diff --git a/Cargo.lock b/Cargo.lock index 6d4c38d8d9d0..0c386b761ab7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2516,6 +2516,7 @@ dependencies = [ "insta", "itertools 0.15.0", "log", + "loom", "num-traits", "parking_lot", "pin-project-lite", @@ -3228,6 +3229,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" @@ -4084,6 +4100,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" @@ -4099,6 +4128,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" @@ -5541,6 +5579,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" @@ -6578,10 +6622,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 efeb58007413..2d027ede7f5b 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 } @@ -322,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 534cea8ea9cb..1f927345d20b 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -105,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/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 5f2274166234..b72e180543f9 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,12 +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 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, - // 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,10 +367,29 @@ 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) { + self.probe_completion.record_batch(non_empty, has_null); + } + + /// 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_completion.saw_null_key_hint() + } + + /// Marks this probe partition as finished, returning `Some` for the last + /// one together with what every partition saw. + /// + /// 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 { + self.probe_completion.report_completed() } } @@ -2980,12 +2993,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, }; @@ -7537,6 +7548,278 @@ 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 { + 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() + } + + /// 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 { + 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( + join: &HashJoinExec, + order: &[usize], + task_ctx: &Arc, + ) -> Result> { + 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 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( + 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 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 + /// 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( + 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/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index b915802ea401..7c1b0d76c2f7 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 000000000000..11fb72d352c6 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/probe_completion.rs @@ -0,0 +1,350 @@ +// 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. +//! +//! 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. Both are model-checked in +//! `loom_tests` instead. + +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 the given atomic types. +/// +/// [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 { + (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. + /// + /// See the [module docs](self) for the invariant this upholds. + #[derive(Debug)] + pub(super) struct ProbeCompletion { + /// Probe partitions that have not finished yet. + running: Counter, + /// Set once any probe partition has seen a row. + saw_row: Flag, + /// Set once any probe partition has seen a NULL join key. + saw_null_key: Flag, + } + + impl ProbeCompletion { + /// Creates the protocol state for `probe_threads` partitions. + pub(super) fn new(probe_threads: usize) -> Self { + Self { + running: Counter::new(probe_threads), + saw_row: Flag::new(false), + saw_null_key: Flag::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 { + // `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), + has_null: self.saw_null_key.load(Ordering::Relaxed), + }) + } + } + }; +} + +define_probe_completion!( + counter: std::sync::atomic::AtomicUsize, + flag: 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`. +/// +/// 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; + use std::sync::atomic::Ordering; + + define_probe_completion!( + counter: loom::sync::atomic::AtomicUsize, + flag: 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 2c1ad9446054..625661188d81 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -21,7 +21,6 @@ //! [`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}; @@ -29,6 +28,7 @@ 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::probe_completion::ProbeSideSummary; use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, SharedBuildAccumulator, }; @@ -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. @@ -1029,24 +1033,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); - } + }; // Every probe partition has finished, so the bitmap is final: snapshot // it once and release the lock for the whole emission phase. @@ -1065,6 +1060,7 @@ impl HashJoinStream { HashJoinStreamState::EmitUnmatchedBuildRows(EmitUnmatchedBuildRowsState { visited, cursor: 0, + probe_summary, }); Ok(StatefulStreamResult::Continue) @@ -1092,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, @@ -1106,6 +1104,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, ); @@ -1114,6 +1113,7 @@ impl HashJoinStream { Some(NullAwareMode::LeftMark { .. }) => { let mark_column = null_aware_left_mark_column( &build_side.left_data, + probe_summary, &left_side, &right_side, ); @@ -1223,12 +1223,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 { @@ -1237,11 +1231,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() } } } @@ -1271,15 +1265,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 @@ -1294,14 +1296,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 @@ -1313,8 +1314,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, ) }