From 84f6d6e3bd5d2fccf1d21c3dbc6b3dd0e9295fd1 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Wed, 7 Oct 2026 01:31:36 +0800 Subject: [PATCH 01/12] 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 64f1817184f8f..176cd4e893ddd 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -64,6 +64,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 a8cb34a6b8be866bfccd597f09d237275fb9bfbe Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 17:10:46 +0800 Subject: [PATCH 02/12] 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 176cd4e893ddd..8fd11453e0378 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -64,7 +64,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. /// @@ -73,7 +72,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 { @@ -86,18 +85,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 { @@ -124,7 +119,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. @@ -195,6 +189,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. @@ -206,6 +213,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, @@ -214,6 +222,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(); @@ -269,6 +278,7 @@ impl SortMergeJoinExec { right_sort_exprs, sort_options, null_equality, + null_aware, projection: None, cache: Arc::new(cache), }) @@ -437,6 +447,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(), @@ -616,6 +627,7 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, + self.null_aware, )? .with_projection(self.projection.clone())?, )), @@ -768,6 +780,7 @@ impl ExecutionPlan for SortMergeJoinExec { self.join_type, self.sort_options.clone(), self.null_equality, + self.null_aware, )?))) } else { try_embed_projection(projection, self) @@ -791,6 +804,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: _, @@ -818,6 +832,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)) @@ -856,6 +871,7 @@ impl ExecutionPlan for SortMergeJoinExec { indices.iter().map(|index| *index as u32).collect() } }, + null_aware, }, )), ), @@ -892,6 +908,7 @@ impl SortMergeJoinExec { sort_options, null_equality, projection, + null_aware, } = &**sort_join; let left = @@ -925,6 +942,9 @@ impl SortMergeJoinExec { *null_equality, "SortMergeJoinExec", )?; + + let null_aware = *null_aware; + let filter = filter .as_ref() .map(|filter| { @@ -959,6 +979,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 f921812ec0619..f0ac15a7519cb 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -238,6 +238,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, @@ -246,6 +247,7 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, + null_aware, ) } @@ -257,6 +259,7 @@ fn join_with_options( sort_options: Vec, null_equality: NullEquality, ) -> Result { + let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -265,6 +268,7 @@ fn join_with_options( join_type, sort_options, null_equality, + null_aware, ) } @@ -277,6 +281,7 @@ fn join_with_filter( sort_options: Vec, null_equality: NullEquality, ) -> Result { + let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -285,6 +290,7 @@ fn join_with_filter( join_type, sort_options, null_equality, + null_aware, ) } @@ -384,6 +390,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, @@ -392,6 +399,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( @@ -419,6 +427,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( @@ -4555,6 +4567,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, @@ -4563,6 +4577,7 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, + null_aware, )?; let stream = join.execute(0, task_ctx)?; @@ -6316,6 +6331,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), @@ -6324,6 +6340,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(); @@ -6365,6 +6382,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), @@ -6373,6 +6392,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(); @@ -6426,6 +6446,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), @@ -6434,6 +6455,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(); @@ -6711,6 +6733,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), @@ -6719,6 +6742,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?; @@ -6816,14 +6840,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?; @@ -7486,6 +7513,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), @@ -7494,6 +7522,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(); @@ -7505,6 +7534,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), @@ -7513,6 +7543,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(); @@ -7544,15 +7575,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]))?; @@ -7616,6 +7651,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), @@ -7624,6 +7661,7 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + null_aware, )?; let swapped = join.swap_inputs()?; @@ -7665,14 +7703,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![]))?; @@ -7755,6 +7797,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, @@ -7763,6 +7806,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 b82efd3b2b583..41aca5ea52fff 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1792,6 +1792,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 01b7f025636b1..3d2b697c828f1 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -27125,6 +27125,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)?; @@ -27154,6 +27157,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() } } @@ -27175,6 +27181,8 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "null_equality", "nullEquality", "projection", + "null_aware", + "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -27187,6 +27195,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 @@ -27216,6 +27225,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)), } } @@ -27243,6 +27253,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 => { @@ -27296,6 +27307,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 { @@ -27307,6 +27324,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 61ba256704f93..4e82bc53cd22c 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2704,6 +2704,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 93390d976a3fec0ac3cd3ccef56b9f074c681a86 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 17:36:11 +0800 Subject: [PATCH 03/12] 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 8fd11453e0378..03b78adfff125 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -377,6 +377,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, @@ -473,6 +478,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 ccf8a217724a34788eb7db87e9fcb986eb612da4 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Fri, 9 Oct 2026 18:12:28 +0800 Subject: [PATCH 04/12] 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 03b78adfff125..278fa2f46e7ca 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -587,6 +587,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 28a3300d09019b6bbe4c72b72d015fc5e97933ea Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 17:27:12 +0800 Subject: [PATCH 05/12] 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 96bb9fa291298..168b3e06a5957 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 @@ -123,6 +123,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, compare_join_arrays}; use crate::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time, @@ -198,6 +199,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 @@ -246,6 +263,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, @@ -290,6 +309,7 @@ impl BitwiseSortMergeJoinStream { schema: SchemaRef, sort_options: Vec, null_equality: NullEquality, + null_aware: Option, outer: SendableRecordBatchStream, inner: SendableRecordBatchStream, on_outer: Vec, @@ -328,6 +348,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, @@ -345,6 +381,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 278fa2f46e7ca..bb34874a2adaa 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -70,14 +70,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, @@ -703,6 +703,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))?; @@ -713,6 +727,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, @@ -1028,6 +1043,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` @@ -1056,6 +1072,7 @@ pub(crate) fn sort_merge_join_stream( schema, sort_options, null_equality, + null_aware, left, right, on_left, @@ -1099,6 +1116,7 @@ pub(crate) fn sort_merge_join_stream( schema, sort_options, null_equality, + null_aware, streamed, buffered, on_streamed, From 0ac32a12dc0e28ed4b05a959ca35c201fdb1f838 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 17:59:58 +0800 Subject: [PATCH 06/12] 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 bb34874a2adaa..5782c046cb9b2 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -213,7 +213,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, @@ -222,7 +221,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(); @@ -278,12 +276,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 { @@ -452,8 +461,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(), @@ -661,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"), @@ -821,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) } @@ -843,6 +854,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 { @@ -853,7 +872,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: _, @@ -881,7 +900,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)) @@ -920,7 +938,6 @@ impl ExecutionPlan for SortMergeJoinExec { indices.iter().map(|index| *index as u32).collect() } }, - null_aware, }, )), ), @@ -957,7 +974,6 @@ impl SortMergeJoinExec { sort_options, null_equality, projection, - null_aware, } = &**sort_join; let left = @@ -992,8 +1008,6 @@ impl SortMergeJoinExec { "SortMergeJoinExec", )?; - let null_aware = *null_aware; - let filter = filter .as_ref() .map(|filter| { @@ -1028,7 +1042,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 f0ac15a7519cb..a1f7f6cc245dc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -238,7 +238,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, @@ -247,7 +246,6 @@ fn join( join_type, sort_options, NullEquality::NullEqualsNothing, - null_aware, ) } @@ -259,7 +257,6 @@ fn join_with_options( sort_options: Vec, null_equality: NullEquality, ) -> Result { - let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -268,7 +265,6 @@ fn join_with_options( join_type, sort_options, null_equality, - null_aware, ) } @@ -281,7 +277,6 @@ fn join_with_filter( sort_options: Vec, null_equality: NullEquality, ) -> Result { - let null_aware = matches!(join_type, LeftAnti); SortMergeJoinExec::try_new( left, right, @@ -290,7 +285,6 @@ fn join_with_filter( join_type, sort_options, null_equality, - null_aware, ) } @@ -390,7 +384,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, @@ -399,7 +392,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( @@ -427,10 +419,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( @@ -4567,8 +4555,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, @@ -4577,7 +4563,6 @@ async fn join_get_stream_and_get_expected( join_type, sort_options, null_equality, - null_aware, )?; let stream = join.execute(0, task_ctx)?; @@ -5446,6 +5431,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))], @@ -5545,6 +5531,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], @@ -5930,6 +5917,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, @@ -6024,6 +6012,7 @@ async fn no_filter_boundary_pending_loses_outer_rows() -> Result<()> { left_schema, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer, @@ -6115,6 +6104,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(), @@ -6235,6 +6225,7 @@ async fn filtered_boundary_pending_outer_rows() -> Result<()> { left_schema, vec![SortOptions::default()], NullEquality::NullEqualsNothing, + None, outer, inner, on_outer, @@ -6331,7 +6322,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), @@ -6340,7 +6330,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(); @@ -6382,8 +6371,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), @@ -6392,7 +6379,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(); @@ -6446,7 +6432,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), @@ -6455,7 +6440,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(); @@ -6733,7 +6717,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), @@ -6742,7 +6725,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?; @@ -6840,17 +6822,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?; @@ -6992,6 +6971,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(), @@ -7513,7 +7493,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), @@ -7522,7 +7501,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(); @@ -7534,7 +7512,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), @@ -7543,7 +7520,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(); @@ -7575,19 +7551,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]))?; @@ -7651,8 +7623,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), @@ -7661,7 +7631,6 @@ async fn swap_inputs_swaps_the_filter() -> Result<()> { join_type, vec![SortOptions::default()], NullEquality::NullEqualsNothing, - null_aware, )?; let swapped = join.swap_inputs()?; @@ -7703,18 +7672,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![]))?; @@ -7797,7 +7762,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, @@ -7806,7 +7770,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 41aca5ea52fff..b82efd3b2b583 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1792,7 +1792,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 3d2b697c828f1..01b7f025636b1 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -27125,9 +27125,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)?; @@ -27157,9 +27154,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() } } @@ -27181,8 +27175,6 @@ impl<'de> serde::Deserialize<'de> for SortMergeJoinExecNode { "null_equality", "nullEquality", "projection", - "null_aware", - "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -27195,7 +27187,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 @@ -27225,7 +27216,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)), } } @@ -27253,7 +27243,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 => { @@ -27307,12 +27296,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 { @@ -27324,7 +27307,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 4e82bc53cd22c..61ba256704f93 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2704,8 +2704,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 80f4948778a446523de296394b856dac274eacca Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sat, 10 Oct 2026 22:32:01 +0800 Subject: [PATCH 07/12] 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 168b3e06a5957..5cca647bda1c6 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 @@ -1141,12 +1141,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 6ea385fd0a44b5c7d95e6cf32f7f35ea855480bc Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Sun, 11 Oct 2026 22:12:36 +0800 Subject: [PATCH 08/12] 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 f1bc734d7292b..624f5e78be5bd 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -73,7 +73,7 @@ use datafusion_common::{ 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; @@ -1781,6 +1781,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_state.config().target_partitions() > 1 && session_state.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 5cca647bda1c6..256e94752b575 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 @@ -132,7 +132,7 @@ use crate::spill::in_progress_spill_file::InProgressSpillFile; use crate::spill::spill_manager::SpillManager; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch}; -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; @@ -199,6 +199,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 @@ -624,7 +650,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 b96492e777e72eef974cc87e8f5b1b58a8de1c20 Mon Sep 17 00:00:00 2001 From: Tiny_Murky Date: Wed, 14 Oct 2026 07:53:01 +0800 Subject: [PATCH 09/12] 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 624f5e78be5bd..4263d05e91c74 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1639,9 +1639,10 @@ impl DefaultPhysicalPlanner { let prefer_hash_join = session_state.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_state.config().target_partitions() > 1 && session_state.config().repartition_joins() && !*null_aware; @@ -1821,7 +1822,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 2b3d2c580b0dd5ad916fe1ef12c841695f3d82c4 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 11:16:18 +0800 Subject: [PATCH 10/12] 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 4263d05e91c74..55c4fb9b362a6 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -1789,6 +1789,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, @@ -1805,12 +1809,17 @@ impl DefaultPhysicalPlanner { && session_state.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 256e94752b575..3388b9ced8463 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 @@ -1215,13 +1215,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 8ef0f065beaea866bf434d154d950eb783a4bf5b Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 14:29:29 +0800 Subject: [PATCH 11/12] 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 912903f31809c..da04ec1d849dd 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -197,17 +197,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, @@ -220,7 +220,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( @@ -482,17 +483,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()), @@ -519,7 +520,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 f41e2e9325437c3d10db2c8eeafcbab5ace7bbf8 Mon Sep 17 00:00:00 2001 From: TinyMurky Date: Wed, 14 Oct 2026 14:29:41 +0800 Subject: [PATCH 12/12] 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 09f648df27702..8428a923615d7 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2819,8 +2819,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; @@ -2839,9 +2839,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); ---- @@ -2851,9 +2850,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