feature(physical-plan): add NullAwareMode to SortMergeJoinExec

datafusion/core/src/physical_planner.rs

@@ -72,7 +72,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;
@@ -1762,6 +1762,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

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,
@@ -131,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;
@@ -198,6 +199,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,
+    subqury_state: SubqueryState,
+) -> Result<BooleanArray> {
+    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
+/// 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 +289,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,
@@ -290,6 +335,7 @@ impl BitwiseSortMergeJoinStream {
         schema: SchemaRef,
         sort_options: Vec<SortOptions>,
         null_equality: NullEquality,
+        null_aware: Option<NullAwareMode>,
         outer: SendableRecordBatchStream,
         inner: SendableRecordBatchStream,
         on_outer: Vec<PhysicalExprRef>,
@@ -328,6 +374,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 +407,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)),
@@ -586,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)?;
@@ -1103,12 +1175,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<RecordBatch, DataFusionError>,
     ) -> 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.