Fix null-aware anti join with shared atomic state for partitioned execution

The previous implementation used per-partition flags to track probe side state,
which caused incorrect results when hash partitioning distributed rows across
multiple partitions. With CollectLeft mode, each output partition only had local
knowledge of its own probe data, not global state.

This commit fixes the issue by:
1. Adding shared AtomicBool flags to JoinLeftData (probe_side_non_empty, probe_side_has_null)
2. All partitions write to and read from these shared atomic flags
3. Ensures global knowledge of probe side state across all partitions

Example of the bug:
- With 16 partitions, NULL rows hash to partition 5, value 2 hashes to partition 12
- Partition 5 sees no probe data (local view: empty)
- Partition 12 sees probe data (local view: non-empty)
- If partition 5 outputs final results, it incorrectly returns NULL rows

With shared atomic state, partition 5 now sees the global truth and correctly
filters NULL rows when probe side is non-empty.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Sonnet 4.5 <[email protected]>

Author: viirya

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -17,7 +17,7 @@
 
 use std::fmt;
 use std::mem::size_of;
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, OnceLock};
 use std::{any::Any, vec};
 
@@ -118,6 +118,11 @@ pub(super) struct JoinLeftData {
     /// Membership testing strategy for filter pushdown
     /// Contains either InList values for small build sides or hash table reference for large build sides
     pub(super) membership: PushdownStrategy,
+    /// Shared atomic flag indicating if any probe partition saw data (for null-aware anti joins)
+    /// This is shared across all probe partitions to provide global knowledge
+    pub(super) probe_side_non_empty: AtomicBool,
+    /// Shared atomic flag indicating if any probe partition saw NULL in join keys (for null-aware anti joins)
+    pub(super) probe_side_has_null: AtomicBool,
 }
 
 impl JoinLeftData {
@@ -1585,6 +1590,8 @@ async fn collect_left_input(
         _reservation: reservation,
         bounds,
         membership,
+        probe_side_non_empty: AtomicBool::new(false),
+        probe_side_has_null: AtomicBool::new(false),
     };
 
     Ok(data)

datafusion/physical-plan/src/joins/hash_join/stream.rs

@@ -21,6 +21,7 @@
 //! [`super::HashJoinExec`]. See comments in [`HashJoinStream`] for more details.
 
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 use std::task::Poll;
 
 use crate::joins::PartitionMode;
@@ -225,14 +226,6 @@ pub(super) struct HashJoinStream {
     output_buffer: Box<BatchCoalescer>,
     /// Whether this is a null-aware anti join
     null_aware: bool,
-    /// Whether the probe side (RIGHT) contains any NULL values in join keys
-    /// Only relevant when null_aware is true.
-    /// For LeftAnti with null-aware semantics, if probe side has NULL, no rows should be output.
-    probe_side_has_null: bool,
-    /// Whether any probe batches were processed (i.e., probe side was non-empty)
-    /// Only relevant when null_aware is true.
-    /// Used to distinguish between empty probe side (should return NULL rows) vs non-empty (should filter NULL rows).
-    probe_side_non_empty: bool,
 }
 
 impl RecordBatchStream for HashJoinStream {
@@ -413,8 +406,6 @@ impl HashJoinStream {
             mode,
             output_buffer,
             null_aware,
-            probe_side_has_null: false,
-            probe_side_non_empty: false,
         }
     }
 
@@ -613,20 +604,21 @@ impl HashJoinStream {
         if self.null_aware {
             // Mark that we've seen a probe batch with actual rows (probe side is non-empty)
             // Only set this if batch has rows - empty batches don't count
+            // Use shared atomic state so all partitions can see this global information
             if state.batch.num_rows() > 0 {
-                self.probe_side_non_empty = true;
+                build_side.left_data.probe_side_non_empty.store(true, Ordering::Relaxed);
             }
 
             // Check if probe side (RIGHT) contains NULL
             // Since null_aware validation ensures single column join, we only check the first column
             let probe_key_column = &state.values[0];
-            if !self.probe_side_has_null && probe_key_column.null_count() > 0 {
-                // Found NULL in probe side - set flag to prevent any output
-                self.probe_side_has_null = true;
+            if probe_key_column.null_count() > 0 {
+                // Found NULL in probe side - set shared flag to prevent any output
+                build_side.left_data.probe_side_has_null.store(true, Ordering::Relaxed);
             }
 
-            // If probe side has NULL (detected in this or previous batch), return empty result
-            if self.probe_side_has_null {
+            // If probe side has NULL (detected in this or any other partition), return empty result
+            if build_side.left_data.probe_side_has_null.load(Ordering::Relaxed) {
                 timer.done();
                 self.state = HashJoinStreamState::FetchProbeBatch;
                 return Ok(StatefulStreamResult::Continue);
@@ -784,14 +776,15 @@ impl HashJoinStream {
             return Ok(StatefulStreamResult::Continue);
         }
 
+        let build_side = self.build_side.try_as_ready()?;
+
         // For null-aware anti join, if probe side had NULL, no rows should be output
-        if self.null_aware && self.probe_side_has_null {
+        // Check shared atomic state to get global knowledge across all partitions
+        if self.null_aware && build_side.left_data.probe_side_has_null.load(Ordering::Relaxed) {
             timer.done();
             self.state = HashJoinStreamState::Completed;
             return Ok(StatefulStreamResult::Continue);
         }
-
-        let build_side = self.build_side.try_as_ready()?;
         if !build_side.left_data.report_probe_completed() {
             self.state = HashJoinStreamState::Completed;
             return Ok(StatefulStreamResult::Continue);
@@ -807,9 +800,10 @@ impl HashJoinStream {
         // For null-aware anti join, filter out LEFT rows with NULL in join keys
         // BUT only if the probe side (RIGHT) was non-empty. If probe side is empty,
         // NULL NOT IN (empty) = TRUE, so NULL rows should be returned.
+        // Use shared atomic state to get global knowledge across all partitions
         if self.null_aware
             && self.join_type == JoinType::LeftAnti
-            && self.probe_side_non_empty
+            && build_side.left_data.probe_side_non_empty.load(Ordering::Relaxed)
         {
             // Since null_aware validation ensures single column join, we only check the first column
             let build_key_column = &build_side.left_data.values()[0];