From 3d586a3e8fbff1d6b01d793dc1c8b1557d42faac Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 11 Oct 2026 10:21:01 -0700 Subject: [PATCH 1/3] feat: build multi-column null-aware NOT IN on the subquery side A multi-column or correlated `NOT IN` plans as a null-aware LeftAnti hash join that JoinSelection never swapped, so the outer side was always the build side and held in memory even when the subquery was small. Let JoinSelection swap every null-aware LeftAnti into a RightAnti. A correlated null-aware RightAnti (more than one key or a join filter) now decides each outer (probe) row while streaming: it reuses the NULL candidate pairing of the correlated LeftAnti (scope maps, the element-wise mismatch check and the join filter), but marks the UNKNOWN probe rows of the current batch instead of build rows in the shared bitmap, and drops them from the anti join output. Co-authored-by: Isaac --- .../null_aware_join/benchmarks/q10.benchmark | 6 +- .../null_aware_join/benchmarks/q12.benchmark | 8 +- .../tests/fuzz_cases/null_aware_join_fuzz.rs | 58 ++- datafusion/core/tests/memory_limit/mod.rs | 1 + .../tests/memory_limit/null_aware_not_in.rs | 103 +++++ .../physical_optimizer/join_selection.rs | 81 +++- .../physical-optimizer/src/join_selection.rs | 12 +- .../physical-plan/src/joins/hash_join/exec.rs | 359 ++++++++++++------ .../src/joins/hash_join/stream.rs | 306 ++++++++++----- .../test_files/null_aware_anti_join.slt | 17 +- .../test_files/null_aware_multi_column.slt | 7 +- 11 files changed, 714 insertions(+), 244 deletions(-) create mode 100644 datafusion/core/tests/memory_limit/null_aware_not_in.rs diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q10.benchmark index 1857ef64d2e15..3ed9ff477c755 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q10.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q10.benchmark @@ -23,9 +23,9 @@ expect_plan HashJoinExec run -- Q10: uncorrelated multi-column NOT IN over nullable keys that hold no NULL. --- A tuple NOT IN always decides UNKNOWN per build row (the outer side), but --- without NULLs there is no candidate pair to check, so this is the floor for --- the multi-column shape. +-- A tuple NOT IN decides UNKNOWN per row (here per outer row, as the join is +-- swapped to build on the subquery), but without NULLs there is no candidate +-- pair to check, so this is the floor for the multi-column shape. SELECT count(*) FROM small_outer o WHERE (o.id_n0, o.z) NOT IN (SELECT i.id_n0, i.z FROM small_inner i); diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q12.benchmark index c126c9c16163c..8197a476984b8 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q12.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q12.benchmark @@ -26,10 +26,10 @@ expect_plan HashJoinExec run -- Q12: uncorrelated multi-column NOT IN with 1% NULL in one outer element, --- the mirror of Q11. The outer side is the build side, so each NULL-valued --- build row is paired with every subquery row and kept as UNKNOWN only when --- the other element matches: the cost grows with (NULL-valued outer rows) x --- (subquery rows). Compare with Q10 and Q11. +-- the mirror of Q11. Each NULL-valued outer row is paired with every subquery +-- row and is UNKNOWN as soon as the other element of one of them matches: the +-- cost grows with (NULL-valued outer rows) x (subquery rows). Compare with Q10 +-- and Q11. SELECT count(*) FROM small_outer o WHERE (o.id_n1, o.z) NOT IN (SELECT i.id_n0, i.z FROM small_inner i); diff --git a/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs b/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs index 4bf7171771570..a4d45ac0fb973 100644 --- a/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs @@ -21,6 +21,10 @@ //! null-aware hash joins, are compared with `NOT EXISTS` / `EXISTS` //! formulations of the same three-valued logic, which do not. Random tables //! with NULLs on both sides are swept across batch sizes and partition counts. +//! +//! The subquery table is sometimes much smaller than the outer table, so +//! `JoinSelection` also swaps `NOT IN` into a null-aware `RightAnti` that +//! builds on the subquery side; the test checks that both plan shapes ran. use std::sync::Arc; @@ -28,16 +32,21 @@ use arrow::array::{Int32Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::util::pretty::pretty_format_batches; use datafusion::datasource::MemTable; +use datafusion::physical_plan::displayable; use datafusion::prelude::{SessionConfig, SessionContext}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; #[tokio::test(flavor = "multi_thread")] async fn null_aware_join_matches_reference() { + let mut plans = NullAwarePlans::default(); for seed in 0..20 { let mut rng = StdRng::seed_from_u64(seed); - let outer = random_table(&mut rng, &["id", "k", "a", "b", "c"]); - let inner = random_table(&mut rng, &["k", "x", "y", "z"]); + let outer = random_table(&mut rng, &["id", "k", "a", "b", "c"], 6); + // Alternate between a subquery table of similar size and a much + // smaller one that `JoinSelection` moves to the build side. + let inner_rows = if seed % 2 == 0 { 6 } else { 2 }; + let inner = random_table(&mut rng, &["k", "x", "y", "z"], inner_rows); for target_partitions in [1, 4] { for batch_size in [1, 3, 8192] { let config = SessionConfig::new() @@ -57,6 +66,7 @@ async fn null_aware_join_matches_reference() { ctx.register_table("o", Arc::clone(&outer) as _).unwrap(); ctx.register_table("i", Arc::clone(&inner) as _).unwrap(); for (query, reference) in cases() { + plans.record(&ctx, &query).await; let actual = run(&ctx, &query).await; let expected = run(&ctx, &reference).await; assert_eq!( @@ -68,6 +78,43 @@ async fn null_aware_join_matches_reference() { } } } + // A fuzz run that never reaches one of the plan shapes would be vacuous. + assert!(plans.left_anti > 0, "{plans:?}"); + assert!(plans.right_anti > 0, "{plans:?}"); + assert!(plans.correlated_right_anti > 0, "{plans:?}"); +} + +/// Counts the null-aware anti join plans the fuzz queries ran. +#[derive(Debug, Default)] +struct NullAwarePlans { + left_anti: usize, + right_anti: usize, + /// `RightAnti` joins with more than one key or a join filter, which + /// decide `NOT IN` per probe row. + correlated_right_anti: usize, +} + +impl NullAwarePlans { + async fn record(&mut self, ctx: &SessionContext, sql: &str) { + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan = displayable(plan.as_ref()).indent(true).to_string(); + for line in plan.lines().filter(|line| line.contains("null_aware")) { + if line.contains("join_type=LeftAnti") { + self.left_anti += 1; + } else if line.contains("join_type=RightAnti") { + self.right_anti += 1; + if line.contains("), (") || line.contains("filter=") { + self.correlated_right_anti += 1; + } + } + } + } } /// Pairs of (query, reference) that must return the same rows. @@ -131,8 +178,9 @@ fn cases() -> Vec<(String, String)> { } /// A table of nullable `Int32` columns with values in `0..3`, split into up -/// to three partitions. A column named `id` holds unique row ids instead. -fn random_table(rng: &mut StdRng, columns: &[&str]) -> Arc { +/// to three partitions of fewer than `max_rows` rows each. A column named `id` +/// holds unique row ids instead. +fn random_table(rng: &mut StdRng, columns: &[&str], max_rows: i32) -> Arc { let schema = Arc::new(Schema::new( columns .iter() @@ -143,7 +191,7 @@ fn random_table(rng: &mut StdRng, columns: &[&str]) -> Arc { let mut next_id = 0; let partitions = (0..rng.random_range(1..4)) .map(|_| { - let rows = rng.random_range(0..6); + let rows = rng.random_range(0..max_rows); let arrays = columns .iter() .map(|name| { diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 2d48ad6790be0..f457864808972 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -23,6 +23,7 @@ use std::sync::{Arc, LazyLock}; #[cfg(feature = "extended_tests")] mod memory_limit_validation; mod nlj_spill_unmatched; +mod null_aware_not_in; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; diff --git a/datafusion/core/tests/memory_limit/null_aware_not_in.rs b/datafusion/core/tests/memory_limit/null_aware_not_in.rs new file mode 100644 index 0000000000000..4bdeb8fe03451 --- /dev/null +++ b/datafusion/core/tests/memory_limit/null_aware_not_in.rs @@ -0,0 +1,103 @@ +// 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. + +//! A multi-column `NOT IN` over a small subquery must build its hash join on +//! the subquery, so a large outer table runs under a memory limit that only +//! fits the subquery side. +//! +//! The null-aware join of a multi-column `NOT IN` used to stay a `LeftAnti` +//! that builds on the outer side, so it ran out of memory where the scalar +//! `NOT IN`, swapped to a `RightAnti` that builds on the subquery, ran. + +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::MemTable; +use datafusion::prelude::*; +use datafusion_common::{Result, assert_contains}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; + +const OUTER_ROWS: i64 = 200_000; +/// Far below the outer side's build state, far above the subquery's. +const MEMORY_LIMIT: usize = 2 * 1024 * 1024; + +/// A table `(a, b)` of `a = v, b = v % 1000` for `v` in `1..=rows`, followed by +/// `extra` rows. +fn table(rows: i64, extra: &[(Option, Option)]) -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + ])); + let a = (1..=rows).map(Some).chain(extra.iter().map(|row| row.0)); + let b = (1..=rows) + .map(|v| Some(v % 1000)) + .chain(extra.iter().map(|row| row.1)); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(a.collect::()), + Arc::new(b.collect::()), + ], + ) + .unwrap(); + // Several batches, so the outer side streams through the probe. + let batches = (0..batch.num_rows()) + .step_by(8192) + .map(|offset| batch.slice(offset, 8192.min(batch.num_rows() - offset))) + .collect(); + Arc::new(MemTable::try_new(schema, vec![batches]).unwrap()) +} + +async fn count_not_in(sql: &str, join_reordering: bool) -> Result { + let config = SessionConfig::new() + .with_target_partitions(1) + .set_bool("datafusion.optimizer.join_reordering", join_reordering); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(MEMORY_LIMIT, 1.0) + .build_arc()?; + let ctx = SessionContext::new_with_config_rt(config, runtime); + ctx.register_table("big", table(OUTER_ROWS, &[]))?; + // A NULL-valued subquery row makes `NOT IN` UNKNOWN for every outer row + // with `b = 7` (and no definite mismatch in `a`). + ctx.register_table("small", table(100, &[(None, Some(7))]))?; + + let batches = ctx.sql(sql).await?.collect().await?; + Ok(batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0)) +} + +#[tokio::test] +async fn multi_column_not_in_builds_on_small_subquery() -> Result<()> { + let sql = "SELECT count(*) FROM big \ + WHERE (a, b) NOT IN (SELECT a, b FROM small)"; + + // Rows 1..=100 match exactly; of the rest, those with `b = 7` are + // UNKNOWN: 199 of the 200 values `v % 1000 = 7`, all but `v = 7`. + assert_eq!(count_not_in(sql, true).await?, OUTER_ROWS - 100 - 199); + + // Building on the outer side does not fit, which shows the limit binds. + let err = count_not_in(sql, false).await.unwrap_err(); + assert_contains!(err.to_string(), "Resources exhausted"); + assert_contains!(err.to_string(), "HashJoinInput"); + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index d4561decc6331..f025e3d17010a 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -43,7 +43,9 @@ use datafusion_physical_optimizer::join_selection::JoinSelection; use datafusion_physical_plan::displayable; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::utils::JoinFilter; -use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode}; +use datafusion_physical_plan::joins::{ + HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode, +}; use datafusion_physical_plan::operator_statistics::{ ClosureStatisticsProvider, StatisticsRegistry, StatisticsResult, }; @@ -607,11 +609,10 @@ async fn test_null_aware_left_anti_respects_disabled_join_reordering() -> Result } /// A filtered null-aware `LeftAnti` (produced by decorrelating a correlated -/// `NOT IN`) must not be swapped to `RightAnti`: the swap is only valid for -/// unfiltered null-aware joins, since a filtered `RightAnti` would apply its -/// build-side NULL short-circuit before the filter. +/// `NOT IN`) swaps to a null-aware `RightAnti` that decides `NOT IN` per probe +/// row; the filter's sides are swapped with the inputs. #[tokio::test] -async fn test_null_aware_left_anti_with_filter_does_not_swap() -> Result<()> { +async fn test_null_aware_left_anti_with_filter_swaps() -> Result<()> { let (big, small) = create_big_and_small(); let join = HashJoinExec::try_new( Arc::clone(&big), @@ -630,16 +631,70 @@ async fn test_null_aware_left_anti_with_filter_does_not_swap() -> Result<()> { let optimized_join = JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?; - let unswapped_join = optimized_join + let swapped_join = optimized_join .downcast_ref::() - .expect("filtered null-aware anti join should not swap"); + .expect("anti join swap should not require a projection"); - assert_eq!(*unswapped_join.join_type(), JoinType::LeftAnti); - assert_eq!(*unswapped_join.partition_mode(), PartitionMode::CollectLeft); - assert!(unswapped_join.null_aware); - assert!(unswapped_join.filter().is_some()); - assert_eq!(unswapped_join.left().schema().field(0).name(), "big_col"); - assert_eq!(unswapped_join.right().schema().field(0).name(), "small_col"); + assert_eq!(*swapped_join.join_type(), JoinType::RightAnti); + assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft); + assert!(swapped_join.null_aware); + let filter = swapped_join.filter().expect("filter should be kept"); + assert_eq!( + filter + .column_indices() + .iter() + .map(|c| c.side) + .collect::>(), + vec![JoinSide::Right, JoinSide::Left] + ); + assert_eq!(swapped_join.left().schema().field(0).name(), "small_col"); + assert_eq!(swapped_join.right().schema().field(0).name(), "big_col"); + + Ok(()) +} + +/// A multi-key null-aware `LeftAnti` (a multi-column `NOT IN`, correlation +/// scope keys, or both) swaps to a `RightAnti` that keeps the number of +/// `NOT IN` value keys and the key order. +#[rstest] +#[case(1)] +#[case(2)] +#[tokio::test] +async fn test_null_aware_multi_key_left_anti_swaps_to_right_anti( + #[case] value_keys: usize, +) -> Result<()> { + let (big, small) = create_big_and_small(); + let big_col = Arc::new(Column::new_with_schema("big_col", &big.schema())?) as _; + let small_col = Arc::new(Column::new_with_schema("small_col", &small.schema())?) as _; + let join = HashJoinExecBuilder::new( + Arc::clone(&big), + Arc::clone(&small), + vec![ + (Arc::clone(&big_col), Arc::clone(&small_col)), + (big_col, small_col), + ], + JoinType::LeftAnti, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .with_null_aware(true) + .with_null_aware_value_keys(value_keys) + .build()?; + let original_schema = join.schema(); + + let optimized_join = + JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?; + let swapped_join = optimized_join + .downcast_ref::() + .expect("anti join swap should not require a projection"); + + assert_eq!(*swapped_join.join_type(), JoinType::RightAnti); + assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft); + assert!(swapped_join.null_aware); + assert_eq!(swapped_join.null_aware_value_keys, value_keys); + assert_eq!(swapped_join.on().len(), 2); + assert_eq!(swapped_join.left().schema().field(0).name(), "small_col"); + assert_eq!(swapped_join.right().schema().field(0).name(), "big_col"); + assert_eq!(swapped_join.schema(), original_schema); Ok(()) } diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 2b342bece7040..eafec278f6277 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -172,14 +172,14 @@ impl PhysicalOptimizerRule for JoinSelection { } } -/// Determines whether it is possible to swap inputs of a hash join - for null-aware joins, we can only swap an uncorrelated `LeftAnti` -/// (a single join key and no filter), because the swapped `RightAnti` has no per-row NULL handling +/// Determines whether it is possible to swap inputs of a hash join. Of the +/// null-aware joins only a `LeftAnti` can be swapped: the resulting `RightAnti` +/// decides `NOT IN` per probe row, including multi-column and correlated +/// (scope keys or a join filter) shapes. A null-aware `LeftMark` has no +/// null-aware `RightMark` counterpart. fn can_swap_hash_join(hash_join: &HashJoinExec) -> bool { hash_join.join_type().supports_swap() - && (!hash_join.null_aware - || (*hash_join.join_type() == JoinType::LeftAnti - && hash_join.on().len() == 1 - && hash_join.filter().is_none())) + && (!hash_join.null_aware || *hash_join.join_type() == JoinType::LeftAnti) } /// Tries to create a [`HashJoinExec`] in [`PartitionMode::CollectLeft`] when possible. diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index de39c1638608a..036b07793934f 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -207,8 +207,9 @@ fn array_map_key_range( /// multi-column key), used by correlated null-aware joins (see /// [`NullAwareMode`]). /// -/// Such rows are UNKNOWN whenever *any* probe row in their correlation scope -/// passes the join filter, so every probe row must be tested against them. +/// Such rows make `NOT IN` UNKNOWN whenever *any* probe row in their +/// correlation scope passes the join filter (unless another element is a +/// definite mismatch), so every probe row must be tested against them. /// Restricting this lookup to the NULL-valued build rows keeps it /// proportional to the number of NULLs instead of enumerating every scope /// match of every probe row. @@ -242,9 +243,10 @@ pub(super) fn null_value_key_mask(value_keys: &[ArrayRef]) -> Option Self::RightAnti, + // `PartitionMode::CollectLeft` is safe because `RightAnti` is + // probe-driven: every probe row sees the whole build side. + (JoinType::RightAnti, PartitionMode::CollectLeft) => Self::RightAnti { + correlated, + value_keys: num_value_keys, + }, (JoinType::LeftMark, _) => Self::LeftMark { correlated, value_keys: num_value_keys, @@ -298,9 +307,6 @@ impl NullAwareMode { } }; match mode { - Self::RightAnti if num_keys != 1 => plan_err!( - "null_aware {join_type} joins only support single column join key, got {num_keys} columns" - ), // Correlated joins share the per-build-row null bitmap across all // probe partitions. Self::LeftMark { .. } @@ -311,35 +317,32 @@ impl NullAwareMode { "null_aware joins require PartitionMode::CollectLeft, got PartitionMode::Partitioned" ) } - Self::RightAnti if has_filter => { - plan_err!("null_aware RightAnti join does not support a join filter") - } _ => Ok(mode), } } - /// Whether this join decides UNKNOWN per build row (see [`NullAwareMode`]). + /// Whether this join decides UNKNOWN per row (see [`NullAwareMode`]). pub(super) fn is_correlated(self) -> bool { - matches!( - self, - Self::LeftAnti { - correlated: true, - .. - } | Self::LeftMark { - correlated: true, - .. - } - ) + match self { + Self::LeftAnti { correlated, .. } + | Self::RightAnti { correlated, .. } + | Self::LeftMark { correlated, .. } => correlated, + } + } + + /// Whether this join records UNKNOWN per build row in the shared + /// null-indices bitmap: a correlated `LeftAnti` or `LeftMark`. + pub(super) fn marks_build_rows(self) -> bool { + self.is_correlated() && !matches!(self, Self::RightAnti { .. }) } /// Number of leading keys that are `NOT IN` value keys; the keys after /// them are correlation scope keys. pub(super) fn value_keys(self) -> usize { match self { - Self::LeftAnti { value_keys, .. } | Self::LeftMark { value_keys, .. } => { - value_keys - } - Self::RightAnti => 1, + Self::LeftAnti { value_keys, .. } + | Self::RightAnti { value_keys, .. } + | Self::LeftMark { value_keys, .. } => value_keys, } } } @@ -390,7 +393,8 @@ pub(super) struct JoinLeftData { /// Tracks the final probe partition and what every partition observed. /// See [`ProbeCompletion`] for the completion invariant. probe_completion: ProbeCompletion, - /// Whether the smaller RightAnti subquery contains a NULL key. + /// Whether the subquery of an uncorrelated null-aware `RightAnti` join + /// contains a NULL key. pub(super) build_side_has_null: bool, // Drop the reservation after the mutable state it also accounts for. build: Arc, @@ -1015,10 +1019,14 @@ pub struct HashJoinExec { /// A join filter holds the non-equality part of a correlated subquery, and /// only the probe rows that pass it take part in the three-valued logic. /// + /// `JoinSelection` may swap a null-aware `LeftAnti` into a `RightAnti` + /// join that builds on the subquery side and decides each outer (probe) + /// row as it streams. + /// /// Key-ordering convention (relied on positionally, not enforced): for a - /// null-aware `LeftAnti` or `LeftMark` join, `on[..V]` are the `NOT IN` - /// value keys and `on[V..]` the correlated equality scope keys, where `V` - /// is [`Self::null_aware_value_keys`]. + /// null-aware join, `on[..V]` are the `NOT IN` value keys and `on[V..]` + /// the correlated equality scope keys, where `V` is + /// [`Self::null_aware_value_keys`]. /// Reordering these keys would silently produce wrong results, which is why /// such joins are pinned to `PartitionMode::CollectLeft` (the only key /// reorderer acts solely on `PartitionMode::Partitioned`). @@ -1033,13 +1041,10 @@ pub struct HashJoinExec { /// one element does not make every comparison UNKNOWN the way a scalar /// NULL does. /// - /// Such a join always decides UNKNOWN per build row, like a correlated - /// scalar `NOT IN`, which has two costs the uncorrelated scalar join does - /// not: the outer side stays the build side, since only a single-key join - /// is swapped, and without correlation keys every NULL-valued row is paired - /// with every row on the other side to check its other elements. See - /// and - /// for the follow-ups. + /// Such a join decides UNKNOWN per row, like a correlated scalar `NOT IN`: + /// without correlation keys every NULL-valued row is paired with every row + /// on the other side to check its other elements. See + /// for the follow-up. pub null_aware_value_keys: usize, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, @@ -3256,9 +3261,13 @@ async fn collect_left_input( let schema = left_stream.schema(); let prepared = mode == BuildMode::Prepared; - // The extra scope maps + null bitmap are only built for correlated - // null-aware joins (see `NullAwareMode`). - let with_null_aware_row_state = null_aware.is_some_and(NullAwareMode::is_correlated); + // The extra scope maps are only built for correlated null-aware joins, and + // the null bitmap only for those that mark build rows (see + // `NullAwareMode`). + let with_null_aware_scope_state = + null_aware.is_some_and(NullAwareMode::is_correlated); + let with_null_aware_row_state = + null_aware.is_some_and(NullAwareMode::marks_build_rows); let null_aware_value_keys = null_aware.map_or(1, NullAwareMode::value_keys); let is_phj_candidate = is_perfect_hash_join_candidate(&on_left, &schema)?; @@ -3459,7 +3468,7 @@ async fn collect_left_input( BooleanBufferBuilder::new(0) }; - let (null_aware_scope_map, null_value_build_rows) = if with_null_aware_row_state { + let (null_aware_scope_map, null_value_build_rows) = if with_null_aware_scope_state { // Null-aware convention: `on_left[..V]` are the value keys and // `on_left[V..]` the (possibly empty) correlation scope keys. let scope_keys = &on_left[null_aware_value_keys..]; @@ -3591,8 +3600,13 @@ async fn collect_left_input( } }; - let build_has_null = null_aware == Some(NullAwareMode::RightAnti) - && !left_values.is_empty() + let build_has_null = matches!( + null_aware, + Some(NullAwareMode::RightAnti { + correlated: false, + .. + }) + ) && !left_values.is_empty() && left_values[0].logical_null_count() > 0; let left_mark_build_key_nulls = if matches!( @@ -10511,13 +10525,13 @@ mod tests { )); } - /// Test that null_aware validation rejects multi-column joins + /// A multi-column null-aware `RightAnti` (a multi-column or correlated + /// `NOT IN` swapped to build on the subquery) is accepted. #[tokio::test] - async fn test_null_aware_validation_multi_column() { + async fn test_null_aware_validation_multi_column_right_anti() { let left = build_table(("a", &vec![1]), ("b", &vec![2]), ("c", &vec![3])); let right = build_table(("x", &vec![1]), ("y", &vec![2]), ("z", &vec![3])); - // Try multi-column join let on = vec![ ( Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, @@ -10529,9 +10543,7 @@ mod tests { ), ]; - // Try to create null-aware right anti join with 2 columns (should fail). - // A multi-column `LeftAnti` is a correlated `NOT IN` and is accepted. - let result = HashJoinExec::try_new( + let join = HashJoinExec::try_new( left, right, on, @@ -10540,14 +10552,15 @@ mod tests { None, PartitionMode::CollectLeft, NullEquality::NullEqualsNothing, - true, // null_aware = true (invalid for multi-column) - ); - - assert!(result.is_err()); - assert!( - result.unwrap_err().to_string().contains( - "null_aware RightAnti joins only support single column join key" - ) + true, + ) + .unwrap(); + assert_eq!( + join.null_aware_mode().unwrap(), + Some(NullAwareMode::RightAnti { + correlated: true, + value_keys: 1, + }) ); } @@ -10568,7 +10581,7 @@ mod tests { ), ]; - // LeftMark (not LeftAnti) so the single-key check is skipped and the partition-mode check is reached. + // A correlated (multi-key) LeftMark requires CollectLeft. let result = HashJoinExec::try_new( left, right, @@ -10590,14 +10603,12 @@ mod tests { ); } - /// A null-aware `RightAnti` short-circuits on the build-side NULL before any - /// filter runs, so the combination must be rejected at construction (and thus - /// via protobuf decoding, which routes through the same builder). A filtered - /// null-aware `LeftAnti`, by contrast, is a valid decorrelated correlated - /// `NOT IN` plan and must still be accepted. + /// A filtered null-aware anti join is a decorrelated correlated `NOT IN` + /// and is accepted both as `LeftAnti` and as the swapped `RightAnti`, which + /// is correlated and so decides `NOT IN` per probe row after the filter. #[tokio::test] - async fn test_null_aware_filter_rejected_only_for_right_anti() { - let make_inputs = || { + async fn test_null_aware_filter_accepted_for_anti_joins() { + for join_type in [JoinType::LeftAnti, JoinType::RightAnti] { let left = build_table_two_cols(("c1", &vec![Some(1)]), ("dummy", &vec![Some(10)])); let right = @@ -10606,44 +10617,20 @@ mod tests { Arc::new(Column::new_with_schema("c1", &left.schema()).unwrap()) as _, Arc::new(Column::new_with_schema("c2", &right.schema()).unwrap()) as _, )]; - (left, right, on) - }; - - // RightAnti + filter is rejected. - let (left, right, on) = make_inputs(); - let result = HashJoinExec::try_new( - left, - right, - on, - Some(prepare_join_filter()), - &JoinType::RightAnti, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - true, - ); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("null_aware RightAnti join does not support a join filter") - ); - - // LeftAnti + filter is accepted (correlated NOT IN decorrelates to this). - let (left, right, on) = make_inputs(); - let result = HashJoinExec::try_new( - left, - right, - on, - Some(prepare_join_filter()), - &JoinType::LeftAnti, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - true, - ); - assert!(result.is_ok()); + let join = HashJoinExec::try_new( + left, + right, + on, + Some(prepare_join_filter()), + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(); + assert!(join.null_aware_mode().unwrap().unwrap().is_correlated()); + } } #[apply(hash_join_exec_configs)] @@ -11271,10 +11258,12 @@ mod tests { }) .collect::>()?; - let join = HashJoinExec::try_new( - left, - right, - on, + // The outer side `left` as the build side (`LeftAnti`), and swapped to + // the probe side (`RightAnti`), must keep the same rows. + let left_anti = HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), None, &JoinType::LeftAnti, None, @@ -11282,9 +11271,14 @@ mod tests { NullEquality::NullEqualsNothing, true, )?; - - let stream = join.execute(0, task_ctx)?; - let batches = common::collect(stream).await?; + let right_anti = swap_for_separate_execution(&left_anti)?; + let batches = + common::collect(left_anti.execute(0, Arc::clone(&task_ctx))?).await?; + let swapped = common::collect(right_anti.execute(0, task_ctx)?).await?; + assert_eq!( + batches_to_sort_string(&batches), + batches_to_sort_string(&swapped) + ); // Scope (1, 1) holds 2, 3, 4 and 5: `(1, 1, 1)` is kept, `(2, 1, 1)` // matches, and both NULL-valued build rows there are UNKNOWN. Scope @@ -11306,6 +11300,141 @@ mod tests { Ok(()) } + /// Swaps a null-aware `LeftAnti` into the `RightAnti` that `JoinSelection` + /// plans, with its own build state so that both can be executed. + fn swap_for_separate_execution( + left_anti: &HashJoinExec, + ) -> Result> { + let swapped = left_anti.swap_inputs(PartitionMode::CollectLeft)?; + let swapped = swapped.downcast_ref::().unwrap(); + Ok(Arc::new(swapped.builder().reset_state().build()?)) + } + + /// Uncorrelated two-column `(a, b) NOT IN (SELECT a, b ...)`, run both as + /// `LeftAnti` (built on the outer side) and as the swapped `RightAnti` + /// (built on the subquery side), which decides each outer row per probe + /// batch. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_multi_column_right_anti(batch_size: usize) -> Result<()> { + type Rows = Vec<(Option, Option)>; + fn table(rows: &Rows) -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.0))), + Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.1))), + ], + ) + .unwrap(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + async fn not_in( + outer: &Rows, + subquery: &Rows, + batch_size: usize, + ) -> Result { + let (outer, subquery) = (table(outer), table(subquery)); + let on = ["a", "b"] + .into_iter() + .map(|name| { + Ok(( + Arc::new(Column::new_with_schema(name, &outer.schema())?) as _, + Arc::new(Column::new_with_schema(name, &subquery.schema())?) as _, + )) + }) + .collect::>()?; + let left_anti = + HashJoinExecBuilder::new(outer, subquery, on, JoinType::LeftAnti) + .with_partition_mode(PartitionMode::CollectLeft) + .with_null_aware(true) + .with_null_aware_value_keys(2) + .build()?; + let right_anti = swap_for_separate_execution(&left_anti)?; + let right_anti_join = right_anti.downcast_ref::().unwrap(); + assert_eq!( + right_anti_join.null_aware_mode()?, + Some(NullAwareMode::RightAnti { + correlated: true, + value_keys: 2, + }) + ); + + let task_ctx = prepare_task_ctx(batch_size, false); + let expected = batches_to_sort_string( + &common::collect(left_anti.execute(0, Arc::clone(&task_ctx))?).await?, + ); + let actual = batches_to_sort_string( + &common::collect(right_anti.execute(0, task_ctx)?).await?, + ); + assert_eq!(expected, actual); + Ok(actual) + } + + let outer = vec![ + (Some(1), Some(2)), + (Some(1), Some(3)), + (Some(5), Some(7)), + (None, Some(9)), + (None, Some(6)), + (None, None), + ]; + // `(1, 2)` matches exactly, `(1, 3)` has no definite mismatch with + // `(NULL, 3)`, nor `(NULL, 6)` with `(5, 6)` and `(NULL, NULL)` with any + // row. `(5, 7)` and `(NULL, 9)` have a definite mismatch with every row. + let subquery = vec![(Some(1), Some(2)), (None, Some(3)), (Some(5), Some(6))]; + let actual = not_in(&outer, &subquery, batch_size).await?; + allow_duplicates! { + assert_snapshot!(actual, @r" + +---+---+ + | a | b | + +---+---+ + | | 9 | + | 5 | 7 | + +---+---+ + "); + } + + // Every subquery row is NULL-valued, so the full-key hash map is empty + // and every outer row is decided by the NULL candidates alone: only + // `(5, 7)` has a definite mismatch with both `(NULL, 3)` and `(1, NULL)`. + let subquery = vec![(None, Some(3)), (Some(1), None)]; + let actual = not_in(&outer, &subquery, batch_size).await?; + allow_duplicates! { + assert_snapshot!(actual, @r" + +---+---+ + | a | b | + +---+---+ + | 5 | 7 | + +---+---+ + "); + } + + // An empty subquery keeps every outer row, NULL-valued ones included. + let actual = not_in(&outer, &vec![], batch_size).await?; + allow_duplicates! { + assert_snapshot!(actual, @r" + +---+---+ + | a | b | + +---+---+ + | | | + | | 6 | + | | 9 | + | 1 | 2 | + | 1 | 3 | + | 5 | 7 | + +---+---+ + "); + } + + Ok(()) + } + /// Correlated null-aware `LeftAnti` whose two scope keys are Float64, with /// `-0.0` on the build side and `0.0` on the probe side. Both scopes must /// still match in every direction: the full key lookup, NULL probe values diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index cc911e14ff0af..369639357dd6c 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -27,6 +27,7 @@ use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::joins::Map; use crate::joins::MapOffset; use crate::joins::PartitionMode; +use crate::joins::SharedBitmapBuilder; use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode, null_value_key_mask}; use crate::joins::hash_join::probe_completion::ProbeSideSummary; use crate::joins::hash_join::shared_bounds::{ @@ -49,9 +50,12 @@ use crate::{ }, }; -use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, UInt32Array, UInt64Array}; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, UInt32Array, + UInt64Array, +}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::{filter, take}; +use arrow::compute::{filter, filter_record_batch, take}; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ @@ -205,6 +209,11 @@ pub(super) struct ProcessProbeBatchState { /// Key comparator for this batch, built on first use and reused by every /// chunk key_comparator: Option, + /// For a correlated null-aware `RightAnti` join, the probe rows whose + /// `NOT IN` is UNKNOWN because of a NULL (see + /// `mark_null_candidates_for_probe_batch`); they are never emitted. `None` + /// when there is no such row; boxed to keep this state small. + unknown_probe_rows: Option>, } impl ProcessProbeBatchState { @@ -879,6 +888,7 @@ impl HashJoinStream { joined_probe_idx: None, matched_probe_idx: None, key_comparator: None, + unknown_probe_rows: None, }); } Some(Err(err)) => return Poll::Ready(Err(err)), @@ -920,31 +930,55 @@ impl HashJoinStream { // Must precede the empty-build-map return below: an all-NULL-key // build side has an empty full-key map but still needs UNKNOWN marks. if mode.is_correlated() && state.offset == (0, None) { + // `RightAnti` marks the rows of this probe batch, the others + // the build rows shared by all probe partitions. + let probe_marks = (!mode.marks_build_rows()).then(|| { + let mut marks = BooleanBufferBuilder::new(state.batch.num_rows()); + marks.append_n(state.batch.num_rows(), false); + SharedBitmapBuilder::new(marks) + }); + let marks = match &probe_marks { + Some(probe_marks) => UnknownMarks::Probe(probe_marks), + None => { + UnknownMarks::Build(build_side.left_data.null_indices_bitmap()) + } + }; mark_null_candidates_for_probe_batch( build_side, state, mode.value_keys(), self.filter.as_ref(), self.join_type, + marks, &self.random_state, self.batch_size, &mut self.null_mark_hashes_buffer, &mut self.null_mark_probe_indices_buffer, &mut self.null_mark_build_indices_buffer, )?; + state.unknown_probe_rows = probe_marks + .map(|marks| marks.into_inner().finish()) + .filter(|marks| marks.count_set_bits() > 0) + .map(Box::new); } } let is_empty = !build_side.left_data.has_matchable_build_rows(); if is_empty { - let result = build_batch_empty_build_side( + let mut result = build_batch_empty_build_side( &self.schema, build_side.left_data.batch(), &state.batch, &self.column_indices, self.join_type, )?; + if let Some(unknown) = &state.unknown_probe_rows { + result = filter_record_batch( + &result, + &BooleanArray::new(!unknown.as_ref(), None), + )?; + } timer.done(); self.output_buffer.push_batch(result)?; self.state = HashJoinStreamState::FetchProbeBatch; @@ -1066,8 +1100,22 @@ impl HashJoinStream { self.right_side_ordered, )?; - if self.null_aware == Some(NullAwareMode::RightAnti) { - right_indices = drop_null_probe_keys(state, right_indices); + match self.null_aware { + Some(NullAwareMode::RightAnti { + correlated: false, .. + }) => { + right_indices = drop_null_probe_keys(state, right_indices); + } + Some(NullAwareMode::RightAnti { + correlated: true, .. + }) => { + if let Some(unknown) = &state.unknown_probe_rows { + right_indices = retain_probe_rows(&right_indices, |idx| { + !unknown.value(idx as usize) + }); + } + } + _ => {} } // Build output batch and push to coalescer @@ -1317,8 +1365,8 @@ fn next_final_indices_chunk( /// Applies the pre-lookup bookkeeping of a null-aware join to a probe batch /// and returns `true` if the batch cannot contribute any output row. /// -/// Under SQL three-valued logic a NULL on either side of `NOT IN` makes the -/// predicate UNKNOWN for every row of the other side: +/// Under SQL three-valued logic a NULL on either side of an uncorrelated scalar +/// `NOT IN` makes the predicate UNKNOWN for every row of the other side: /// - `RightAnti` emits probe rows, so a NULL build key silences every batch. /// - `LeftAnti` emits build rows in the final stage, so a NULL probe key seen /// by any partition silences everything; the shared flags carry this to the @@ -1331,10 +1379,15 @@ fn null_aware_skip_probe_batch( left_data: &JoinLeftData, ) -> bool { match mode { - NullAwareMode::RightAnti => left_data.build_side_has_null, - // Correlated joins decide UNKNOWN per build row instead, in + NullAwareMode::RightAnti { + correlated: false, .. + } => left_data.build_side_has_null, + // Correlated joins decide UNKNOWN per row instead, in // `mark_null_candidates_for_probe_batch`. - NullAwareMode::LeftAnti { + NullAwareMode::RightAnti { + correlated: true, .. + } + | NullAwareMode::LeftAnti { correlated: true, .. } | NullAwareMode::LeftMark { @@ -1378,17 +1431,27 @@ fn drop_null_probe_keys( .or_else(|| probe_key.logical_nulls()); // we only need this copy if there are NULLs match mask.filter(|m| m.null_count() > 0) { - Some(mask) => UInt32Array::from_iter_values( - right_indices - .values() - .iter() - .copied() - .filter(|idx| !mask.is_null(*idx as usize)), - ), + Some(mask) => { + retain_probe_rows(&right_indices, |idx| !mask.is_null(idx as usize)) + } None => right_indices, } } +/// Keeps the probe rows of `right_indices` for which `keep` returns true. +fn retain_probe_rows( + right_indices: &UInt32Array, + keep: impl Fn(u32) -> bool, +) -> UInt32Array { + UInt32Array::from_iter_values( + right_indices + .values() + .iter() + .copied() + .filter(|idx| keep(*idx)), + ) +} + /// Final-stage rules of a null-aware `LeftAnti` join, evaluated by the last /// probe partition. /// @@ -1465,21 +1528,51 @@ fn null_aware_left_mark_column( ) } -/// Records which build rows of a correlated null-aware join are UNKNOWN -/// candidates for this probe batch. +/// Where a correlated null-aware join records the rows whose `NOT IN` is +/// UNKNOWN, see [`mark_null_candidates_for_probe_batch`]. +#[derive(Clone, Copy)] +enum UnknownMarks<'a> { + /// `LeftAnti` and `LeftMark` mark build rows, in the null-indices bitmap + /// shared by all probe partitions. + Build(&'a SharedBitmapBuilder), + /// `RightAnti` marks the rows of the current probe batch. + Probe(&'a SharedBitmapBuilder), +} + +impl<'a> UnknownMarks<'a> { + fn bitmap(self) -> &'a SharedBitmapBuilder { + match self { + Self::Build(bitmap) | Self::Probe(bitmap) => bitmap, + } + } + + /// The row of a (build, probe) candidate pair that this join marks. + fn marked_row(self, build_idx: u64, probe_idx: u32) -> usize { + match self { + Self::Build(_) => build_idx as usize, + Self::Probe(_) => probe_idx as usize, + } + } +} + +/// Records which rows of a correlated null-aware join are UNKNOWN candidates +/// for this probe batch: the build rows for `LeftAnti` and `LeftMark`, the +/// probe rows for `RightAnti` (see [`UnknownMarks`]). /// /// Key layout: `on[..V]` are the `NOT IN` value keys, `on[V..]` the (possibly /// empty) correlation scope keys (see `HashJoinExec::null_aware`). A row is -/// NULL-valued when any of its value keys is NULL. An unmatched build row's -/// `NOT IN` is UNKNOWN instead of TRUE (its mark is NULL instead of FALSE) -/// when either: -/// 1. it is NULL-valued and a probe row in its correlation scope passes the -/// join filter, or -/// 2. a NULL-valued probe row in its correlation scope passes the join filter, +/// NULL-valued when any of its value keys is NULL. A (build, probe) pair makes +/// the `NOT IN` of the marked row UNKNOWN instead of TRUE (a `LeftMark` mark +/// NULL instead of FALSE) when either: +/// 1. the build row is NULL-valued, or +/// 2. the probe row is NULL-valued, /// -/// and, for a multi-column value key, the two rows' value tuples are not a -/// definite mismatch: every element pair is equal or involves a NULL. With a -/// single value key a NULL on either side already rules out a mismatch. +/// the two rows are in the same correlation scope, the pair passes the join +/// filter, and, for a multi-column value key, the two rows' value tuples are +/// not a definite mismatch: every element pair is equal or involves a NULL. +/// With a single value key a NULL on either side already rules out a mismatch. +/// An exact match, which makes `NOT IN` FALSE, is found by the main hash lookup +/// instead and takes precedence over these marks. /// /// Case 1 pairs the NULL-valued build rows with all probe rows; case 2 pairs /// all build rows with the NULL-valued probe rows. Scope keys narrow these @@ -1487,10 +1580,10 @@ fn null_aware_left_mark_column( /// The value tuples and then the join filter, if any, decide which candidates /// count. /// -/// A build row stays UNKNOWN once it is marked, so candidates whose build row -/// is already marked are skipped, and the join filter is not evaluated for -/// them. Without scope keys this also ends the pairing as soon as no unmarked -/// build row is left. +/// A row stays UNKNOWN once it is marked, so candidates whose marked row is +/// already marked are skipped, and the join filter is not evaluated for them. +/// Without scope keys this also ends the pairing as soon as no unmarked row is +/// left. #[expect(clippy::too_many_arguments)] fn mark_null_candidates_for_probe_batch( build_side: &BuildSideReadyState, @@ -1498,6 +1591,7 @@ fn mark_null_candidates_for_probe_batch( num_value_keys: usize, filter: Option<&JoinFilter>, join_type: JoinType, + marks: UnknownMarks<'_>, random_state: &RandomState, batch_size: usize, hashes_buffer: &mut Vec, @@ -1521,10 +1615,10 @@ fn mark_null_candidates_for_probe_batch( let (probe_value_keys, probe_scope_values) = state.values.split_at(num_value_keys); // Keeps the candidate pairs whose value tuples are not a definite mismatch - // and that pass the join filter, and marks their build rows as UNKNOWN. + // and that pass the join filter, and marks their rows as UNKNOWN. let mut mark = |build_indices: UInt64Array, probe_indices: UInt32Array| { let (build_indices, probe_indices) = - retain_unmarked(left_data, build_indices, probe_indices); + retain_unmarked(marks, build_indices, probe_indices); let (build_indices, probe_indices) = if num_value_keys > 1 { retain_value_mismatch_free( build_value_keys, @@ -1538,33 +1632,32 @@ fn mark_null_candidates_for_probe_batch( if build_indices.is_empty() { return Ok(()); } - let build_indices = match filter { - Some(filter) => { - apply_join_filter_to_indices( - left_data.batch(), - &state.batch, - build_indices, - probe_indices, - filter, - JoinSide::Left, - None, - join_type, - )? - .0 - } - None => build_indices, + let (build_indices, probe_indices) = match filter { + Some(filter) => apply_join_filter_to_indices( + left_data.batch(), + &state.batch, + build_indices, + probe_indices, + filter, + JoinSide::Left, + None, + join_type, + )?, + None => (build_indices, probe_indices), }; if !build_indices.is_empty() { - let mut null_bitmap = left_data.null_indices_bitmap().lock(); - for build_idx in build_indices.values() { - null_bitmap.set_bit(*build_idx as usize, true); + let mut bitmap = marks.bitmap().lock(); + for (build_idx, probe_idx) in + build_indices.values().iter().zip(probe_indices.values()) + { + bitmap.set_bit(marks.marked_row(*build_idx, *probe_idx), true); } } Ok(()) }; - // Case 1: NULL-valued build rows are UNKNOWN as soon as a probe row in - // their correlation scope passes the checks in `mark`. + // Case 1: NULL-valued build rows paired with the probe rows in their + // correlation scope. if let Some(null_rows) = null_value_build_rows { match &null_rows.scope_map { Some(scope_map) => { @@ -1595,7 +1688,7 @@ fn mark_null_candidates_for_probe_batch( } None => { for_each_unmarked_cross_product( - left_data, + marks, null_rows.build_indices.values().iter().copied(), 0..state.batch.num_rows() as u32, batch_size, @@ -1605,8 +1698,8 @@ fn mark_null_candidates_for_probe_batch( } } - // Case 2: NULL-valued probe rows make every build row in their correlation - // scope that passes the checks in `mark` UNKNOWN. + // Case 2: NULL-valued probe rows paired with the build rows in their + // correlation scope. if let Some(null_mask) = probe_null_value_mask { let null_probe_rows = UInt32Array::from_iter_values( null_mask.values().set_indices().map(|i| i as u32), @@ -1648,7 +1741,7 @@ fn mark_null_candidates_for_probe_batch( } None => { for_each_unmarked_cross_product( - left_data, + marks, 0..left_data.batch().num_rows() as u64, null_probe_rows.values().iter().copied(), batch_size, @@ -1759,75 +1852,108 @@ fn for_each_scope_match( Ok(()) } -/// Removes the candidate pairs whose build row is already marked UNKNOWN. +/// Removes the candidate pairs whose marked row is already marked UNKNOWN. fn retain_unmarked( - left_data: &JoinLeftData, + marks: UnknownMarks<'_>, build_indices: UInt64Array, probe_indices: UInt32Array, ) -> (UInt64Array, UInt32Array) { - let bitmap = left_data.null_indices_bitmap().lock(); - let is_unmarked = |build_idx: &u64| !bitmap.get_bit(*build_idx as usize); - if build_indices.values().iter().all(is_unmarked) { + let bitmap = marks.bitmap().lock(); + let is_unmarked = |(build_idx, probe_idx): (&u64, &u32)| { + !bitmap.get_bit(marks.marked_row(*build_idx, *probe_idx)) + }; + if build_indices + .values() + .iter() + .zip(probe_indices.values().iter()) + .all(is_unmarked) + { return (build_indices, probe_indices); } let (build, probe): (Vec, Vec) = build_indices .values() .iter() .zip(probe_indices.values().iter()) - .filter(|(build_idx, _)| is_unmarked(build_idx)) + .filter(|pair| is_unmarked(*pair)) .unzip(); (build.into(), probe.into()) } -/// Calls `f` with the pairs of `build_rows` x `probe_rows` whose build row is -/// not marked UNKNOWN, as chunks of at most `batch_size` pairs. -/// -/// `f` marks build rows, so the unmarked build rows are found again after each -/// chunk. The pairing stops when no unmarked build row is left. +/// Calls `f` with the pairs of `build_rows` x `probe_rows` whose marked row +/// (see [`UnknownMarks`]) is not marked UNKNOWN, as chunks of at most +/// `batch_size` pairs. fn for_each_unmarked_cross_product( - left_data: &JoinLeftData, + marks: UnknownMarks<'_>, build_rows: impl Iterator, probe_rows: impl Iterator, batch_size: usize, mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, ) -> Result<()> { - let retain_unmarked_rows = |rows: &mut Vec| { - let bitmap = left_data.null_indices_bitmap().lock(); - rows.retain(|idx| !bitmap.get_bit(*idx as usize)); + match marks { + UnknownMarks::Build(bitmap) => for_each_unmarked_pair( + bitmap, + build_rows.collect(), + probe_rows, + batch_size, + |build, probe| f(build.into(), probe.into()), + ), + UnknownMarks::Probe(bitmap) => for_each_unmarked_pair( + bitmap, + probe_rows.collect(), + build_rows, + batch_size, + |probe, build| f(build.into(), probe.into()), + ), + } +} + +/// Calls `f` with the pairs of `marked_rows` x `other_rows` whose row in +/// `marked_rows` is not set in `bitmap`, as chunks of at most `batch_size` +/// pairs. +/// +/// `f` sets bits in `bitmap`, so the unmarked rows are found again after each +/// chunk. The pairing stops when no unmarked row is left. +fn for_each_unmarked_pair, O: Copy>( + bitmap: &SharedBitmapBuilder, + mut marked_rows: Vec, + other_rows: impl Iterator, + batch_size: usize, + mut f: impl FnMut(Vec, Vec) -> Result<()>, +) -> Result<()> { + let retain_unmarked_rows = |rows: &mut Vec| { + let bitmap = bitmap.lock(); + rows.retain(|idx| !bitmap.get_bit((*idx).into() as usize)); }; - let mut build_rows: Vec = build_rows.collect(); - retain_unmarked_rows(&mut build_rows); + retain_unmarked_rows(&mut marked_rows); - let mut build_chunk = Vec::with_capacity(batch_size); - let mut probe_chunk = Vec::with_capacity(batch_size); + let mut marked_chunk = Vec::with_capacity(batch_size); + let mut other_chunk = Vec::with_capacity(batch_size); let mut marks_since_refresh = false; - for probe_row in probe_rows { + for other_row in other_rows { // Refresh only after a chunk was sent, so the cost of the refresh // stays proportional to the pairs already evaluated. if marks_since_refresh { - retain_unmarked_rows(&mut build_rows); + retain_unmarked_rows(&mut marked_rows); marks_since_refresh = false; } - if build_rows.is_empty() { + if marked_rows.is_empty() { break; } - for build_row in &build_rows { - build_chunk.push(*build_row); - probe_chunk.push(probe_row); - if build_chunk.len() == batch_size { + for marked_row in &marked_rows { + marked_chunk.push(*marked_row); + other_chunk.push(other_row); + if marked_chunk.len() == batch_size { f( - std::mem::replace(&mut build_chunk, Vec::with_capacity(batch_size)) - .into(), - std::mem::replace(&mut probe_chunk, Vec::with_capacity(batch_size)) - .into(), + std::mem::replace(&mut marked_chunk, Vec::with_capacity(batch_size)), + std::mem::replace(&mut other_chunk, Vec::with_capacity(batch_size)), )?; marks_since_refresh = true; } } } - if !build_chunk.is_empty() { - f(build_chunk.into(), probe_chunk.into())?; + if !marked_chunk.is_empty() { + f(marked_chunk, other_chunk)?; } Ok(()) } diff --git a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt index 916819fbb9d89..29159f257795b 100644 --- a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt @@ -411,7 +411,9 @@ NULL 0 # Results under a batch size sweep are in null_aware_multi_column.slt. A tuple # NOT IN is a null-aware join whose leading keys are the tuple elements -# (`value_keys`), followed by any correlation keys. +# (`value_keys`), followed by any correlation keys. Like the scalar NOT IN, +# JoinSelection swaps it into a null-aware RightAnti that builds on the +# subquery side when that side is not larger, keeping the key order. statement ok CREATE TABLE multi_col_outer(k INT, a INT, b INT) AS VALUES @@ -436,7 +438,7 @@ logical_plan 06)------SubqueryAlias: i 07)--------TableScan: multi_col_inner projection=[k, x, y] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(a@1, x@0), (b@2, y@1), (k@0, k@2)], null_aware(value_keys=2) +01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(x@0, a@1), (y@1, b@2), (k@2, k@0)], null_aware(value_keys=2) 02)--DataSourceExec: partitions=1, partition_sizes=[1] 03)--DataSourceExec: partitions=1, partition_sizes=[1] @@ -455,10 +457,10 @@ logical_plan 06)------Projection: multi_col_inner.x, CAST(multi_col_inner.y AS Int64) 07)--------TableScan: multi_col_inner projection=[x, y] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(a@1, x@0), (__correlated_sq_1_value_1@3, multi_col_inner.y@1)], projection=[k@0, a@1, b@2], null_aware(value_keys=2) -02)--ProjectionExec: expr=[k@0 as k, a@1 as a, b@2 as b, 2 as __correlated_sq_1_value_1] +01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(x@0, a@1), (multi_col_inner.y@1, __correlated_sq_1_value_1@3)], projection=[k@0, a@1, b@2], null_aware(value_keys=2) +02)--ProjectionExec: expr=[x@0 as x, CAST(y@1 AS Int64) as multi_col_inner.y] 03)----DataSourceExec: partitions=1, partition_sizes=[1] -04)--ProjectionExec: expr=[x@0 as x, CAST(y@1 AS Int64) as multi_col_inner.y] +04)--ProjectionExec: expr=[k@0 as k, a@1 as a, b@2 as b, 2 as __correlated_sq_1_value_1] 05)----DataSourceExec: partitions=1, partition_sizes=[1] # As a nullable value, the tuple NOT IN is a null-aware mark join. @@ -990,7 +992,8 @@ NULL 0 # The shape must stay a single anti join that carries the correlation as a # filter. The alternative lowering materializes the mark of three joins, which -# is correct but far slower. +# is correct but far slower. JoinSelection may still swap it into a null-aware +# RightAnti, which applies the filter before deciding each outer row. query TT EXPLAIN SELECT id, g FROM nai_res_og WHERE nai_res_og.id NOT IN (SELECT i.id FROM nai_res_ig i WHERE nai_res_og.g > 0); ---- @@ -1001,7 +1004,7 @@ logical_plan 04)----SubqueryAlias: i 05)------TableScan: nai_res_ig projection=[id] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)], filter=g@0 > 0, null_aware +01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(id@0, id@0)], filter=g@0 > 0, null_aware 02)--DataSourceExec: partitions=1, partition_sizes=[1] 03)--DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/sqllogictest/test_files/null_aware_multi_column.slt b/datafusion/sqllogictest/test_files/null_aware_multi_column.slt index 917d1caec972f..96dd612a42426 100644 --- a/datafusion/sqllogictest/test_files/null_aware_multi_column.slt +++ b/datafusion/sqllogictest/test_files/null_aware_multi_column.slt @@ -20,11 +20,16 @@ # mismatch and at least one of them involves a NULL: `(NULL, 8) = (1, 2)` is # `UNKNOWN AND FALSE`, which is FALSE, while `(NULL, 2) = (1, 2)` is UNKNOWN. # The batch size sweep exercises the chunked UNKNOWN candidate pairing of the -# null-aware hash join. Plans are checked in null_aware_anti_join.slt. +# null-aware hash join. With join reordering on, JoinSelection swaps a NOT IN +# whose subquery is not larger into a null-aware RightAnti that builds on the +# subquery and decides each outer row as it streams; with it off, the join +# stays a LeftAnti that builds on the outer side and decides per build row. +# Plans are checked in null_aware_anti_join.slt. # Matrix rules: no EXPLAIN, no in-file SET of a swept knob, rowsort every # multi-row query. # configMatrix: datafusion.execution.batch_size=1,2,100,8192 +# configMatrix: datafusion.optimizer.join_reordering=true,false statement ok CREATE TABLE mc_outer(a INT, b INT, value TEXT) AS VALUES From fa12ec3983d013e3489b0dfaf4e72dd2f5b96a02 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 11 Oct 2026 10:33:00 -0700 Subject: [PATCH 2/3] test: update correlated NOT IN plan in subquery_projection.slt Co-authored-by: Isaac --- .../sqllogictest/test_files/subquery_projection.slt | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/datafusion/sqllogictest/test_files/subquery_projection.slt b/datafusion/sqllogictest/test_files/subquery_projection.slt index f86de505fe342..f2845c38788e6 100644 --- a/datafusion/sqllogictest/test_files/subquery_projection.slt +++ b/datafusion/sqllogictest/test_files/subquery_projection.slt @@ -403,8 +403,9 @@ statement ok DROP TABLE t2; # A non-equality correlation leaves one key and a residual join filter. The -# null-aware `LeftAnti` hash join applies the residual filter when it looks -# for UNKNOWN rows. For `k = 1` the key is NULL and the correlated subquery +# null-aware anti hash join (here swapped to a `RightAnti` that builds on the +# smaller subquery) applies the residual filter when it looks for UNKNOWN +# rows. For `k = 1` the key is NULL and the correlated subquery # result is empty, and `NULL NOT IN ()` is TRUE. Both rows are # correct. statement ok @@ -438,10 +439,10 @@ logical_plan 05)------Projection: CAST(rb.k AS Int64), rb.z 06)--------TableScan: rb projection=[k, z] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(nullif(ra.k,Int64(1))@2, rb.k@0)], filter=z@1 > z@0, projection=[k@0], null_aware -02)--ProjectionExec: expr=[k@0 as k, z@1 as z, nullif(CAST(k@0 AS Int64), 1) as nullif(ra.k,Int64(1))] +01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(rb.k@0, nullif(ra.k,Int64(1))@2)], filter=z@1 > z@0, projection=[k@0], null_aware +02)--ProjectionExec: expr=[CAST(k@0 AS Int64) as rb.k, z@1 as z] 03)----DataSourceExec: partitions=1, partition_sizes=[1] -04)--ProjectionExec: expr=[CAST(k@0 AS Int64) as rb.k, z@1 as z] +04)--ProjectionExec: expr=[k@0 as k, z@1 as z, nullif(CAST(k@0 AS Int64), 1) as nullif(ra.k,Int64(1))] 05)----DataSourceExec: partitions=1, partition_sizes=[1] statement ok From 785332e66e47f0d599154fc638d8adaa015c7cb2 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 15:57:24 -0700 Subject: [PATCH 3/3] Address review: skip the per-batch UNKNOWN bitmap, trim overlapping tests - A correlated null-aware RightAnti allocates the per-batch UNKNOWN bitmap only when the build side has a NULL-valued row or the probe batch has a NULL in a value key, and process_probe_batch matches on the join mode. - Rename `UnknownMarks::marked_row` to `unknown_row` and the build state flags to `with_null_candidate_state` / `with_build_unknown_bitmap`. - Move the all-NULL-valued subquery case to null_aware_multi_column.slt and drop the overlapping exec.rs test; fold the construction tests into one parametrized test and the swap tests into one rstest. - The fuzz walks the plan for HashJoinExec instead of matching the displayed plan, and plans each query once. - Move the memory-limit test to null_aware_memory_limit.slt. Co-authored-by: Isaac --- .../tests/fuzz_cases/null_aware_join_fuzz.rs | 57 ++-- datafusion/core/tests/memory_limit/mod.rs | 1 - .../tests/memory_limit/null_aware_not_in.rs | 103 ------- .../physical_optimizer/join_selection.rs | 142 +++------- .../physical-plan/src/joins/hash_join/exec.rs | 256 ++++-------------- .../src/joins/hash_join/stream.rs | 140 +++++----- .../test_files/null_aware_memory_limit.slt | 67 +++++ .../test_files/null_aware_multi_column.slt | 17 ++ 8 files changed, 288 insertions(+), 495 deletions(-) delete mode 100644 datafusion/core/tests/memory_limit/null_aware_not_in.rs create mode 100644 datafusion/sqllogictest/test_files/null_aware_memory_limit.slt diff --git a/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs b/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs index a4d45ac0fb973..c76e4d1a83e71 100644 --- a/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/null_aware_join_fuzz.rs @@ -31,8 +31,10 @@ use std::sync::Arc; use arrow::array::{Int32Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::util::pretty::pretty_format_batches; +use datafusion::common::JoinType; use datafusion::datasource::MemTable; -use datafusion::physical_plan::displayable; +use datafusion::physical_plan::joins::HashJoinExec; +use datafusion::physical_plan::{ExecutionPlan, collect}; use datafusion::prelude::{SessionConfig, SessionContext}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -66,9 +68,9 @@ async fn null_aware_join_matches_reference() { ctx.register_table("o", Arc::clone(&outer) as _).unwrap(); ctx.register_table("i", Arc::clone(&inner) as _).unwrap(); for (query, reference) in cases() { - plans.record(&ctx, &query).await; - let actual = run(&ctx, &query).await; - let expected = run(&ctx, &reference).await; + let (actual, plan) = run(&ctx, &query).await; + plans.record(plan.as_ref()); + let (expected, _) = run(&ctx, &reference).await; assert_eq!( actual, expected, "seed={seed} target_partitions={target_partitions} \ @@ -95,25 +97,25 @@ struct NullAwarePlans { } impl NullAwarePlans { - async fn record(&mut self, ctx: &SessionContext, sql: &str) { - let plan = ctx - .sql(sql) - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - let plan = displayable(plan.as_ref()).indent(true).to_string(); - for line in plan.lines().filter(|line| line.contains("null_aware")) { - if line.contains("join_type=LeftAnti") { - self.left_anti += 1; - } else if line.contains("join_type=RightAnti") { - self.right_anti += 1; - if line.contains("), (") || line.contains("filter=") { - self.correlated_right_anti += 1; + /// Counts the null-aware anti joins in `plan` and its descendants. + fn record(&mut self, plan: &dyn ExecutionPlan) { + if let Some(join) = plan.downcast_ref::() + && join.null_aware + { + match join.join_type() { + JoinType::LeftAnti => self.left_anti += 1, + JoinType::RightAnti => { + self.right_anti += 1; + if join.on().len() > 1 || join.filter().is_some() { + self.correlated_right_anti += 1; + } } + _ => {} } } + for child in plan.children() { + self.record(child.as_ref()); + } } } @@ -215,9 +217,16 @@ fn random_table(rng: &mut StdRng, columns: &[&str], max_rows: i32) -> Arc Vec { - let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); +/// The rows of `sql`, sorted, and the physical plan that produced them. +async fn run(ctx: &SessionContext, sql: &str) -> (Vec, Arc) { + let plan = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); let formatted = pretty_format_batches(&batches).unwrap().to_string(); let mut rows: Vec = formatted .lines() @@ -226,5 +235,5 @@ async fn run(ctx: &SessionContext, sql: &str) -> Vec { .map(String::from) .collect(); rows.sort(); - rows + (rows, plan) } diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 7f7d7215b1c1b..0faa4a64cf983 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -23,7 +23,6 @@ use std::sync::{Arc, LazyLock}; #[cfg(feature = "extended_tests")] mod memory_limit_validation; mod nlj_spill_unmatched; -mod null_aware_not_in; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; diff --git a/datafusion/core/tests/memory_limit/null_aware_not_in.rs b/datafusion/core/tests/memory_limit/null_aware_not_in.rs deleted file mode 100644 index 4bdeb8fe03451..0000000000000 --- a/datafusion/core/tests/memory_limit/null_aware_not_in.rs +++ /dev/null @@ -1,103 +0,0 @@ -// 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. - -//! A multi-column `NOT IN` over a small subquery must build its hash join on -//! the subquery, so a large outer table runs under a memory limit that only -//! fits the subquery side. -//! -//! The null-aware join of a multi-column `NOT IN` used to stay a `LeftAnti` -//! that builds on the outer side, so it ran out of memory where the scalar -//! `NOT IN`, swapped to a `RightAnti` that builds on the subquery, ran. - -use std::sync::Arc; - -use arrow::array::{Int64Array, RecordBatch}; -use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::MemTable; -use datafusion::prelude::*; -use datafusion_common::{Result, assert_contains}; -use datafusion_execution::runtime_env::RuntimeEnvBuilder; - -const OUTER_ROWS: i64 = 200_000; -/// Far below the outer side's build state, far above the subquery's. -const MEMORY_LIMIT: usize = 2 * 1024 * 1024; - -/// A table `(a, b)` of `a = v, b = v % 1000` for `v` in `1..=rows`, followed by -/// `extra` rows. -fn table(rows: i64, extra: &[(Option, Option)]) -> Arc { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, true), - Field::new("b", DataType::Int64, true), - ])); - let a = (1..=rows).map(Some).chain(extra.iter().map(|row| row.0)); - let b = (1..=rows) - .map(|v| Some(v % 1000)) - .chain(extra.iter().map(|row| row.1)); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(a.collect::()), - Arc::new(b.collect::()), - ], - ) - .unwrap(); - // Several batches, so the outer side streams through the probe. - let batches = (0..batch.num_rows()) - .step_by(8192) - .map(|offset| batch.slice(offset, 8192.min(batch.num_rows() - offset))) - .collect(); - Arc::new(MemTable::try_new(schema, vec![batches]).unwrap()) -} - -async fn count_not_in(sql: &str, join_reordering: bool) -> Result { - let config = SessionConfig::new() - .with_target_partitions(1) - .set_bool("datafusion.optimizer.join_reordering", join_reordering); - let runtime = RuntimeEnvBuilder::new() - .with_memory_limit(MEMORY_LIMIT, 1.0) - .build_arc()?; - let ctx = SessionContext::new_with_config_rt(config, runtime); - ctx.register_table("big", table(OUTER_ROWS, &[]))?; - // A NULL-valued subquery row makes `NOT IN` UNKNOWN for every outer row - // with `b = 7` (and no definite mismatch in `a`). - ctx.register_table("small", table(100, &[(None, Some(7))]))?; - - let batches = ctx.sql(sql).await?.collect().await?; - Ok(batches[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0)) -} - -#[tokio::test] -async fn multi_column_not_in_builds_on_small_subquery() -> Result<()> { - let sql = "SELECT count(*) FROM big \ - WHERE (a, b) NOT IN (SELECT a, b FROM small)"; - - // Rows 1..=100 match exactly; of the rest, those with `b = 7` are - // UNKNOWN: 199 of the 200 values `v % 1000 = 7`, all but `v = 7`. - assert_eq!(count_not_in(sql, true).await?, OUTER_ROWS - 100 - 199); - - // Building on the outer side does not fit, which shows the limit binds. - let err = count_not_in(sql, false).await.unwrap_err(); - assert_contains!(err.to_string(), "Resources exhausted"); - assert_contains!(err.to_string(), "HashJoinInput"); - - Ok(()) -} diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index fe0392a1150be..f10b33469e052 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -471,29 +471,41 @@ async fn test_join_with_swap_semi() { } } +/// A null-aware `LeftAnti` swaps to a null-aware `RightAnti` that builds on +/// the smaller subquery side. Besides the scalar `NOT IN`, this covers a +/// multi-key one (a multi-column `NOT IN`, correlation scope keys, or both), +/// which keeps the number of `NOT IN` value keys and the key order, and a +/// filtered one (a decorrelated correlated `NOT IN`), whose filter sides are +/// swapped with the inputs. #[rstest] -#[case(PartitionMode::CollectLeft)] -#[case(PartitionMode::Auto)] -#[case(PartitionMode::Partitioned)] +#[case::collect_left(PartitionMode::CollectLeft, 1, 1, false)] +#[case::auto(PartitionMode::Auto, 1, 1, false)] +#[case::partitioned(PartitionMode::Partitioned, 1, 1, false)] +#[case::filter(PartitionMode::CollectLeft, 1, 1, true)] +#[case::scope_key(PartitionMode::CollectLeft, 2, 1, false)] +#[case::multi_column(PartitionMode::CollectLeft, 2, 2, false)] #[tokio::test] async fn test_null_aware_left_anti_swaps_to_right_anti( #[case] partition_mode: PartitionMode, + #[case] num_keys: usize, + #[case] value_keys: usize, + #[case] with_filter: bool, ) -> Result<()> { let (big, small) = create_big_and_small(); - let join = HashJoinExec::try_new( + let big_col = Arc::new(Column::new_with_schema("big_col", &big.schema())?) as _; + let small_col = Arc::new(Column::new_with_schema("small_col", &small.schema())?) as _; + let on = vec![(big_col, small_col); num_keys]; + let join = HashJoinExecBuilder::new( Arc::clone(&big), Arc::clone(&small), - vec![( - Arc::new(Column::new_with_schema("big_col", &big.schema())?), - Arc::new(Column::new_with_schema("small_col", &small.schema())?), - )], - None, - &JoinType::LeftAnti, - None, - partition_mode, - NullEquality::NullEqualsNothing, - true, - )?; + on, + JoinType::LeftAnti, + ) + .with_filter(if with_filter { nl_join_filter() } else { None }) + .with_partition_mode(partition_mode) + .with_null_aware(true) + .with_null_aware_value_keys(value_keys) + .build()?; let original_schema = join.schema(); let optimized_join = @@ -505,9 +517,18 @@ async fn test_null_aware_left_anti_swaps_to_right_anti( assert_eq!(*swapped_join.join_type(), JoinType::RightAnti); assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft); assert!(swapped_join.null_aware); + assert_eq!(swapped_join.null_aware_value_keys, value_keys); + assert_eq!(swapped_join.on().len(), num_keys); assert_eq!(swapped_join.left().schema().field(0).name(), "small_col"); assert_eq!(swapped_join.right().schema().field(0).name(), "big_col"); assert_eq!(swapped_join.schema(), original_schema); + let filter_sides = swapped_join + .filter() + .map(|filter| filter.column_indices().iter().map(|c| c.side).collect()); + assert_eq!( + filter_sides, + with_filter.then(|| vec![JoinSide::Right, JoinSide::Left]) + ); Ok(()) } @@ -619,97 +640,6 @@ async fn test_null_aware_left_anti_respects_disabled_join_reordering() -> Result Ok(()) } -/// A filtered null-aware `LeftAnti` (produced by decorrelating a correlated -/// `NOT IN`) swaps to a null-aware `RightAnti` that decides `NOT IN` per probe -/// row; the filter's sides are swapped with the inputs. -#[tokio::test] -async fn test_null_aware_left_anti_with_filter_swaps() -> Result<()> { - let (big, small) = create_big_and_small(); - let join = HashJoinExec::try_new( - Arc::clone(&big), - Arc::clone(&small), - vec![( - Arc::new(Column::new_with_schema("big_col", &big.schema())?), - Arc::new(Column::new_with_schema("small_col", &small.schema())?), - )], - nl_join_filter(), - &JoinType::LeftAnti, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - true, - )?; - - let optimized_join = - JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?; - let swapped_join = optimized_join - .downcast_ref::() - .expect("anti join swap should not require a projection"); - - assert_eq!(*swapped_join.join_type(), JoinType::RightAnti); - assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft); - assert!(swapped_join.null_aware); - let filter = swapped_join.filter().expect("filter should be kept"); - assert_eq!( - filter - .column_indices() - .iter() - .map(|c| c.side) - .collect::>(), - vec![JoinSide::Right, JoinSide::Left] - ); - assert_eq!(swapped_join.left().schema().field(0).name(), "small_col"); - assert_eq!(swapped_join.right().schema().field(0).name(), "big_col"); - - Ok(()) -} - -/// A multi-key null-aware `LeftAnti` (a multi-column `NOT IN`, correlation -/// scope keys, or both) swaps to a `RightAnti` that keeps the number of -/// `NOT IN` value keys and the key order. -#[rstest] -#[case(1)] -#[case(2)] -#[tokio::test] -async fn test_null_aware_multi_key_left_anti_swaps_to_right_anti( - #[case] value_keys: usize, -) -> Result<()> { - let (big, small) = create_big_and_small(); - let big_col = Arc::new(Column::new_with_schema("big_col", &big.schema())?) as _; - let small_col = Arc::new(Column::new_with_schema("small_col", &small.schema())?) as _; - let join = HashJoinExecBuilder::new( - Arc::clone(&big), - Arc::clone(&small), - vec![ - (Arc::clone(&big_col), Arc::clone(&small_col)), - (big_col, small_col), - ], - JoinType::LeftAnti, - ) - .with_partition_mode(PartitionMode::CollectLeft) - .with_null_aware(true) - .with_null_aware_value_keys(value_keys) - .build()?; - let original_schema = join.schema(); - - let optimized_join = - JoinSelection::new().optimize(Arc::new(join), &ConfigOptions::new())?; - let swapped_join = optimized_join - .downcast_ref::() - .expect("anti join swap should not require a projection"); - - assert_eq!(*swapped_join.join_type(), JoinType::RightAnti); - assert_eq!(*swapped_join.partition_mode(), PartitionMode::CollectLeft); - assert!(swapped_join.null_aware); - assert_eq!(swapped_join.null_aware_value_keys, value_keys); - assert_eq!(swapped_join.on().len(), 2); - assert_eq!(swapped_join.left().schema().field(0).name(), "small_col"); - assert_eq!(swapped_join.right().schema().field(0).name(), "big_col"); - assert_eq!(swapped_join.schema(), original_schema); - - Ok(()) -} - #[tokio::test] async fn test_join_swap_uses_context_statistics_registry() -> Result<()> { let left: Arc = Arc::new(StatisticsExec::new( diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 2be17a09fb6df..2188135fed8ff 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3261,12 +3261,13 @@ async fn collect_left_input( let schema = left_stream.schema(); let prepared = mode == BuildMode::Prepared; - // The extra scope maps are only built for correlated null-aware joins, and - // the null bitmap only for those that mark build rows (see - // `NullAwareMode`). - let with_null_aware_scope_state = - null_aware.is_some_and(NullAwareMode::is_correlated); - let with_null_aware_row_state = + // Correlated null-aware joins (see `NullAwareMode`), multi-column value + // keys without scope keys included, pair NULL-valued rows with the other + // side to find UNKNOWN rows: they need the NULL-valued build rows and, with + // scope keys, the scope map. Of those, only the ones that decide build rows + // also need the per-build-row UNKNOWN bitmap. + let with_null_candidate_state = null_aware.is_some_and(NullAwareMode::is_correlated); + let with_build_unknown_bitmap = null_aware.is_some_and(NullAwareMode::marks_build_rows); let null_aware_value_keys = null_aware.map_or(1, NullAwareMode::value_keys); @@ -3462,13 +3463,13 @@ async fn collect_left_input( BooleanBufferBuilder::new(0) }; - let null_indices_bitmap = if with_null_aware_row_state { + let null_indices_bitmap = if with_build_unknown_bitmap { allocate_bitmap()? } else { BooleanBufferBuilder::new(0) }; - let (null_aware_scope_map, null_value_build_rows) = if with_null_aware_scope_state { + let (null_aware_scope_map, null_value_build_rows) = if with_null_candidate_state { // Null-aware convention: `on_left[..V]` are the value keys and // `on_left[V..]` the (possibly empty) correlation scope keys. let scope_keys = &on_left[null_aware_value_keys..]; @@ -10526,45 +10527,6 @@ mod tests { )); } - /// A multi-column null-aware `RightAnti` (a multi-column or correlated - /// `NOT IN` swapped to build on the subquery) is accepted. - #[tokio::test] - async fn test_null_aware_validation_multi_column_right_anti() { - let left = build_table(("a", &vec![1]), ("b", &vec![2]), ("c", &vec![3])); - let right = build_table(("x", &vec![1]), ("y", &vec![2]), ("z", &vec![3])); - - let on = vec![ - ( - Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, - Arc::new(Column::new_with_schema("x", &right.schema()).unwrap()) as _, - ), - ( - Arc::new(Column::new_with_schema("b", &left.schema()).unwrap()) as _, - Arc::new(Column::new_with_schema("y", &right.schema()).unwrap()) as _, - ), - ]; - - let join = HashJoinExec::try_new( - left, - right, - on, - None, - &JoinType::RightAnti, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - true, - ) - .unwrap(); - assert_eq!( - join.null_aware_mode().unwrap(), - Some(NullAwareMode::RightAnti { - correlated: true, - value_keys: 1, - }) - ); - } - /// Test that null_aware validation rejects PartitionMode::Partitioned (must be CollectLeft). #[tokio::test] async fn test_null_aware_validation_partitioned() { @@ -10604,34 +10566,55 @@ mod tests { ); } - /// A filtered null-aware anti join is a decorrelated correlated `NOT IN` - /// and is accepted both as `LeftAnti` and as the swapped `RightAnti`, which - /// is correlated and so decides `NOT IN` per probe row after the filter. - #[tokio::test] - async fn test_null_aware_filter_accepted_for_anti_joins() { - for join_type in [JoinType::LeftAnti, JoinType::RightAnti] { - let left = - build_table_two_cols(("c1", &vec![Some(1)]), ("dummy", &vec![Some(10)])); - let right = - build_table_two_cols(("c2", &vec![Some(1)]), ("dummy", &vec![Some(100)])); - let on = vec![( - Arc::new(Column::new_with_schema("c1", &left.schema()).unwrap()) as _, - Arc::new(Column::new_with_schema("c2", &right.schema()).unwrap()) as _, - )]; - let join = HashJoinExec::try_new( - left, - right, - on, - Some(prepare_join_filter()), - &join_type, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - true, - ) - .unwrap(); - assert!(join.null_aware_mode().unwrap().unwrap().is_correlated()); + /// A correlated null-aware anti join, with more than one key or a join + /// filter (a decorrelated correlated `NOT IN`), is accepted both as + /// `LeftAnti` and as the `RightAnti` that `JoinSelection` swaps it into. + #[rstest] + #[case::right_anti_multi_key(JoinType::RightAnti, true, false)] + #[case::right_anti_filter(JoinType::RightAnti, false, true)] + #[case::left_anti_filter(JoinType::LeftAnti, false, true)] + fn test_null_aware_correlated_anti_join_accepted( + #[case] join_type: JoinType, + #[case] multi_key: bool, + #[case] with_filter: bool, + ) -> Result<()> { + let left = + build_table_two_cols(("c1", &vec![Some(1)]), ("dummy", &vec![Some(10)])); + let right = + build_table_two_cols(("c2", &vec![Some(1)]), ("dummy", &vec![Some(100)])); + let key = |l: &str, r: &str| -> Result<(PhysicalExprRef, PhysicalExprRef)> { + Ok(( + Arc::new(Column::new_with_schema(l, &left.schema())?) as _, + Arc::new(Column::new_with_schema(r, &right.schema())?) as _, + )) + }; + let mut on = vec![key("c1", "c2")?]; + if multi_key { + on.push(key("dummy", "dummy")?); } + let join = HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on, + with_filter.then(prepare_join_filter), + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + let expected = match join_type { + JoinType::RightAnti => NullAwareMode::RightAnti { + correlated: true, + value_keys: 1, + }, + _ => NullAwareMode::LeftAnti { + correlated: true, + value_keys: 1, + }, + }; + assert_eq!(join.null_aware_mode()?, Some(expected)); + Ok(()) } #[apply(hash_join_exec_configs)] @@ -11311,131 +11294,6 @@ mod tests { Ok(Arc::new(swapped.builder().reset_state().build()?)) } - /// Uncorrelated two-column `(a, b) NOT IN (SELECT a, b ...)`, run both as - /// `LeftAnti` (built on the outer side) and as the swapped `RightAnti` - /// (built on the subquery side), which decides each outer row per probe - /// batch. - #[apply(hash_join_exec_configs)] - #[tokio::test] - async fn test_null_aware_multi_column_right_anti(batch_size: usize) -> Result<()> { - type Rows = Vec<(Option, Option)>; - fn table(rows: &Rows) -> Arc { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.0))), - Arc::new(Int32Array::from_iter(rows.iter().map(|r| r.1))), - ], - ) - .unwrap(); - TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() - } - - async fn not_in( - outer: &Rows, - subquery: &Rows, - batch_size: usize, - ) -> Result { - let (outer, subquery) = (table(outer), table(subquery)); - let on = ["a", "b"] - .into_iter() - .map(|name| { - Ok(( - Arc::new(Column::new_with_schema(name, &outer.schema())?) as _, - Arc::new(Column::new_with_schema(name, &subquery.schema())?) as _, - )) - }) - .collect::>()?; - let left_anti = - HashJoinExecBuilder::new(outer, subquery, on, JoinType::LeftAnti) - .with_partition_mode(PartitionMode::CollectLeft) - .with_null_aware(true) - .with_null_aware_value_keys(2) - .build()?; - let right_anti = swap_for_separate_execution(&left_anti)?; - let right_anti_join = right_anti.downcast_ref::().unwrap(); - assert_eq!( - right_anti_join.null_aware_mode()?, - Some(NullAwareMode::RightAnti { - correlated: true, - value_keys: 2, - }) - ); - - let task_ctx = prepare_task_ctx(batch_size, false); - let expected = batches_to_sort_string( - &common::collect(left_anti.execute(0, Arc::clone(&task_ctx))?).await?, - ); - let actual = batches_to_sort_string( - &common::collect(right_anti.execute(0, task_ctx)?).await?, - ); - assert_eq!(expected, actual); - Ok(actual) - } - - let outer = vec![ - (Some(1), Some(2)), - (Some(1), Some(3)), - (Some(5), Some(7)), - (None, Some(9)), - (None, Some(6)), - (None, None), - ]; - // `(1, 2)` matches exactly, `(1, 3)` has no definite mismatch with - // `(NULL, 3)`, nor `(NULL, 6)` with `(5, 6)` and `(NULL, NULL)` with any - // row. `(5, 7)` and `(NULL, 9)` have a definite mismatch with every row. - let subquery = vec![(Some(1), Some(2)), (None, Some(3)), (Some(5), Some(6))]; - let actual = not_in(&outer, &subquery, batch_size).await?; - allow_duplicates! { - assert_snapshot!(actual, @r" - +---+---+ - | a | b | - +---+---+ - | | 9 | - | 5 | 7 | - +---+---+ - "); - } - - // Every subquery row is NULL-valued, so the full-key hash map is empty - // and every outer row is decided by the NULL candidates alone: only - // `(5, 7)` has a definite mismatch with both `(NULL, 3)` and `(1, NULL)`. - let subquery = vec![(None, Some(3)), (Some(1), None)]; - let actual = not_in(&outer, &subquery, batch_size).await?; - allow_duplicates! { - assert_snapshot!(actual, @r" - +---+---+ - | a | b | - +---+---+ - | 5 | 7 | - +---+---+ - "); - } - - // An empty subquery keeps every outer row, NULL-valued ones included. - let actual = not_in(&outer, &vec![], batch_size).await?; - allow_duplicates! { - assert_snapshot!(actual, @r" - +---+---+ - | a | b | - +---+---+ - | | | - | | 6 | - | | 9 | - | 1 | 2 | - | 1 | 3 | - | 5 | 7 | - +---+---+ - "); - } - - Ok(()) - } - /// Correlated null-aware `LeftAnti` whose two scope keys are Float64, with /// `-0.0` on the build side and `0.0` on the probe side. Both scopes must /// still match in every direction: the full key lookup, NULL probe values diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 72f4c3d7c080e..964baa58aae7f 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -940,36 +940,50 @@ impl HashJoinStream { // Must precede the empty-build-map return below: an all-NULL-key // build side has an empty full-key map but still needs UNKNOWN marks. if mode.is_correlated() && state.is_first_chunk() { - // `RightAnti` marks the rows of this probe batch, the others - // the build rows shared by all probe partitions. - let probe_marks = (!mode.marks_build_rows()).then(|| { - let mut marks = BooleanBufferBuilder::new(state.batch.num_rows()); - marks.append_n(state.batch.num_rows(), false); - SharedBitmapBuilder::new(marks) - }); - let marks = match &probe_marks { - Some(probe_marks) => UnknownMarks::Probe(probe_marks), - None => { - UnknownMarks::Build(build_side.left_data.null_indices_bitmap()) - } + let mut mark = |marks| { + mark_null_candidates_for_probe_batch( + build_side, + state, + mode.value_keys(), + self.filter.as_ref(), + self.join_type, + marks, + &self.random_state, + self.batch_size, + &mut self.null_mark_hashes_buffer, + &mut self.null_mark_probe_indices_buffer, + &mut self.null_mark_build_indices_buffer, + ) }; - mark_null_candidates_for_probe_batch( - build_side, - state, - mode.value_keys(), - self.filter.as_ref(), - self.join_type, - marks, - &self.random_state, - self.batch_size, - &mut self.null_mark_hashes_buffer, - &mut self.null_mark_probe_indices_buffer, - &mut self.null_mark_build_indices_buffer, - )?; - state.unknown_probe_rows = probe_marks - .map(|marks| marks.into_inner().finish()) - .filter(|marks| marks.count_set_bits() > 0) - .map(Box::new); + match mode { + // `RightAnti` decides the rows of this probe batch. Without + // a NULL-valued row on either side none of them can be + // UNKNOWN, so the bitmap is only allocated when one exists. + NullAwareMode::RightAnti { value_keys, .. } => { + let may_be_unknown = + build_side.left_data.null_value_build_rows().is_some() + || state.values[..value_keys] + .iter() + .any(|values| values.logical_null_count() > 0); + if may_be_unknown { + let num_rows = state.batch.num_rows(); + let mut unknown = BooleanBufferBuilder::new(num_rows); + unknown.append_n(num_rows, false); + let unknown = SharedBitmapBuilder::new(unknown); + mark(UnknownMarks::Probe(&unknown))?; + let unknown = unknown.into_inner().finish(); + state.unknown_probe_rows = + (unknown.count_set_bits() > 0).then(|| Box::new(unknown)); + } + } + // The others decide the build rows shared by all probe + // partitions. + NullAwareMode::LeftAnti { .. } | NullAwareMode::LeftMark { .. } => { + mark(UnknownMarks::Build( + build_side.left_data.null_indices_bitmap(), + ))?; + } + } } } @@ -1556,8 +1570,9 @@ impl<'a> UnknownMarks<'a> { } } - /// The row of a (build, probe) candidate pair that this join marks. - fn marked_row(self, build_idx: u64, probe_idx: u32) -> usize { + /// The row of a (build, probe) candidate pair whose `NOT IN` the pair + /// decides, and that is marked when the pair makes it UNKNOWN. + fn unknown_row(self, build_idx: u64, probe_idx: u32) -> usize { match self { Self::Build(_) => build_idx as usize, Self::Probe(_) => probe_idx as usize, @@ -1571,9 +1586,10 @@ impl<'a> UnknownMarks<'a> { /// /// Key layout: `on[..V]` are the `NOT IN` value keys, `on[V..]` the (possibly /// empty) correlation scope keys (see `HashJoinExec::null_aware`). A row is -/// NULL-valued when any of its value keys is NULL. A (build, probe) pair makes -/// the `NOT IN` of the marked row UNKNOWN instead of TRUE (a `LeftMark` mark -/// NULL instead of FALSE) when either: +/// NULL-valued when any of its value keys is NULL. Each (build, probe) pair +/// decides one of its rows (see [`UnknownMarks::unknown_row`]), and makes that +/// row's `NOT IN` UNKNOWN instead of TRUE (a `LeftMark` mark NULL instead of +/// FALSE) when either: /// 1. the build row is NULL-valued, or /// 2. the probe row is NULL-valued, /// @@ -1590,10 +1606,10 @@ impl<'a> UnknownMarks<'a> { /// The value tuples and then the join filter, if any, decide which candidates /// count. /// -/// A row stays UNKNOWN once it is marked, so candidates whose marked row is -/// already marked are skipped, and the join filter is not evaluated for them. -/// Without scope keys this also ends the pairing as soon as no unmarked row is -/// left. +/// A row stays UNKNOWN once it is marked, so pairs that decide a row already +/// known to be UNKNOWN are skipped, and the join filter is not evaluated for +/// them. Without scope keys this also ends the pairing as soon as every row +/// that can be decided is UNKNOWN. #[expect(clippy::too_many_arguments)] fn mark_null_candidates_for_probe_batch( build_side: &BuildSideReadyState, @@ -1660,7 +1676,7 @@ fn mark_null_candidates_for_probe_batch( for (build_idx, probe_idx) in build_indices.values().iter().zip(probe_indices.values()) { - bitmap.set_bit(marks.marked_row(*build_idx, *probe_idx), true); + bitmap.set_bit(marks.unknown_row(*build_idx, *probe_idx), true); } } Ok(()) @@ -1862,7 +1878,7 @@ fn for_each_scope_match( Ok(()) } -/// Removes the candidate pairs whose marked row is already marked UNKNOWN. +/// Removes the candidate pairs that decide a row already marked UNKNOWN. fn retain_unmarked( marks: UnknownMarks<'_>, build_indices: UInt64Array, @@ -1870,7 +1886,7 @@ fn retain_unmarked( ) -> (UInt64Array, UInt32Array) { let bitmap = marks.bitmap().lock(); let is_unmarked = |(build_idx, probe_idx): (&u64, &u32)| { - !bitmap.get_bit(marks.marked_row(*build_idx, *probe_idx)) + !bitmap.get_bit(marks.unknown_row(*build_idx, *probe_idx)) }; if build_indices .values() @@ -1889,9 +1905,9 @@ fn retain_unmarked( (build.into(), probe.into()) } -/// Calls `f` with the pairs of `build_rows` x `probe_rows` whose marked row -/// (see [`UnknownMarks`]) is not marked UNKNOWN, as chunks of at most -/// `batch_size` pairs. +/// Calls `f` with the pairs of `build_rows` x `probe_rows` that decide a row +/// (see [`UnknownMarks::unknown_row`]) not yet marked UNKNOWN, as chunks of at +/// most `batch_size` pairs. fn for_each_unmarked_cross_product( marks: UnknownMarks<'_>, build_rows: impl Iterator, @@ -1917,53 +1933,53 @@ fn for_each_unmarked_cross_product( } } -/// Calls `f` with the pairs of `marked_rows` x `other_rows` whose row in -/// `marked_rows` is not set in `bitmap`, as chunks of at most `batch_size` +/// Calls `f` with the pairs of `decided_rows` x `other_rows` whose row in +/// `decided_rows` is not set in `bitmap`, as chunks of at most `batch_size` /// pairs. /// -/// `f` sets bits in `bitmap`, so the unmarked rows are found again after each -/// chunk. The pairing stops when no unmarked row is left. -fn for_each_unmarked_pair, O: Copy>( +/// `f` sets bits in `bitmap`, so the unset rows are found again after each +/// chunk. The pairing stops when every row of `decided_rows` is set. +fn for_each_unmarked_pair, O: Copy>( bitmap: &SharedBitmapBuilder, - mut marked_rows: Vec, + mut decided_rows: Vec, other_rows: impl Iterator, batch_size: usize, - mut f: impl FnMut(Vec, Vec) -> Result<()>, + mut f: impl FnMut(Vec, Vec) -> Result<()>, ) -> Result<()> { - let retain_unmarked_rows = |rows: &mut Vec| { + let retain_undecided_rows = |rows: &mut Vec| { let bitmap = bitmap.lock(); rows.retain(|idx| !bitmap.get_bit((*idx).into() as usize)); }; - retain_unmarked_rows(&mut marked_rows); + retain_undecided_rows(&mut decided_rows); - let mut marked_chunk = Vec::with_capacity(batch_size); + let mut decided_chunk = Vec::with_capacity(batch_size); let mut other_chunk = Vec::with_capacity(batch_size); let mut marks_since_refresh = false; for other_row in other_rows { // Refresh only after a chunk was sent, so the cost of the refresh // stays proportional to the pairs already evaluated. if marks_since_refresh { - retain_unmarked_rows(&mut marked_rows); + retain_undecided_rows(&mut decided_rows); marks_since_refresh = false; } - if marked_rows.is_empty() { + if decided_rows.is_empty() { break; } - for marked_row in &marked_rows { - marked_chunk.push(*marked_row); + for decided_row in &decided_rows { + decided_chunk.push(*decided_row); other_chunk.push(other_row); - if marked_chunk.len() == batch_size { + if decided_chunk.len() == batch_size { f( - std::mem::replace(&mut marked_chunk, Vec::with_capacity(batch_size)), + std::mem::replace(&mut decided_chunk, Vec::with_capacity(batch_size)), std::mem::replace(&mut other_chunk, Vec::with_capacity(batch_size)), )?; marks_since_refresh = true; } } } - if !marked_chunk.is_empty() { - f(marked_chunk, other_chunk)?; + if !decided_chunk.is_empty() { + f(decided_chunk, other_chunk)?; } Ok(()) } diff --git a/datafusion/sqllogictest/test_files/null_aware_memory_limit.slt b/datafusion/sqllogictest/test_files/null_aware_memory_limit.slt new file mode 100644 index 0000000000000..ec0f16f271189 --- /dev/null +++ b/datafusion/sqllogictest/test_files/null_aware_memory_limit.slt @@ -0,0 +1,67 @@ +# 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. + +# A multi-column NOT IN over a small subquery builds its null-aware hash join +# on the subquery, so a large outer table runs under a memory limit that only +# fits the subquery side. The join used to stay a LeftAnti that builds on the +# outer side and ran out of memory where the scalar NOT IN ran. + +statement ok +CREATE TABLE na_mem_big AS +SELECT v AS a, v % 1000 AS b FROM generate_series(1, 200000) AS t(v); + +# A NULL-valued subquery row makes NOT IN UNKNOWN for every outer row with +# `b = 7`. +statement ok +CREATE TABLE na_mem_small AS +SELECT v AS a, v % 1000 AS b FROM generate_series(1, 100) AS t(v) +UNION ALL +SELECT NULL, 7; + +# Far below the outer side's build state, far above the subquery's. +statement ok +SET datafusion.runtime.memory_limit = '2M' + +# Rows 1..=100 match exactly; of the rest, those with `b = 7` are UNKNOWN: 199 +# of the 200 values `v % 1000 = 7`, all but `v = 7`. +query I +SELECT count(*) FROM na_mem_big +WHERE (a, b) NOT IN (SELECT a, b FROM na_mem_small); +---- +199701 + +# Building on the outer side does not fit, which shows the limit binds. +statement ok +SET datafusion.optimizer.join_reordering = false + +query error Resources exhausted +SELECT count(*) FROM na_mem_big +WHERE (a, b) NOT IN (SELECT a, b FROM na_mem_small); + +statement ok +RESET datafusion.optimizer.join_reordering + +statement ok +DROP TABLE na_mem_big; + +statement ok +DROP TABLE na_mem_small; + +# Setting a runtime option rebuilds the session state from the existing one, +# which turns this option off because the default catalog already exists. +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema diff --git a/datafusion/sqllogictest/test_files/null_aware_multi_column.slt b/datafusion/sqllogictest/test_files/null_aware_multi_column.slt index 96dd612a42426..ebf33f199154f 100644 --- a/datafusion/sqllogictest/test_files/null_aware_multi_column.slt +++ b/datafusion/sqllogictest/test_files/null_aware_multi_column.slt @@ -52,6 +52,11 @@ CREATE TABLE mc_inner_no_null(x INT, y INT) AS VALUES statement ok CREATE TABLE mc_empty(x INT, y INT); +statement ok +CREATE TABLE mc_inner_all_null(x INT, y INT) AS VALUES +(NULL, 4), +(1, NULL); + ############# ## NULL in the subquery ############# @@ -66,6 +71,15 @@ WHERE (a, b) NOT IN (SELECT x, y FROM mc_inner_null); 5 6 z NULL 8 w +# Every subquery row holds a NULL, so no row can match a whole tuple and only +# the NULL-aware checks decide. Only (5, 6) has a definite mismatch with both +# (NULL, 4) and (1, NULL). +query IIT rowsort +SELECT * FROM mc_outer +WHERE (a, b) NOT IN (SELECT x, y FROM mc_inner_all_null); +---- +5 6 z + ############# ## NULL only in the outer tuple ############# @@ -548,6 +562,9 @@ DROP TABLE mc_inner_no_null; statement ok DROP TABLE mc_empty; +statement ok +DROP TABLE mc_inner_all_null; + statement ok DROP TABLE mc3_outer;