From 2d49af70c8350d227adc9815efd4a3ff4ec8adce Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Tue, 8 Sep 2026 21:13:25 +0800
Subject: [PATCH 1/6] refactor(hash_join): enhance null-aware join logic with
 probe summary tracking

---
 .../physical-plan/src/joins/hash_join/exec.rs | 212 +++++++++++++++++-
 .../src/joins/hash_join/stream.rs             |  67 +++---
 2 files changed, 233 insertions(+), 46 deletions(-)

diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 24b70a22e37e5..8dec40301ae8a 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -308,11 +308,15 @@ 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/mark 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,
+    /// Shared flag set once any probe partition saw a row (null-aware anti/mark joins).
+    ///
+    /// Private on purpose: the final stage must read it through
+    /// [`Self::report_probe_completed`], which orders it after every
+    /// partition's stores.
+    probe_side_non_empty: AtomicBool,
+    /// Shared flag set once any probe partition saw a NULL join key (null-aware anti/mark joins).
+    /// Private for the same reason as `probe_side_non_empty`.
+    probe_side_has_null: AtomicBool,
 
     // For RightAnti joins, where the build side is a smaller subquery, truthy if has null for the single join key
     pub(super) build_side_has_null: bool,
@@ -373,11 +377,54 @@ impl JoinLeftData {
         &self.membership
     }
 
-    /// Decrements the counter of running threads, and returns `true`
-    /// if caller is the last running thread
-    pub(super) fn report_probe_completed(&self) -> bool {
-        self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
+    /// Records what a probe partition saw in one batch, for the null-aware
+    /// rules evaluated in the final stage.
+    pub(super) fn record_probe_batch(&self, non_empty: bool, has_null: bool) {
+        // Relaxed is enough: `report_probe_completed` orders these stores
+        // before the last partition's reads.
+        if non_empty {
+            self.probe_side_non_empty.store(true, Ordering::Relaxed);
+        }
+        if has_null {
+            self.probe_side_has_null.store(true, Ordering::Relaxed);
+        }
+    }
+
+    /// Whether some probe partition has already recorded a NULL join key.
+    ///
+    /// Only an early-exit hint for the probe phase: it may lag behind sibling
+    /// partitions. Final-stage decisions must use the [`ProbeSideSummary`]
+    /// returned by [`Self::report_probe_completed`] instead.
+    pub(super) fn probe_side_has_null_hint(&self) -> bool {
+        self.probe_side_has_null.load(Ordering::Relaxed)
     }
+
+    /// Decrements the counter of running probe partitions. Returns `Some` for
+    /// the last one, together with the shared probe-side flags.
+    ///
+    /// This is the synchronization point between probe partitions: the
+    /// `AcqRel` decrement publishes everything a finishing partition wrote to
+    /// the last partition, and the summary is read only after it. Handing the
+    /// flags out here, rather than exposing them, keeps the final stage from
+    /// reading them before its own decrement, which could miss a NULL that a
+    /// sibling partition records between the read and the decrement.
+    pub(super) fn report_probe_completed(&self) -> Option<ProbeSideSummary> {
+        let is_last = self.probe_threads_counter.fetch_sub(1, Ordering::AcqRel) == 1;
+        is_last.then(|| ProbeSideSummary {
+            has_null: self.probe_side_has_null.load(Ordering::Relaxed),
+            non_empty: self.probe_side_non_empty.load(Ordering::Relaxed),
+        })
+    }
+}
+
+/// What every probe partition together saw, as observed by the last partition
+/// to finish. Only obtainable from [`JoinLeftData::report_probe_completed`].
+#[derive(Debug, Clone, Copy)]
+pub(super) struct ProbeSideSummary {
+    /// Some probe partition saw a row.
+    pub(super) non_empty: bool,
+    /// Some probe partition saw a NULL join key.
+    pub(super) has_null: bool,
 }
 
 /// Helps to build [`HashJoinExec`].
@@ -7360,6 +7407,153 @@ mod tests {
         Ok(())
     }
 
+    /// Builds a two-partition probe side for the cross-partition null-aware
+    /// tests: partition 0 holds the only NULL key, partition 1 holds none.
+    fn build_two_partition_probe_with_null_in_partition_0() -> Arc<dyn ExecutionPlan> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("c2", DataType::Int32, true),
+            Field::new("dummy", DataType::Int32, true),
+        ]));
+        let partition_0 = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![Some(1), None])),
+                Arc::new(Int32Array::from(vec![Some(100), Some(400)])),
+            ],
+        )
+        .unwrap();
+        let partition_1 = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(Int32Array::from(vec![Some(2)])),
+                Arc::new(Int32Array::from(vec![Some(200)])),
+            ],
+        )
+        .unwrap();
+        TestMemoryExec::try_new_exec(
+            &[vec![partition_0], vec![partition_1]],
+            schema,
+            None,
+        )
+        .unwrap()
+    }
+
+    /// Drains the probe partitions of `join` one after another in the given
+    /// order and returns everything they emitted.
+    async fn collect_partitions_in_order(
+        join: &HashJoinExec,
+        order: &[usize],
+        task_ctx: &Arc<TaskContext>,
+    ) -> Result<Vec<RecordBatch>> {
+        let mut batches = vec![];
+        for &partition in order {
+            let stream = join.execute(partition, Arc::clone(task_ctx))?;
+            batches.extend(common::collect(stream).await?);
+        }
+        Ok(batches)
+    }
+
+    /// A NULL probe key silences a null-aware anti join even when the
+    /// partition that saw it is not the one emitting the final build rows.
+    ///
+    /// `CollectLeft` with several probe partitions is what the planner
+    /// produces for `NOT IN`, and only the last partition to finish emits the
+    /// build rows, reading the shared NULL flag set by its siblings. Both
+    /// finishing orders must produce no rows.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_anti_join_probe_null_in_other_partition(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        for order in [[0, 1], [1, 0]] {
+            let left = build_table_two_cols(
+                ("c1", &vec![Some(1), Some(2), Some(3), Some(4)]),
+                ("dummy", &vec![Some(10), Some(20), Some(30), Some(40)]),
+            );
+            let right = build_two_partition_probe_with_null_in_partition_0();
+
+            let on = vec![(
+                Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+                Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+            )];
+
+            let join = HashJoinExec::try_new(
+                left,
+                right,
+                on,
+                None,
+                &JoinType::LeftAnti,
+                None,
+                PartitionMode::CollectLeft,
+                NullEquality::NullEqualsNothing,
+                true, // null_aware = true
+            )?;
+
+            let batches = collect_partitions_in_order(&join, &order, &task_ctx).await?;
+            let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+            assert_eq!(
+                rows, 0,
+                "probe order {order:?} emitted rows although a probe partition saw NULL"
+            );
+        }
+        Ok(())
+    }
+
+    /// The null-aware mark join counterpart of
+    /// [`test_null_aware_anti_join_probe_null_in_other_partition`]: an
+    /// unmatched build row must get an UNKNOWN (NULL) mark when the NULL probe
+    /// key was seen by a sibling partition, whichever partition finishes last.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_probe_null_in_other_partition(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        for order in [[0, 1], [1, 0]] {
+            let left = build_table_two_cols(
+                ("c1", &vec![Some(1), Some(4)]),
+                ("dummy", &vec![Some(10), Some(40)]),
+            );
+            let right = build_two_partition_probe_with_null_in_partition_0();
+
+            let on = vec![(
+                Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+                Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+            )];
+
+            let join = HashJoinExec::try_new(
+                left,
+                right,
+                on,
+                None,
+                &JoinType::LeftMark,
+                None,
+                PartitionMode::CollectLeft,
+                NullEquality::NullEqualsNothing,
+                true, // null_aware = true
+            )?;
+
+            let batches = collect_partitions_in_order(&join, &order, &task_ctx).await?;
+
+            // c1=1 matches probe row 1 (mark true); c1=4 is unmatched and the
+            // probe side had a NULL, so its mark is UNKNOWN.
+            allow_duplicates! {
+                assert_snapshot!(batches_to_sort_string(&batches), @r"
+                +----+-------+------+
+                | c1 | dummy | mark |
+                +----+-------+------+
+                | 1  | 10    | true |
+                | 4  | 40    |      |
+                +----+-------+------+
+                ");
+            }
+        }
+        Ok(())
+    }
+
     /// Test null-aware anti join when build side (left) contains NULL keys
     /// Expected: rows with NULL keys should not be output
     #[apply(hash_join_exec_configs)]
diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs
index ff08828e94262..932a6b8018d9e 100644
--- a/datafusion/physical-plan/src/joins/hash_join/stream.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs
@@ -21,14 +21,13 @@
 //! [`super::HashJoinExec`]. See comments in [`HashJoinStream`] for more details.
 
 use std::sync::Arc;
-use std::sync::atomic::Ordering;
 use std::task::Poll;
 
 use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus};
 use crate::joins::Map;
 use crate::joins::MapOffset;
 use crate::joins::PartitionMode;
-use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode};
+use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode, ProbeSideSummary};
 use crate::joins::hash_join::shared_bounds::{
     PartitionBounds, PartitionBuildData, SharedBuildAccumulator,
 };
@@ -987,24 +986,15 @@ impl HashJoinStream {
 
         let build_side = self.build_side.try_as_ready()?;
 
-        // For null-aware anti join, if probe side had NULL, no rows should be output
-        // Check shared atomic state to get global knowledge across all partitions
-        if self.null_aware == Some(NullAwareMode::LeftAnti)
-            && build_side
-                .left_data
-                .probe_side_has_null
-                .load(Ordering::Relaxed)
-        {
+        // Only the last probe partition to finish emits the build-side rows,
+        // and only it receives the shared probe-side summary (see
+        // `JoinLeftData::report_probe_completed` for why the flags are not
+        // readable any other way here).
+        let Some(probe_summary) = build_side.left_data.report_probe_completed() else {
             timer.done();
             self.state = HashJoinStreamState::Completed;
             return Ok(StatefulStreamResult::Continue);
-        }
-
-        if !build_side.left_data.report_probe_completed() {
-            timer.done();
-            self.state = HashJoinStreamState::Completed;
-            return Ok(StatefulStreamResult::Continue);
-        }
+        };
 
         // use the global left bitmap to produce the left indices and right indices
         let (left_side, right_side) = get_final_indices_from_shared_bitmap(
@@ -1019,6 +1009,7 @@ impl HashJoinStream {
             Some(NullAwareMode::LeftAnti) => {
                 let (left_side, right_side) = null_aware_left_anti_final_indices(
                     &build_side.left_data,
+                    probe_summary,
                     left_side,
                     right_side,
                 );
@@ -1027,6 +1018,7 @@ impl HashJoinStream {
             Some(NullAwareMode::LeftMark { .. }) => {
                 let mark_column = null_aware_left_mark_column(
                     &build_side.left_data,
+                    probe_summary,
                     &left_side,
                     &right_side,
                 );
@@ -1087,12 +1079,6 @@ fn null_aware_skip_probe_batch(
     match mode {
         NullAwareMode::RightAnti => left_data.build_side_has_null,
         NullAwareMode::LeftAnti | NullAwareMode::LeftMark { .. } => {
-            // Only batches with rows count: `NULL NOT IN (empty)` is TRUE.
-            if state.batch.num_rows() > 0 {
-                left_data
-                    .probe_side_non_empty
-                    .store(true, Ordering::Relaxed);
-            }
             // `on[0]` is the `NOT IN` value key for both modes.
             let probe_key_column = &state.values[0];
             let probe_has_null = match mode {
@@ -1101,11 +1087,11 @@ fn null_aware_skip_probe_batch(
                 }
                 _ => probe_key_column.null_count() > 0,
             };
-            if probe_has_null {
-                left_data.probe_side_has_null.store(true, Ordering::Relaxed);
-            }
-            mode == NullAwareMode::LeftAnti
-                && left_data.probe_side_has_null.load(Ordering::Relaxed)
+            // Only batches with rows count: `NULL NOT IN (empty)` is TRUE.
+            left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null);
+            // Best-effort early exit; the final stage re-checks the flag
+            // through `report_probe_completed`.
+            mode == NullAwareMode::LeftAnti && left_data.probe_side_has_null_hint()
         }
     }
 }
@@ -1135,15 +1121,23 @@ fn drop_null_probe_keys(
     }
 }
 
-/// Final-stage rule of a null-aware `LeftAnti` join: a NULL build key means
-/// `NULL NOT IN (probe)`, which is UNKNOWN (row dropped) unless the probe side
-/// was empty, where it is TRUE (row kept).
+/// Final-stage rules of a null-aware `LeftAnti` join, evaluated by the last
+/// probe partition from what every partition together saw:
+/// - a NULL probe key seen by any partition makes `build.key NOT IN (probe)`
+///   UNKNOWN for every build row, so nothing is emitted;
+/// - otherwise a NULL build key means `NULL NOT IN (probe)`, which is UNKNOWN
+///   (row dropped) unless the probe side was empty, where it is TRUE (row
+///   kept).
 fn null_aware_left_anti_final_indices(
     left_data: &JoinLeftData,
+    probe_summary: ProbeSideSummary,
     left_side: UInt64Array,
     right_side: UInt32Array,
 ) -> (UInt64Array, UInt32Array) {
-    if !left_data.probe_side_non_empty.load(Ordering::Relaxed) {
+    if probe_summary.has_null {
+        return (UInt64Array::new_null(0), UInt32Array::new_null(0));
+    }
+    if !probe_summary.non_empty {
         return (left_side, right_side);
     }
     // null_aware validation ensures a single join key
@@ -1158,14 +1152,13 @@ fn null_aware_left_anti_final_indices(
 }
 
 /// Builds the nullable mark column of a null-aware `LeftMark` join from the
-/// final indices and the shared NULL-tracking state.
+/// final indices and what every probe partition together saw.
 fn null_aware_left_mark_column(
     left_data: &JoinLeftData,
+    probe_summary: ProbeSideSummary,
     left_side: &UInt64Array,
     right_side: &UInt32Array,
 ) -> ArrayRef {
-    let probe_side_has_null = left_data.probe_side_has_null.load(Ordering::Relaxed);
-    let probe_side_non_empty = left_data.probe_side_non_empty.load(Ordering::Relaxed);
     let build_key_column = &left_data.values()[0];
     // Correlated joins precomputed the UNKNOWN decision per build row.
     let null_indices_bitmap = left_data
@@ -1177,8 +1170,8 @@ fn null_aware_left_mark_column(
         right_side,
         build_key_column.as_ref(),
         null_indices_bitmap.as_deref(),
-        probe_side_has_null,
-        probe_side_non_empty,
+        probe_summary.has_null,
+        probe_summary.non_empty,
     )
 }
 

From 6f6b16d59a9fe4115e1ad3d21a9ad05efabf015b Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Wed, 9 Sep 2026 22:22:04 +0800
Subject: [PATCH 2/6] feat(hash_join): implement probe completion protocol for
 null-aware joins

---
 Cargo.lock                                    |  48 +++
 Cargo.toml                                    |   1 +
 datafusion/physical-plan/Cargo.toml           |   4 +
 .../physical-plan/src/joins/hash_join/exec.rs |  73 ++--
 .../physical-plan/src/joins/hash_join/mod.rs  |   1 +
 .../src/joins/hash_join/probe_completion.rs   | 344 ++++++++++++++++++
 .../src/joins/hash_join/stream.rs             |   3 +-
 7 files changed, 422 insertions(+), 52 deletions(-)
 create mode 100644 datafusion/physical-plan/src/joins/hash_join/probe_completion.rs

diff --git a/Cargo.lock b/Cargo.lock
index b29cca419a037..d7ab6a8789da2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2511,6 +2511,7 @@ dependencies = [
  "insta",
  "itertools 0.15.0",
  "log",
+ "loom",
  "num-traits",
  "parking_lot",
  "pin-project-lite",
@@ -3223,6 +3224,21 @@ dependencies = [
  "prost-build",
 ]
 
+[[package]]
+name = "generator"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3b854b0e584ead1a33f18b2fcad7cf7be18b3875c78816b753639aa501513ae"
+dependencies = [
+ "cc",
+ "cfg-if",
+ "libc",
+ "log",
+ "rustversion",
+ "windows-link",
+ "windows-result",
+]
+
 [[package]]
 name = "generic-array"
 version = "0.14.7"
@@ -4079,6 +4095,19 @@ version = "0.4.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6"
 
+[[package]]
+name = "loom"
+version = "0.7.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca"
+dependencies = [
+ "cfg-if",
+ "generator",
+ "scoped-tls",
+ "tracing",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "lru-slab"
 version = "0.1.2"
@@ -4094,6 +4123,15 @@ dependencies = [
  "twox-hash",
 ]
 
+[[package]]
+name = "matchers"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
+dependencies = [
+ "regex-automata",
+]
+
 [[package]]
 name = "matchit"
 version = "0.8.4"
@@ -5536,6 +5574,12 @@ dependencies = [
  "syn 2.0.119",
 ]
 
+[[package]]
+name = "scoped-tls"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294"
+
 [[package]]
 name = "scopeguard"
 version = "1.2.0"
@@ -6573,10 +6617,14 @@ version = "0.3.23"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
 dependencies = [
+ "matchers",
  "nu-ansi-term",
+ "once_cell",
+ "regex-automata",
  "sharded-slab",
  "smallvec",
  "thread_local",
+ "tracing",
  "tracing-core",
  "tracing-log",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 66bd816945908..437379fcf6ab9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -175,6 +175,7 @@ itertools = "0.15"
 itoa = "1.0"
 liblzma = { version = "0.4.6", features = ["static"] }
 log = "^0.4"
+loom = "0.7"
 memchr = "2.8.1"
 num-traits = { version = "0.2" }
 object_store = { version = "0.13.2", default-features = false }
diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml
index 0aa22653b44ee..6649295f03866 100644
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
@@ -96,6 +96,10 @@ criterion = { workspace = true, features = ["async_futures"] }
 datafusion-functions-aggregate = { workspace = true }
 datafusion-functions-window = { workspace = true }
 insta = { workspace = true }
+# Model checker for the hash join probe-completion protocol; see
+# joins/hash_join/probe_completion.rs. Test-only, so it stays out of the
+# dependency tree of anything that links DataFusion.
+loom = { workspace = true }
 rand = { workspace = true }
 rstest = { workspace = true }
 rstest_reuse = "0.7.0"
diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 8dec40301ae8a..1a8463c6cd5c8 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -18,7 +18,6 @@
 use std::collections::HashSet;
 use std::fmt;
 use std::mem::size_of;
-use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::{Arc, OnceLock};
 use std::vec;
 
@@ -33,6 +32,7 @@ use crate::filter_pushdown::{
 use crate::joins::Map;
 use crate::joins::array_map::ArrayMap;
 use crate::joins::hash_join::inlist_builder::build_struct_inlist_values;
+use crate::joins::hash_join::probe_completion::{ProbeCompletion, ProbeSideSummary};
 use crate::joins::hash_join::shared_bounds::{
     ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator,
 };
@@ -293,9 +293,9 @@ pub(super) struct JoinLeftData {
     visited_indices_bitmap: SharedBitmapBuilder,
     /// Shared bitmap builder for null marks
     null_indices_bitmap: SharedBitmapBuilder,
-    /// Counter of running probe-threads, potentially
-    /// able to update `visited_indices_bitmap`
-    probe_threads_counter: AtomicUsize,
+    /// Tracks which probe partition finishes last and what the partitions
+    /// collectively saw. See [`ProbeCompletion`] for the invariant it upholds.
+    probe_completion: ProbeCompletion,
     /// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it.
     /// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage.
     /// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption.
@@ -308,16 +308,6 @@ 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 flag set once any probe partition saw a row (null-aware anti/mark joins).
-    ///
-    /// Private on purpose: the final stage must read it through
-    /// [`Self::report_probe_completed`], which orders it after every
-    /// partition's stores.
-    probe_side_non_empty: AtomicBool,
-    /// Shared flag set once any probe partition saw a NULL join key (null-aware anti/mark joins).
-    /// Private for the same reason as `probe_side_non_empty`.
-    probe_side_has_null: AtomicBool,
-
     // For RightAnti joins, where the build side is a smaller subquery, truthy if has null for the single join key
     pub(super) build_side_has_null: bool,
 }
@@ -380,14 +370,7 @@ impl JoinLeftData {
     /// Records what a probe partition saw in one batch, for the null-aware
     /// rules evaluated in the final stage.
     pub(super) fn record_probe_batch(&self, non_empty: bool, has_null: bool) {
-        // Relaxed is enough: `report_probe_completed` orders these stores
-        // before the last partition's reads.
-        if non_empty {
-            self.probe_side_non_empty.store(true, Ordering::Relaxed);
-        }
-        if has_null {
-            self.probe_side_has_null.store(true, Ordering::Relaxed);
-        }
+        self.probe_completion.record_batch(non_empty, has_null);
     }
 
     /// Whether some probe partition has already recorded a NULL join key.
@@ -396,37 +379,20 @@ impl JoinLeftData {
     /// partitions. Final-stage decisions must use the [`ProbeSideSummary`]
     /// returned by [`Self::report_probe_completed`] instead.
     pub(super) fn probe_side_has_null_hint(&self) -> bool {
-        self.probe_side_has_null.load(Ordering::Relaxed)
+        self.probe_completion.saw_null_key_hint()
     }
 
-    /// Decrements the counter of running probe partitions. Returns `Some` for
-    /// the last one, together with the shared probe-side flags.
+    /// Marks this probe partition as finished, returning `Some` for the last
+    /// one together with what every partition saw.
     ///
-    /// This is the synchronization point between probe partitions: the
-    /// `AcqRel` decrement publishes everything a finishing partition wrote to
-    /// the last partition, and the summary is read only after it. Handing the
-    /// flags out here, rather than exposing them, keeps the final stage from
-    /// reading them before its own decrement, which could miss a NULL that a
-    /// sibling partition records between the read and the decrement.
+    /// The summary is reachable only through this call, which is what stops
+    /// the final stage from reading the shared flags before its own
+    /// decrement. See [`ProbeCompletion`] for why that ordering matters.
     pub(super) fn report_probe_completed(&self) -> Option<ProbeSideSummary> {
-        let is_last = self.probe_threads_counter.fetch_sub(1, Ordering::AcqRel) == 1;
-        is_last.then(|| ProbeSideSummary {
-            has_null: self.probe_side_has_null.load(Ordering::Relaxed),
-            non_empty: self.probe_side_non_empty.load(Ordering::Relaxed),
-        })
+        self.probe_completion.report_completed()
     }
 }
 
-/// What every probe partition together saw, as observed by the last partition
-/// to finish. Only obtainable from [`JoinLeftData::report_probe_completed`].
-#[derive(Debug, Clone, Copy)]
-pub(super) struct ProbeSideSummary {
-    /// Some probe partition saw a row.
-    pub(super) non_empty: bool,
-    /// Some probe partition saw a NULL join key.
-    pub(super) has_null: bool,
-}
-
 /// Helps to build [`HashJoinExec`].
 ///
 /// Builder can be created from an existing [`HashJoinExec`] using [`From::from`].
@@ -3024,12 +2990,10 @@ async fn collect_left_input(
         values: left_values,
         visited_indices_bitmap: Mutex::new(visited_indices_bitmap),
         null_indices_bitmap: Mutex::new(null_indices_bitmap),
-        probe_threads_counter: AtomicUsize::new(probe_threads_count),
+        probe_completion: ProbeCompletion::new(probe_threads_count),
         _reservation: reservation,
         bounds,
         membership,
-        probe_side_non_empty: AtomicBool::new(false),
-        probe_side_has_null: AtomicBool::new(false),
         build_side_has_null: build_has_null,
     };
 
@@ -7458,8 +7422,13 @@ mod tests {
     ///
     /// `CollectLeft` with several probe partitions is what the planner
     /// produces for `NOT IN`, and only the last partition to finish emits the
-    /// build rows, reading the shared NULL flag set by its siblings. Both
-    /// finishing orders must produce no rows.
+    /// build rows, reading what its siblings recorded. Both finishing orders
+    /// must produce no rows.
+    ///
+    /// These partitions are drained sequentially, so this covers the
+    /// cross-partition plumbing, not the concurrent race that motivated
+    /// [`ProbeCompletion`]. The interleavings are model-checked in
+    /// `probe_completion::loom_tests` instead.
     #[apply(hash_join_exec_configs)]
     #[tokio::test]
     async fn test_null_aware_anti_join_probe_null_in_other_partition(
@@ -7505,6 +7474,8 @@ mod tests {
     /// [`test_null_aware_anti_join_probe_null_in_other_partition`]: an
     /// unmatched build row must get an UNKNOWN (NULL) mark when the NULL probe
     /// key was seen by a sibling partition, whichever partition finishes last.
+    ///
+    /// Sequentially drained, with the same caveat as that test.
     #[apply(hash_join_exec_configs)]
     #[tokio::test]
     async fn test_null_aware_left_mark_probe_null_in_other_partition(
diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs
index b915802ea4015..7c1b0d76c2f73 100644
--- a/datafusion/physical-plan/src/joins/hash_join/mod.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs
@@ -23,5 +23,6 @@ pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState
 mod exec;
 mod inlist_builder;
 mod partitioned_hash_eval;
+mod probe_completion;
 mod shared_bounds;
 mod stream;
diff --git a/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
new file mode 100644
index 0000000000000..bc063dcdc4bad
--- /dev/null
+++ b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
@@ -0,0 +1,344 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Completion protocol shared by the probe partitions of a hash join.
+//!
+//! The probe partitions of one `HashJoinExec` run concurrently over a shared
+//! build side. Two facts have to cross partition boundaries:
+//!
+//! - which partition emits the build-side rows in the final stage, namely the
+//!   last one to finish, and
+//! - what the partitions collectively saw, which the null-aware `NOT IN` rules
+//!   need: did any partition see a row, and did any see a NULL join key.
+//!
+//! [`ProbeCompletion`] owns both, because reading the second without ordering
+//! it against the first produced wrong results. A partition that read the NULL
+//! flag *before* decrementing the counter could observe `false`, have a
+//! sibling record a NULL and finish, then become the last partition itself and
+//! emit build rows that `NOT IN` must suppress.
+//!
+//! The invariant is therefore: **whichever caller observes itself to be last
+//! sees every fact recorded by every other partition.** It rests on two
+//! properties, and the type is shaped to keep both.
+//!
+//! 1. *Order.* The facts are reachable only through the value returned by
+//!    [`ProbeCompletion::report_completed`], so no caller can read them before
+//!    its own decrement.
+//! 2. *Visibility.* That decrement is `AcqRel`, so the last partition
+//!    synchronizes with the release of every partition that finished before
+//!    it. Were the counter `Relaxed`, nothing would order a sibling's store
+//!    before the final load even when the decrement happened first.
+//!
+//! # Testing
+//!
+//! Neither property can be pinned down by an ordinary concurrent test. The
+//! window is a few instructions wide, and on x86 a `Relaxed` decrement lowers
+//! to the same instruction as an `AcqRel` one, so a stress test would pass on
+//! the broken version. The invariant is instead model-checked with [loom],
+//! which enumerates the thread interleavings *and* the store visibility the
+//! memory model permits, in `loom_tests` below.
+//!
+//! Because loom can only explore its own atomic types, the protocol is written
+//! once in [`define_probe_completion`] and instantiated twice: over
+//! [`std::sync::atomic`] for the real join, and over `loom::sync::atomic` for
+//! the model. Both instantiations share this single copy of the orderings and
+//! the call sequence, so weakening the decrement to `Relaxed`, or reading the
+//! facts before it, fails the model. That is what makes these tests regression
+//! coverage rather than a restatement of the fix.
+//!
+//! [loom]: https://docs.rs/loom
+
+use std::sync::atomic::Ordering;
+
+/// What every probe partition together saw, as observed by the last one to
+/// finish.
+///
+/// Obtainable only from [`ProbeCompletion::report_completed`], which is what
+/// stops the final stage from reading the shared state too early.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) struct ProbeSideSummary {
+    /// Some probe partition saw a row. An entirely empty probe side is not the
+    /// same as one that matched nothing: `NULL NOT IN (empty)` is TRUE.
+    pub(super) non_empty: bool,
+    /// Some probe partition saw a NULL join key.
+    pub(super) has_null: bool,
+}
+
+/// Defines [`ProbeCompletion`] over a given pair of atomic types.
+///
+/// The indirection exists so the loom model below can instantiate the very
+/// same logic over loom's atomics. See the [module docs](self); do not add a
+/// second copy of these orderings anywhere.
+macro_rules! define_probe_completion {
+    ($atomic_usize:ty, $atomic_bool:ty) => {
+        /// Tracks how many probe partitions are still running, together with
+        /// the null-aware facts they contribute.
+        ///
+        /// See the [module docs](self) for the invariant this upholds.
+        #[derive(Debug)]
+        pub(super) struct ProbeCompletion {
+            /// Probe partitions that have not finished yet.
+            running: $atomic_usize,
+            /// Set once any probe partition has seen a row.
+            saw_row: $atomic_bool,
+            /// Set once any probe partition has seen a NULL join key.
+            saw_null_key: $atomic_bool,
+        }
+
+        impl ProbeCompletion {
+            /// Creates the protocol state for `probe_threads` partitions.
+            pub(super) fn new(probe_threads: usize) -> Self {
+                Self {
+                    running: <$atomic_usize>::new(probe_threads),
+                    saw_row: <$atomic_bool>::new(false),
+                    saw_null_key: <$atomic_bool>::new(false),
+                }
+            }
+
+            /// Records what one probe partition saw in one batch.
+            ///
+            /// `Relaxed` suffices: the `AcqRel` decrement in
+            /// [`Self::report_completed`] publishes these stores to the last
+            /// partition.
+            pub(super) fn record_batch(&self, non_empty: bool, has_null: bool) {
+                if non_empty {
+                    self.saw_row.store(true, Ordering::Relaxed);
+                }
+                if has_null {
+                    self.saw_null_key.store(true, Ordering::Relaxed);
+                }
+            }
+
+            /// Whether some partition has already recorded a NULL join key.
+            ///
+            /// A hint for skipping work during the probe phase, which may lag
+            /// behind sibling partitions. It must never decide what the final
+            /// stage emits; use the [`ProbeSideSummary`] from
+            /// [`Self::report_completed`] for that.
+            pub(super) fn saw_null_key_hint(&self) -> bool {
+                self.saw_null_key.load(Ordering::Relaxed)
+            }
+
+            /// Marks the calling partition finished, returning `Some` only for
+            /// the last one, together with what all partitions saw.
+            ///
+            /// The `AcqRel` decrement publishes this partition's own stores
+            /// and acquires those of the partitions that finished earlier, so
+            /// the summary handed to the last caller is complete.
+            pub(super) fn report_completed(&self) -> Option<ProbeSideSummary> {
+                let was_last = self.running.fetch_sub(1, Ordering::AcqRel) == 1;
+                was_last.then(|| ProbeSideSummary {
+                    non_empty: self.saw_row.load(Ordering::Relaxed),
+                    has_null: self.saw_null_key.load(Ordering::Relaxed),
+                })
+            }
+        }
+    };
+}
+
+define_probe_completion!(
+    std::sync::atomic::AtomicUsize,
+    std::sync::atomic::AtomicBool
+);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The sequential contract: only the last caller is handed the summary,
+    /// and it carries what the earlier partitions recorded.
+    ///
+    /// The concurrent invariant is covered by `loom_tests` instead.
+    #[test]
+    fn only_the_last_partition_receives_the_summary() {
+        let completion = ProbeCompletion::new(3);
+
+        completion.record_batch(true, false);
+        assert_eq!(completion.report_completed(), None);
+
+        completion.record_batch(true, true);
+        assert_eq!(completion.report_completed(), None);
+
+        assert_eq!(
+            completion.report_completed(),
+            Some(ProbeSideSummary {
+                non_empty: true,
+                has_null: true,
+            })
+        );
+    }
+
+    /// A probe side that never saw a row leaves both facts clear, which is how
+    /// `NULL NOT IN (empty)` stays TRUE.
+    #[test]
+    fn an_untouched_probe_side_reports_nothing_seen() {
+        let completion = ProbeCompletion::new(1);
+
+        assert_eq!(
+            completion.report_completed(),
+            Some(ProbeSideSummary {
+                non_empty: false,
+                has_null: false,
+            })
+        );
+    }
+
+    /// The hint may lag, but it must never report a NULL that was not
+    /// recorded.
+    #[test]
+    fn the_hint_reflects_recorded_nulls() {
+        let completion = ProbeCompletion::new(1);
+        assert!(!completion.saw_null_key_hint());
+
+        completion.record_batch(true, false);
+        assert!(!completion.saw_null_key_hint());
+
+        completion.record_batch(true, true);
+        assert!(completion.saw_null_key_hint());
+    }
+}
+
+/// Model-checked concurrency tests for the protocol.
+///
+/// These instantiate the protocol over loom's atomics (see the
+/// [module docs](self)) and let loom enumerate the interleavings and store
+/// visibility the memory model allows. They fail if the decrement in
+/// `report_completed` is weakened to `Relaxed`.
+#[cfg(test)]
+mod loom_tests {
+    use super::ProbeSideSummary;
+    use loom::sync::Arc;
+    use std::sync::atomic::Ordering;
+
+    define_probe_completion!(
+        loom::sync::atomic::AtomicUsize,
+        loom::sync::atomic::AtomicBool
+    );
+
+    /// The invariant behind the fix: whichever partition observes itself to be
+    /// last sees the NULL its sibling recorded.
+    ///
+    /// This is the regression test for the wrong-results bug. One partition
+    /// records a NULL and finishes while the other finishes having recorded
+    /// nothing, and the last one must not conclude that no NULL was seen.
+    #[test]
+    fn the_last_partition_observes_a_sibling_null() {
+        loom::model(|| {
+            let completion = Arc::new(ProbeCompletion::new(2));
+
+            let recorder = {
+                let completion = Arc::clone(&completion);
+                loom::thread::spawn(move || {
+                    completion.record_batch(true, true);
+                    completion.report_completed()
+                })
+            };
+
+            let observer = completion.report_completed();
+            let recorded = recorder.join().unwrap();
+
+            let summary = match (recorded, observer) {
+                (Some(summary), None) | (None, Some(summary)) => summary,
+                (Some(_), Some(_)) => panic!("two partitions both finished last"),
+                (None, None) => panic!("no partition finished last"),
+            };
+
+            assert!(
+                summary.has_null,
+                "the last partition missed the NULL its sibling recorded"
+            );
+            assert!(
+                summary.non_empty,
+                "the last partition missed the row its sibling recorded"
+            );
+        });
+    }
+
+    /// With both partitions recording, the summary must be the union of what
+    /// they saw, whichever finishes last.
+    #[test]
+    fn the_summary_unions_what_every_partition_recorded() {
+        loom::model(|| {
+            let completion = Arc::new(ProbeCompletion::new(2));
+
+            let other = {
+                let completion = Arc::clone(&completion);
+                loom::thread::spawn(move || {
+                    // Rows, but no NULL key.
+                    completion.record_batch(true, false);
+                    completion.report_completed()
+                })
+            };
+
+            // A NULL key in an otherwise empty batch.
+            completion.record_batch(false, true);
+            let here = completion.report_completed();
+            let there = other.join().unwrap();
+
+            let summary = here.or(there).expect("some partition must finish last");
+            assert_eq!(
+                summary,
+                ProbeSideSummary {
+                    non_empty: true,
+                    has_null: true,
+                },
+                "the summary lost a fact one of the partitions recorded"
+            );
+        });
+    }
+
+    /// Exactly one partition may take the final stage. Two would emit the
+    /// build side twice; none would drop it.
+    #[test]
+    fn exactly_one_partition_finishes_last() {
+        loom::model(|| {
+            let completion = Arc::new(ProbeCompletion::new(2));
+
+            let other = {
+                let completion = Arc::clone(&completion);
+                loom::thread::spawn(move || completion.report_completed().is_some())
+            };
+
+            let here = completion.report_completed().is_some();
+            let there = other.join().unwrap();
+
+            assert!(here ^ there, "exactly one partition must finish last");
+        });
+    }
+
+    /// The probe-phase hint never invents a NULL: if no partition recorded
+    /// one, no observation of the hint may report one.
+    #[test]
+    fn the_hint_never_reports_an_unrecorded_null() {
+        loom::model(|| {
+            let completion = Arc::new(ProbeCompletion::new(2));
+
+            let other = {
+                let completion = Arc::clone(&completion);
+                loom::thread::spawn(move || {
+                    completion.record_batch(true, false);
+                    completion.saw_null_key_hint()
+                })
+            };
+
+            let here = completion.saw_null_key_hint();
+            let there = other.join().unwrap();
+
+            assert!(!here && !there, "the hint reported an unrecorded NULL");
+        });
+    }
+}
diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs
index 932a6b8018d9e..1c89e3a2c71f7 100644
--- a/datafusion/physical-plan/src/joins/hash_join/stream.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs
@@ -27,7 +27,8 @@ use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus};
 use crate::joins::Map;
 use crate::joins::MapOffset;
 use crate::joins::PartitionMode;
-use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode, ProbeSideSummary};
+use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode};
+use crate::joins::hash_join::probe_completion::ProbeSideSummary;
 use crate::joins::hash_join::shared_bounds::{
     PartitionBounds, PartitionBuildData, SharedBuildAccumulator,
 };

From 64fa4fef5ffd7314d29e57f1b5ea71c8b870a7c8 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Wed, 9 Sep 2026 22:28:59 +0800
Subject: [PATCH 3/6] test(hash_join): add tests for logical NULL handling in
 dictionary-encoded joins

---
 .../physical-plan/src/joins/hash_join/exec.rs | 118 ++++++++++++++++++
 1 file changed, 118 insertions(+)

diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs
index 1a8463c6cd5c8..51ab0de913f04 100644
--- a/datafusion/physical-plan/src/joins/hash_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs
@@ -7402,6 +7402,70 @@ mod tests {
         .unwrap()
     }
 
+    /// Builds the dictionary-encoded twin of
+    /// [`build_two_partition_probe_with_null_in_partition_0`]: the same two
+    /// probe partitions, but the join key is a dictionary and partition 0's
+    /// NULL exists only logically.
+    ///
+    /// Every dictionary key is a physically valid index; one of them points at
+    /// a NULL dictionary value. The array therefore reports
+    /// `null_count() == 0` while `logical_null_count() > 0`, which is the case
+    /// that a physical-NULL check silently misses.
+    fn build_two_partition_dict_probe_with_logical_null_in_partition_0()
+    -> Arc<dyn ExecutionPlan> {
+        let dict_type =
+            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("c2", dict_type, true),
+            Field::new("dummy", DataType::Int32, true),
+        ]));
+
+        // Dictionary values: [1, NULL]; keys: [0, 1] => logical [1, NULL].
+        let with_logical_null: ArrayRef = Arc::new(DictionaryArray::new(
+            Int32Array::from(vec![0, 1]),
+            Arc::new(Int32Array::from(vec![Some(1), None])),
+        ));
+        assert_eq!(
+            with_logical_null.null_count(),
+            0,
+            "the probe NULL must be logical only, or this test stops covering \
+             the logical_null_count path"
+        );
+        assert!(
+            with_logical_null.logical_null_count() > 0,
+            "the probe key must carry a logical NULL"
+        );
+
+        let partition_0 = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                with_logical_null,
+                Arc::new(Int32Array::from(vec![Some(100), Some(400)])),
+            ],
+        )
+        .unwrap();
+
+        // Dictionary values: [2]; keys: [0] => logical [2], no NULL anywhere.
+        let partition_1 = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(DictionaryArray::new(
+                    Int32Array::from(vec![0]),
+                    Arc::new(Int32Array::from(vec![Some(2)])),
+                )),
+                Arc::new(Int32Array::from(vec![Some(200)])),
+            ],
+        )
+        .unwrap();
+
+        TestMemoryExec::try_new_exec(
+            &[vec![partition_0], vec![partition_1]],
+            schema,
+            None,
+        )
+        .unwrap()
+    }
+
     /// Drains the probe partitions of `join` one after another in the given
     /// order and returns everything they emitted.
     async fn collect_partitions_in_order(
@@ -7470,6 +7534,60 @@ mod tests {
         Ok(())
     }
 
+    /// The dictionary counterpart of
+    /// [`test_null_aware_anti_join_probe_null_in_other_partition`], where the
+    /// silencing probe NULL exists only logically.
+    ///
+    /// A dictionary key that points at a NULL dictionary value has
+    /// `null_count() == 0`, so only the `logical_null_count()` check in
+    /// `null_aware_skip_probe_batch` records it. This pins that check to the
+    /// cross-partition summary: the partition holding the logical NULL is not
+    /// the one that emits the build rows, in either drain order.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_anti_join_probe_logical_null_in_other_partition(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        for order in [[0, 1], [1, 0]] {
+            // Build keys are dictionary-encoded too, and NULL-free.
+            let left = build_table_dict_key(
+                "c1",
+                vec![Some(1), Some(2), Some(3), Some(4)],
+                vec![0, 1, 2, 3],
+                "dummy",
+                vec![Some(10), Some(20), Some(30), Some(40)],
+            );
+            let right = build_two_partition_dict_probe_with_logical_null_in_partition_0();
+
+            let on = vec![(
+                Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+                Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+            )];
+
+            let join = HashJoinExec::try_new(
+                left,
+                right,
+                on,
+                None,
+                &JoinType::LeftAnti,
+                None,
+                PartitionMode::CollectLeft,
+                NullEquality::NullEqualsNothing,
+                true, // null_aware = true
+            )?;
+
+            let batches = collect_partitions_in_order(&join, &order, &task_ctx).await?;
+            let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+            assert_eq!(
+                rows, 0,
+                "probe order {order:?} emitted rows although a probe partition saw a logical NULL"
+            );
+        }
+        Ok(())
+    }
+
     /// The null-aware mark join counterpart of
     /// [`test_null_aware_anti_join_probe_null_in_other_partition`]: an
     /// unmatched build row must get an UNKNOWN (NULL) mark when the NULL probe

From f673b67a3223018f29609e2f4d64d8b76140eebe Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Wed, 9 Sep 2026 22:50:05 +0800
Subject: [PATCH 4/6] refactor(hash_join): simplify probe completion macro and
 improve documentation

---
 .../src/joins/hash_join/probe_completion.rs   | 62 +++++++++----------
 1 file changed, 30 insertions(+), 32 deletions(-)

diff --git a/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
index bc063dcdc4bad..eb00bc0cd6b9b 100644
--- a/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
@@ -43,24 +43,10 @@
 //!    it. Were the counter `Relaxed`, nothing would order a sibling's store
 //!    before the final load even when the decrement happened first.
 //!
-//! # Testing
-//!
-//! Neither property can be pinned down by an ordinary concurrent test. The
+//! Neither property can be pinned down by an ordinary concurrent test: the
 //! window is a few instructions wide, and on x86 a `Relaxed` decrement lowers
-//! to the same instruction as an `AcqRel` one, so a stress test would pass on
-//! the broken version. The invariant is instead model-checked with [loom],
-//! which enumerates the thread interleavings *and* the store visibility the
-//! memory model permits, in `loom_tests` below.
-//!
-//! Because loom can only explore its own atomic types, the protocol is written
-//! once in [`define_probe_completion`] and instantiated twice: over
-//! [`std::sync::atomic`] for the real join, and over `loom::sync::atomic` for
-//! the model. Both instantiations share this single copy of the orderings and
-//! the call sequence, so weakening the decrement to `Relaxed`, or reading the
-//! facts before it, fails the model. That is what makes these tests regression
-//! coverage rather than a restatement of the fix.
-//!
-//! [loom]: https://docs.rs/loom
+//! to the same instruction as an `AcqRel` one. Both are model-checked in
+//! `loom_tests` instead.
 
 use std::sync::atomic::Ordering;
 
@@ -78,13 +64,23 @@ pub(super) struct ProbeSideSummary {
     pub(super) has_null: bool,
 }
 
-/// Defines [`ProbeCompletion`] over a given pair of atomic types.
+/// Defines [`ProbeCompletion`] over the given atomic types.
 ///
-/// The indirection exists so the loom model below can instantiate the very
-/// same logic over loom's atomics. See the [module docs](self); do not add a
-/// second copy of these orderings anywhere.
+/// [loom] can only explore its own atomics, so the protocol is written once
+/// here and instantiated twice: over [`std::sync::atomic`] for the join, and
+/// over `loom::sync::atomic` for the model in `loom_tests`. Both share this
+/// single copy of the orderings and the call sequence, which is what lets the
+/// model speak for the real thing. Keep it that way: a second copy of these
+/// orderings would not be covered.
+///
+/// [loom]: https://docs.rs/loom
 macro_rules! define_probe_completion {
-    ($atomic_usize:ty, $atomic_bool:ty) => {
+    (counter: $counter:ty, flag: $flag:ty) => {
+        /// Counts the probe partitions that are still running.
+        type Counter = $counter;
+        /// One fact, set by any partition and read by the last one.
+        type Flag = $flag;
+
         /// Tracks how many probe partitions are still running, together with
         /// the null-aware facts they contribute.
         ///
@@ -92,20 +88,20 @@ macro_rules! define_probe_completion {
         #[derive(Debug)]
         pub(super) struct ProbeCompletion {
             /// Probe partitions that have not finished yet.
-            running: $atomic_usize,
+            running: Counter,
             /// Set once any probe partition has seen a row.
-            saw_row: $atomic_bool,
+            saw_row: Flag,
             /// Set once any probe partition has seen a NULL join key.
-            saw_null_key: $atomic_bool,
+            saw_null_key: Flag,
         }
 
         impl ProbeCompletion {
             /// Creates the protocol state for `probe_threads` partitions.
             pub(super) fn new(probe_threads: usize) -> Self {
                 Self {
-                    running: <$atomic_usize>::new(probe_threads),
-                    saw_row: <$atomic_bool>::new(false),
-                    saw_null_key: <$atomic_bool>::new(false),
+                    running: Counter::new(probe_threads),
+                    saw_row: Flag::new(false),
+                    saw_null_key: Flag::new(false),
                 }
             }
 
@@ -140,6 +136,8 @@ macro_rules! define_probe_completion {
             /// and acquires those of the partitions that finished earlier, so
             /// the summary handed to the last caller is complete.
             pub(super) fn report_completed(&self) -> Option<ProbeSideSummary> {
+                // `fetch_sub` returns the value from *before* the decrement,
+                // so `1` means this call is the one that took it to zero.
                 let was_last = self.running.fetch_sub(1, Ordering::AcqRel) == 1;
                 was_last.then(|| ProbeSideSummary {
                     non_empty: self.saw_row.load(Ordering::Relaxed),
@@ -151,8 +149,8 @@ macro_rules! define_probe_completion {
 }
 
 define_probe_completion!(
-    std::sync::atomic::AtomicUsize,
-    std::sync::atomic::AtomicBool
+    counter: std::sync::atomic::AtomicUsize,
+    flag: std::sync::atomic::AtomicBool
 );
 
 #[cfg(test)]
@@ -225,8 +223,8 @@ mod loom_tests {
     use std::sync::atomic::Ordering;
 
     define_probe_completion!(
-        loom::sync::atomic::AtomicUsize,
-        loom::sync::atomic::AtomicBool
+        counter: loom::sync::atomic::AtomicUsize,
+        flag: loom::sync::atomic::AtomicBool
     );
 
     /// The invariant behind the fix: whichever partition observes itself to be

From 2a643360415d7641c49a78a016104137f5f6d3a1 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Thu, 10 Sep 2026 20:25:00 +0800
Subject: [PATCH 5/6] fix(hash_join): carry probe summary into chunked final
 build-row emission

Merging main's chunked emission of final build rows split the final stage
into `prepare_unmatched_build_rows` and `emit_unmatched_build_rows`, leaving
the null-aware post-processing in the second one while `probe_summary` was
still bound in the first. `report_probe_completed` hands the summary out
exactly once, so capture it alongside the bitmap snapshot in
`EmitUnmatchedBuildRowsState` and read it back per chunk.
---
 datafusion/physical-plan/src/joins/hash_join/stream.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs
index 7e42a45d049f6..625661188d81c 100644
--- a/datafusion/physical-plan/src/joins/hash_join/stream.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs
@@ -210,6 +210,10 @@ pub(super) struct EmitUnmatchedBuildRowsState {
     visited: BooleanBuffer,
     /// Index of the next build row to examine
     cursor: usize,
+    /// What every probe partition together saw, needed by the null-aware
+    /// post-processing of each chunk. Captured with the bitmap snapshot
+    /// because [`JoinLeftData::report_probe_completed`] hands it out once.
+    probe_summary: ProbeSideSummary,
 }
 
 /// Lifecycle of this partition's build-data report to the shared coordinator.
@@ -1056,6 +1060,7 @@ impl HashJoinStream {
             HashJoinStreamState::EmitUnmatchedBuildRows(EmitUnmatchedBuildRowsState {
                 visited,
                 cursor: 0,
+                probe_summary,
             });
 
         Ok(StatefulStreamResult::Continue)
@@ -1083,6 +1088,8 @@ impl HashJoinStream {
 
         let build_side = self.build_side.try_as_ready()?;
 
+        let probe_summary = state.probe_summary;
+
         // use the global left bitmap to produce the left indices and right indices
         let (left_side, right_side) = next_final_indices_chunk(
             &state.visited,

From 98346799752db567301828afb4672ea56285f5d8 Mon Sep 17 00:00:00 2001
From: Jay Zhan <jayzhan211@gmail.com>
Date: Thu, 10 Sep 2026 22:24:47 +0800
Subject: [PATCH 6/6] test(hash_join): keep loom out of the ordinary build

`loom` was a plain dev-dependency, so it joined every `cargo test
--workspace` resolve. It is not inert there: loom pulls in
tracing-subscriber, whose env-filter turns on `regex-automata`'s
`dfa-build`/`dfa-search` features for the whole workspace build. Those
change which engine the shared `regex` selects, and the new path costs
enough extra stack to abort unrelated deep-recursion tests --
`sql::unparser::test_tpch_unparser_roundtrip` overflowed its stack in CI
on every run since the dependency was added.

Measured on `core_integration`, varying only whether loom is in the
resolve (`RUST_MIN_STACK` at which the test passes):

    without loom   480K ok   1M ok     2M ok     4M ok
    with loom      480K ok   1M FAIL   2M FAIL   4M ok

Declaring loom under `cfg(datafusion_loom)` leaves the ordinary build
byte-identical to one that never mentioned it. The cfg is namespaced
because a bare `loom` is also read by tokio, which would switch tokio
into its own loom build and drop `tokio::fs`.

Run the model checks with:

    RUSTFLAGS="--cfg datafusion_loom" \
        cargo test -p datafusion-physical-plan --lib loom_tests
---
 Cargo.toml                                    |  5 ++++
 datafusion/physical-plan/Cargo.toml           | 23 +++++++++++++++----
 .../src/joins/hash_join/probe_completion.rs   | 10 +++++++-
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 01672a2b39be9..2d027ede7f5b1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -323,6 +323,11 @@ unexpected_cfgs = { level = "warn", check-cfg = [
     'cfg(datafusion_coop, values("tokio", "tokio_fallback", "per_stream"))',
     "cfg(coverage)",
     "cfg(coverage_nightly)",
+    # Enables the loom model-checked concurrency tests; see
+    # datafusion/physical-plan/src/joins/hash_join/probe_completion.rs.
+    # Namespaced rather than a bare `loom`, which tokio also reads and which
+    # would switch tokio itself into its loom build.
+    "cfg(datafusion_loom)",
 ] }
 unused_qualifications = "deny"
 
diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml
index 614a6437160f0..1f927345d20ba 100644
--- a/datafusion/physical-plan/Cargo.toml
+++ b/datafusion/physical-plan/Cargo.toml
@@ -96,10 +96,6 @@ criterion = { workspace = true, features = ["async_futures"] }
 datafusion-functions-aggregate = { workspace = true }
 datafusion-functions-window = { workspace = true }
 insta = { workspace = true }
-# Model checker for the hash join probe-completion protocol; see
-# joins/hash_join/probe_completion.rs. Test-only, so it stays out of the
-# dependency tree of anything that links DataFusion.
-loom = { workspace = true }
 rand = { workspace = true }
 rstest = { workspace = true }
 rstest_reuse = "0.7.0"
@@ -109,6 +105,25 @@ tokio = { workspace = true, features = [
     "parking_lot",
 ] }
 
+# Model checker for the hash join probe-completion protocol; see
+# joins/hash_join/probe_completion.rs.
+#
+# Gated on `--cfg datafusion_loom` rather than declared a plain dev-dependency: loom pulls
+# in tracing-subscriber, which turns on `regex-automata`'s DFA features for
+# every crate in the workspace build. That is not inert. It changes which
+# engine the shared `regex` picks and costs enough extra stack to overflow
+# unrelated deep-recursion tests (`sql::unparser::test_tpch_unparser_roundtrip`
+# aborts in CI). Keeping loom behind its own cfg leaves the ordinary build
+# untouched.
+#
+# Run the model checks with:
+#   RUSTFLAGS="--cfg datafusion_loom" cargo test -p datafusion-physical-plan --lib loom_tests
+#
+# The cfg is namespaced because a bare `loom` is also read by tokio, which
+# would switch tokio into its own loom build and drop `tokio::fs`.
+[target.'cfg(datafusion_loom)'.dev-dependencies]
+loom = { workspace = true }
+
 [[bench]]
 harness = false
 name = "partial_ordering"
diff --git a/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
index eb00bc0cd6b9b..11fb72d352c6a 100644
--- a/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
+++ b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs
@@ -216,7 +216,15 @@ mod tests {
 /// [module docs](self)) and let loom enumerate the interleavings and store
 /// visibility the memory model allows. They fail if the decrement in
 /// `report_completed` is weakened to `Relaxed`.
-#[cfg(test)]
+///
+/// Compiled only under `--cfg datafusion_loom`, which is also what pulls the `loom`
+/// dev-dependency into the build (see `datafusion-physical-plan`'s
+/// `Cargo.toml` for why it is kept out of the ordinary one). Run them with:
+///
+/// ```text
+/// RUSTFLAGS="--cfg datafusion_loom" cargo test -p datafusion-physical-plan --lib loom_tests
+/// ```
+#[cfg(all(test, datafusion_loom))]
 mod loom_tests {
     use super::ProbeSideSummary;
     use loom::sync::Arc;