feature(physical-plan): add NullAwareMode to SortMergeJoinExec

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;
@@ -1721,17 +1721,18 @@ 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;
 
-                // 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 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.
                 if *null_aware && join_on.is_empty() {
                     return plan_err!(
@@ -1863,16 +1864,44 @@ 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
+                {
+                    // 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,
+                            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
                     && !*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,
@@ -1884,7 +1913,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

datafusion/core/tests/physical_optimizer/test_utils.rs

@@ -176,6 +176,7 @@ pub fn create_test_schema3() -> Result<SchemaRef> {
     Ok(schema)
 }
 
+/// Creates a sort-merge join with null-aware semantics disabled by default.
 pub fn sort_merge_join_exec(
     left: Arc<dyn ExecutionPlan>,
     right: Arc<dyn ExecutionPlan>,

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::<SortMergeJoinExec>()
-    {
+    } else if let Some(origin_exec) = plan.downcast_ref::<SortMergeJoinExec>() {
+        let SortMergeJoinExec {
+            left,
+            right,
+            on,
+            filter,
+            join_type,
+            sort_options,
+            null_equality,
+            ..
+        } = origin_exec;
         let join_constructor = |new_conditions: (
             Vec<(PhysicalExprRef, PhysicalExprRef)>,
             Vec<SortOptions>,
@@ -208,6 +208,7 @@ pub fn adjust_input_keys_ordering(
                 new_conditions.1,
                 *null_equality,
             )
+            .and_then(|join| join.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::<SortMergeJoinExec>()
-    {
+    } else if let Some(origin_exec) = plan.downcast_ref::<SortMergeJoinExec>() {
+        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()),
@@ -507,6 +508,7 @@ pub fn reorder_join_keys_to_inputs(
                 new_sort_options,
                 *null_equality,
             )
+            .and_then(|join| join.with_null_aware(origin_exec.null_aware()))
             .map(|smj| Arc::new(smj) as _);
         }
     }
@@ -687,6 +689,7 @@ fn reorder_join_keys_to_range_inputs(
             new_sort_options,
             join.null_equality,
         )?
+        .with_null_aware(join.null_aware())?
         .with_projection(join.projection.clone())?;
         return Ok(Arc::new(reordered) as _);
     }

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,
@@ -120,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;
@@ -187,6 +188,48 @@ 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,
+    subquery_state: SubqueryState,
+) -> Result<BooleanArray> {
+    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.
+            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
+/// 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 +278,8 @@ pub(crate) struct BitwiseSortMergeJoinStream {
     filter: Option<JoinFilter>,
     sort_options: Vec<SortOptions>,
     null_equality: NullEquality,
+    null_aware: Option<NullAwareMode>,
+    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,
@@ -274,11 +319,20 @@ 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,
         sort_options: Vec<SortOptions>,
         null_equality: NullEquality,
+        null_aware: Option<NullAwareMode>,
         outer: SendableRecordBatchStream,
         inner: SendableRecordBatchStream,
         on_outer: Vec<PhysicalExprRef>,
@@ -334,6 +388,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)),
@@ -575,7 +631,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)?;
@@ -1092,12 +1156,56 @@ impl BitwiseSortMergeJoinStream {
         }
     }
 
+    /// Initializes the inner input's empty/NULL summary for null-aware joins.
+    async fn initialize_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 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_key_is_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<RecordBatch, DataFusionError>,
     ) -> Result<()> {
+        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 semantics)
+                return Ok(());
+            }
+            Some(NullAwareMode::LeftAnti) => {}
+        }
+
         // 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.