From e18fcbd93a842394a132eca8d73a6d4cb8a66239 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Wed, 7 Oct 2026 01:31:36 +0800 Subject: [PATCH 01/25] feature(physical-plan): add NullAwareMode to SortMergeJoinExec --- .../src/joins/sort_merge_join/exec.rs | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 6b345c86d903c..8507c08648156 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -51,6 +51,67 @@ use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; use futures::StreamExt; + +/// Null-aware (`NOT IN`) semantics of a soft merge join, derived from +/// [`SortMergeJoinExec::null_aware`] and the join type. +/// +/// Only these combinations are legal (see [`Self::try_new`]), so the +/// stream matches on this instead of re-checking `null_aware && join_type == ..`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum NullAwareMode { + /// `left.key NOT IN (right.key)`: emits left rows. + LeftAnti +} + +impl NullAwareMode { + /// Validates that `null_aware` may be set for this join and returns its mode. + pub(super) fn try_new( + join_type: JoinType, + num_keys: usize, + has_filter: bool, + null_equality: NullEquality, + sort_options: &[SortOptions], + ) -> Result { + if num_keys != 1 || has_filter { + return plan_err!( + "null-aware SortMergeJoin requires one key and no filter" + ); + } + + // NULL keys must not match each other. + // The null-aware execution path handles UNKNOWN results separately. + // UNKNOWN happens when `WHERE x NOT IN (NULL)` + if null_equality != NullEquality::NullEqualsNothing { + return plan_err!( + "null-aware SortMergeJoin requires NullEqualsNothing" + ); + } + + let mode = match join_type { + JoinType::LeftAnti => Self::LeftAnti, + _ => { + return plan_err!( + "null-aware SortMergeJoin only supports LeftAnti, got {join_type}" + ); + } + }; + + match mode { + Self::LeftAnti => { + // With a single right-side partition and NULLS FIRST, + // the first row reveals whether the subquery contains NULL. + if sort_options.len() != 1 || !sort_options[0].nulls_first { + return plan_err!( + "null-aware {join_type} SortMergeJoin requires NULLS FIRST ordering" + ); + } + Ok(mode) + } + } + } +} + + /// Join execution plan that executes equi-join predicates on multiple partitions using Sort-Merge /// join algorithm and applies an optional filter post join. Can be used to join arbitrarily large /// inputs where one or both of the inputs don't fit in the available memory. From 7c9359154dd72944cdac35e7f982893573c20193 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 17:10:46 +0800 Subject: [PATCH 02/25] chore(physical-plain): add null_aware to SortMergeJoinExec --- .../src/joins/sort_merge_join/exec.rs | 39 +++++++++++---- .../src/joins/sort_merge_join/tests.rs | 50 +++++++++++++++++-- .../proto-models/proto/datafusion.proto | 1 + .../proto-models/src/generated/pbjson.rs | 18 +++++++ .../proto-models/src/generated/prost.rs | 2 + 5 files changed, 98 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 8507c08648156..9985fdcb90f7a 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -51,7 +51,6 @@ use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; use futures::StreamExt; - /// Null-aware (`NOT IN`) semantics of a soft merge join, derived from /// [`SortMergeJoinExec::null_aware`] and the join type. /// @@ -60,7 +59,7 @@ use futures::StreamExt; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum NullAwareMode { /// `left.key NOT IN (right.key)`: emits left rows. - LeftAnti + LeftAnti, } impl NullAwareMode { @@ -73,18 +72,14 @@ impl NullAwareMode { sort_options: &[SortOptions], ) -> Result { if num_keys != 1 || has_filter { - return plan_err!( - "null-aware SortMergeJoin requires one key and no filter" - ); + return plan_err!("null-aware SortMergeJoin requires one key and no filter"); } // NULL keys must not match each other. // The null-aware execution path handles UNKNOWN results separately. // UNKNOWN happens when `WHERE x NOT IN (NULL)` if null_equality != NullEquality::NullEqualsNothing { - return plan_err!( - "null-aware SortMergeJoin requires NullEqualsNothing" - ); + return plan_err!("null-aware SortMergeJoin requires NullEqualsNothing"); } let mode = match join_type { @@ -111,7 +106,6 @@ impl NullAwareMode { } } - /// Join execution plan that executes equi-join predicates on multiple partitions using Sort-Merge /// join algorithm and applies an optional filter post join. Can be used to join arbitrarily large /// inputs where one or both of the inputs don't fit in the available memory. @@ -182,6 +176,19 @@ pub struct SortMergeJoinExec { pub sort_options: Vec, /// Defines the null equality for the join. pub null_equality: NullEquality, + /// Flag to indicate if this join uses null-aware equality semantics. + /// + /// Currently supported only for uncorrelated `LeftAnti` joins with one + /// join key, no join filter, and [`NullEquality::NullEqualsNothing`]. + /// + /// Both inputs must have a single partition and be sorted with NULLS FIRST. + /// This lets the join determine whether the entire right input is empty + /// or contains a NULL key before emitting any left rows. + /// + /// An empty right input preserves all left rows, including NULL keys. + /// A right input containing a NULL key produces no rows. + /// Otherwise, only unmatched left rows with non-NULL keys are emitted. + pub null_aware: bool, /// The columns of `schema` to emit, in order. `None` emits all of them. pub projection: Option>, /// Cache holding plan properties like equivalences, output partitioning etc. @@ -193,6 +200,7 @@ impl SortMergeJoinExec { /// The inputs are sorted using `sort_options` are applied to the columns in the `on` /// # Error /// This function errors when it is not possible to join the left and right sides on keys `on`. + #[expect(clippy::too_many_arguments)] pub fn try_new( left: Arc, right: Arc, @@ -201,6 +209,7 @@ impl SortMergeJoinExec { join_type: JoinType, sort_options: Vec, null_equality: NullEquality, + null_aware: bool, ) -> Result { let left_schema = left.schema(); let right_schema = right.schema(); @@ -256,6 +265,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, + null_aware, projection: None, cache: Arc::new(cache), }) @@ -424,6 +434,7 @@ impl SortMergeJoinExec { self.join_type().swap(), self.sort_options.clone(), self.null_equality, + self.null_aware, )? .with_projection(swap_join_projection( left.schema().fields().len(), @@ -603,6 +614,7 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, + self.null_aware, )? .with_projection(self.projection.clone())?, )), @@ -755,6 +767,7 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, + self.null_aware, )?))) } else { try_embed_projection(projection, self) @@ -778,6 +791,7 @@ impl ExecutionPlan for SortMergeJoinExec { join_type, sort_options, null_equality, + null_aware, projection, // derived from the children's schemas by `try_new` on decode schema: _, @@ -805,6 +819,7 @@ impl ExecutionPlan for SortMergeJoinExec { let join_type = crate::joins::proto::join_type_to_proto(*join_type); let null_equality = crate::joins::proto::null_equality_to_proto(*null_equality); + let null_aware = *null_aware; let filter = filter .as_ref() .map(|filter| crate::joins::proto::join_filter_to_proto(filter, ctx)) @@ -843,6 +858,7 @@ impl ExecutionPlan for SortMergeJoinExec { indices.iter().map(|index| *index as u32).collect() } }, + null_aware, }, )), ), @@ -879,6 +895,7 @@ impl SortMergeJoinExec { sort_options, null_equality, projection, + null_aware, } = &**sort_join; let left = @@ -912,6 +929,9 @@ impl SortMergeJoinExec { *null_equality, "SortMergeJoinExec", )?; + + let null_aware = *null_aware; + let filter = filter .as_ref() .map(|filter| { @@ -946,6 +966,7 @@ impl SortMergeJoinExec { join_type, sort_options, null_equality, + null_aware, )? .with_projection(projection)?, )) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 207113a7338db..e111c420d4d64 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -225,6 +225,7 @@ fn join( join_type: JoinType, ) -> Result { let sort_options = vec![SortOptions::default(); on.len()]; + let null_aware = is_join_type_null_aware(join_type); SortMergeJoinExec::try_new( left, right, @@ -233,6 +234,7 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, + null_aware, ) } @@ -244,6 +246,7 @@ fn join_with_options( sort_options: Vec, null_equality: NullEquality, ) -> Result { + let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -252,6 +255,7 @@ fn join_with_options( join_type, sort_options, null_equality, + null_aware, ) } @@ -264,6 +268,7 @@ fn join_with_filter( sort_options: Vec, null_equality: NullEquality, ) -> Result { + let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -272,6 +277,7 @@ fn join_with_filter( join_type, sort_options, null_equality, + null_aware, ) } @@ -371,6 +377,7 @@ fn join_and_projection_for_pushdown( Arc::new(Column::new("b1", 1)) as _, Arc::new(Column::new("b2", 1)) as _, )]; + let null_aware = false; let join = Arc::new(SortMergeJoinExec::try_new( left, right, @@ -379,6 +386,7 @@ fn join_and_projection_for_pushdown( Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )?); let input: Arc = Arc::clone(&join) as _; let projection = ProjectionExec::try_new( @@ -406,6 +414,10 @@ fn join_and_projection_for_pushdown( Ok((join, projection)) } +fn is_join_type_null_aware(join_type: JoinType) -> bool { + matches!(join_type, LeftAnti) +} + /// `c1 < c2`, referencing the last column of each child. fn filter_for_pushdown() -> JoinFilter { JoinFilter::new( @@ -4579,6 +4591,8 @@ async fn join_get_stream_and_get_expected( prepare_record_batches_for_cmp(batches) }; + let null_aware = is_join_type_null_aware(LeftAnti); + let join = SortMergeJoinExec::try_new( left, right, @@ -4587,6 +4601,7 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, + null_aware, )?; let stream = join.execute(0, task_ctx)?; @@ -6486,6 +6501,7 @@ async fn bitwise_spill_with_filter() -> Result<()> { .with_runtime(Arc::clone(&runtime)), ); + let null_aware = is_join_type_null_aware(LeftAnti); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6494,6 +6510,7 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -6535,6 +6552,8 @@ async fn bitwise_spill_with_filter() -> Result<()> { let task_ctx_no_spill = Arc::new( TaskContext::default().with_session_config(session_config.clone()), ); + + let null_aware = is_join_type_null_aware(LeftAnti); let join_no_spill = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6543,6 +6562,7 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -6596,6 +6616,7 @@ async fn bitwise_filtered_no_spill() -> Result<()> { .with_runtime(Arc::clone(&runtime)), ); + let null_aware = is_join_type_null_aware(LeftAnti); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6604,6 +6625,7 @@ async fn bitwise_filtered_no_spill() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join.execute(0, task_ctx)?; let err = common::collect(stream).await.unwrap_err(); @@ -6881,6 +6903,7 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { (&left, &right) }; + let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( Arc::clone(join_left), Arc::clone(join_right), @@ -6889,6 +6912,7 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -6986,14 +7010,17 @@ async fn bitwise_spill_null_key_group() -> Result<()> { .with_runtime(runtime), ); + let join_type = LeftSemi; + let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, on, Some(filter), - LeftSemi, + join_type, sort_options, NullEquality::NullEqualsNull, + null_aware, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -7656,6 +7683,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { let task_ctx = Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7664,6 +7692,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -7675,6 +7704,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { ); let task_ctx_no_spill = Arc::new(TaskContext::default()); + let null_aware = is_join_type_null_aware(join_type); let join_no_spill = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7683,6 +7713,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + null_aware, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -7714,15 +7745,19 @@ async fn swap_inputs_swaps_the_projection() -> Result<()> { Arc::new(Column::new("a1", 0)) as _, Arc::new(Column::new("a2", 0)) as _, )]; + + let join_type = Inner; + let null_aware = is_join_type_null_aware(join_type); // One column from each side, in an order that tells the two sides apart. let join = SortMergeJoinExec::try_new( left, right, on, None, - Inner, + join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )? .with_projection(Some(vec![4, 2]))?; @@ -7786,6 +7821,8 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { for join_type in [ Inner, Left, Right, Full, LeftSemi, LeftAnti, RightSemi, RightAnti, ] { + let null_aware = is_join_type_null_aware(join_type); + let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7794,6 +7831,7 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )?; let swapped = join.swap_inputs()?; @@ -7835,14 +7873,18 @@ async fn an_empty_projection_keeps_the_rows() -> Result<()> { Arc::new(Column::new("a1", 0)) as _, Arc::new(Column::new("a2", 0)) as _, )]; + + let join_type = Inner; + let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, on, None, - Inner, + join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )? .with_projection(Some(vec![]))?; @@ -7925,6 +7967,7 @@ async fn collect_streamed_keys( Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)])), ); + let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, @@ -7933,6 +7976,7 @@ async fn collect_streamed_keys( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )?; // A small batch size keeps the gate firing often enough to interleave the diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 6c4ef8f0a4b2f..16338bc7a3df0 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1835,6 +1835,7 @@ message SortMergeJoinExecNode { repeated SortExprNode sort_options = 6; datafusion_common.NullEquality null_equality = 7; repeated uint32 projection = 8; + bool null_aware = 9; } message PiecewiseMergeJoinExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 725cceeb75970..34f2a293ae816 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -27666,6 +27666,9 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { len += 1; } + if self.null_aware { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.SortMergeJoinExecNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -27695,6 +27698,9 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { struct_ser.serialize_field("projection", &self.projection)?; } + if self.null_aware { + struct_ser.serialize_field("nullAware", &self.null_aware)?; + } struct_ser.end() } } @@ -27716,6 +27722,8 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "null_equality", "nullEquality", "projection", + "null_aware", + "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -27728,6 +27736,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { SortOptions, NullEquality, Projection, + NullAware, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -27757,6 +27766,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "sortOptions" | "sort_options" => Ok(GeneratedField::SortOptions), "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), "projection" => Ok(GeneratedField::Projection), + "nullAware" | "null_aware" => Ok(GeneratedField::NullAware), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -27784,6 +27794,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { let mut sort_options__ = None; let mut null_equality__ = None; let mut projection__ = None; + let mut null_aware__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -27837,6 +27848,12 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { .into_iter().map(|x| x.0).collect()) ; } + GeneratedField::NullAware => { + if null_aware__.is_some() { + return Err(serde::de::Error::duplicate_field("nullAware")); + } + null_aware__ = Some(map_.next_value()?); + } } } Ok(SortMergeJoinExecNode { @@ -27848,6 +27865,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { sort_options: sort_options__.unwrap_or_default(), null_equality: null_equality__.unwrap_or_default(), projection: projection__.unwrap_or_default(), + null_aware: null_aware__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1034d0443dbce..1583244ea7e5a 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2760,6 +2760,8 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, #[prost(uint32, repeated, tag = "8")] pub projection: ::prost::alloc::vec::Vec, + #[prost(bool, tag = "9")] + pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PiecewiseMergeJoinExecNode { From 61e3508eae0c1a60e9730a00b266f48149c41071 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 17:36:11 +0800 Subject: [PATCH 03/25] chore(physical-plan): SortMergeJoin add null_aware_mode --- .../src/joins/sort_merge_join/exec.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 9985fdcb90f7a..ea8a28c211ee9 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -364,6 +364,11 @@ impl SortMergeJoinExec { self.null_equality } + /// Null aware + pub fn null_aware(&self) -> bool { + self.null_aware + } + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. fn compute_properties( left: &Arc, @@ -460,6 +465,23 @@ impl SortMergeJoinExec { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) } } + + /// The null-aware semantics of this join, if [`Self::null_aware`] is set. + /// + /// Errors if `null_aware` is set on a join that does not support it. + pub(super) fn null_aware_mode(&self) -> Result> { + self.null_aware + .then(|| { + NullAwareMode::try_new( + self.join_type, + self.on.len(), + self.filter.is_some(), + self.null_equality(), + self.sort_options(), + ) + }) + .transpose() + } } impl EmbeddedProjection for SortMergeJoinExec { From db9a428c4059a1777f6c318470cbdb1eaaa8264f Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 18:12:28 +0800 Subject: [PATCH 04/25] chore(physical-play): SortMergeJoinExec add SinglePartition to input_distribution_requirements when null_aware is true --- .../physical-plan/src/joins/sort_merge_join/exec.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index ea8a28c211ee9..c011858e3a7cd 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -574,6 +574,18 @@ impl ExecutionPlan for SortMergeJoinExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { + // Null-aware NOT IN needs facts about the entire right input. + // Require a single partition on both sides to use the existing + // sort-merge join's partition pairing without cross-partition coordination. + // + // TODO: Allow multiple left partitions by sharing the right-side summary + // and providing each left partition with the right input for matching. + if self.null_aware { + return InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]); + } let (left_expr, right_expr) = self .on .iter() From 1bdd52df6a317050b1b5c030a3af5979553eaa21 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 17:27:12 +0800 Subject: [PATCH 05/25] chore(physical-plan): checkpoint null-aware sort-merge join stream setup Pass the validated null-aware mode to the bitwise stream and initialize the subquery state. Validate single-partition inputs at execution time. --- .../joins/sort_merge_join/bitwise_stream.rs | 38 +++++++++++++++++++ .../src/joins/sort_merge_join/exec.rs | 22 ++++++++++- 2 files changed, 58 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 5c1dd58d51c5f..1a80d765c6170 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -110,6 +110,7 @@ use std::cmp::Ordering; use std::sync::Arc; use crate::EmptyRecordBatchStream; +use crate::joins::sort_merge_join::exec::NullAwareMode; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time, @@ -187,6 +188,22 @@ fn find_key_group_end(cmp: &JoinKeyComparator, from: usize, len: usize) -> usize lo } +/// State of the right-side subquery used by a null-aware `LeftAnti` join. +/// +/// Determined before emitting left rows, using a single right partition +/// sorted with NULLS FIRST. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SubqueryState { + /// The right input has not been inspected yet. + Pending, + /// The right input contains no rows. + Empty, + /// The right input contains rows but no NULL keys. + NonEmptyNoNull, + /// The right input contains at least one NULL key. + HasNull, +} + /// Sort-Merge join stream for Semi/Anti/Mark joins. /// /// Named "bitwise" because it tracks outer-row matches via a per-batch @@ -235,6 +252,8 @@ pub(crate) struct BitwiseSortMergeJoinStream { filter: Option, sort_options: Vec, null_equality: NullEquality, + null_aware: Option, + subquery_state: SubqueryState, // Decomposed from JoinType: when RightSemi/RightAnti, outer=right, // inner=left, so we swap sides when building the filter batch. outer_is_left: bool, @@ -279,6 +298,7 @@ impl BitwiseSortMergeJoinStream { schema: SchemaRef, sort_options: Vec, null_equality: NullEquality, + null_aware: Option, outer: SendableRecordBatchStream, inner: SendableRecordBatchStream, on_outer: Vec, @@ -317,6 +337,22 @@ impl BitwiseSortMergeJoinStream { let peak_mem_used = MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); + if let Some(mode) = null_aware { + let validated = NullAwareMode::try_new( + join_type, + on_outer.len(), + filter.is_some(), + null_equality, + &sort_options, + )?; + + if mode != validated { + return internal_err!( + "null-aware mode does not match the join configuration" + ); + } + } + let mut state = Self { join_type, outer, @@ -334,6 +370,8 @@ impl BitwiseSortMergeJoinStream { filter, sort_options, null_equality, + null_aware, + subquery_state: SubqueryState::Pending, outer_is_left, coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size) .with_biggest_coalesce_batch_size(Some(batch_size / 2)), diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index c011858e3a7cd..f5b8cba5fb8ce 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -57,14 +57,14 @@ use futures::StreamExt; /// Only these combinations are legal (see [`Self::try_new`]), so the /// stream matches on this instead of re-checking `null_aware && join_type == ..`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum NullAwareMode { +pub(crate) enum NullAwareMode { /// `left.key NOT IN (right.key)`: emits left rows. LeftAnti, } impl NullAwareMode { /// Validates that `null_aware` may be set for this join and returns its mode. - pub(super) fn try_new( + pub(crate) fn try_new( join_type: JoinType, num_keys: usize, has_filter: bool, @@ -690,6 +690,20 @@ impl ExecutionPlan for SortMergeJoinExec { "Invalid SortMergeJoinExec, partition count mismatch {left_partitions}!={right_partitions},\ consider using RepartitionExec" ); + + let null_aware = self.null_aware_mode()?; + + // Null-aware execution requires a single partition on both sides. + // Validate this at execution time because callers may bypass the + // optimizer or modify the public null_aware flag after planning. + if null_aware.is_some() + && (partition != 0 || left_partitions != 1 || right_partitions != 1) + { + return internal_err!( + "null-aware SortMergeJoin requires partition 0 and single-partition inputs" + ); + } + // execute children plans let left = self.left.execute(partition, Arc::clone(&context))?; let right = self.right.execute(partition, Arc::clone(&context))?; @@ -700,6 +714,7 @@ impl ExecutionPlan for SortMergeJoinExec { schema: Arc::clone(&self.schema), sort_options: self.sort_options.clone(), null_equality: self.null_equality, + null_aware, left, right, on_left, @@ -1015,6 +1030,7 @@ pub(crate) struct SortMergeJoinInputs { /// Sort options of the join keys, one per key, that both inputs are sorted with pub(crate) sort_options: Vec, pub(crate) null_equality: NullEquality, + pub(crate) null_aware: Option, /// Left input, sorted on `on_left` with `sort_options` pub(crate) left: SendableRecordBatchStream, /// Right input, sorted on `on_right` with `sort_options` @@ -1043,6 +1059,7 @@ pub(crate) fn sort_merge_join_stream( schema, sort_options, null_equality, + null_aware, left, right, on_left, @@ -1086,6 +1103,7 @@ pub(crate) fn sort_merge_join_stream( schema, sort_options, null_equality, + null_aware, streamed, buffered, on_streamed, From b1fa1225f5d1975c851558aa89183201272f6ef8 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 17:59:58 +0800 Subject: [PATCH 06/25] fix(physical-plan): preserve sort-merge join API compatibility Restore the seven-argument SortMergeJoinExec::try_new constructor and enable null-aware semantics through the validated with_null_aware builder. Preserve the flag when rebuilding plans. Restore the original protobuf schema and generated types. Reject serialization of null-aware sort-merge joins rather than silently losing their semantics. Restore ordinary join test configuration and update direct stream callers to pass no null-aware mode. --- .../src/joins/sort_merge_join/exec.rs | 57 ++++++++++++------- .../src/joins/sort_merge_join/tests.rs | 57 ++++--------------- .../proto-models/proto/datafusion.proto | 1 - .../proto-models/src/generated/pbjson.rs | 18 ------ .../proto-models/src/generated/prost.rs | 2 - 5 files changed, 45 insertions(+), 90 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index f5b8cba5fb8ce..ebc7d9736561d 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -200,7 +200,6 @@ impl SortMergeJoinExec { /// The inputs are sorted using `sort_options` are applied to the columns in the `on` /// # Error /// This function errors when it is not possible to join the left and right sides on keys `on`. - #[expect(clippy::too_many_arguments)] pub fn try_new( left: Arc, right: Arc, @@ -209,7 +208,6 @@ impl SortMergeJoinExec { join_type: JoinType, sort_options: Vec, null_equality: NullEquality, - null_aware: bool, ) -> Result { let left_schema = left.schema(); let right_schema = right.schema(); @@ -265,12 +263,23 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, - null_aware, + null_aware: false, projection: None, cache: Arc::new(cache), }) } + /// Enables or disables null-aware (`NOT IN`) semantics. + /// + /// Enabling requires an uncorrelated `LeftAnti` join with one key, no + /// filter, [`NullEquality::NullEqualsNothing`], and NULLS FIRST ordering. + /// Both inputs must have a single partition when executed. + pub fn with_null_aware(mut self, null_aware: bool) -> Result { + self.null_aware = null_aware; + self.null_aware_mode()?; + Ok(self) + } + /// Returns this join emitting only the columns in `projection`, in that order. /// The indices address the join's own schema, before any projection. pub fn with_projection(&self, projection: Option>) -> Result { @@ -439,8 +448,8 @@ impl SortMergeJoinExec { self.join_type().swap(), self.sort_options.clone(), self.null_equality, - self.null_aware, )? + .with_null_aware(self.null_aware)? .with_projection(swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), @@ -648,8 +657,8 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, - self.null_aware, )? + .with_null_aware(self.null_aware)? .with_projection(self.projection.clone())?, )), _ => internal_err!("SortMergeJoin wrong number of children"), @@ -808,16 +817,18 @@ impl ExecutionPlan for SortMergeJoinExec { self.filter().as_ref(), &column_indices, )? { - Ok(Some(Arc::new(SortMergeJoinExec::try_new( - Arc::new(projected_left_child), - Arc::new(projected_right_child), - join_on, - join_filter, - self.join_type, - self.sort_options.clone(), - self.null_equality, - self.null_aware, - )?))) + Ok(Some(Arc::new( + SortMergeJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_on, + join_filter, + self.join_type, + self.sort_options.clone(), + self.null_equality, + )? + .with_null_aware(self.null_aware)?, + ))) } else { try_embed_projection(projection, self) } @@ -830,6 +841,14 @@ impl ExecutionPlan for SortMergeJoinExec { ) -> Result> { use datafusion_proto_models::protobuf; + // The existing protobuf node cannot represent null-aware semantics. + // Reject encoding rather than silently decoding as an ordinary anti join. + if self.null_aware { + return datafusion_common::not_impl_err!( + "Serializing null-aware SortMergeJoinExec is not supported" + ); + } + // Destructure exhaustively (no `..`) so that a newly added field is a // compile error here instead of being silently left out of the proto. let Self { @@ -840,7 +859,7 @@ impl ExecutionPlan for SortMergeJoinExec { join_type, sort_options, null_equality, - null_aware, + null_aware: _, projection, // derived from the children's schemas by `try_new` on decode schema: _, @@ -868,7 +887,6 @@ impl ExecutionPlan for SortMergeJoinExec { let join_type = crate::joins::proto::join_type_to_proto(*join_type); let null_equality = crate::joins::proto::null_equality_to_proto(*null_equality); - let null_aware = *null_aware; let filter = filter .as_ref() .map(|filter| crate::joins::proto::join_filter_to_proto(filter, ctx)) @@ -907,7 +925,6 @@ impl ExecutionPlan for SortMergeJoinExec { indices.iter().map(|index| *index as u32).collect() } }, - null_aware, }, )), ), @@ -944,7 +961,6 @@ impl SortMergeJoinExec { sort_options, null_equality, projection, - null_aware, } = &**sort_join; let left = @@ -979,8 +995,6 @@ impl SortMergeJoinExec { "SortMergeJoinExec", )?; - let null_aware = *null_aware; - let filter = filter .as_ref() .map(|filter| { @@ -1015,7 +1029,6 @@ impl SortMergeJoinExec { join_type, sort_options, null_equality, - null_aware, )? .with_projection(projection)?, )) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index e111c420d4d64..1106e7c7c7be1 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -225,7 +225,6 @@ fn join( join_type: JoinType, ) -> Result { let sort_options = vec![SortOptions::default(); on.len()]; - let null_aware = is_join_type_null_aware(join_type); SortMergeJoinExec::try_new( left, right, @@ -234,7 +233,6 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, - null_aware, ) } @@ -246,7 +244,6 @@ fn join_with_options( sort_options: Vec, null_equality: NullEquality, ) -> Result { - let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -255,7 +252,6 @@ fn join_with_options( join_type, sort_options, null_equality, - null_aware, ) } @@ -268,7 +264,6 @@ fn join_with_filter( sort_options: Vec, null_equality: NullEquality, ) -> Result { - let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -277,7 +272,6 @@ fn join_with_filter( join_type, sort_options, null_equality, - null_aware, ) } @@ -377,7 +371,6 @@ fn join_and_projection_for_pushdown( Arc::new(Column::new("b1", 1)) as _, Arc::new(Column::new("b2", 1)) as _, )]; - let null_aware = false; let join = Arc::new(SortMergeJoinExec::try_new( left, right, @@ -386,7 +379,6 @@ fn join_and_projection_for_pushdown( Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )?); let input: Arc = Arc::clone(&join) as _; let projection = ProjectionExec::try_new( @@ -414,10 +406,6 @@ fn join_and_projection_for_pushdown( Ok((join, projection)) } -fn is_join_type_null_aware(join_type: JoinType) -> bool { - matches!(join_type, LeftAnti) -} - /// `c1 < c2`, referencing the last column of each child. fn filter_for_pushdown() -> JoinFilter { JoinFilter::new( @@ -4591,8 +4579,6 @@ async fn join_get_stream_and_get_expected( prepare_record_batches_for_cmp(batches) }; - let null_aware = is_join_type_null_aware(LeftAnti); - let join = SortMergeJoinExec::try_new( left, right, @@ -4601,7 +4587,6 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, - null_aware, )?; let stream = join.execute(0, task_ctx)?; @@ -5616,6 +5601,7 @@ async fn bitwise_emits_completed_batches_before_pending_outer( Arc::clone(&schema), vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, vec![Arc::new(Column::new("key", 1))], @@ -5715,6 +5701,7 @@ fn join_time_test_join( outer_schema, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, vec![Arc::new(Column::new("b1", 1)) as PhysicalExprRef], @@ -6100,6 +6087,7 @@ async fn filter_buffer_pending_loses_inner_rows() -> Result<()> { left_schema, // output schema = outer schema for semi vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer, @@ -6194,6 +6182,7 @@ async fn no_filter_boundary_pending_loses_outer_rows() -> Result<()> { left_schema, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer, @@ -6285,6 +6274,7 @@ async fn no_filter_boundary_pending_with_unmatched_prefix() -> Result<()> { Arc::clone(&left_schema), vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer.clone(), @@ -6405,6 +6395,7 @@ async fn filtered_boundary_pending_outer_rows() -> Result<()> { left_schema, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer, @@ -6501,7 +6492,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { .with_runtime(Arc::clone(&runtime)), ); - let null_aware = is_join_type_null_aware(LeftAnti); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6510,7 +6500,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -6552,8 +6541,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { let task_ctx_no_spill = Arc::new( TaskContext::default().with_session_config(session_config.clone()), ); - - let null_aware = is_join_type_null_aware(LeftAnti); let join_no_spill = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6562,7 +6549,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -6616,7 +6602,6 @@ async fn bitwise_filtered_no_spill() -> Result<()> { .with_runtime(Arc::clone(&runtime)), ); - let null_aware = is_join_type_null_aware(LeftAnti); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -6625,7 +6610,6 @@ async fn bitwise_filtered_no_spill() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join.execute(0, task_ctx)?; let err = common::collect(stream).await.unwrap_err(); @@ -6903,7 +6887,6 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { (&left, &right) }; - let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( Arc::clone(join_left), Arc::clone(join_right), @@ -6912,7 +6895,6 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -7010,17 +6992,14 @@ async fn bitwise_spill_null_key_group() -> Result<()> { .with_runtime(runtime), ); - let join_type = LeftSemi; - let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, on, Some(filter), - join_type, + LeftSemi, sort_options, NullEquality::NullEqualsNull, - null_aware, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -7162,6 +7141,7 @@ async fn spill_filtered_boundary_loses_outer_rows() -> Result<()> { Arc::clone(&left_schema), vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer.clone(), @@ -7683,7 +7663,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { let task_ctx = Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); - let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7692,7 +7671,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -7704,7 +7682,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { ); let task_ctx_no_spill = Arc::new(TaskContext::default()); - let null_aware = is_join_type_null_aware(join_type); let join_no_spill = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7713,7 +7690,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - null_aware, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -7745,19 +7721,15 @@ async fn swap_inputs_swaps_the_projection() -> Result<()> { Arc::new(Column::new("a1", 0)) as _, Arc::new(Column::new("a2", 0)) as _, )]; - - let join_type = Inner; - let null_aware = is_join_type_null_aware(join_type); // One column from each side, in an order that tells the two sides apart. let join = SortMergeJoinExec::try_new( left, right, on, None, - join_type, + Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )? .with_projection(Some(vec![4, 2]))?; @@ -7821,8 +7793,6 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { for join_type in [ Inner, Left, Right, Full, LeftSemi, LeftAnti, RightSemi, RightAnti, ] { - let null_aware = is_join_type_null_aware(join_type); - let join = SortMergeJoinExec::try_new( Arc::clone(&left), Arc::clone(&right), @@ -7831,7 +7801,6 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )?; let swapped = join.swap_inputs()?; @@ -7873,18 +7842,14 @@ async fn an_empty_projection_keeps_the_rows() -> Result<()> { Arc::new(Column::new("a1", 0)) as _, Arc::new(Column::new("a2", 0)) as _, )]; - - let join_type = Inner; - let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, on, None, - join_type, + Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )? .with_projection(Some(vec![]))?; @@ -7967,7 +7932,6 @@ async fn collect_streamed_keys( Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)])), ); - let null_aware = is_join_type_null_aware(join_type); let join = SortMergeJoinExec::try_new( left, right, @@ -7976,7 +7940,6 @@ async fn collect_streamed_keys( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )?; // A small batch size keeps the gate firing often enough to interleave the diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 16338bc7a3df0..6c4ef8f0a4b2f 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1835,7 +1835,6 @@ message SortMergeJoinExecNode { repeated SortExprNode sort_options = 6; datafusion_common.NullEquality null_equality = 7; repeated uint32 projection = 8; - bool null_aware = 9; } message PiecewiseMergeJoinExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 34f2a293ae816..725cceeb75970 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -27666,9 +27666,6 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { len += 1; } - if self.null_aware { - len += 1; - } let mut struct_ser = serializer.serialize_struct("datafusion.SortMergeJoinExecNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -27698,9 +27695,6 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { struct_ser.serialize_field("projection", &self.projection)?; } - if self.null_aware { - struct_ser.serialize_field("nullAware", &self.null_aware)?; - } struct_ser.end() } } @@ -27722,8 +27716,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "null_equality", "nullEquality", "projection", - "null_aware", - "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -27736,7 +27728,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { SortOptions, NullEquality, Projection, - NullAware, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -27766,7 +27757,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "sortOptions" | "sort_options" => Ok(GeneratedField::SortOptions), "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), "projection" => Ok(GeneratedField::Projection), - "nullAware" | "null_aware" => Ok(GeneratedField::NullAware), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -27794,7 +27784,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { let mut sort_options__ = None; let mut null_equality__ = None; let mut projection__ = None; - let mut null_aware__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -27848,12 +27837,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { .into_iter().map(|x| x.0).collect()) ; } - GeneratedField::NullAware => { - if null_aware__.is_some() { - return Err(serde::de::Error::duplicate_field("nullAware")); - } - null_aware__ = Some(map_.next_value()?); - } } } Ok(SortMergeJoinExecNode { @@ -27865,7 +27848,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { sort_options: sort_options__.unwrap_or_default(), null_equality: null_equality__.unwrap_or_default(), projection: projection__.unwrap_or_default(), - null_aware: null_aware__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1583244ea7e5a..1034d0443dbce 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2760,8 +2760,6 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, #[prost(uint32, repeated, tag = "8")] pub projection: ::prost::alloc::vec::Vec, - #[prost(bool, tag = "9")] - pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PiecewiseMergeJoinExecNode { From e6c78bbd8e2c58c5a321d03dfb39e3679d6260da Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 22:32:01 +0800 Subject: [PATCH 07/25] feat(physical-plan): initialize null-aware sort-merge join state Classify the inner input as empty, containing NULL keys, or nonempty without NULL keys before the merge scan. Retain the first nonempty batch at offset zero for matching. Finish null-aware LeftAnti streams without emitting rows when the inner input contains a NULL key. --- .../joins/sort_merge_join/bitwise_stream.rs | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 1a80d765c6170..74d9fac86925d 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -1130,12 +1130,55 @@ impl BitwiseSortMergeJoinStream { } } + /// Initializes the inner input's empty/NULL summary for null-aware joins. + async fn initialized_null_aware(&mut self) -> Result<()> { + if self.subquery_state != SubqueryState::Pending { + return Ok(()); + } + + match self.null_aware { + None => return Ok(()), + Some(NullAwareMode::LeftAnti) => {} + } + + if !self.next_inner_batch().await? { + self.subquery_state = SubqueryState::Empty; + return Ok(()); + } + + // Keep the first nonempty inner batch at offset 0 for the merge scan. + // The complete right input need to be sorted with NULLS FIRST. + let first_batch_has_null = self.inner_key_arrays[0] + .logical_nulls() + .is_some_and(|null_keys| null_keys.is_null(0)); + + self.subquery_state = if first_batch_has_null { + SubqueryState::HasNull + } else { + SubqueryState::NonEmptyNoNull + }; + + Ok(()) + } + /// Main loop: a classic merge-scan over the two sorted inputs, emitting /// output batches as they complete. async fn join( &mut self, emitter: &mut TryEmitter, ) -> Result<()> { + self.initialized_null_aware().await?; + match self.null_aware { + None => {} + Some(NullAwareMode::LeftAnti) => { + // NOT IN cannot be TRUE for any outer row when the inner + // input contains a NULL key. (following NOT IN semetic) + if self.subquery_state == SubqueryState::HasNull { + return Ok(()); + } + } + } + // The `has_current_*` / `has_completed_batch` fast paths keep async // state machinery out of the per-key-group hot path; the awaiting // helpers are only entered at batch boundaries. From 80d183dd005bbfec4113a46c6d2c7bfaaa8f4311 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sun, 11 Oct 2026 22:12:36 +0800 Subject: [PATCH 08/25] feat(physical-plan): route null-aware left anti joins to sort merge Choose null-aware SortMergeJoinExec for single-key LeftAnti joins without a filter when hash joins are not preferred. Require NullEqualsNothing and request ASC NULLS FIRST ordering. Apply an outer-key validity mask to unmatched rows when the inner input is nonempty without NULL keys. Use logical_nulls to account for dictionary values that are NULL. Preserve unmatched rows when the inner input is empty. --- datafusion/core/src/physical_planner.rs | 21 +++++++++- .../joins/sort_merge_join/bitwise_stream.rs | 38 ++++++++++++++++++- 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 198e5bb7412b2..6ed8010d900f7 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -65,7 +65,7 @@ use datafusion_common::{ internal_err, not_impl_err, plan_err, }; use datafusion_common::{ - TableReference, assert_eq_or_internal_err, assert_or_internal_err, + NullEquality, TableReference, assert_eq_or_internal_err, assert_or_internal_err, }; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::memory::MemorySourceConfig; @@ -1863,6 +1863,25 @@ impl DefaultPhysicalPlanner { None, )?) } + } else if *null_aware + && !prefer_hash_join + && *join_type == JoinType::LeftAnti + && join_on.len() == 1 + && join_filter.is_none() + && *null_equality == NullEquality::NullEqualsNothing + { + Arc::new( + SortMergeJoinExec::try_new( + physical_left, + physical_right, + join_on, + join_filter, + *join_type, + vec![SortOptions::new(false, true)], // ASC NULLS FIRST + *null_equality, + )? + .with_null_aware(true)?, + ) } else if session.config().target_partitions() > 1 && session.config().repartition_joins() && !prefer_hash_join diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 74d9fac86925d..3ad61f059cc78 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -121,7 +121,7 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::array::{ Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch, RecordBatchOptions, }; -use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not}; +use arrow::compute::{BatchCoalescer, SortOptions, and, filter_record_batch, not}; use arrow::datatypes::SchemaRef; use arrow::util::bit_chunk_iterator::UnalignedBitChunk; use arrow::util::bit_util::apply_bitwise_binary_op; @@ -188,6 +188,32 @@ fn find_key_group_end(cmp: &JoinKeyComparator, from: usize, len: usize) -> usize lo } +/// Applies SQL NOT IN semantics to the unmatched outer rows. +fn null_aware_left_anti_selection( + unmatched: BooleanArray, + outer_key: &dyn Array, + subqury_state: SubqueryState, +) -> Result { + match subqury_state { + SubqueryState::Empty => Ok(unmatched), + SubqueryState::NonEmptyNoNull => match outer_key.logical_nulls() { + // All outer keys are non-NULL, so unmatched rows can be kept. + None => Ok(unmatched), + Some(nulls) => { + // With a nonempty RHS containing no NULL keys, NOT IN keeps only + // unmatched outer rows whose keys are non-NULL. Combine the + // unmatched mask with the outer key validity bitmap to exclude + // NULL keys, for which NOT IN evaluates to UNKNOWN. + let valids = BooleanArray::new(nulls.inner().clone(), None); + and(&unmatched, &valids).map_err(DataFusionError::from) + } + }, + SubqueryState::Pending | SubqueryState::HasNull => { + internal_err!("Cannot emit null-aware anti rows before RHS validation") + } + } +} + /// State of the right-side subquery used by a null-aware `LeftAnti` join. /// /// Determined before emitting left rows, using a single right partition @@ -613,7 +639,15 @@ impl BitwiseSortMergeJoinStream { } } JoinType::LeftAnti | JoinType::RightAnti => { - let selection = not(&BooleanArray::new(matched_buf, None))?; + let unmatched = not(&BooleanArray::new(matched_buf, None))?; + let selection = match self.null_aware { + None => unmatched, + Some(NullAwareMode::LeftAnti) => null_aware_left_anti_selection( + unmatched, + self.outer_key_arrays[0].as_ref(), + self.subquery_state, + )?, + }; let filtered = filter_record_batch(&batch, &selection)?; if filtered.num_rows() > 0 { self.coalescer.push_batch(filtered)?; From 9de0382d176ac5ee34fa4a1f4984c720f4681b99 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Wed, 14 Oct 2026 07:53:01 +0800 Subject: [PATCH 09/25] docs(core): clarify null-aware join planner comments Distinguish CollectLeft null-aware hash joins from the supported single-partition LeftAnti sort-merge path. Clarify that remaining null-aware joins retain CollectLeft mode. --- datafusion/core/src/physical_planner.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 6ed8010d900f7..3d5271bb49362 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1721,9 +1721,10 @@ impl DefaultPhysicalPlanner { let prefer_hash_join = session.config_options().optimizer.prefer_hash_join; - // Null-aware joins are pinned to CollectLeft hash joins (see - // `HashJoinExec::null_aware`): never repartition them, and - // never route them to the sort-merge path below. + // Null-aware hash joins require CollectLeft for global NULL tracking, + // so they must not be hash-repartitioned. + // Supported null-aware LeftAnti joins can instead use the + // single-partition sort-merge path below. let can_repartition_join = session.config().target_partitions() > 1 && session.config().repartition_joins() && !*null_aware; @@ -1903,7 +1904,7 @@ impl DefaultPhysicalPlanner { *null_equality, )?) } else { - // Null-aware joins need global probe-side state, so keep + // Remaining null-aware joins need global probe-side state, so keep // them in CollectLeft mode. let partition_mode = if can_repartition_join { PartitionMode::Auto From e692ff82367fd01eab3dd3799e1532664487d806 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 11:16:18 +0800 Subject: [PATCH 10/25] chore(smj): Clarify routing and fix match lint Document the supported null-aware LeftAnti routing and HashJoin fallback. Use a match guard for the RHS NULL early exit without changing its behavior. --- datafusion/core/src/physical_planner.rs | 19 ++++++++++++++----- .../joins/sort_merge_join/bitwise_stream.rs | 9 +++++---- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 3d5271bb49362..881ffd4f55f51 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1871,6 +1871,10 @@ impl DefaultPhysicalPlanner { && join_filter.is_none() && *null_equality == NullEquality::NullEqualsNothing { + // Use null-aware sort-merge join for single-key LeftAnti joins + // without a filter. Single-partition inputs and NULLS FIRST ordering + // let the stream determine whether the entire RHS is empty or + // contains NULL before emitting any rows. Arc::new( SortMergeJoinExec::try_new( physical_left, @@ -1887,12 +1891,17 @@ impl DefaultPhysicalPlanner { && session.config().repartition_joins() && !prefer_hash_join && !*null_aware - // Null-aware joins (e.g. `NOT IN` with a nullable subquery) must - // use the CollectLeft HashJoin below: SortMergeJoinExec does not - // implement null-aware anti-join semantics and would return wrong - // results when the right side contains a null join key. + // Null-aware joins (e.g. `NOT IN` with a nullable subquery) that do not + // satisfy the conditions above must use the CollectLeft HashJoin below. + // SortMergeJoinExec currently supports null-aware semantics only for + // single-key LeftAnti joins without a filter, using NullEqualsNothing. + // Routing other null-aware joins through the ordinary sort-merge path + // would not preserve NOT IN semantics, including when the subquery + // contains a NULL key. { - // Use SortMergeJoin if hash join is not preferred + // Use ordinary sort-merge join when hash join is not preferred. + // Null-aware joins not handled above fall back to CollectLeft + // HashJoin because their semantics are not supported by sort-merge join. let join_on_len = join_on.len(); Arc::new(SortMergeJoinExec::try_new( physical_left, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 3ad61f059cc78..ba0c946490ce5 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -1204,13 +1204,14 @@ impl BitwiseSortMergeJoinStream { self.initialized_null_aware().await?; match self.null_aware { None => {} - Some(NullAwareMode::LeftAnti) => { + Some(NullAwareMode::LeftAnti) + if self.subquery_state == SubqueryState::HasNull => + { // NOT IN cannot be TRUE for any outer row when the inner // input contains a NULL key. (following NOT IN semetic) - if self.subquery_state == SubqueryState::HasNull { - return Ok(()); - } + return Ok(()); } + Some(NullAwareMode::LeftAnti) => {} } // The `has_current_*` / `has_completed_batch` fast paths keep async From 69fa0cbebf937a179d2d009bff700f2fe9f66fa4 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 14:29:29 +0800 Subject: [PATCH 11/25] fix(optimizer): Preserve null-aware SMJ mode Carry the original null-aware flag through both sort-merge join key reordering constructors so rebuilt joins retain NOT IN semantics. --- .../enforce_distribution.rs | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index b847e195c07b3..e29d7005f393e 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -184,17 +184,17 @@ pub fn adjust_input_keys_ordering( requirements.children[1].data = shift_right_required(&requirements.data, left_columns_len) .unwrap_or_default(); - } else if let Some(SortMergeJoinExec { - left, - right, - on, - filter, - join_type, - sort_options, - null_equality, - .. - }) = plan.downcast_ref::() - { + } else if let Some(origin_exec) = plan.downcast_ref::() { + let SortMergeJoinExec { + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + .. + } = origin_exec; let join_constructor = |new_conditions: ( Vec<(PhysicalExprRef, PhysicalExprRef)>, Vec, @@ -207,7 +207,8 @@ pub fn adjust_input_keys_ordering( *join_type, new_conditions.1, *null_equality, - ) + )? + .with_null_aware(origin_exec.null_aware()) .map(|e| Arc::new(e) as _) }; return reorder_partitioned_join_keys( @@ -469,17 +470,17 @@ pub fn reorder_join_keys_to_inputs( .build_exec(); } } - } else if let Some(SortMergeJoinExec { - left, - right, - on, - filter, - join_type, - sort_options, - null_equality, - .. - }) = plan.downcast_ref::() - { + } else if let Some(origin_exec) = plan.downcast_ref::() { + let SortMergeJoinExec { + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + .. + } = origin_exec; let (join_keys, positions) = reorder_current_join_keys( extract_join_keys(on), Some(left.output_partitioning()), @@ -506,7 +507,8 @@ pub fn reorder_join_keys_to_inputs( *join_type, new_sort_options, *null_equality, - ) + )? + .with_null_aware(origin_exec.null_aware()) .map(|smj| Arc::new(smj) as _); } } From 45abd6d08d15e59fba7b9d97720ae61876428536 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 14:29:41 +0800 Subject: [PATCH 12/25] test(smj): Update null-aware join plan expectation Expect SortMergeJoinExec for eligible NOT IN queries when hash join is not preferred. Update the regression comments to reflect null-aware sort-merge join support. --- datafusion/sqllogictest/test_files/subquery.slt | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 2ec240ac67f24..93a9bf62d2a97 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3187,8 +3187,8 @@ DROP TABLE sq_count_orders; # NULL the predicate is never TRUE, so the query must return zero rows. This must # hold regardless of the chosen physical join operator. Previously, with # prefer_hash_join = false and multiple partitions, the planner routed the -# null-aware anti join to SortMergeJoin (which is not null-aware) and returned -# wrong results; null-aware anti joins must use the CollectLeft HashJoin. +# null-aware anti join to an ordinary SortMergeJoin and returned wrong results. +# The null-aware SortMergeJoin path must preserve the same NOT IN semantics. statement ok set datafusion.optimizer.prefer_hash_join = false; @@ -3207,9 +3207,8 @@ query I SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null) ORDER BY x; ---- -# The null-aware anti join must be planned as a CollectLeft HashJoinExec even with -# prefer_hash_join = false: SortMergeJoinExec is not null-aware and must not be used. -# This test has no usable size statistics, so JoinSelection keeps LeftAnti. +# With prefer_hash_join = false, this single-key, unfiltered null-aware LeftAnti +# join uses SortMergeJoinExec with sorted, single-partition inputs. query TT EXPLAIN SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null); ---- @@ -3219,9 +3218,11 @@ logical_plan 03)--SubqueryAlias: __correlated_sq_1 04)----TableScan: nia_right_with_null projection=[y] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(x@0, y@0)], null_aware -02)--DataSourceExec: partitions=1, partition_sizes=[1] -03)--DataSourceExec: partitions=1, partition_sizes=[1] +01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)] +02)--SortExec: expr=[x@0 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortExec: expr=[y@0 ASC], preserve_partitioning=[false] +05)----DataSourceExec: partitions=1, partition_sizes=[1] # Subquery has no NULL -> NOT IN behaves like a normal anti join. query I From 0e0f128c9633ac62df18ae7524ce19234119d302 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 17 Oct 2026 16:32:55 +0800 Subject: [PATCH 13/25] docs: edit comment of physical_planner of null aware and join check --- datafusion/core/src/physical_planner.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 881ffd4f55f51..5b6b6ba91beae 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1729,10 +1729,10 @@ impl DefaultPhysicalPlanner { && session.config().repartition_joins() && !*null_aware; - // Only `HashJoinExec` implements null-aware semantics, and it - // needs equi-join keys to do so. Without them the join would be - // planned as a nested loop (or piecewise merge) join, which - // silently ignores the flag and returns wrong results for + // Only `HashJoinExec` and the null-aware `LeftAnti` path in `SortMergeJoinExec` + // implements null-aware semantics, and thoes need equi-join keys to do so. + // Without them the join would be planned as a nested loop (or piecewise merge) + // join, which silently ignores the flag and returns wrong results for // `NOT IN` over a nullable subquery. Fail loudly instead. if *null_aware && join_on.is_empty() { return plan_err!( From cf80ccebf8dfed90078902cbbc443ca2f1d98c7f Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 17 Oct 2026 16:42:10 +0800 Subject: [PATCH 14/25] chores(physical-plan): fix typo --- .../src/joins/sort_merge_join/bitwise_stream.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index ba0c946490ce5..69062395dd6ce 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -1165,7 +1165,7 @@ impl BitwiseSortMergeJoinStream { } /// Initializes the inner input's empty/NULL summary for null-aware joins. - async fn initialized_null_aware(&mut self) -> Result<()> { + async fn initialize_null_aware(&mut self) -> Result<()> { if self.subquery_state != SubqueryState::Pending { return Ok(()); } @@ -1181,12 +1181,12 @@ impl BitwiseSortMergeJoinStream { } // Keep the first nonempty inner batch at offset 0 for the merge scan. - // The complete right input need to be sorted with NULLS FIRST. - let first_batch_has_null = self.inner_key_arrays[0] + // The complete right input must be sorted with NULLS FIRST. + let first_key_is_null = self.inner_key_arrays[0] .logical_nulls() .is_some_and(|null_keys| null_keys.is_null(0)); - self.subquery_state = if first_batch_has_null { + self.subquery_state = if first_key_is_null { SubqueryState::HasNull } else { SubqueryState::NonEmptyNoNull @@ -1201,14 +1201,14 @@ impl BitwiseSortMergeJoinStream { &mut self, emitter: &mut TryEmitter, ) -> Result<()> { - self.initialized_null_aware().await?; + self.initialize_null_aware().await?; match self.null_aware { None => {} Some(NullAwareMode::LeftAnti) if self.subquery_state == SubqueryState::HasNull => { // NOT IN cannot be TRUE for any outer row when the inner - // input contains a NULL key. (following NOT IN semetic) + // input contains a NULL key. (following NOT IN semetics) return Ok(()); } Some(NullAwareMode::LeftAnti) => {} From b8901ec96eaadcd67111b3dc6b3617b157085a8f Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 17 Oct 2026 17:51:03 +0800 Subject: [PATCH 15/25] feat(physical-plan): show null-aware sort-merge joins in EXPLAIN Display the null_aware marker in default, verbose, and tree plan formats when enabled. Update the existing EXPLAIN expectation and add a tree-format SQLLogicTest case for null-aware LeftAnti joins. --- .../src/joins/sort_merge_join/exec.rs | 12 +++++++- .../sqllogictest/test_files/subquery.slt | 29 ++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index ebc7d9736561d..a8c8b2d9ffc55 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -515,6 +515,11 @@ impl DisplayAs for SortMergeJoinExec { } else { "" }; + let display_null_aware = if self.null_aware() { + ", null_aware" + } else { + "" + }; let display_projections = match &self.projection { Some(projection) => format!( ", projection=[{}]", @@ -532,7 +537,7 @@ impl DisplayAs for SortMergeJoinExec { }; write!( f, - "{}: join_type={:?}, on=[{}]{}{}{}", + "{}: join_type={:?}, on=[{}]{}{}{}{}", Self::static_name(), self.join_type, on, @@ -541,6 +546,7 @@ impl DisplayAs for SortMergeJoinExec { |f| format!(", filter={}", f.expression()) ), display_null_equality, + display_null_aware, display_projections, ) } @@ -563,6 +569,10 @@ impl DisplayAs for SortMergeJoinExec { writeln!(f, "NullsEqual: true")?; } + if self.null_aware() { + writeln!(f, "null_aware")?; + } + Ok(()) } } diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 93a9bf62d2a97..6f61892a018df 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3218,12 +3218,39 @@ logical_plan 03)--SubqueryAlias: __correlated_sq_1 04)----TableScan: nia_right_with_null projection=[y] physical_plan -01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)] +01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)], null_aware 02)--SortExec: expr=[x@0 ASC], preserve_partitioning=[false] 03)----DataSourceExec: partitions=1, partition_sizes=[1] 04)--SortExec: expr=[y@0 ASC], preserve_partitioning=[false] 05)----DataSourceExec: partitions=1, partition_sizes=[1] +# Tree-formatted plans must also distinguish null-aware joins from ordinary anti joins. +query TT +EXPLAIN FORMAT TREE SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null); +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortMergeJoinExec │ +03)│ -------------------- │ +04)│ join_type: LeftAnti │ +05)│ ├──────────────┐ +06)│ null_aware │ │ +07)│ │ │ +08)│ on: (x = y) │ │ +09)└─────────────┬─────────────┘ │ +10)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +11)│ SortExec ││ SortExec │ +12)│ -------------------- ││ -------------------- │ +13)│ x@0 ASC ││ y@0 ASC │ +14)└─────────────┬─────────────┘└─────────────┬─────────────┘ +15)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +16)│ DataSourceExec ││ DataSourceExec │ +17)│ -------------------- ││ -------------------- │ +18)│ bytes: 112 ││ bytes: 176 │ +19)│ format: memory ││ format: memory │ +20)│ rows: 4 ││ rows: 2 │ +21)└───────────────────────────┘└───────────────────────────┘ + # Subquery has no NULL -> NOT IN behaves like a normal anti join. query I SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_no_null) ORDER BY x; From b179fd65913b104d5988b899621a622404313fca Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 17 Oct 2026 20:45:25 +0800 Subject: [PATCH 16/25] docs(physical-plan): fix typo, change thoes to those --- datafusion/core/src/physical_planner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 5b6b6ba91beae..5ef99206f533a 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1730,7 +1730,7 @@ impl DefaultPhysicalPlanner { && !*null_aware; // Only `HashJoinExec` and the null-aware `LeftAnti` path in `SortMergeJoinExec` - // implements null-aware semantics, and thoes need equi-join keys to do so. + // implements null-aware semantics, and those need equi-join keys to do so. // Without them the join would be planned as a nested loop (or piecewise merge) // join, which silently ignores the flag and returns wrong results for // `NOT IN` over a nullable subquery. Fail loudly instead. From 25351db5fc3e9f2ec9218e24e572c4974031deba Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Mon, 19 Oct 2026 07:37:47 +0800 Subject: [PATCH 17/25] test(smj): Cover null-aware NOT IN semantics Add SQL regression coverage for null-aware sort-merge joins when prefer_hash_join is false. Verify that an empty RHS preserves NULL outer keys and duplicate rows, an RHS containing NULL suppresses all rows, and a nonempty RHS without NULL excludes NULL and matched outer keys. Use a four-partition RHS to verify that NULL detection covers the entire input. Assert null-aware SMJ plans with target_partitions set to 1 or 4 and with join repartitioning enabled or disabled. Verify that unsupported correlated anti joins retain the CollectLeft HashJoin fallback and that null-aware mark joins preserve UNKNOWN results when hash joins are not preferred. Update the related test comments. Refs #25546 --- .../test_files/null_aware_anti_join.slt | 23 ++- .../test_files/null_aware_mark_join.slt | 17 +- .../sqllogictest/test_files/subquery.slt | 194 +++++++++++++++++- 3 files changed, 223 insertions(+), 11 deletions(-) diff --git a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt index f28d0022f91fc..3acc55fa76fdf 100644 --- a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt @@ -611,8 +611,8 @@ query I SELECT id FROM nai_outer WHERE id > 5 AND id NOT IN (SELECT id FROM nai_inner) ORDER BY id; ---- -# Same query under SortMergeJoin + multiple partitions: null-aware joins must -# be planned as a CollectLeft HashJoin, not a plain anti SortMergeJoin. +# When hash joins are not preferred, the supported uncorrelated case uses a null-aware +# SortMergeJoin, which must still preserve the subquery NULL. statement ok SET datafusion.optimizer.prefer_hash_join = false; @@ -771,6 +771,22 @@ NULL # The projected value column must be the FIRST equi-join key: a null-aware hash # join reads `on[0]` as the `NOT IN` value key and `on[1..]` as correlation # scope keys, positionally. Giving the correlation that slot returns wrong rows. +# Correlated null-aware joins must keep the CollectLeft HashJoin fallback even +# when hash joins are not preferred: the single-key SMJ path cannot handle scopes. +statement ok +SET datafusion.optimizer.prefer_hash_join = false; + +query I +SELECT z FROM naconst_corr_t3 WHERE 1 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z = naconst_corr_t3.z) ORDER BY z; +---- +7 +NULL + +query I +SELECT z FROM naconst_corr_t3 WHERE 2 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z = naconst_corr_t3.z) ORDER BY z; +---- +NULL + query TT EXPLAIN SELECT z FROM naconst_corr_t3 WHERE 1 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z = naconst_corr_t3.z); ---- @@ -789,6 +805,9 @@ physical_plan 04)--ProjectionExec: expr=[CAST(id@0 AS Int64) as naconst_corr_t4.id, z@1 as z] 05)----DataSourceExec: partitions=1, partition_sizes=[1] +statement ok +SET datafusion.optimizer.prefer_hash_join = true; + statement ok DROP TABLE naconst_corr_t3; diff --git a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt index fcbc929948fec..d657ffa60289e 100644 --- a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt @@ -387,8 +387,9 @@ c d e -# Null-aware mark joins are pinned to CollectLeft hash joins (mark joins have no -# sort-merge path); flipping prefer_hash_join off must not drop null awareness. +# Null-aware mark joins are pinned to CollectLeft hash joins because their +# null-aware semantics have no sort-merge path. Flipping prefer_hash_join off +# must not drop null awareness. statement ok set datafusion.optimizer.prefer_hash_join = false; @@ -412,6 +413,18 @@ physical_plan 04)------DataSourceExec: partitions=1, partition_sizes=[1] 05)------DataSourceExec: partitions=1, partition_sizes=[1] +# Execute under the same setting: matching keys are FALSE for NOT IN, while +# unmatched keys and the NULL outer key are UNKNOWN because the RHS contains NULL. +query IT rowsort +SELECT id, value +FROM outer_table +WHERE (id NOT IN (SELECT id FROM inner_table_with_null)) IS NULL; +---- +1 a +3 c +4 d +NULL e + #################################### ## Nested loop mark join with NULLs #################################### diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 6f61892a018df..4b553cc9aa64e 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -3196,6 +3196,12 @@ set datafusion.optimizer.prefer_hash_join = false; statement ok CREATE TABLE nia_left(x INT) AS VALUES (1), (2), (3), (4); +statement ok +CREATE TABLE nia_left_with_null(x INT) AS VALUES (NULL), (1), (2), (3), (3), (4); + +statement ok +CREATE TABLE nia_right_with_empty(y INT); + statement ok CREATE TABLE nia_right_with_null(y INT) AS VALUES (2), (NULL); @@ -3207,6 +3213,46 @@ query I SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null) ORDER BY x; ---- +# Subquery is empty -> NOT IN keeps every left row. +query I +SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_empty) ORDER BY x; +---- +1 +2 +3 +4 + +# Subquery has no NULL -> NOT IN behaves like a normal anti join. +query I +SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_no_null) ORDER BY x; +---- +1 +3 + +# A NULL in the RHS also suppresses a NULL left key. +query I +SELECT x FROM nia_left_with_null WHERE x NOT IN (SELECT y FROM nia_right_with_null) ORDER BY x; +---- + +# An empty RHS preserves NULL left keys and duplicate rows. +query I +SELECT x FROM nia_left_with_null WHERE x NOT IN (SELECT y FROM nia_right_with_empty) ORDER BY x NULLS FIRST; +---- +NULL +1 +2 +3 +3 +4 + +# A nonempty RHS without NULL excludes NULL left keys but preserves duplicates. +query I +SELECT x FROM nia_left_with_null WHERE x NOT IN (SELECT y FROM nia_right_no_null) ORDER BY x; +---- +1 +3 +3 + # With prefer_hash_join = false, this single-key, unfiltered null-aware LeftAnti # join uses SortMergeJoinExec with sorted, single-partition inputs. query TT @@ -3251,16 +3297,15 @@ physical_plan 20)│ rows: 4 ││ rows: 2 │ 21)└───────────────────────────┘└───────────────────────────┘ -# Subquery has no NULL -> NOT IN behaves like a normal anti join. -query I -SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_no_null) ORDER BY x; ----- -1 -3 - statement ok DROP TABLE nia_left; +statement ok +DROP TABLE nia_left_with_null; + +statement ok +DROP TABLE nia_right_with_empty; + statement ok DROP TABLE nia_right_with_null; @@ -3270,6 +3315,141 @@ DROP TABLE nia_right_no_null; statement ok reset datafusion.optimizer.prefer_hash_join; +# A NULL in one RHS partition must affect every outer row. + +statement ok +SET datafusion.optimizer.prefer_hash_join = false; + +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +SET datafusion.optimizer.repartition_joins = true; + +statement ok +CREATE TABLE nia_mp_left(x INT) AS +VALUES (NULL), (1), (2), (3), (3); + +# Each UNION ALL branch produces one partition. +# CTAS preserves the partitioned result in the memory table. +statement ok +CREATE TABLE nia_mp_right AS +SELECT CAST(10 AS INT) AS y +UNION ALL +SELECT CAST(20 AS INT) AS y +UNION ALL +SELECT CAST(30 AS INT) AS y +UNION ALL +SELECT CAST(NULL AS INT) AS y; + +# Verify that the RHS really has four nonempty partitions. +# partition_sizes counts batches, not rows. +query TT +EXPLAIN SELECT y FROM nia_mp_right; +---- +logical_plan TableScan: nia_mp_right projection=[y] +physical_plan DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + +# No left key matches 10, 20, or 30. +# The NULL in the remaining RHS partition must still suppress every row. +query I +SELECT x +FROM nia_mp_left +WHERE x NOT IN (SELECT y FROM nia_mp_right) +ORDER BY x; +---- + +# The optimizer must bring all RHS partitions into the null-aware join. +query TT +EXPLAIN SELECT x FROM nia_mp_left WHERE x NOT IN (SELECT y FROM nia_mp_right); +---- +logical_plan +01)LeftAnti Join: nia_mp_left.x = __correlated_sq_1.y null_aware +02)--TableScan: nia_mp_left projection=[x] +03)--SubqueryAlias: __correlated_sq_1 +04)----TableScan: nia_mp_right projection=[y] +physical_plan +01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)], null_aware +02)--SortExec: expr=[x@0 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortPreservingMergeExec: [y@0 ASC] +05)----SortExec: expr=[y@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + +# Disabling join repartitioning must preserve the same semantics. +statement ok +SET datafusion.optimizer.repartition_joins = false; + +query I +SELECT x +FROM nia_mp_left +WHERE x NOT IN (SELECT y FROM nia_mp_right) +ORDER BY x; +---- + +# Disabling hash repartitioning must not bypass the single-partition requirement. +query TT +EXPLAIN SELECT x FROM nia_mp_left WHERE x NOT IN (SELECT y FROM nia_mp_right); +---- +logical_plan +01)LeftAnti Join: nia_mp_left.x = __correlated_sq_1.y null_aware +02)--TableScan: nia_mp_left projection=[x] +03)--SubqueryAlias: __correlated_sq_1 +04)----TableScan: nia_mp_right projection=[y] +physical_plan +01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)], null_aware +02)--SortExec: expr=[x@0 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortPreservingMergeExec: [y@0 ASC] +05)----SortExec: expr=[y@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + +# Requesting one target partition must also preserve the result. +statement ok +SET datafusion.execution.target_partitions = 1; + +query I +SELECT x +FROM nia_mp_left +WHERE x NOT IN (SELECT y FROM nia_mp_right) +ORDER BY x; +---- + +query TT +EXPLAIN +SELECT x +FROM nia_mp_left +WHERE x NOT IN (SELECT y FROM nia_mp_right); +---- +logical_plan +01)LeftAnti Join: nia_mp_left.x = __correlated_sq_1.y null_aware +02)--TableScan: nia_mp_left projection=[x] +03)--SubqueryAlias: __correlated_sq_1 +04)----TableScan: nia_mp_right projection=[y] +physical_plan +01)SortMergeJoinExec: join_type=LeftAnti, on=[(x@0, y@0)], null_aware +02)--SortExec: expr=[x@0 ASC], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)--SortPreservingMergeExec: [y@0 ASC] +05)----SortExec: expr=[y@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1] + +statement ok +DROP TABLE nia_mp_left; + +statement ok +DROP TABLE nia_mp_right; + +statement ok +RESET datafusion.optimizer.prefer_hash_join; + +statement ok +RESET datafusion.optimizer.repartition_joins; + +# Restore the SLT runner's target partition count. +statement ok +SET datafusion.execution.target_partitions = 4; + # Regression test for #24139: CTE inside a correlated scalar subquery causing a schema error statement ok CREATE TABLE metrics (host VARCHAR, bytes_sent DOUBLE, ts BIGINT) AS VALUES From b4ff8b8b5507622f7f17240c22d82ead4fd43a04 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Mon, 19 Oct 2026 15:51:28 +0800 Subject: [PATCH 18/25] test(smj): Cover null-aware executor boundaries Add five executor tests for null-aware LeftAnti sort-merge joins. Cover empty inputs, empty batches before RHS NULLs, startup prefetch across key groups, draining remaining left batches, and dictionary value NULLs. Use row IDs to verify that NULL and duplicate rows are retained or removed according to NOT IN semantics. Exercise batch sizes 1 and 2 at key-group boundaries and check dictionaries with valid indices pointing to NULL values and different index mappings on each side. Reuse the existing join and input helpers, with small fixture helpers modeled on the null-aware HashJoin tests. Execute plans directly so the optimizer cannot simplify away empty-input cases. Refs #25546 --- .../src/joins/sort_merge_join/tests.rs | 182 ++++++++++++++++++ 1 file changed, 182 insertions(+) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 1106e7c7c7be1..e4fc698f95d04 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -1761,6 +1761,188 @@ async fn join_left_anti() -> Result<()> { Ok(()) } +mod null_aware_left_anti_tests { + use super::*; + use arrow::array::{Array, ArrayRef, DictionaryArray}; + use arrow::datatypes::Int32Type; + + fn batch(keys: ArrayRef, row_ids: &[i32]) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", keys.data_type().clone(), true), + Field::new("row_id", DataType::Int32, false), + ])); + Ok(RecordBatch::try_new( + schema, + vec![keys, Arc::new(Int32Array::from(row_ids.to_vec()))], + )?) + } + + fn int_batch(keys: Vec>, row_ids: &[i32]) -> Result { + batch(Arc::new(Int32Array::from(keys)), row_ids) + } + + async fn collect_ids( + left_batches: Vec, + right_batches: Vec, + batch_size: usize, + ) -> Result> { + let left = build_table_from_batches(left_batches); + // Unlike build_table_from_batches, this must also accept zero batches. + let right = TestMemoryExec::try_new_exec(&[right_batches], left.schema(), None)?; + let on: JoinOn = vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )]; + // Direct execution does not insert sorts. Fixtures are globally sorted + // ASC NULLS FIRST, including across batch boundaries. + let join = join_with_options( + left, + right, + on, + LeftAnti, + vec![SortOptions::new(false, true)], + NullEquality::NullEqualsNothing, + )? + .with_null_aware(true)?; + let context = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(batch_size)), + ); + let batches = common::collect(join.execute(0, context)?).await?; + let mut row_ids = Vec::new(); + for batch in batches { + let ids = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("Int32 row IDs"); + row_ids.extend(ids.values().iter().copied()); + } + Ok(row_ids) + } + + #[tokio::test] + async fn empty_rhs_keeps_nulls_and_duplicates() -> Result<()> { + let left = vec![ + int_batch(vec![None, Some(1), Some(2)], &[0, 1, 2])?, + int_batch(vec![Some(3), Some(3)], &[3, 4])?, + ]; + // EOF without any batch and EOF after empty batches are both empty RHS. + for right in [ + vec![], + vec![int_batch(vec![], &[])?, int_batch(vec![], &[])?], + ] { + let actual = collect_ids(left.clone(), right, 2).await?; + assert_eq!(actual, vec![0, 1, 2, 3, 4]); + } + Ok(()) + } + + #[tokio::test] + async fn rhs_null_after_empty_batch_returns_no_rows() -> Result<()> { + let left = vec![ + int_batch(vec![None, Some(1)], &[0, 1])?, + int_batch(vec![Some(2), Some(3)], &[2, 3])?, + ]; + // Cover an all-NULL RHS and a NULL followed by a non-NULL batch. + for right in [ + vec![int_batch(vec![], &[])?, int_batch(vec![None], &[10])?], + vec![ + int_batch(vec![], &[])?, + int_batch(vec![None], &[10])?, + int_batch(vec![Some(2)], &[11])?, + ], + ] { + let actual = collect_ids(left.clone(), right, 2).await?; + assert!(actual.is_empty(), "unexpected rows: {actual:?}"); + } + Ok(()) + } + + #[tokio::test] + async fn startup_preserves_first_rhs_batch_across_key_groups() -> Result<()> { + let left = vec![ + int_batch(vec![None, Some(1), Some(2)], &[0, 1, 2])?, + int_batch(vec![], &[])?, + int_batch(vec![Some(2), Some(3), Some(3)], &[3, 4, 5])?, + int_batch(vec![Some(4), Some(5)], &[6, 7])?, + ]; + let right = vec![ + int_batch(vec![], &[])?, + int_batch(vec![Some(2)], &[10])?, + int_batch(vec![], &[])?, + int_batch(vec![Some(4)], &[11])?, + int_batch(vec![Some(4)], &[12])?, + ]; + for batch_size in [1, 2] { + let actual = collect_ids(left.clone(), right.clone(), batch_size).await?; + // Keep 1, 3, 3, 5; both occurrences of 2 must match the first RHS batch. + assert_eq!(actual, vec![1, 4, 5, 7], "batch size {batch_size}"); + } + Ok(()) + } + + #[tokio::test] + async fn exhausted_rhs_drains_remaining_left_batches() -> Result<()> { + let left = vec![ + int_batch(vec![None, Some(1)], &[0, 1])?, + int_batch(vec![Some(2)], &[2])?, + int_batch(vec![Some(3), Some(3)], &[3, 4])?, + int_batch(vec![Some(4), Some(5)], &[5, 6])?, + ]; + let right = vec![int_batch(vec![Some(2)], &[10])?]; + let actual = collect_ids(left, right, 2).await?; + // Exclude the initial NULL and retain unmatched rows after RHS EOF, + // including duplicates and the final partial output batch. + assert_eq!(actual, vec![1, 3, 4, 5, 6]); + Ok(()) + } + + // Like HashJoin's build_table_dict_key, keep all dictionary indices valid + // so NULL dictionary values exercise logical rather than physical nulls. + fn dictionary_batch( + keys: Vec, + values: Vec>, + row_ids: &[i32], + ) -> Result { + let keys = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(Int32Array::from(values)), + )?; + batch(Arc::new(keys), row_ids) + } + + #[tokio::test] + async fn dictionary_value_nulls_follow_not_in_semantics() -> Result<()> { + let left = dictionary_batch( + vec![0, 1, 2, 3, 3], + vec![None, Some(1), Some(2), Some(3)], + &[0, 1, 2, 3, 4], + )?; + assert_eq!(left.column(0).null_count(), 0); + assert_eq!(left.column(0).logical_null_count(), 1); + + // The same logical key (2) has a different dictionary index on each side. + let rhs_with_null = dictionary_batch(vec![0, 1], vec![None, Some(2)], &[10, 11])?; + assert_eq!(rhs_with_null.column(0).null_count(), 0); + assert_eq!(rhs_with_null.column(0).logical_null_count(), 1); + + let cases = [ + (vec![], vec![0, 1, 2, 3, 4]), + ( + vec![dictionary_batch(vec![0], vec![Some(2)], &[10])?], + vec![1, 3, 4], + ), + (vec![rhs_with_null], vec![]), + ]; + for (right, expected) in cases { + let actual = collect_ids(vec![left.clone()], right, 2).await?; + assert_eq!(actual, expected); + } + Ok(()) + } +} + #[tokio::test] async fn join_right_anti_one_one() -> Result<()> { let left = build_table( From 441f5b9665fbf3d2a53ac23c1c85c5f87cb2f312 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 24 Oct 2026 16:30:31 +0800 Subject: [PATCH 19/25] test(smj): unify batch builders and clarify null-aware coverage Share nullable two-column and dictionary batch builders across sort-merge join tests, and give the LeftAnti result helper an explicit name. Split RHS NULL coverage by nullable and non-null left keys, and document the startup and batch-boundary cases. --- .../src/joins/sort_merge_join/tests.rs | 369 +++++++++++++----- 1 file changed, 272 insertions(+), 97 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index e4fc698f95d04..15263374c4e8b 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -24,19 +24,19 @@ use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::sorts::sort::SortExec; use crate::spill::spill_manager::SpillManager; use crate::test::TestMemoryExec; +use crate::test::build_table_i32; use crate::test::exec::{BarrierExec, PanicExec}; -use crate::test::{build_table_i32, build_table_i32_two_cols}; use crate::{ExecutionPlan, RecordBatchStream, common}; use crate::{ expressions::Column, joins::sort_merge_join::filter::get_corrected_filter_mask, joins::sort_merge_join::materializing_stream::JoinedRecordBatches, }; use arrow::array::{ - BinaryArray, BooleanArray, Date32Array, Date64Array, FixedSizeBinaryArray, - Int32Array, RecordBatch, UInt64Array, + ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, DictionaryArray, + FixedSizeBinaryArray, Int32Array, RecordBatch, UInt64Array, }; use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch}; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use arrow_ord::sort::SortColumn; use arrow_schema::SchemaRef; use bytes::Bytes; @@ -209,11 +209,47 @@ pub fn build_table_i32_nullable( TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } -pub fn build_table_two_cols( - a: (&str, &Vec), - b: (&str, &Vec), +/// Builds a record batch with two nullable columns. +fn build_batch_two_cols(a: (&str, ArrayRef), b: (&str, ArrayRef)) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new(a.0, a.1.data_type().clone(), true), + Field::new(b.0, b.1.data_type().clone(), true), + ])); + Ok(RecordBatch::try_new(schema, vec![a.1, b.1])?) +} + +/// Builds a two-column batch whose join key is a dictionary. +/// +/// All dictionary indices are physically valid. NULL dictionary values +/// exercise logical NULL handling. +fn build_batch_dict_key( + key_name: &str, + dict_values: Vec>, + keys: Vec, + row_id_name: &str, + row_ids: Vec, +) -> Result { + let dict = DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::new(Int32Array::from(dict_values)), + )?; + + build_batch_two_cols( + (key_name, Arc::new(dict)), + (row_id_name, Arc::new(Int32Array::from(row_ids))), + ) +} + +/// Builds a single-batch table with two nullable Int32 columns. +fn build_table_two_cols( + a: (&str, &[Option]), + b: (&str, &[Option]), ) -> Arc { - let batch = build_table_i32_two_cols(a, b); + let batch = build_batch_two_cols( + (a.0, Arc::new(Int32Array::from(a.1.to_vec()))), + (b.0, Arc::new(Int32Array::from(b.1.to_vec()))), + ) + .unwrap(); let schema = batch.schema(); TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } @@ -1246,8 +1282,8 @@ async fn join_right_different_columns_count_with_filter() -> Result<()> { ); let right = build_table_two_cols( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the left ); let on = vec![( @@ -1298,8 +1334,8 @@ async fn join_left_different_columns_count_with_filter() -> Result<()> { // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![4, 5, 6]), // 6 does not exist on the right + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1358,8 +1394,8 @@ async fn join_left_with_filter_columns_right_before_left() -> Result<()> { // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![4, 5, 6]), // 6 does not exist on the right + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1418,7 +1454,10 @@ async fn join_right_with_filter_columns_right_before_left() -> Result<()> { // from t2 // right join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 - let left = build_table_two_cols(("a2", &vec![10, 20, 30]), ("b1", &vec![4, 5, 6])); + let left = build_table_two_cols( + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), + ); let right = build_table( ("a1", &vec![1, 21, 3]), // 20(t2.a2) > 21(t1.a1) is false @@ -1550,8 +1589,8 @@ async fn join_left_mark_different_columns_count_with_filter() -> Result<()> { // left mark join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![4, 5, 6]), // 6 does not exist on the right + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1652,8 +1691,8 @@ async fn join_right_mark_different_columns_count_with_filter() -> Result<()> { ); let right = build_table_two_cols( - ("a2", &vec![10, 20, 30]), - ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the left ); let on = vec![( @@ -1761,27 +1800,13 @@ async fn join_left_anti() -> Result<()> { Ok(()) } -mod null_aware_left_anti_tests { +mod null_aware_tests { use super::*; - use arrow::array::{Array, ArrayRef, DictionaryArray}; - use arrow::datatypes::Int32Type; - - fn batch(keys: ArrayRef, row_ids: &[i32]) -> Result { - let schema = Arc::new(Schema::new(vec![ - Field::new("key", keys.data_type().clone(), true), - Field::new("row_id", DataType::Int32, false), - ])); - Ok(RecordBatch::try_new( - schema, - vec![keys, Arc::new(Int32Array::from(row_ids.to_vec()))], - )?) - } - - fn int_batch(keys: Vec>, row_ids: &[i32]) -> Result { - batch(Arc::new(Int32Array::from(keys)), row_ids) - } + use arrow::array::Array; - async fn collect_ids( + /// Runs a null-aware LeftAnti join and collects the surviving left-side row IDs. + /// Expects `key` as the first column and `row_id` as the second column. + async fn collect_left_anti_ids( left_batches: Vec, right_batches: Vec, batch_size: usize, @@ -1824,58 +1849,181 @@ mod null_aware_left_anti_tests { #[tokio::test] async fn empty_rhs_keeps_nulls_and_duplicates() -> Result<()> { let left = vec![ - int_batch(vec![None, Some(1), Some(2)], &[0, 1, 2])?, - int_batch(vec![Some(3), Some(3)], &[3, 4])?, + build_batch_two_cols( + ( + "key", + Arc::new(Int32Array::from(vec![None, Some(1), Some(2)])), + ), + ("row_id", Arc::new(Int32Array::from(vec![0, 1, 2]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(3), Some(3)]))), + ("row_id", Arc::new(Int32Array::from(vec![3, 4]))), + )?, ]; // EOF without any batch and EOF after empty batches are both empty RHS. for right in [ vec![], - vec![int_batch(vec![], &[])?, int_batch(vec![], &[])?], + vec![ + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + ], ] { - let actual = collect_ids(left.clone(), right, 2).await?; + let actual = collect_left_anti_ids(left.clone(), right, 2).await?; assert_eq!(actual, vec![0, 1, 2, 3, 4]); } Ok(()) } #[tokio::test] - async fn rhs_null_after_empty_batch_returns_no_rows() -> Result<()> { - let left = vec![ - int_batch(vec![None, Some(1)], &[0, 1])?, - int_batch(vec![Some(2), Some(3)], &[2, 3])?, - ]; + async fn rhs_null_after_empty_batch_with_null_lhs_returns_no_rows() -> Result<()> { + let left = vec![build_batch_two_cols( + ( + "key", + Arc::new(Int32Array::from(vec![None, Some(1), Some(2)])), + ), + ("row_id", Arc::new(Int32Array::from(vec![0, 1, 2]))), + )?]; // Cover an all-NULL RHS and a NULL followed by a non-NULL batch. for right in [ - vec![int_batch(vec![], &[])?, int_batch(vec![None], &[10])?], vec![ - int_batch(vec![], &[])?, - int_batch(vec![None], &[10])?, - int_batch(vec![Some(2)], &[11])?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![None]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?, + ], + vec![ + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![None]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![11]))), + )?, ], ] { - let actual = collect_ids(left.clone(), right, 2).await?; + let actual = collect_left_anti_ids(left.clone(), right, 2).await?; assert!(actual.is_empty(), "unexpected rows: {actual:?}"); } Ok(()) } + #[tokio::test] + async fn rhs_null_after_empty_batch_without_null_lhs_returns_no_rows() -> Result<()> { + let left = vec![build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(1), Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![0, 1]))), + )?]; + // Cover an all-NULL RHS and a NULL followed by a non-NULL batch. + for right in [ + vec![ + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![None]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?, + ], + vec![ + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![None]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![11]))), + )?, + ], + ] { + let actual = collect_left_anti_ids(left.clone(), right, 2).await?; + assert!(actual.is_empty(), "unexpected rows: {actual:?}"); + } + Ok(()) + } + + // Initialization reads the first nonempty RHS batch to check for NULL keys. + // The merge scan must reuse that batch: skipping its sole key (2) would + // incorrectly retain both left rows with key 2. + // + // Left key 2 and right key 4 span batch boundaries, with empty batches + // interspersed to exercise skipping them without treating them as EOF. + // Both left rows with key 3 must survive. Run with batch sizes 1 and 2 to + // verify the same result when output is emitted in small batches. + // + // Also covers a nonempty RHS without NULL keys: retain only unmatched, + // non-NULL left keys, preserving duplicate rows. #[tokio::test] async fn startup_preserves_first_rhs_batch_across_key_groups() -> Result<()> { let left = vec![ - int_batch(vec![None, Some(1), Some(2)], &[0, 1, 2])?, - int_batch(vec![], &[])?, - int_batch(vec![Some(2), Some(3), Some(3)], &[3, 4, 5])?, - int_batch(vec![Some(4), Some(5)], &[6, 7])?, + build_batch_two_cols( + ( + "key", + Arc::new(Int32Array::from(vec![None, Some(1), Some(2)])), + ), + ("row_id", Arc::new(Int32Array::from(vec![0, 1, 2]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ( + "key", + Arc::new(Int32Array::from(vec![Some(2), Some(3), Some(3)])), + ), + ("row_id", Arc::new(Int32Array::from(vec![3, 4, 5]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(4), Some(5)]))), + ("row_id", Arc::new(Int32Array::from(vec![6, 7]))), + )?, ]; let right = vec![ - int_batch(vec![], &[])?, - int_batch(vec![Some(2)], &[10])?, - int_batch(vec![], &[])?, - int_batch(vec![Some(4)], &[11])?, - int_batch(vec![Some(4)], &[12])?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(Vec::>::new()))), + ("row_id", Arc::new(Int32Array::from(Vec::::new()))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(4)]))), + ("row_id", Arc::new(Int32Array::from(vec![11]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(4)]))), + ("row_id", Arc::new(Int32Array::from(vec![12]))), + )?, ]; for batch_size in [1, 2] { - let actual = collect_ids(left.clone(), right.clone(), batch_size).await?; + let actual = + collect_left_anti_ids(left.clone(), right.clone(), batch_size).await?; // Keep 1, 3, 3, 5; both occurrences of 2 must match the first RHS batch. assert_eq!(actual, vec![1, 4, 5, 7], "batch size {batch_size}"); } @@ -1885,58 +2033,73 @@ mod null_aware_left_anti_tests { #[tokio::test] async fn exhausted_rhs_drains_remaining_left_batches() -> Result<()> { let left = vec![ - int_batch(vec![None, Some(1)], &[0, 1])?, - int_batch(vec![Some(2)], &[2])?, - int_batch(vec![Some(3), Some(3)], &[3, 4])?, - int_batch(vec![Some(4), Some(5)], &[5, 6])?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![None, Some(1)]))), + ("row_id", Arc::new(Int32Array::from(vec![0, 1]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![2]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(3), Some(3)]))), + ("row_id", Arc::new(Int32Array::from(vec![3, 4]))), + )?, + build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(4), Some(5)]))), + ("row_id", Arc::new(Int32Array::from(vec![5, 6]))), + )?, ]; - let right = vec![int_batch(vec![Some(2)], &[10])?]; - let actual = collect_ids(left, right, 2).await?; + let right = vec![build_batch_two_cols( + ("key", Arc::new(Int32Array::from(vec![Some(2)]))), + ("row_id", Arc::new(Int32Array::from(vec![10]))), + )?]; + let actual = collect_left_anti_ids(left, right, 2).await?; // Exclude the initial NULL and retain unmatched rows after RHS EOF, // including duplicates and the final partial output batch. assert_eq!(actual, vec![1, 3, 4, 5, 6]); Ok(()) } - // Like HashJoin's build_table_dict_key, keep all dictionary indices valid - // so NULL dictionary values exercise logical rather than physical nulls. - fn dictionary_batch( - keys: Vec, - values: Vec>, - row_ids: &[i32], - ) -> Result { - let keys = DictionaryArray::::try_new( - Int32Array::from(keys), - Arc::new(Int32Array::from(values)), - )?; - batch(Arc::new(keys), row_ids) - } - #[tokio::test] async fn dictionary_value_nulls_follow_not_in_semantics() -> Result<()> { - let left = dictionary_batch( - vec![0, 1, 2, 3, 3], + let left = build_batch_dict_key( + "key", vec![None, Some(1), Some(2), Some(3)], - &[0, 1, 2, 3, 4], + vec![0, 1, 2, 3, 3], + "row_id", + vec![0, 1, 2, 3, 4], )?; assert_eq!(left.column(0).null_count(), 0); assert_eq!(left.column(0).logical_null_count(), 1); // The same logical key (2) has a different dictionary index on each side. - let rhs_with_null = dictionary_batch(vec![0, 1], vec![None, Some(2)], &[10, 11])?; + let rhs_with_null = build_batch_dict_key( + "key", + vec![None, Some(2)], + vec![0, 1], + "row_id", + vec![10, 11], + )?; assert_eq!(rhs_with_null.column(0).null_count(), 0); assert_eq!(rhs_with_null.column(0).logical_null_count(), 1); let cases = [ (vec![], vec![0, 1, 2, 3, 4]), ( - vec![dictionary_batch(vec![0], vec![Some(2)], &[10])?], + vec![build_batch_dict_key( + "key", + vec![Some(2)], + vec![0], + "row_id", + vec![10], + )?], vec![1, 3, 4], ), (vec![rhs_with_null], vec![]), ]; for (right, expected) in cases { - let actual = collect_ids(vec![left.clone()], right, 2).await?; + let actual = collect_left_anti_ids(vec![left.clone()], right, 2).await?; assert_eq!(actual, expected); } Ok(()) @@ -1950,7 +2113,10 @@ async fn join_right_anti_one_one() -> Result<()> { ("b1", &vec![4, 5, 5]), ("c1", &vec![7, 8, 8]), ); - let right = build_table_two_cols(("a2", &vec![10, 20, 30]), ("b1", &vec![4, 5, 6])); + let right = build_table_two_cols( + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), + ); let on = vec![( Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, @@ -2002,7 +2168,10 @@ async fn join_right_anti_two_two() -> Result<()> { ("b1", &vec![4, 5, 5]), ("c1", &vec![7, 8, 8]), ); - let right = build_table_two_cols(("a2", &vec![10, 20, 30]), ("b1", &vec![4, 5, 6])); + let right = build_table_two_cols( + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), Some(5), Some(6)]), + ); let on = vec![ ( Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, @@ -2113,7 +2282,10 @@ async fn join_right_anti_two_with_filter() -> Result<()> { #[tokio::test] async fn join_right_anti_filtered_with_mismatched_columns() -> Result<()> { - let left = build_table_two_cols(("a1", &vec![31, 31]), ("b1", &vec![32, 33])); + let left = build_table_two_cols( + ("a1", &vec![Some(31), Some(31)]), + ("b1", &vec![Some(32), Some(33)]), + ); let right = build_table( ("a2", &vec![31, 31]), ("b2", &vec![32, 35]), @@ -3049,8 +3221,14 @@ async fn join_full_multiple_batches() -> Result<()> { async fn join_full_null_filter_result() -> Result<()> { // Left: (a, b) all non-null, sorted on a let left = build_table_two_cols( - ("a1", &vec![1, 1, 2, 2, 3, 3]), - ("b1", &vec![1, 2, 1, 2, 1, 2]), + ( + "a1", + &vec![Some(1), Some(1), Some(2), Some(2), Some(3), Some(3)], + ), + ( + "b1", + &vec![Some(1), Some(2), Some(1), Some(2), Some(1), Some(2)], + ), ); // Right: (a, b) with b nullable, sorted on a @@ -5553,10 +5731,7 @@ async fn assert_key_batch_boundaries( )?; } let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from(ids)) as arrow::array::ArrayRef, - ), + ("id", Arc::new(Int32Array::from(ids)) as ArrayRef), ("key", keys), ])?; Ok(build_table_from_batches( From c3f8d79749929a547d24aa67f3e2b11d78c812ca Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 24 Oct 2026 18:29:17 +0800 Subject: [PATCH 20/25] feat(physical-plan)!: serialize null-aware sort-merge joins Accept null_aware in SortMergeJoinExec::try_new, validate supported modes, and preserve the flag when rebuilding joins. Encode and decode the flag in protobuf and JSON, regenerate models, and cover roundtrips with projections. BREAKING CHANGE: SortMergeJoinExec::try_new requires an eighth null_aware argument; SortMergeJoinExecNode gains a public null_aware field. --- datafusion/core/src/physical_planner.rs | 23 ++++---- datafusion/core/tests/execution/coop.rs | 1 + datafusion/core/tests/fuzz_cases/join_fuzz.rs | 2 + .../enforce_distribution.rs | 2 + .../physical_optimizer/ensure_requirements.rs | 1 + .../physical_optimizer/projection_pushdown.rs | 1 + .../tests/physical_optimizer/test_utils.rs | 1 + .../enforce_distribution.rs | 9 ++-- .../physical-plan/benches/sort_merge_join.rs | 1 + .../src/joins/sort_merge_join/exec.rs | 52 ++++++++++--------- .../src/joins/sort_merge_join/tests.rs | 23 ++++++-- .../proto-models/proto/datafusion.proto | 1 + .../proto-models/src/generated/pbjson.rs | 18 +++++++ .../proto-models/src/generated/prost.rs | 2 + datafusion/proto/tests/cases/plans/joins.rs | 52 +++++++++++++++++++ 15 files changed, 146 insertions(+), 43 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 5ef99206f533a..4978761232917 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1875,18 +1875,16 @@ impl DefaultPhysicalPlanner { // without a filter. Single-partition inputs and NULLS FIRST ordering // let the stream determine whether the entire RHS is empty or // contains NULL before emitting any rows. - Arc::new( - SortMergeJoinExec::try_new( - physical_left, - physical_right, - join_on, - join_filter, - *join_type, - vec![SortOptions::new(false, true)], // ASC NULLS FIRST - *null_equality, - )? - .with_null_aware(true)?, - ) + Arc::new(SortMergeJoinExec::try_new( + physical_left, + physical_right, + join_on, + join_filter, + *join_type, + vec![SortOptions::new(false, true)], // ASC NULLS FIRST + *null_equality, + true, + )?) } else if session.config().target_partitions() > 1 && session.config().repartition_joins() && !prefer_hash_join @@ -1911,6 +1909,7 @@ impl DefaultPhysicalPlanner { *join_type, vec![SortOptions::default(); join_on_len], *null_equality, + false, )?) } else { // Remaining null-aware joins need global probe-side state, so keep diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs index eda649770ff5c..48ead7a17ac34 100644 --- a/datafusion/core/tests/execution/coop.rs +++ b/datafusion/core/tests/execution/coop.rs @@ -367,6 +367,7 @@ async fn sort_merge_join_yields( JoinType::Inner, vec![inf1.properties().eq_properties.output_ordering().unwrap()[0].options], NullEquality::NullEqualsNull, + false, )?); query_yields(join, session_ctx.task_ctx()).await diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index c98e0b40dc7f9..438fcb4b08367 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -949,6 +949,7 @@ impl JoinFuzzTestCase { self.join_type, vec![SortOptions::default(); self.on_columns().len()], NullEquality::NullEqualsNothing, + false, ) .unwrap(), ) @@ -1363,6 +1364,7 @@ async fn test_filtered_join_spill_fuzz() { *join_type, vec![SortOptions::default(); on.len()], NullEquality::NullEqualsNothing, + false, ) .unwrap(), ); diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index e728a5512e330..384572e8500b4 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1413,6 +1413,7 @@ fn range_join_key_alignment_sort_merge_join() -> Result<()> { JoinType::Inner, vec![desc, asc], NullEquality::NullEqualsNothing, + false, )? .with_projection(Some(vec![0, 5]))?, ); @@ -1610,6 +1611,7 @@ fn range_join_key_alignment_keeps_streamable_sort_merge_join_keys() -> Result<() JoinType::Inner, vec![SortOptions::default(); 2], NullEquality::NullEqualsNothing, + false, )?)) }; diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 86adb307db75b..0dc66abb4241b 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -1241,6 +1241,7 @@ fn test_idempotent_sort_merge_join() { nulls_first: false, }], NullEquality::NullEqualsNothing, + false, ) .expect("SortMergeJoinExec creation failed"), ); diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 0da2f759e3265..3fc3042b8edff 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -1876,6 +1876,7 @@ fn test_sort_merge_join_interleaved_projection_embeds() -> Result<()> { JoinType::Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )?); let projection: Arc = Arc::new(ProjectionExec::try_new( vec![ diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index b93281def20fb..9443091dc1387 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -191,6 +191,7 @@ pub fn sort_merge_join_exec( *join_type, vec![SortOptions::default(); join_on.len()], NullEquality::NullEqualsNothing, + false, ) .unwrap(), ) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index e29d7005f393e..49725b5d8adcb 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -207,8 +207,8 @@ pub fn adjust_input_keys_ordering( *join_type, new_conditions.1, *null_equality, - )? - .with_null_aware(origin_exec.null_aware()) + origin_exec.null_aware(), + ) .map(|e| Arc::new(e) as _) }; return reorder_partitioned_join_keys( @@ -507,8 +507,8 @@ pub fn reorder_join_keys_to_inputs( *join_type, new_sort_options, *null_equality, - )? - .with_null_aware(origin_exec.null_aware()) + origin_exec.null_aware(), + ) .map(|smj| Arc::new(smj) as _); } } @@ -688,6 +688,7 @@ fn reorder_join_keys_to_range_inputs( join.join_type, new_sort_options, join.null_equality, + join.null_aware(), )? .with_projection(join.projection.clone())?; return Ok(Arc::new(reordered) as _); diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index c46ad9500b6ba..e665daa9f6268 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -145,6 +145,7 @@ fn do_join_with_filter( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, ) .unwrap(); diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index a8c8b2d9ffc55..63af68029f646 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -200,6 +200,7 @@ impl SortMergeJoinExec { /// The inputs are sorted using `sort_options` are applied to the columns in the `on` /// # Error /// This function errors when it is not possible to join the left and right sides on keys `on`. + #[expect(clippy::too_many_arguments)] pub fn try_new( left: Arc, right: Arc, @@ -208,7 +209,17 @@ impl SortMergeJoinExec { join_type: JoinType, sort_options: Vec, null_equality: NullEquality, + null_aware: bool, ) -> Result { + if null_aware { + NullAwareMode::try_new( + join_type, + on.len(), + filter.is_some(), + null_equality, + &sort_options, + )?; + } let left_schema = left.schema(); let right_schema = right.schema(); @@ -263,7 +274,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, - null_aware: false, + null_aware, projection: None, cache: Arc::new(cache), }) @@ -448,8 +459,8 @@ impl SortMergeJoinExec { self.join_type().swap(), self.sort_options.clone(), self.null_equality, + self.null_aware, )? - .with_null_aware(self.null_aware)? .with_projection(swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), @@ -667,8 +678,8 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, + self.null_aware, )? - .with_null_aware(self.null_aware)? .with_projection(self.projection.clone())?, )), _ => internal_err!("SortMergeJoin wrong number of children"), @@ -827,18 +838,16 @@ impl ExecutionPlan for SortMergeJoinExec { self.filter().as_ref(), &column_indices, )? { - Ok(Some(Arc::new( - SortMergeJoinExec::try_new( - Arc::new(projected_left_child), - Arc::new(projected_right_child), - join_on, - join_filter, - self.join_type, - self.sort_options.clone(), - self.null_equality, - )? - .with_null_aware(self.null_aware)?, - ))) + Ok(Some(Arc::new(SortMergeJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_on, + join_filter, + self.join_type, + self.sort_options.clone(), + self.null_equality, + self.null_aware, + )?))) } else { try_embed_projection(projection, self) } @@ -851,14 +860,6 @@ impl ExecutionPlan for SortMergeJoinExec { ) -> Result> { use datafusion_proto_models::protobuf; - // The existing protobuf node cannot represent null-aware semantics. - // Reject encoding rather than silently decoding as an ordinary anti join. - if self.null_aware { - return datafusion_common::not_impl_err!( - "Serializing null-aware SortMergeJoinExec is not supported" - ); - } - // Destructure exhaustively (no `..`) so that a newly added field is a // compile error here instead of being silently left out of the proto. let Self { @@ -869,7 +870,7 @@ impl ExecutionPlan for SortMergeJoinExec { join_type, sort_options, null_equality, - null_aware: _, + null_aware, projection, // derived from the children's schemas by `try_new` on decode schema: _, @@ -921,6 +922,7 @@ impl ExecutionPlan for SortMergeJoinExec { filter, sort_options, null_equality: null_equality.into(), + null_aware: *null_aware, // Proto3 `repeated` cannot distinguish `None` from // `Some(vec![])`. `Some(vec![])` (reachable via // `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`) @@ -970,6 +972,7 @@ impl SortMergeJoinExec { filter, sort_options, null_equality, + null_aware, projection, } = &**sort_join; @@ -1039,6 +1042,7 @@ impl SortMergeJoinExec { join_type, sort_options, null_equality, + *null_aware, )? .with_projection(projection)?, )) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 15263374c4e8b..f12852469357c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -269,6 +269,7 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, + false, ) } @@ -288,6 +289,7 @@ fn join_with_options( join_type, sort_options, null_equality, + false, ) } @@ -308,6 +310,7 @@ fn join_with_filter( join_type, sort_options, null_equality, + false, ) } @@ -415,6 +418,7 @@ fn join_and_projection_for_pushdown( Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )?); let input: Arc = Arc::clone(&join) as _; let projection = ProjectionExec::try_new( @@ -1820,15 +1824,16 @@ mod null_aware_tests { )]; // Direct execution does not insert sorts. Fixtures are globally sorted // ASC NULLS FIRST, including across batch boundaries. - let join = join_with_options( + let join = SortMergeJoinExec::try_new( left, right, on, + None, LeftAnti, vec![SortOptions::new(false, true)], NullEquality::NullEqualsNothing, - )? - .with_null_aware(true)?; + true, + )?; let context = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(batch_size)), @@ -4947,6 +4952,7 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, + false, )?; let stream = join.execute(0, task_ctx)?; @@ -6857,6 +6863,7 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -6906,6 +6913,7 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -6967,6 +6975,7 @@ async fn bitwise_filtered_no_spill() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join.execute(0, task_ctx)?; let err = common::collect(stream).await.unwrap_err(); @@ -7252,6 +7261,7 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -7357,6 +7367,7 @@ async fn bitwise_spill_null_key_group() -> Result<()> { LeftSemi, sort_options, NullEquality::NullEqualsNull, + false, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -8028,6 +8039,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -8047,6 +8059,7 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, + false, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -8087,6 +8100,7 @@ async fn swap_inputs_swaps_the_projection() -> Result<()> { Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )? .with_projection(Some(vec![4, 2]))?; @@ -8158,6 +8172,7 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )?; let swapped = join.swap_inputs()?; @@ -8207,6 +8222,7 @@ async fn an_empty_projection_keeps_the_rows() -> Result<()> { Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )? .with_projection(Some(vec![]))?; @@ -8297,6 +8313,7 @@ async fn collect_streamed_keys( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )?; // A small batch size keeps the gate firing often enough to interleave the diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 6c4ef8f0a4b2f..16338bc7a3df0 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1835,6 +1835,7 @@ message SortMergeJoinExecNode { repeated SortExprNode sort_options = 6; datafusion_common.NullEquality null_equality = 7; repeated uint32 projection = 8; + bool null_aware = 9; } message PiecewiseMergeJoinExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 725cceeb75970..34f2a293ae816 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -27666,6 +27666,9 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { len += 1; } + if self.null_aware { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.SortMergeJoinExecNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -27695,6 +27698,9 @@ impl serde::Serialize for SortMergeJoinExecNode { if !self.projection.is_empty() { struct_ser.serialize_field("projection", &self.projection)?; } + if self.null_aware { + struct_ser.serialize_field("nullAware", &self.null_aware)?; + } struct_ser.end() } } @@ -27716,6 +27722,8 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "null_equality", "nullEquality", "projection", + "null_aware", + "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -27728,6 +27736,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { SortOptions, NullEquality, Projection, + NullAware, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -27757,6 +27766,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "sortOptions" | "sort_options" => Ok(GeneratedField::SortOptions), "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), "projection" => Ok(GeneratedField::Projection), + "nullAware" | "null_aware" => Ok(GeneratedField::NullAware), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -27784,6 +27794,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { let mut sort_options__ = None; let mut null_equality__ = None; let mut projection__ = None; + let mut null_aware__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -27837,6 +27848,12 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { .into_iter().map(|x| x.0).collect()) ; } + GeneratedField::NullAware => { + if null_aware__.is_some() { + return Err(serde::de::Error::duplicate_field("nullAware")); + } + null_aware__ = Some(map_.next_value()?); + } } } Ok(SortMergeJoinExecNode { @@ -27848,6 +27865,7 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { sort_options: sort_options__.unwrap_or_default(), null_equality: null_equality__.unwrap_or_default(), projection: projection__.unwrap_or_default(), + null_aware: null_aware__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1034d0443dbce..1583244ea7e5a 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2760,6 +2760,8 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, #[prost(uint32, repeated, tag = "8")] pub projection: ::prost::alloc::vec::Vec, + #[prost(bool, tag = "9")] + pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PiecewiseMergeJoinExecNode { diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 98397e85b0b3c..94b7770ed1fb4 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -415,6 +415,7 @@ fn roundtrip_sort_merge_join_with_projection() -> Result<()> { JoinType::Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + false, )? .with_projection(projection.clone())?, ), @@ -429,6 +430,55 @@ fn roundtrip_sort_merge_join_with_projection() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_sort_merge_join_null_aware() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let on = vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )]; + + for null_aware in [false, true] { + // Preserve the flag together with absent, empty, and nonempty projections. + for projection in [None, Some(vec![]), Some(vec![0])] { + let plan: Arc = Arc::new( + SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + on.clone(), + None, + JoinType::LeftAnti, + vec![SortOptions::new(false, true)], + NullEquality::NullEqualsNothing, + null_aware, + )? + .with_projection(projection.clone())?, + ); + let decoded = roundtrip_test_and_return( + Arc::clone(&plan), + &ctx, + &codec, + &proto_converter, + )?; + let decoded = decoded.downcast_ref::().unwrap(); + assert_eq!(decoded.null_aware(), null_aware); + assert_eq!(decoded.projection, projection); + + #[cfg(feature = "json")] + { + let decoded = super::roundtrip_test_json_and_return(plan, &ctx)?; + let decoded = decoded.downcast_ref::().unwrap(); + assert_eq!(decoded.null_aware(), null_aware); + assert_eq!(decoded.projection, projection); + } + } + } + Ok(()) +} + #[test] fn roundtrip_sort_merge_join() -> Result<()> { let ctx = SessionContext::new(); @@ -494,6 +544,7 @@ fn roundtrip_sort_merge_join() -> Result<()> { join_type, sort_options.clone(), null_equality, + false, )?), &ctx, &codec, @@ -502,6 +553,7 @@ fn roundtrip_sort_merge_join() -> Result<()> { let result = result.downcast_ref::().unwrap(); assert_eq!(result.join_type(), join_type); assert_eq!(result.null_equality(), null_equality); + assert!(!result.null_aware()); assert_eq!(result.sort_options(), sort_options); assert_eq!( result.filter().as_ref().map(|f| f.column_indices()), From 9833de06f737b62cd43376cf860ae30c5b4cafae Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sun, 25 Oct 2026 12:46:24 +0800 Subject: [PATCH 21/25] test(smj): replace temporary vectors with slices Pass array slices to build_table_two_cols to avoid unnecessary Vec allocations and fix the 20 clippy::useless_vec errors. --- .../src/joins/sort_merge_join/tests.rs | 40 +++++++++---------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index f12852469357c..f1f747abb8a01 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -1286,8 +1286,8 @@ async fn join_right_different_columns_count_with_filter() -> Result<()> { ); let right = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the left + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), // 6 does not exist on the left ); let on = vec![( @@ -1338,8 +1338,8 @@ async fn join_left_different_columns_count_with_filter() -> Result<()> { // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1398,8 +1398,8 @@ async fn join_left_with_filter_columns_right_before_left() -> Result<()> { // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1459,8 +1459,8 @@ async fn join_right_with_filter_columns_right_before_left() -> Result<()> { // right join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), ); let right = build_table( @@ -1593,8 +1593,8 @@ async fn join_left_mark_different_columns_count_with_filter() -> Result<()> { // left mark join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 let left = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the right + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), // 6 does not exist on the right ); let right = build_table( @@ -1695,8 +1695,8 @@ async fn join_right_mark_different_columns_count_with_filter() -> Result<()> { ); let right = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), // 6 does not exist on the left + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), // 6 does not exist on the left ); let on = vec![( @@ -2119,8 +2119,8 @@ async fn join_right_anti_one_one() -> Result<()> { ("c1", &vec![7, 8, 8]), ); let right = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), ); let on = vec![( Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, @@ -2174,8 +2174,8 @@ async fn join_right_anti_two_two() -> Result<()> { ("c1", &vec![7, 8, 8]), ); let right = build_table_two_cols( - ("a2", &vec![Some(10), Some(20), Some(30)]), - ("b1", &vec![Some(4), Some(5), Some(6)]), + ("a2", &[Some(10), Some(20), Some(30)]), + ("b1", &[Some(4), Some(5), Some(6)]), ); let on = vec![ ( @@ -2288,8 +2288,8 @@ async fn join_right_anti_two_with_filter() -> Result<()> { #[tokio::test] async fn join_right_anti_filtered_with_mismatched_columns() -> Result<()> { let left = build_table_two_cols( - ("a1", &vec![Some(31), Some(31)]), - ("b1", &vec![Some(32), Some(33)]), + ("a1", &[Some(31), Some(31)]), + ("b1", &[Some(32), Some(33)]), ); let right = build_table( ("a2", &vec![31, 31]), @@ -3228,11 +3228,11 @@ async fn join_full_null_filter_result() -> Result<()> { let left = build_table_two_cols( ( "a1", - &vec![Some(1), Some(1), Some(2), Some(2), Some(3), Some(3)], + &[Some(1), Some(1), Some(2), Some(2), Some(3), Some(3)], ), ( "b1", - &vec![Some(1), Some(2), Some(1), Some(2), Some(1), Some(2)], + &[Some(1), Some(2), Some(1), Some(2), Some(1), Some(2)], ), ); From fa829ceee2c0c964cb5d157082d29ec134adac8e Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Thu, 29 Oct 2026 22:41:37 +0800 Subject: [PATCH 22/25] refactor: remove redundant null-aware sort merge validation --- .../joins/sort_merge_join/bitwise_stream.rs | 28 +++++++------------ .../src/joins/sort_merge_join/exec.rs | 4 +-- 2 files changed, 12 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index 69062395dd6ce..c0d39b54bd4c0 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -192,9 +192,9 @@ fn find_key_group_end(cmp: &JoinKeyComparator, from: usize, len: usize) -> usize fn null_aware_left_anti_selection( unmatched: BooleanArray, outer_key: &dyn Array, - subqury_state: SubqueryState, + subquery_state: SubqueryState, ) -> Result { - match subqury_state { + match subquery_state { SubqueryState::Empty => Ok(unmatched), SubqueryState::NonEmptyNoNull => match outer_key.logical_nulls() { // All outer keys are non-NULL, so unmatched rows can be kept. @@ -319,6 +319,14 @@ pub(crate) struct BitwiseSortMergeJoinStream { } impl BitwiseSortMergeJoinStream { + /// Creates a sort-merge join stream for semi, anti, and mark joins. + /// + /// # Preconditions + /// + /// If `null_aware` is `Some`, the mode and join configuration must satisfy + /// the requirements of [`NullAwareMode::try_new`]. + /// `SortMergeJoinExec::execute` performs this validation before creating + /// the stream. #[expect(clippy::too_many_arguments)] pub fn try_new( schema: SchemaRef, @@ -363,22 +371,6 @@ impl BitwiseSortMergeJoinStream { let peak_mem_used = MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); - if let Some(mode) = null_aware { - let validated = NullAwareMode::try_new( - join_type, - on_outer.len(), - filter.is_some(), - null_equality, - &sort_options, - )?; - - if mode != validated { - return internal_err!( - "null-aware mode does not match the join configuration" - ); - } - } - let mut state = Self { join_type, outer, diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 63af68029f646..70d4256843b12 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -51,7 +51,7 @@ use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; use futures::StreamExt; -/// Null-aware (`NOT IN`) semantics of a soft merge join, derived from +/// Null-aware (`NOT IN`) semantics of a sort merge join, derived from /// [`SortMergeJoinExec::null_aware`] and the join type. /// /// Only these combinations are legal (see [`Self::try_new`]), so the @@ -384,7 +384,7 @@ impl SortMergeJoinExec { self.null_equality } - /// Null aware + /// Return whether this join uses null-aware (NOT IN) semantics pub fn null_aware(&self) -> bool { self.null_aware } From 76a6be74158faba67ff0a8f2790bc3c44f01e3d1 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 30 Oct 2026 00:16:22 +0800 Subject: [PATCH 23/25] docs: explain null-aware sort merge join protobuf compatibility --- .../proto-models/proto/datafusion.proto | 5 ++ .../library-user-guide/upgrading/56.0.0.md | 48 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 16338bc7a3df0..3fde7c45b4f28 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1835,6 +1835,11 @@ message SortMergeJoinExecNode { repeated SortExprNode sort_options = 6; datafusion_common.NullEquality null_equality = 7; repeated uint32 projection = 8; + // Enables null-aware (`NOT IN`) semantics. + // (See `SortMergeJoinExec::null_aware`) + // + // Currently supported only for uncorrelated `LeftAnti` joins with one + // join key and no join filter. bool null_aware = 9; } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 2b11283bd09a9..b2912abfeda82 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1044,6 +1044,54 @@ reader decodes a payload without the field (value `0`) as a scalar `NOT IN`. A ignores the unknown field and would read the extra tuple elements as correlation keys, returning wrong results. +### Null-aware sort-merge joins add a protobuf field + +`SortMergeJoinExec` now supports null-aware (`NOT IN`) semantics for +uncorrelated, single-key `LeftAnti` joins without a join filter. +The existing seven-argument `SortMergeJoinExec::try_new` continues +to create a non-null-aware join. Enable null-aware semantics with +`.with_null_aware(true)?`: + +```rust,ignore +let exec = SortMergeJoinExec::try_new( + left, + right, + on, + None, + JoinType::LeftAnti, + vec![SortOptions::new(false, true)], // ASC NULLS FIRST + NullEquality::NullEqualsNothing, +)? +.with_null_aware(true)?; +``` + +Both inputs must have a single partition when executed. + +The generated protobuf `SortMergeJoinExecNode` type has a new +`null_aware` field. Exhaustive struct literals must initialize it; +`false` preserves the previous behavior: + +```rust,ignore +let node = SortMergeJoinExecNode { + left, + right, + on, + join_type, + filter, + sort_options, + null_equality, + projection, + null_aware: false, +}; +``` + +For binary protobuf plans, an older sort-merge join payload without this +field decodes with `null_aware = false`. A 55.x reader must not consume +a new sort-merge join plan with `null_aware = true`: it ignores the +unknown field and executes a plain anti join, which can return incorrect +results. Upgrade all plan readers before sending them null-aware +sort-merge join plans. + ### `FunctionalDependence` has a new `null_equality` field `FunctionalDependence` now records the NULL semantics its dependency holds under. From 6fd44a206afcc79421b2975a2e07be9cda95084b Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Mon, 2 Nov 2026 14:36:59 +0800 Subject: [PATCH 24/25] fix(physical-plan): Preserve sort-merge join API Restore the seven-argument SortMergeJoinExec::try_new constructor and validate null-aware semantics through with_null_aware when planning, rebuilding, or decoding joins. Keep the additive protobuf field. Cover rejection of unsupported null-aware joins during proto decoding. Regenerate the Prost field documentation, fix the semantics typo, and remove redundant ArrayRef qualifications reported by Clippy. --- datafusion/core/src/physical_planner.rs | 23 +++++----- datafusion/core/tests/execution/coop.rs | 1 - datafusion/core/tests/fuzz_cases/join_fuzz.rs | 2 - .../enforce_distribution.rs | 2 - .../physical_optimizer/ensure_requirements.rs | 1 - .../physical_optimizer/projection_pushdown.rs | 1 - .../tests/physical_optimizer/test_utils.rs | 1 - .../enforce_distribution.rs | 6 +-- .../physical-plan/benches/sort_merge_join.rs | 1 - .../joins/sort_merge_join/bitwise_stream.rs | 2 +- .../src/joins/sort_merge_join/exec.rs | 44 ++++++++---------- .../src/joins/sort_merge_join/tests.rs | 28 +++-------- .../proto-models/src/generated/prost.rs | 5 ++ datafusion/proto/tests/cases/plans/joins.rs | 46 +++++++++++++++++-- 14 files changed, 89 insertions(+), 74 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4978761232917..5ef99206f533a 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1875,16 +1875,18 @@ impl DefaultPhysicalPlanner { // without a filter. Single-partition inputs and NULLS FIRST ordering // let the stream determine whether the entire RHS is empty or // contains NULL before emitting any rows. - Arc::new(SortMergeJoinExec::try_new( - physical_left, - physical_right, - join_on, - join_filter, - *join_type, - vec![SortOptions::new(false, true)], // ASC NULLS FIRST - *null_equality, - true, - )?) + Arc::new( + SortMergeJoinExec::try_new( + physical_left, + physical_right, + join_on, + join_filter, + *join_type, + vec![SortOptions::new(false, true)], // ASC NULLS FIRST + *null_equality, + )? + .with_null_aware(true)?, + ) } else if session.config().target_partitions() > 1 && session.config().repartition_joins() && !prefer_hash_join @@ -1909,7 +1911,6 @@ impl DefaultPhysicalPlanner { *join_type, vec![SortOptions::default(); join_on_len], *null_equality, - false, )?) } else { // Remaining null-aware joins need global probe-side state, so keep diff --git a/datafusion/core/tests/execution/coop.rs b/datafusion/core/tests/execution/coop.rs index 48ead7a17ac34..eda649770ff5c 100644 --- a/datafusion/core/tests/execution/coop.rs +++ b/datafusion/core/tests/execution/coop.rs @@ -367,7 +367,6 @@ async fn sort_merge_join_yields( JoinType::Inner, vec![inf1.properties().eq_properties.output_ordering().unwrap()[0].options], NullEquality::NullEqualsNull, - false, )?); query_yields(join, session_ctx.task_ctx()).await diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 438fcb4b08367..c98e0b40dc7f9 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -949,7 +949,6 @@ impl JoinFuzzTestCase { self.join_type, vec![SortOptions::default(); self.on_columns().len()], NullEquality::NullEqualsNothing, - false, ) .unwrap(), ) @@ -1364,7 +1363,6 @@ async fn test_filtered_join_spill_fuzz() { *join_type, vec![SortOptions::default(); on.len()], NullEquality::NullEqualsNothing, - false, ) .unwrap(), ); diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 384572e8500b4..e728a5512e330 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -1413,7 +1413,6 @@ fn range_join_key_alignment_sort_merge_join() -> Result<()> { JoinType::Inner, vec![desc, asc], NullEquality::NullEqualsNothing, - false, )? .with_projection(Some(vec![0, 5]))?, ); @@ -1611,7 +1610,6 @@ fn range_join_key_alignment_keeps_streamable_sort_merge_join_keys() -> Result<() JoinType::Inner, vec![SortOptions::default(); 2], NullEquality::NullEqualsNothing, - false, )?)) }; diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 0dc66abb4241b..86adb307db75b 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -1241,7 +1241,6 @@ fn test_idempotent_sort_merge_join() { nulls_first: false, }], NullEquality::NullEqualsNothing, - false, ) .expect("SortMergeJoinExec creation failed"), ); diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 3fc3042b8edff..0da2f759e3265 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -1876,7 +1876,6 @@ fn test_sort_merge_join_interleaved_projection_embeds() -> Result<()> { JoinType::Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )?); let projection: Arc = Arc::new(ProjectionExec::try_new( vec![ diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 9443091dc1387..b93281def20fb 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -191,7 +191,6 @@ pub fn sort_merge_join_exec( *join_type, vec![SortOptions::default(); join_on.len()], NullEquality::NullEqualsNothing, - false, ) .unwrap(), ) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 49725b5d8adcb..5e176e7140fb6 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -207,8 +207,8 @@ pub fn adjust_input_keys_ordering( *join_type, new_conditions.1, *null_equality, - origin_exec.null_aware(), ) + .and_then(|join| join.with_null_aware(origin_exec.null_aware())) .map(|e| Arc::new(e) as _) }; return reorder_partitioned_join_keys( @@ -507,8 +507,8 @@ pub fn reorder_join_keys_to_inputs( *join_type, new_sort_options, *null_equality, - origin_exec.null_aware(), ) + .and_then(|join| join.with_null_aware(origin_exec.null_aware())) .map(|smj| Arc::new(smj) as _); } } @@ -688,8 +688,8 @@ fn reorder_join_keys_to_range_inputs( join.join_type, new_sort_options, join.null_equality, - join.null_aware(), )? + .with_null_aware(join.null_aware())? .with_projection(join.projection.clone())?; return Ok(Arc::new(reordered) as _); } diff --git a/datafusion/physical-plan/benches/sort_merge_join.rs b/datafusion/physical-plan/benches/sort_merge_join.rs index e665daa9f6268..c46ad9500b6ba 100644 --- a/datafusion/physical-plan/benches/sort_merge_join.rs +++ b/datafusion/physical-plan/benches/sort_merge_join.rs @@ -145,7 +145,6 @@ fn do_join_with_filter( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, ) .unwrap(); diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index c0d39b54bd4c0..e313594d6b07c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -1200,7 +1200,7 @@ impl BitwiseSortMergeJoinStream { if self.subquery_state == SubqueryState::HasNull => { // NOT IN cannot be TRUE for any outer row when the inner - // input contains a NULL key. (following NOT IN semetics) + // input contains a NULL key. (following NOT IN semantics) return Ok(()); } Some(NullAwareMode::LeftAnti) => {} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 70d4256843b12..1cdb506b562ca 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -198,9 +198,11 @@ pub struct SortMergeJoinExec { impl SortMergeJoinExec { /// Tries to create a new [SortMergeJoinExec]. /// The inputs are sorted using `sort_options` are applied to the columns in the `on` + /// + /// Null-aware semantics are disabled by default. Enable them with + /// [`Self::with_null_aware`]. /// # Error /// This function errors when it is not possible to join the left and right sides on keys `on`. - #[expect(clippy::too_many_arguments)] pub fn try_new( left: Arc, right: Arc, @@ -209,17 +211,7 @@ impl SortMergeJoinExec { join_type: JoinType, sort_options: Vec, null_equality: NullEquality, - null_aware: bool, ) -> Result { - if null_aware { - NullAwareMode::try_new( - join_type, - on.len(), - filter.is_some(), - null_equality, - &sort_options, - )?; - } let left_schema = left.schema(); let right_schema = right.schema(); @@ -274,7 +266,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, - null_aware, + null_aware: false, projection: None, cache: Arc::new(cache), }) @@ -459,8 +451,8 @@ impl SortMergeJoinExec { self.join_type().swap(), self.sort_options.clone(), self.null_equality, - self.null_aware, )? + .with_null_aware(self.null_aware)? .with_projection(swap_join_projection( left.schema().fields().len(), right.schema().fields().len(), @@ -678,8 +670,8 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, - self.null_aware, )? + .with_null_aware(self.null_aware)? .with_projection(self.projection.clone())?, )), _ => internal_err!("SortMergeJoin wrong number of children"), @@ -838,16 +830,18 @@ impl ExecutionPlan for SortMergeJoinExec { self.filter().as_ref(), &column_indices, )? { - Ok(Some(Arc::new(SortMergeJoinExec::try_new( - Arc::new(projected_left_child), - Arc::new(projected_right_child), - join_on, - join_filter, - self.join_type, - self.sort_options.clone(), - self.null_equality, - self.null_aware, - )?))) + Ok(Some(Arc::new( + SortMergeJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_on, + join_filter, + self.join_type, + self.sort_options.clone(), + self.null_equality, + )? + .with_null_aware(self.null_aware)?, + ))) } else { try_embed_projection(projection, self) } @@ -1042,8 +1036,8 @@ impl SortMergeJoinExec { join_type, sort_options, null_equality, - *null_aware, )? + .with_null_aware(*null_aware)? .with_projection(projection)?, )) } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index f1f747abb8a01..babc9851a985c 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -269,7 +269,6 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, - false, ) } @@ -289,7 +288,6 @@ fn join_with_options( join_type, sort_options, null_equality, - false, ) } @@ -310,7 +308,6 @@ fn join_with_filter( join_type, sort_options, null_equality, - false, ) } @@ -418,7 +415,6 @@ fn join_and_projection_for_pushdown( Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )?); let input: Arc = Arc::clone(&join) as _; let projection = ProjectionExec::try_new( @@ -1832,8 +1828,8 @@ mod null_aware_tests { LeftAnti, vec![SortOptions::new(false, true)], NullEquality::NullEqualsNothing, - true, - )?; + )? + .with_null_aware(true)?; let context = Arc::new( TaskContext::default() .with_session_config(SessionConfig::new().with_batch_size(batch_size)), @@ -4952,7 +4948,6 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, - false, )?; let stream = join.execute(0, task_ctx)?; @@ -5682,7 +5677,7 @@ async fn join_key_type_batch_boundaries( #[values(1, 2)] input_batch_size: usize, ) -> Result<()> { // Single-digit keys so Utf8 dictionary values sort like the integers. - let keys = |unmatched_key| -> Result { + let keys = |unmatched_key| -> Result { let keys = arrow::array::Int64Array::from(vec![ None, Some(1), @@ -5716,15 +5711,15 @@ async fn join_key_type_batch_boundaries( /// batches. Row 0 must be null, row 4 must differ between the sides, and all /// remaining rows must match. async fn assert_key_batch_boundaries( - left_keys: arrow::array::ArrayRef, - right_keys: arrow::array::ArrayRef, + left_keys: ArrayRef, + right_keys: ArrayRef, join_type: JoinType, null_equality: NullEquality, descending: bool, input_batch_size: usize, ) -> Result<()> { let key_type = left_keys.data_type().clone(); - let input = |mut keys: arrow::array::ArrayRef| -> Result> { + let input = |mut keys: ArrayRef| -> Result> { let mut ids: Vec = (0..keys.len() as i32).collect(); if descending { ids.reverse(); @@ -6863,7 +6858,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -6913,7 +6907,6 @@ async fn bitwise_spill_with_filter() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -6975,7 +6968,6 @@ async fn bitwise_filtered_no_spill() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join.execute(0, task_ctx)?; let err = common::collect(stream).await.unwrap_err(); @@ -7261,7 +7253,6 @@ async fn bitwise_multi_spill_inner_key_group() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -7367,7 +7358,6 @@ async fn bitwise_spill_null_key_group() -> Result<()> { LeftSemi, sort_options, NullEquality::NullEqualsNull, - false, )?; let stream = join.execute(0, task_ctx)?; let batches = common::collect(stream).await?; @@ -8039,7 +8029,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join.execute(0, task_ctx)?; let spilled_result = common::collect(stream).await.unwrap(); @@ -8059,7 +8048,6 @@ async fn bitwise_spill_pending_stream() -> Result<()> { join_type, sort_options.clone(), NullEquality::NullEqualsNothing, - false, )?; let stream = join_no_spill.execute(0, task_ctx_no_spill)?; let no_spill_result = common::collect(stream).await.unwrap(); @@ -8100,7 +8088,6 @@ async fn swap_inputs_swaps_the_projection() -> Result<()> { Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )? .with_projection(Some(vec![4, 2]))?; @@ -8172,7 +8159,6 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )?; let swapped = join.swap_inputs()?; @@ -8222,7 +8208,6 @@ async fn an_empty_projection_keeps_the_rows() -> Result<()> { Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )? .with_projection(Some(vec![]))?; @@ -8313,7 +8298,6 @@ async fn collect_streamed_keys( join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )?; // A small batch size keeps the gate firing often enough to interleave the diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 1583244ea7e5a..142e560baea77 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2760,6 +2760,11 @@ pub struct SortMergeJoinExecNode { pub null_equality: i32, #[prost(uint32, repeated, tag = "8")] pub projection: ::prost::alloc::vec::Vec, + /// Enables null-aware (`NOT IN`) semantics. + /// (See `SortMergeJoinExec::null_aware`) + /// + /// Currently supported only for uncorrelated `LeftAnti` joins with one + /// join key and no join filter. #[prost(bool, tag = "9")] pub null_aware: bool, } diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 94b7770ed1fb4..7f764888db074 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -415,7 +415,6 @@ fn roundtrip_sort_merge_join_with_projection() -> Result<()> { JoinType::Inner, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - false, )? .with_projection(projection.clone())?, ), @@ -453,8 +452,8 @@ fn roundtrip_sort_merge_join_null_aware() -> Result<()> { JoinType::LeftAnti, vec![SortOptions::new(false, true)], NullEquality::NullEqualsNothing, - null_aware, )? + .with_null_aware(null_aware)? .with_projection(projection.clone())?, ); let decoded = roundtrip_test_and_return( @@ -479,6 +478,48 @@ fn roundtrip_sort_merge_join_null_aware() -> Result<()> { Ok(()) } +#[test] +fn sort_merge_join_decode_validates_null_aware() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let plan = Arc::new(SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::new(EmptyExec::new(schema)), + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + JoinType::Inner, + vec![SortOptions::new(false, true)], + NullEquality::NullEqualsNothing, + )?); + let bytes = + physical_plan_to_bytes_with_proto_converter(plan, &codec, &proto_converter)?; + let mut node = protobuf::PhysicalPlanNode::decode(bytes.as_ref()) + .expect("encoded physical plan must decode"); + let Some(PhysicalPlanType::SortMergeJoin(join)) = node.physical_plan_type.as_mut() + else { + panic!("expected a SortMergeJoin node"); + }; + // The ordinary Inner join is valid, but enabling null awareness is not. + join.null_aware = true; + let err = physical_plan_from_bytes_with_proto_converter( + &node.encode_to_vec(), + ctx.task_ctx().as_ref(), + &codec, + &proto_converter, + ) + .expect_err("decoding must validate null-aware join restrictions"); + assert!( + err.to_string() + .contains("null-aware SortMergeJoin only supports LeftAnti") + ); + Ok(()) +} + #[test] fn roundtrip_sort_merge_join() -> Result<()> { let ctx = SessionContext::new(); @@ -544,7 +585,6 @@ fn roundtrip_sort_merge_join() -> Result<()> { join_type, sort_options.clone(), null_equality, - false, )?), &ctx, &codec, From 9545a0661d52aa1f76e99601cb31919c361a4d8a Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Tue, 3 Nov 2026 18:19:39 +0800 Subject: [PATCH 25/25] docs: clarify null-aware join test coverage Document that the SMJ test helper disables null-aware semantics by default. Explain why protobuf decoding must reject Inner joins with null_aware enabled and give examples of valid configurations. --- .../tests/physical_optimizer/test_utils.rs | 1 + datafusion/proto/tests/cases/plans/joins.rs | 18 +++++++++++++++++- 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index b93281def20fb..146e4a0d7c54a 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -176,6 +176,7 @@ pub fn create_test_schema3() -> Result { Ok(schema) } +/// Creates a sort-merge join with null-aware semantics disabled by default. pub fn sort_merge_join_exec( left: Arc, right: Arc, diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 7f764888db074..db0fb6b85dd72 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -478,12 +478,20 @@ fn roundtrip_sort_merge_join_null_aware() -> Result<()> { Ok(()) } +/// Null-aware SortMergeJoinExec currently supports only LeftAnti joins. +/// Checks that decoding rejects an unsupported Inner join whose protobuf flag +/// enables null-aware semantics. +/// The flag is changed after serialization to simulate invalid input from a +/// producer. This guards against bypassing builder validation during decoding +/// now that `try_new` no longer takes the `null_aware` argument. #[test] fn sort_merge_join_decode_validates_null_aware() -> Result<()> { let ctx = SessionContext::new(); let codec = DefaultPhysicalExtensionCodec {}; let proto_converter = DefaultPhysicalProtoConverter {}; let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + + // join with no null_aware let plan = Arc::new(SortMergeJoinExec::try_new( Arc::new(EmptyExec::new(Arc::clone(&schema))), Arc::new(EmptyExec::new(schema)), @@ -496,15 +504,23 @@ fn sort_merge_join_decode_validates_null_aware() -> Result<()> { vec![SortOptions::new(false, true)], NullEquality::NullEqualsNothing, )?); + + // serialize let bytes = physical_plan_to_bytes_with_proto_converter(plan, &codec, &proto_converter)?; + + // unserialize should be able to decode back to PhysicalPlanNode let mut node = protobuf::PhysicalPlanNode::decode(bytes.as_ref()) .expect("encoded physical plan must decode"); let Some(PhysicalPlanType::SortMergeJoin(join)) = node.physical_plan_type.as_mut() else { panic!("expected a SortMergeJoin node"); }; - // The ordinary Inner join is valid, but enabling null awareness is not. + // This makes the protobuf describe an JoinType::Inner join with null_aware = true, + // which is invalid because null-aware SortMergeJoinExec currently supports + // only JoinType::LeftAnti joins. + // With the other settings used here, LeftAnti + null_aware = true would be + // valid; keeping JoinType::Inner + null_aware = false would also be valid. join.null_aware = true; let err = physical_plan_from_bytes_with_proto_converter( &node.encode_to_vec(),