From ac05cb90aab01431e0a8089c8e369929b0a82749 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:29:30 -0500 Subject: [PATCH 1/5] fix: apply the join filter when null-aware hash joins mark NOT IN rows UNKNOWN A null-aware LeftAnti join with a join filter ignored the filter for NULL keys: one NULL probe key removed every build row, even when the filter excluded that NULL row for every build row. A null-aware LeftAnti join with correlation scope keys failed to plan. Treat a null-aware LeftAnti or LeftMark join as correlated when it has scope keys or a join filter. Correlated joins record the UNKNOWN decision per build row in the null-indices bitmap: the candidate (build, probe) pairs come from the scope map, or from all pairs when there are no scope keys, and the join filter decides which pairs count. The LeftAnti final stage drops the rows marked UNKNOWN. JoinSelection only swaps an uncorrelated null-aware LeftAnti. Co-Authored-By: Claude Opus 5 --- .../physical-optimizer/src/join_selection.rs | 4 +- .../physical-plan/src/joins/hash_join/exec.rs | 195 ++++++----- .../src/joins/hash_join/stream.rs | 328 ++++++++++++------ 3 files changed, 336 insertions(+), 191 deletions(-) diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index f11e8612de370..2b342bece7040 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -172,11 +172,13 @@ impl PhysicalOptimizerRule for JoinSelection { } } -/// Determines whether it is possible to swap inputs of a hash join - for null-aware joins, we can only swap `LeftAnti` with no filters +/// 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 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())) } diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 7b9e701119ef4..64c462507fcbf 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -192,19 +192,20 @@ fn try_create_array_map( Ok(Some((array_map, batch, left_values))) } -/// Correlation-scope hash map over only the build rows whose scalar `NOT IN` -/// value key is NULL, used by correlated null-aware `LeftMark` joins. +/// The build rows whose scalar `NOT IN` value key is NULL, used by correlated +/// null-aware joins (see [`NullAwareMode`]). /// -/// Such rows produce a NULL (UNKNOWN) mark whenever *any* probe row shares -/// their correlation scope, so every probe row must be tested against them. -/// Restricting this map to the NULL-valued build rows keeps that lookup +/// 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. +/// 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. -pub(super) struct NullValueScopeMap { +pub(super) struct NullValueBuildRows { /// Hash table keyed by the correlation scope values of the NULL-valued /// build rows. Stored positions index into `scope_values`/`build_indices`, - /// not the full build batch. - pub(super) map: Box, + /// not the full build batch. `None` when the join has no correlation + /// scope keys, so every probe row is in scope. + pub(super) scope_map: Option>, /// Correlation scope key values of the NULL-valued build rows. pub(super) scope_values: Vec, /// Maps positions in `map`/`scope_values` back to row indices in the full @@ -215,19 +216,23 @@ pub(super) struct NullValueScopeMap { /// Null-aware (`NOT IN`) semantics of a hash join, derived from /// [`HashJoinExec::null_aware`] and the join type. /// -/// Only these three combinations are legal (see [`Self::try_new`]), so the +/// Only these combinations are legal (see [`Self::try_new`]), so the /// stream matches on this instead of re-checking `null_aware && join_type == ..`. +/// +/// A `correlated` join has correlation scope keys (`on[1..]`, see +/// [`HashJoinExec::null_aware`]) or a join filter, or both. A NULL then makes +/// `NOT IN` UNKNOWN only for the build rows whose scope and filter keep that +/// NULL, so the join records the decision per build row in the null-indices +/// bitmap instead of in shared probe-side flags. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) enum NullAwareMode { - /// Uncorrelated `build.key NOT IN (probe.key)`: emits build rows, and - /// none of them once any probe key is NULL. - LeftAnti, + /// `build.key NOT IN (probe.key)`: emits build rows. When uncorrelated, + /// none of them are emitted once any probe key is NULL. + LeftAnti { correlated: bool }, /// Uncorrelated `probe.key NOT IN (build.key)`: emits probe rows, and /// none of them once any build key is NULL. RightAnti, - /// `NOT IN` as a nullable mark column on the build rows. `correlated` - /// means `on[1..]` are correlation scope keys (see - /// [`HashJoinExec::null_aware`]). + /// `NOT IN` as a nullable mark column on the build rows. LeftMark { correlated: bool }, } @@ -239,13 +244,12 @@ impl NullAwareMode { num_keys: usize, has_filter: bool, ) -> Result { + let correlated = num_keys > 1 || has_filter; let mode = match (join_type, partition_mode) { - (JoinType::LeftAnti, _) => Self::LeftAnti, + (JoinType::LeftAnti, _) => Self::LeftAnti { correlated }, // `PartitionMode::CollectLeft` is safe because `RightAnti` is probe-driven (JoinType::RightAnti, PartitionMode::CollectLeft) => Self::RightAnti, - (JoinType::LeftMark, _) => Self::LeftMark { - correlated: num_keys > 1, - }, + (JoinType::LeftMark, _) => Self::LeftMark { correlated }, _ => { return plan_err!( "null_aware can only be true for LeftAnti joins and RightAnti joins with `CollectLeft` `PartitionMode`, or LeftMark joins, got {join_type} with {partition_mode}" @@ -253,10 +257,14 @@ impl NullAwareMode { } }; match mode { - Self::LeftAnti | Self::RightAnti if num_keys != 1 => plan_err!( + Self::RightAnti if num_keys != 1 => plan_err!( "null_aware {join_type} joins only support single column join key, got {num_keys} columns" ), - Self::LeftMark { .. } if partition_mode == PartitionMode::Partitioned => { + // Correlated joins share the per-build-row null bitmap across all + // probe partitions. + Self::LeftMark { .. } | Self::LeftAnti { correlated: true } + if partition_mode == PartitionMode::Partitioned => + { plan_err!( "null_aware joins require PartitionMode::CollectLeft, got PartitionMode::Partitioned" ) @@ -267,6 +275,14 @@ impl NullAwareMode { _ => Ok(mode), } } + + /// Whether this join decides UNKNOWN per build row (see [`NullAwareMode`]). + pub(super) fn is_correlated(self) -> bool { + matches!( + self, + Self::LeftAnti { correlated: true } | Self::LeftMark { correlated: true } + ) + } } /// HashTable and input data for the left (build side) of a join @@ -274,17 +290,16 @@ pub(super) struct JoinLeftData { /// The hash table with indices into `batch` /// Arc is used to allow sharing with SharedBuildAccumulator for hash map pushdown pub(super) map: Arc, - /// Hash table over correlated scope keys for scalar null-aware mark joins. + /// Hash table over correlated scope keys for correlated null-aware joins. /// - /// For null-aware `LeftMark`, key 0 is the scalar `NOT IN` value key and - /// keys 1..N are correlated equality scope keys. This map covers all build - /// rows and is probed only with NULL-valued probe rows; the complementary - /// direction uses `null_value_scope_map`. - null_aware_mark_scope_map: Option>, - /// Scope map restricted to the build rows whose value key is NULL (see - /// [`NullValueScopeMap`]). `None` when the build side has no NULL value - /// keys. - null_value_scope_map: Option, + /// Key 0 is the scalar `NOT IN` value key and keys 1..N are correlated + /// equality scope keys. This map covers all build rows and is probed only + /// with NULL-valued probe rows; the complementary direction uses + /// `null_value_build_rows`. `None` when there are no scope keys. + null_aware_scope_map: Option>, + /// The build rows whose value key is NULL (see [`NullValueBuildRows`]). + /// `None` when the build side has no NULL value keys. + null_value_build_rows: Option, /// The input rows for the build side batch: RecordBatch, /// The build side on expressions values @@ -318,12 +333,12 @@ impl JoinLeftData { &self.map } - pub(super) fn null_aware_mark_scope_map(&self) -> Option<&dyn JoinHashMapType> { - self.null_aware_mark_scope_map.as_deref() + pub(super) fn null_aware_scope_map(&self) -> Option<&dyn JoinHashMapType> { + self.null_aware_scope_map.as_deref() } - pub(super) fn null_value_scope_map(&self) -> Option<&NullValueScopeMap> { - self.null_value_scope_map.as_ref() + pub(super) fn null_value_build_rows(&self) -> Option<&NullValueBuildRows> { + self.null_value_build_rows.as_ref() } /// returns a reference to the build side batch @@ -878,13 +893,16 @@ pub struct HashJoinExec { /// Flag to indicate if this join uses null-aware equality semantics. /// /// Set for the physical lowering of scalar `NOT IN` subqueries (producing - /// `JoinType::LeftAnti` when uncorrelated or `JoinType::LeftMark` when - /// correlated). When `true`, NULLs in the join keys follow SQL `NOT IN` - /// three-valued logic rather than ordinary equi-join semantics. + /// `JoinType::LeftAnti` at the top level of a filter or `JoinType::LeftMark` + /// inside a larger expression). When `true`, NULLs in the join keys follow + /// SQL `NOT IN` three-valued logic rather than ordinary equi-join semantics. + /// 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. /// /// Key-ordering convention (relied on positionally, not enforced): for a - /// null-aware `LeftMark` join with more than one key, `on[0]` is the scalar - /// `NOT IN` value key and `on[1..N]` are the correlated equality scope keys. + /// null-aware `LeftAnti` or `LeftMark` join with more than one key, `on[0]` + /// is the scalar `NOT IN` value key and `on[1..N]` are the correlated + /// equality scope 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`). @@ -2858,11 +2876,8 @@ async fn collect_left_input( let schema = left_stream.schema(); // The extra scope maps + null bitmap are only built for correlated - // null-aware LeftMark joins (`on_left[1..]` are correlation scope keys). - let with_null_aware_mark_state = matches!( - null_aware, - Some(NullAwareMode::LeftMark { correlated: true }) - ); + // null-aware joins (see `NullAwareMode`). + let with_null_aware_row_state = null_aware.is_some_and(NullAwareMode::is_correlated); let is_phj_candidate = is_perfect_hash_join_candidate(&on_left, &schema)?; @@ -2992,42 +3007,42 @@ async fn collect_left_input( BooleanBufferBuilder::new(0) }; - let null_indices_bitmap = if with_null_aware_mark_state { + let null_indices_bitmap = if with_null_aware_row_state { allocate_bitmap()? } else { BooleanBufferBuilder::new(0) }; - let (null_aware_mark_scope_map, null_value_scope_map) = if with_null_aware_mark_state - { - // Null-aware `LeftMark` convention: `on_left[0]` is the value key and - // `on_left[1..]` the scope keys, so the scope map needs more than one key. - debug_assert!( - on_left.len() > 1, - "null-aware LeftMark needs on_left[0]=value, on_left[1..]=scope, got {} key(s)", - on_left.len() - ); - // Scope-only NULL marking uses a HashMap (the primary join map may use - // ArrayMap for full-key matches, but scope keys have arbitrary shape). - let mut scope_map = new_join_hashmap(num_rows, &mut reservation, &metrics)?; - - let mut hashes_buffer = vec![0; batch.num_rows()]; - update_hash( - &on_left[1..], - &batch, - &mut *scope_map, - 0, - &random_state, - &mut hashes_buffer, - 0, - true, - NullEquality::NullEqualsNothing, - )?; + let (null_aware_scope_map, null_value_build_rows) = if with_null_aware_row_state { + // Null-aware convention: `on_left[0]` is the value key and + // `on_left[1..]` the (possibly empty) correlation scope keys. + let scope_keys = &on_left[1..]; + let scope_map = if scope_keys.is_empty() { + None + } else { + // Scope-only NULL marking uses a HashMap (the primary join map may + // use ArrayMap for full-key matches, but scope keys have arbitrary + // shape). + let mut scope_map = new_join_hashmap(num_rows, &mut reservation, &metrics)?; + + let mut hashes_buffer = vec![0; batch.num_rows()]; + update_hash( + scope_keys, + &batch, + &mut *scope_map, + 0, + &random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNothing, + )?; + Some(scope_map) + }; - // Build the dedicated scope map over the NULL-valued build rows (see - // `NullValueScopeMap`). + // Collect the NULL-valued build rows (see `NullValueBuildRows`). let value_key = &left_values[0]; - let null_value_scope_map = if value_key.null_count() > 0 { + let null_value_build_rows = if value_key.logical_null_count() > 0 { let null_mask = arrow::compute::is_null(value_key.as_ref())?; let build_indices = UInt64Array::from_iter_values( null_mask.values().set_indices().map(|i| i as u64), @@ -3048,14 +3063,19 @@ async fn collect_left_input( reservation.try_grow(retained_size)?; metrics.build_mem_used.add(retained_size); - let null_rows = build_indices.len(); - let mut map = new_join_hashmap(null_rows, &mut reservation, &metrics)?; - let mut hashes_buffer = vec![0; null_rows]; - create_hashes(&scope_values, &random_state, &mut hashes_buffer)?; - map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0); + let scope_map = if scope_values.is_empty() { + None + } else { + let null_rows = build_indices.len(); + let mut map = new_join_hashmap(null_rows, &mut reservation, &metrics)?; + let mut hashes_buffer = vec![0; null_rows]; + create_hashes(&scope_values, &random_state, &mut hashes_buffer)?; + map.update_from_iter(Box::new(hashes_buffer.iter().enumerate().rev()), 0); + Some(map) + }; - Some(NullValueScopeMap { - map, + Some(NullValueBuildRows { + scope_map, scope_values, build_indices, }) @@ -3063,7 +3083,7 @@ async fn collect_left_input( None }; - (Some(scope_map), null_value_scope_map) + (scope_map, null_value_build_rows) } else { (None, None) }; @@ -3106,8 +3126,8 @@ async fn collect_left_input( let data = JoinLeftData { map, - null_aware_mark_scope_map, - null_value_scope_map, + null_aware_scope_map, + null_value_build_rows, batch, values: left_values, visited_indices_bitmap: Mutex::new(visited_indices_bitmap), @@ -8825,13 +8845,14 @@ mod tests { ), ]; - // Try to create null-aware anti join with 2 columns (should fail) + // 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( left, right, on, None, - &JoinType::LeftAnti, + &JoinType::RightAnti, None, PartitionMode::CollectLeft, NullEquality::NullEqualsNothing, @@ -8841,7 +8862,7 @@ mod tests { assert!(result.is_err()); assert!( result.unwrap_err().to_string().contains( - "null_aware LeftAnti joins only support single column join key" + "null_aware RightAnti joins only support single column join key" ) ); } diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 7b70b2bf3b318..de6de146cddd1 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -37,7 +37,6 @@ use crate::stream::EmptyRecordBatchStream; use crate::{ RecordBatchStream, SendableRecordBatchStream, handle_state, hash_utils::create_hashes, - joins::SharedBitmapBuilder, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, @@ -581,8 +580,8 @@ impl HashJoinStream { hashes_buffer, probe_indices_buffer: Vec::with_capacity(batch_size), build_indices_buffer: Vec::with_capacity(batch_size), - // Left unallocated: only correlated null-aware LeftMark joins ever - // use these, and they grow them on first use. + // Left unallocated: only correlated null-aware joins ever use + // these, and they grow them on first use. null_mark_hashes_buffer: Vec::new(), null_mark_probe_indices_buffer: Vec::new(), null_mark_build_indices_buffer: Vec::new(), @@ -837,29 +836,24 @@ impl HashJoinStream { let timer = self.join_metrics.join_time.timer(); if let Some(mode) = self.null_aware { - if null_aware_skip_probe_batch( - mode, - state, - &build_side.left_data, - self.filter.is_some(), - ) { + if null_aware_skip_probe_batch(mode, state, &build_side.left_data) { timer.done(); self.state = HashJoinStreamState::FetchProbeBatch; return Ok(StatefulStreamResult::Continue); } - // For correlated null-aware LeftMark, record this batch's UNKNOWN + // For correlated null-aware joins, record this batch's UNKNOWN // candidates once, before the first chunked lookup // (offset == (0, None)). // // 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 matches!(mode, NullAwareMode::LeftMark { correlated: true }) - && state.offset == (0, None) - { + if mode.is_correlated() && state.offset == (0, None) { mark_null_candidates_for_probe_batch( build_side, state, + self.filter.as_ref(), + self.join_type, &self.random_state, self.batch_size, &mut self.null_mark_hashes_buffer, @@ -1126,18 +1120,20 @@ impl HashJoinStream { // Null-aware joins post-process the build rows under SQL three-valued // logic; see the helpers for the rules. let (left_side, right_side, mark_column) = match self.null_aware { - Some(NullAwareMode::LeftAnti) => { + Some(NullAwareMode::LeftAnti { correlated }) => { let (left_side, right_side) = null_aware_left_anti_final_indices( &build_side.left_data, + correlated, probe_summary, left_side, right_side, ); (left_side, right_side, None) } - Some(NullAwareMode::LeftMark { .. }) => { + Some(NullAwareMode::LeftMark { correlated }) => { let mark_column = null_aware_left_mark_column( &build_side.left_data, + correlated, probe_summary, &left_side, &right_side, @@ -1243,24 +1239,28 @@ fn null_aware_skip_probe_batch( mode: NullAwareMode, state: &ProcessProbeBatchState, left_data: &JoinLeftData, - has_filter: bool, ) -> bool { match mode { NullAwareMode::RightAnti => left_data.build_side_has_null, - NullAwareMode::LeftAnti | NullAwareMode::LeftMark { .. } => { + // Correlated joins decide UNKNOWN per build row instead, in + // `mark_null_candidates_for_probe_batch`. + NullAwareMode::LeftAnti { correlated: true } + | NullAwareMode::LeftMark { correlated: true } => false, + NullAwareMode::LeftAnti { correlated: false } + | NullAwareMode::LeftMark { correlated: false } => { // `on[0]` is the `NOT IN` value key for both modes. let probe_key_column = &state.values[0]; - let probe_has_null = match mode { - NullAwareMode::LeftAnti if !has_filter => { - probe_key_column.logical_null_count() > 0 - } - _ => probe_key_column.null_count() > 0, + let is_anti = matches!(mode, NullAwareMode::LeftAnti { .. }); + let probe_has_null = if is_anti { + probe_key_column.logical_null_count() > 0 + } else { + probe_key_column.null_count() > 0 }; // Only batches with rows count: `NULL NOT IN (empty)` is TRUE. left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null); // Best-effort early exit; the final stage re-checks the flag // through `report_probe_completed`. - mode == NullAwareMode::LeftAnti && left_data.probe_side_has_null_hint() + is_anti && left_data.probe_side_has_null_hint() } } } @@ -1291,7 +1291,12 @@ fn drop_null_probe_keys( } /// Final-stage rules of a null-aware `LeftAnti` join, evaluated by the last -/// probe partition from what every partition together saw: +/// probe partition. +/// +/// A correlated join (see `NullAwareMode`) drops the unmatched build rows +/// whose `NOT IN` was marked UNKNOWN in the null-indices bitmap. +/// +/// Otherwise the rules use what every partition together saw: /// - a NULL probe key seen by any partition makes `build.key NOT IN (probe)` /// UNKNOWN for every build row, so nothing is emitted; /// - otherwise a NULL build key means `NULL NOT IN (probe)`, which is UNKNOWN @@ -1299,10 +1304,23 @@ fn drop_null_probe_keys( /// kept). fn null_aware_left_anti_final_indices( left_data: &JoinLeftData, + correlated: bool, probe_summary: ProbeSideSummary, left_side: UInt64Array, right_side: UInt32Array, ) -> (UInt64Array, UInt32Array) { + if correlated { + let null_indices_bitmap = left_data.null_indices_bitmap().lock(); + let left_side = UInt64Array::from_iter_values( + left_side + .values() + .iter() + .copied() + .filter(|idx| !null_indices_bitmap.get_bit(*idx as usize)), + ); + let right_side = UInt32Array::new_null(left_side.len()); + return (left_side, right_side); + } if probe_summary.has_null { return (UInt64Array::new_null(0), UInt32Array::new_null(0)); } @@ -1324,16 +1342,14 @@ fn null_aware_left_anti_final_indices( /// final indices and what every probe partition together saw. fn null_aware_left_mark_column( left_data: &JoinLeftData, + correlated: bool, probe_summary: ProbeSideSummary, left_side: &UInt64Array, right_side: &UInt32Array, ) -> ArrayRef { let build_key_column = &left_data.values()[0]; // Correlated joins precomputed the UNKNOWN decision per build row. - let null_indices_bitmap = left_data - .null_aware_mark_scope_map() - .is_some() - .then(|| left_data.null_indices_bitmap().lock()); + let null_indices_bitmap = correlated.then(|| left_data.null_indices_bitmap().lock()); build_null_aware_left_mark_column( left_side, right_side, @@ -1344,104 +1360,184 @@ fn null_aware_left_mark_column( ) } -/// Records which build rows of a correlated null-aware `LeftMark` join are -/// UNKNOWN candidates for this probe batch. +/// Records which build rows of a correlated null-aware join are UNKNOWN +/// candidates for this probe batch. /// -/// Key layout: `on[0]` is the `NOT IN` value key, `on[1..]` the correlation -/// scope keys (see `HashJoinExec::null_aware`). A build row's mark must be -/// NULL (SQL UNKNOWN) instead of FALSE when it is unmatched and either: -/// 1. its value key is NULL and any probe row shares its correlation scope, or -/// 2. some probe row in its correlation scope has a NULL value key. +/// Key layout: `on[0]` is the `NOT IN` value key, `on[1..]` the (possibly +/// empty) correlation scope keys (see `HashJoinExec::null_aware`). An +/// unmatched build row's `NOT IN` is UNKNOWN instead of TRUE (its mark is NULL +/// instead of FALSE) when either: +/// 1. its value key is NULL and any probe row in its correlation scope passes +/// the join filter, or +/// 2. some probe row in its correlation scope with a NULL value key passes the +/// join filter. /// -/// Case 1 probes the build-side NULL-value scope map with all probe rows; -/// case 2 probes the full scope map with only the NULL-valued probe rows. +/// 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 +/// pairs through a hash lookup; without scope keys every pair is a candidate. +/// The join filter, if any, then decides which candidates count. +#[expect(clippy::too_many_arguments)] fn mark_null_candidates_for_probe_batch( build_side: &BuildSideReadyState, state: &ProcessProbeBatchState, + filter: Option<&JoinFilter>, + join_type: JoinType, random_state: &RandomState, batch_size: usize, hashes_buffer: &mut Vec, probe_indices_buffer: &mut Vec, build_indices_buffer: &mut Vec, ) -> Result<()> { - let Some(scope_map) = build_side.left_data.null_aware_mark_scope_map() else { + let left_data = &build_side.left_data; + let null_value_build_rows = left_data.null_value_build_rows(); + let probe_value_key = &state.values[0]; + let probe_has_null_values = probe_value_key.logical_null_count() > 0; + if null_value_build_rows.is_none() && !probe_has_null_values { return Ok(()); - }; + } debug_assert_eq!( - build_side.left_data.values().len(), + left_data.values().len(), state.values.len(), "build/probe key counts must match" ); - debug_assert!(state.values.len() > 1, "keys must be [value, scope..]"); - - let probe_value_key = &state.values[0]; - let build_scope_values = &build_side.left_data.values()[1..]; + let build_scope_values = &left_data.values()[1..]; let probe_scope_values = &state.values[1..]; - let null_value_scope_map = build_side.left_data.null_value_scope_map(); - let probe_has_null_values = probe_value_key.null_count() > 0; - if null_value_scope_map.is_none() && !probe_has_null_values { - return Ok(()); - } + // Keeps the candidate pairs that pass the join filter and marks their + // build rows as UNKNOWN. + let mut mark = |build_indices: UInt64Array, probe_indices: UInt32Array| { + 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, + }; + 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); + } + } + Ok(()) + }; // Case 1: build rows with a NULL value key are UNKNOWN as soon as any - // probe row shares their correlation scope. - if let Some(null_value_scope_map) = null_value_scope_map { - hashes_buffer.clear(); - hashes_buffer.resize(state.batch.num_rows(), 0); - create_hashes(probe_scope_values, random_state, hashes_buffer)?; - - scan_scope_matches_into_bitmap( - null_value_scope_map.map.as_ref(), - &null_value_scope_map.scope_values, - probe_scope_values, - hashes_buffer, - batch_size, - probe_indices_buffer, - build_indices_buffer, - // The map indexes only the NULL-valued build rows; translate its - // positions back to row indices in the full build batch. - |position| null_value_scope_map.build_indices.value(position as usize), - build_side.left_data.null_indices_bitmap(), - )?; + // probe row in their correlation scope passes the filter. + if let Some(null_rows) = null_value_build_rows { + match &null_rows.scope_map { + Some(scope_map) => { + hashes_buffer.clear(); + hashes_buffer.resize(state.batch.num_rows(), 0); + create_hashes(probe_scope_values, random_state, hashes_buffer)?; + + for_each_scope_match( + scope_map.as_ref(), + &null_rows.scope_values, + probe_scope_values, + hashes_buffer, + batch_size, + probe_indices_buffer, + build_indices_buffer, + |positions, probe_indices| { + // The map indexes only the NULL-valued build rows; + // translate its positions back to build row indices. + let build_indices = UInt64Array::from_iter_values( + positions + .values() + .iter() + .map(|p| null_rows.build_indices.value(*p as usize)), + ); + mark(build_indices, probe_indices) + }, + )?; + } + None => { + let probe_rows = + UInt32Array::from_iter_values(0..state.batch.num_rows() as u32); + for_each_cross_product( + &null_rows.build_indices, + &probe_rows, + batch_size, + &mut mark, + )?; + } + } } // Case 2: NULL-valued probe rows make every build row in their correlation - // scope an UNKNOWN candidate. + // scope that passes the filter an UNKNOWN candidate. if probe_has_null_values { let null_mask = arrow::compute::is_null(probe_value_key.as_ref())?; - let probe_null_scope_values = probe_scope_values - .iter() - .map(|values| Ok(arrow::compute::filter(values.as_ref(), &null_mask)?)) - .collect::>>()?; - - hashes_buffer.clear(); - hashes_buffer.resize(null_mask.true_count(), 0); - create_hashes(&probe_null_scope_values, random_state, hashes_buffer)?; + let null_probe_rows = UInt32Array::from_iter_values( + null_mask.values().set_indices().map(|i| i as u32), + ); - scan_scope_matches_into_bitmap( - scope_map, - build_scope_values, - &probe_null_scope_values, - hashes_buffer, - batch_size, - probe_indices_buffer, - build_indices_buffer, - |position| position, - build_side.left_data.null_indices_bitmap(), - )?; + match left_data.null_aware_scope_map() { + Some(scope_map) => { + let probe_null_scope_values = probe_scope_values + .iter() + .map(|values| { + Ok(arrow::compute::filter(values.as_ref(), &null_mask)?) + }) + .collect::>>()?; + + hashes_buffer.clear(); + hashes_buffer.resize(null_probe_rows.len(), 0); + create_hashes(&probe_null_scope_values, random_state, hashes_buffer)?; + + for_each_scope_match( + scope_map, + build_scope_values, + &probe_null_scope_values, + hashes_buffer, + batch_size, + probe_indices_buffer, + build_indices_buffer, + |build_indices, positions| { + // The lookup ran over only the NULL-valued probe rows; + // translate its positions back to probe row indices. + let probe_indices = UInt32Array::from_iter_values( + positions + .values() + .iter() + .map(|p| null_probe_rows.value(*p as usize)), + ); + mark(build_indices, probe_indices) + }, + )?; + } + None => { + let build_rows = + UInt64Array::from_iter_values(0..left_data.batch().num_rows() as u64); + for_each_cross_product( + &build_rows, + &null_probe_rows, + batch_size, + &mut mark, + )?; + } + } } Ok(()) } -/// Scans all correlation-scope matches between `build_scope_values` and -/// `probe_scope_values` and sets the bit of every matched build row in -/// `null_bitmap`, translating matched map positions through -/// `map_position_to_build_row`. +/// Calls `f` with all correlation-scope matches between `build_scope_values` +/// and `probe_scope_values`, as chunks of at most `batch_size` pairs of +/// (position in `build_scope_values`, position in `probe_scope_values`). #[expect(clippy::too_many_arguments)] -fn scan_scope_matches_into_bitmap( +fn for_each_scope_match( scope_map: &dyn JoinHashMapType, build_scope_values: &[ArrayRef], probe_scope_values: &[ArrayRef], @@ -1449,12 +1545,11 @@ fn scan_scope_matches_into_bitmap( batch_size: usize, probe_indices_buffer: &mut Vec, build_indices_buffer: &mut Vec, - map_position_to_build_row: impl Fn(u64) -> u64, - null_bitmap: &SharedBitmapBuilder, + mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, ) -> Result<()> { let mut offset = (0, None); loop { - let (build_indices, _probe_indices, next_offset) = lookup_join_hashmap( + let (build_indices, probe_indices, next_offset) = lookup_join_hashmap( scope_map, build_scope_values, probe_scope_values, @@ -1468,13 +1563,7 @@ fn scan_scope_matches_into_bitmap( )?; if !build_indices.is_empty() { - let mut null_bitmap = null_bitmap.lock(); - - for build_idx in build_indices.iter() { - let build_idx = build_idx - .expect("scope lookup should produce non-null build indices"); - null_bitmap.set_bit(map_position_to_build_row(build_idx) as usize, true); - } + f(build_indices, probe_indices)?; } let Some(next_offset) = next_offset else { @@ -1486,6 +1575,39 @@ fn scan_scope_matches_into_bitmap( Ok(()) } +/// Calls `f` with every pair of `build_rows` x `probe_rows`, as chunks of at +/// most `batch_size` pairs. +fn for_each_cross_product( + build_rows: &UInt64Array, + probe_rows: &UInt32Array, + batch_size: usize, + mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, +) -> Result<()> { + let chunk_size = batch_size + .max(1) + .min(build_rows.len().saturating_mul(probe_rows.len())); + let mut build_chunk = Vec::with_capacity(chunk_size); + let mut probe_chunk = Vec::with_capacity(chunk_size); + for probe_row in probe_rows.values() { + for build_row in build_rows.values() { + build_chunk.push(*build_row); + probe_chunk.push(*probe_row); + if build_chunk.len() == chunk_size { + f( + std::mem::replace(&mut build_chunk, Vec::with_capacity(chunk_size)) + .into(), + std::mem::replace(&mut probe_chunk, Vec::with_capacity(chunk_size)) + .into(), + )?; + } + } + } + if !build_chunk.is_empty() { + f(build_chunk.into(), probe_chunk.into())?; + } + Ok(()) +} + impl Stream for HashJoinStream { type Item = Result; From 61ce30669e73e13dd4343ed2ce2271967c1e6df8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:29:30 -0500 Subject: [PATCH 2/5] fix: plan NOT IN mark joins with a non-equality correlation as null-aware The hash join now applies the join filter when it marks UNKNOWN rows, so a NOT IN mark join no longer needs to fall back to a non null-aware join when a non-equality correlation stays behind as a join filter. The fallback gave FALSE instead of NULL, so NOT (x IN (...)) returned extra rows. Co-Authored-By: Claude Opus 5 --- .../src/decorrelate_predicate_subquery.rs | 38 ++++--------------- .../sqllogictest/test_files/subquery.slt | 10 ++--- 2 files changed, 12 insertions(+), 36 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 0ad8b44c40def..8ea4dbd0451c6 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -21,7 +21,6 @@ use std::ops::Deref; use std::sync::Arc; use crate::decorrelate::PullUpCorrelatedExpr; -use crate::extract_equijoin_predicate::split_eq_and_noneq_join_predicate; use crate::optimizer::ApplyOrder; use crate::utils::replace_qualified_name; use crate::{OptimizerConfig, OptimizerRule}; @@ -584,25 +583,12 @@ fn build_join( sub_query_alias.clone() }; - let mark_filter_is_hashable_only = - if join_type == JoinType::LeftMark && in_predicate_opt.is_some() { - let (_, residual_filter) = split_eq_and_noneq_join_predicate( - join_filter.clone(), - left.schema(), - right_projected.schema(), - )?; - residual_filter.is_none() - } else { - false - }; - // For scalar NOT IN mark joins, propagate null-aware semantics into the - // nullable mark column when the predicate can be implemented by hash keys. - // Non-equality correlated filters stay on the legacy path because hash join - // execution cannot mark UNKNOWN candidates for residual predicates. + // nullable mark column. A non-equality correlation stays behind as a + // join filter, which the hash join also applies when it decides + // whether a NULL makes the mark UNKNOWN. let null_aware = join_type == JoinType::LeftMark && in_predicate_opt.is_some() - && mark_filter_is_hashable_only && join_keys_may_be_null( &join_filter, left.schema(), @@ -756,18 +742,6 @@ mod tests { plan.inputs().into_iter().any(has_null_aware_left_mark_join) } - fn has_non_null_aware_left_mark_join(plan: &LogicalPlan) -> bool { - if let LogicalPlan::Join(join) = plan - && join.join_type == JoinType::LeftMark - { - return !join.null_aware; - } - - plan.inputs() - .into_iter() - .any(has_non_null_aware_left_mark_join) - } - fn optimize_with_decorrelate(plan: LogicalPlan) -> Result { let optimizer = crate::Optimizer::with_rules(vec![Arc::new( DecorrelatePredicateSubquery::new(), @@ -1556,8 +1530,10 @@ mod tests { Ok(()) } + /// A non-equality correlation stays behind as a join filter, which the hash + /// join applies when it marks UNKNOWN rows, so the mark is still null-aware. #[test] - fn correlated_not_in_mark_join_is_not_null_aware_for_residual_filter() -> Result<()> { + fn correlated_not_in_mark_join_is_null_aware_for_residual_filter() -> Result<()> { let outer_scan = nullable_scalar_mark_scan("outer_t")?; let inner_scan = nullable_scalar_mark_scan("inner_t")?; @@ -1576,7 +1552,7 @@ mod tests { let optimized = optimize_with_decorrelate(plan)?; assert!( - has_non_null_aware_left_mark_join(&optimized), + has_null_aware_left_mark_join(&optimized), "{}", optimized.display_indent_schema() ); diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 626ef60762b91..0dad17ccfb687 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1235,7 +1235,7 @@ where t1.t1_id > 40 or t1.t1_id in (select t2.t2_id from t2 where t1.t1_int > 0) logical_plan 01)Projection: t1.t1_id, t1.t1_name, t1.t1_int 02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_1.mark -03)----LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) +03)----LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) null_aware 04)------TableScan: t1 projection=[t1_id, t1_name, t1_int] 05)------SubqueryAlias: __correlated_sq_1 06)--------TableScan: t2 projection=[t2_id] @@ -1262,7 +1262,7 @@ where t1.t1_id = 11 or t1.t1_id + 12 not in (select t2.t2_id + 1 from t2 where t logical_plan 01)Projection: t1.t1_id, t1.t1_name, t1.t1_int 02)--Filter: t1.t1_id = Int32(11) OR NOT __correlated_sq_1.mark -03)----LeftMark Join: CAST(t1.t1_id AS Int64) + Int64(12) = __correlated_sq_1.t2.t2_id + Int64(1) Filter: t1.t1_int > Int32(0) +03)----LeftMark Join: CAST(t1.t1_id AS Int64) + Int64(12) = __correlated_sq_1.t2.t2_id + Int64(1) Filter: t1.t1_int > Int32(0) null_aware 04)------TableScan: t1 projection=[t1_id, t1_name, t1_int] 05)------SubqueryAlias: __correlated_sq_1 06)--------Projection: CAST(t2.t2_id AS Int64) + Int64(1) @@ -1402,13 +1402,13 @@ logical_plan 01)Projection: t1.t1_name, t1.t1_id 02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_1.mark 03)----Projection: t1.t1_id, t1.t1_name, __correlated_sq_1.mark -04)------LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) +04)------LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) null_aware 05)--------TableScan: t1 projection=[t1_id, t1_name, t1_int] 06)--------SubqueryAlias: __correlated_sq_1 07)----------TableScan: t2 projection=[t2_id] physical_plan 01)FilterExec: t1_id@0 > 40 OR mark@2, projection=[t1_name@1, t1_id@0] -02)--HashJoinExec: mode=CollectLeft, join_type=RightMark, on=[(t2_id@0, t1_id@0)], filter=t1_int@0 > 0, projection=[t1_id@0, t1_name@1, mark@3] +02)--HashJoinExec: mode=CollectLeft, join_type=LeftMark, on=[(t1_id@0, t2_id@0)], filter=t1_int@0 > 0, projection=[t1_id@0, t1_name@1, mark@3], null_aware 03)----DataSourceExec: partitions=1, partition_sizes=[2] 04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 05)------DataSourceExec: partitions=1, partition_sizes=[2] @@ -1497,7 +1497,7 @@ where t1.t1_id in (select t3.t3_id from t3) and (t1.t1_id > 40 or t1.t1_id in (s logical_plan 01)Projection: t1.t1_id, t1.t1_name, t1.t1_int 02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_2.mark -03)----LeftMark Join: t1.t1_id = __correlated_sq_2.t2_id Filter: t1.t1_int > Int32(0) +03)----LeftMark Join: t1.t1_id = __correlated_sq_2.t2_id Filter: t1.t1_int > Int32(0) null_aware 04)------LeftSemi Join: t1.t1_id = __correlated_sq_1.t3_id 05)--------TableScan: t1 projection=[t1_id, t1_name, t1_int] 06)--------SubqueryAlias: __correlated_sq_1 From 9926d7557305ee56457c6fe9380c8bf6dd00479f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:29:30 -0500 Subject: [PATCH 3/5] test: correlated NOT IN with a non-equality correlation Adds sqllogictest regression tests for https://github.com/apache/datafusion/issues/25336 (expected results checked with DuckDB and PostgreSQL) and HashJoinExec unit tests for null-aware LeftAnti and LeftMark joins that have a join filter and no scope keys. Co-Authored-By: Claude Opus 5 --- .../physical-plan/src/joins/hash_join/exec.rs | 145 +++++++++++++++ .../test_files/null_aware_anti_join.slt | 165 ++++++++++++++++++ .../test_files/null_aware_mark_join.slt | 124 +++++++++++++ 3 files changed, 434 insertions(+) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 64c462507fcbf..2abbeee86a766 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -9158,6 +9158,151 @@ mod tests { Ok(()) } + /// `left.z > right.z` over the second column of two-column tables: the + /// non-equality correlation of + /// `id NOT IN (SELECT r.id FROM r WHERE r.z < l.z)`. + fn prepare_second_column_gt_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("z", DataType::Int32, true), + Field::new("z", DataType::Int32, true), + ]); + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("z", 0)), + Operator::Gt, + Arc::new(Column::new("z", 1)), + )) as Arc; + + JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ) + } + + /// Build and probe sides of a null-aware join whose only correlation is + /// the non-equality filter from [`prepare_second_column_gt_filter`]. + /// + /// For each build row, the probe rows with a smaller `z` form its + /// subquery result: + /// - `(1, 10)` and `(2, 20)`: `{1, NULL}` + /// - `(NULL, 30)`: `{1, NULL}` + /// - `(4, 40)`: `{1, 4, NULL}` + /// - `(NULL, 1)` and `(5, 1)`: empty + /// + /// The probe row `(NULL, 50)` never passes the filter. + fn build_null_aware_filter_only_inputs() + -> (Arc, Arc, JoinOn) { + let left = build_table_two_cols( + ("id", &vec![Some(1), Some(2), None, Some(4), None, Some(5)]), + ( + "z", + &vec![Some(10), Some(20), Some(30), Some(40), Some(1), Some(1)], + ), + ); + let right = build_table_two_cols( + ("id", &vec![Some(1), None, Some(4), None]), + ("z", &vec![Some(5), Some(50), Some(35), Some(2)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("id", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("id", &right.schema()).unwrap()) as _, + )]; + (left, right, on) + } + + /// Null-aware `LeftAnti` with a join filter and no correlation scope keys. + /// + /// A NULL on either side only makes `NOT IN` UNKNOWN for the build rows + /// where the filter keeps the NULL, so the NULLs must not remove every row. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_left_anti_filter_only(batch_size: usize) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + let (left, right, on) = build_null_aware_filter_only_inputs(); + + let join = HashJoinExec::try_new( + left, + right, + on, + Some(prepare_second_column_gt_filter()), + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // Only the rows with an empty subquery result are TRUE. + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+---+ + | id | z | + +----+---+ + | | 1 | + | 5 | 1 | + +----+---+ + "); + } + + Ok(()) + } + + /// Null-aware `LeftMark` with a join filter and no correlation scope keys. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_left_mark_filter_only(batch_size: usize) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + let (left, right, on) = build_null_aware_filter_only_inputs(); + + let join = HashJoinExec::try_new( + left, + right, + on, + Some(prepare_second_column_gt_filter()), + &JoinType::LeftMark, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // `(1, 10)` and `(4, 40)` match (true); `(2, 20)` and `(NULL, 30)` + // keep the NULL probe row (UNKNOWN); `(NULL, 1)` and `(5, 1)` have an + // empty subquery result (false). + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-------+ + | id | z | mark | + +----+----+-------+ + | | 1 | false | + | | 30 | | + | 1 | 10 | true | + | 2 | 20 | | + | 4 | 40 | true | + | 5 | 1 | false | + +----+----+-------+ + "); + } + + Ok(()) + } + #[test] fn test_lr_is_preserved() { assert_eq!(lr_is_preserved(JoinType::Inner), (true, true)); diff --git a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt index 8023684ac3ee0..1944c1aa988e0 100644 --- a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt @@ -666,3 +666,168 @@ DROP TABLE naconst_t1; statement ok DROP TABLE naconst_t2; + +## Correlated NOT IN with a non-equality correlation +## https://github.com/apache/datafusion/issues/25336 +############# + +# The non-equality correlation stays behind as a residual join filter. A NULL +# subquery value only makes NOT IN UNKNOWN for the outer rows where the +# residual keeps that NULL row. Expected results are verified with DuckDB and +# PostgreSQL. + +statement ok +CREATE TABLE nai_res_t1(id INT, z INT) AS VALUES (1,10), (2,20), (NULL,30), (4,40); + +statement ok +CREATE TABLE nai_res_t2(id INT, z INT) AS VALUES (1,5), (NULL,50); + +# The NULL row (z = 50) never passes `t2.z < t1.z`, so it does not affect the result. +query I +SELECT id FROM nai_res_t1 WHERE id NOT IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z < nai_res_t1.z) ORDER BY id; +---- +2 +4 + +query I +SELECT id FROM nai_res_t1 WHERE NOT (id IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z < nai_res_t1.z)) ORDER BY id; +---- +2 +4 + +query I +SELECT id FROM nai_res_t1 WHERE id IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z < nai_res_t1.z) ORDER BY id; +---- +1 + +# Uncorrelated residuals are pushed into the subquery. +query I +SELECT id FROM nai_res_t1 WHERE id NOT IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z < 40) ORDER BY id; +---- +2 +4 + +query I +SELECT id FROM nai_res_t1 WHERE id NOT IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z < 100) ORDER BY id; +---- + +# The NULL row passes `t2.z > t1.z` for every outer row, so every row is UNKNOWN. +query I +SELECT id FROM nai_res_t1 WHERE id NOT IN (SELECT nai_res_t2.id FROM nai_res_t2 WHERE nai_res_t2.z > nai_res_t1.z) ORDER BY id; +---- + +statement ok +CREATE TABLE nai_res_outer(id INT, z INT, g INT) AS VALUES +(1, 10, 1), +(2, 20, 1), +(NULL, 30, 1), +(4, 40, 2), +(NULL, 1, 2), +(5, 1, 3); + +statement ok +CREATE TABLE nai_res_inner(id INT, z INT, g INT) AS VALUES +(1, 5, 1), +(NULL, 50, 1), +(4, 35, 2), +(NULL, 2, 2); + +# Per outer row, the residual decides which inner rows (NULL or not) are in +# the subquery. A NULL outer value is TRUE only when that set is empty. +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.z < nai_res_outer.z); +---- +5 1 +NULL 1 + +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.z > nai_res_outer.z); +---- + +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.z < nai_res_outer.z AND i.z > 3); +---- +2 20 +5 1 +NULL 1 + +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id + 0 NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.z < nai_res_outer.z); +---- +5 1 +NULL 1 + +# The residual references the subquery value itself. +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.id + i.z > nai_res_outer.z); +---- +1 10 +2 20 +4 40 +5 1 + +# Equality and non-equality correlation together. +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.g = nai_res_outer.g AND i.z < nai_res_outer.z); +---- +2 20 +5 1 +NULL 1 + +# Equality correlation only. +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.g = nai_res_outer.g); +---- +5 1 + +# NOT EXISTS uses two-valued logic and must not change. +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE NOT EXISTS (SELECT 1 FROM nai_res_inner i WHERE i.id = nai_res_outer.id AND i.z < nai_res_outer.z); +---- +2 20 +5 1 +NULL 1 +NULL 30 + +# Output batches of one row split the unmatched build rows and the candidate +# pairs checked against the residual filter into many chunks. +statement ok +SET datafusion.execution.batch_size = 1; + +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.z < nai_res_outer.z); +---- +5 1 +NULL 1 + +query II rowsort +SELECT id, z FROM nai_res_outer +WHERE id NOT IN (SELECT i.id FROM nai_res_inner i WHERE i.g = nai_res_outer.g AND i.z < nai_res_outer.z); +---- +2 20 +5 1 +NULL 1 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +DROP TABLE nai_res_t1; + +statement ok +DROP TABLE nai_res_t2; + +statement ok +DROP TABLE nai_res_outer; + +statement ok +DROP TABLE nai_res_inner; diff --git a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt index dfaa4f23cb72e..27faf4aa9ba34 100644 --- a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt @@ -637,3 +637,127 @@ DROP TABLE nmconst_t1; statement ok DROP TABLE nmconst_t2; + +########################################################## +## Correlated NOT IN mark join with a non-equality correlation +## https://github.com/apache/datafusion/issues/25336 +########################################################## + +# The non-equality correlation stays behind as a residual join filter, so the +# mark must be NULL (UNKNOWN) only when the residual keeps a NULL on either +# side of the comparison. Expected results are verified with DuckDB. + +statement ok +CREATE TABLE nam_res_outer(id INT, z INT, g INT) AS VALUES +(1, 10, 1), +(2, 20, 1), +(NULL, 30, 1), +(4, 40, 2), +(NULL, 1, 2), +(5, 1, 3); + +statement ok +CREATE TABLE nam_res_inner(id INT, z INT, g INT) AS VALUES +(1, 5, 1), +(NULL, 50, 1), +(4, 35, 2), +(NULL, 2, 2); + +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z)) IS NULL; +---- +2 20 +NULL 30 + +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z)) IS TRUE; +---- +5 1 +NULL 1 + +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z)) IS FALSE; +---- +1 10 +4 40 + +# `NOT mark` must stay NULL for UNKNOWN rows instead of turning into TRUE. +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE NOT (id IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z)) OR id = 4; +---- +4 40 +5 1 +NULL 1 + +# Positive IN goes through the same mark. +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id IN (SELECT i.id FROM nam_res_inner i WHERE i.z > nam_res_outer.z)) IS NULL; +---- +1 10 +2 20 +4 40 +5 1 +NULL 1 +NULL 30 + +# Equality and non-equality correlation together. +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.g = nam_res_outer.g AND i.z < nam_res_outer.z)) IS NULL; +---- +NULL 30 + +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.g = nam_res_outer.g AND i.z < nam_res_outer.z)) IS TRUE; +---- +2 20 +5 1 +NULL 1 + +# Equality correlation only (already null-aware; control). +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.g = nam_res_outer.g)) IS NULL; +---- +2 20 +NULL 1 +NULL 30 + +# In a SELECT list, the mark column shows TRUE, FALSE and NULL directly. +query IIB rowsort +SELECT id, z, id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z) +FROM nam_res_outer; +---- +1 10 false +2 20 NULL +4 40 false +5 1 true +NULL 1 true +NULL 30 NULL + +# Output batches of one row split the candidate pairs checked against the +# residual filter into many chunks. +statement ok +SET datafusion.execution.batch_size = 1; + +query II rowsort +SELECT id, z FROM nam_res_outer +WHERE (id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < nam_res_outer.z)) IS NULL; +---- +2 20 +NULL 30 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +DROP TABLE nam_res_outer; + +statement ok +DROP TABLE nam_res_inner; From f6d74b48949134bf5f3f0815b60621c5ca520ea4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:06:41 -0500 Subject: [PATCH 4/5] bench: add correctness asserts to null_aware_join Q05-Q08, pin Q08 null_aware These asserts fail on main (apache/datafusion#25336) and pass with the fix. Q08 is a null-aware mark join once the fix lands. Co-Authored-By: Claude Opus 5 --- .../null_aware_join/benchmarks/q05.benchmark | 14 +++++++++++ .../null_aware_join/benchmarks/q06.benchmark | 14 +++++++++++ .../null_aware_join/benchmarks/q07.benchmark | 16 +++++++++++++ .../null_aware_join/benchmarks/q08.benchmark | 23 +++++++++++++++++++ 4 files changed, 67 insertions(+) diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark index 2d0a04077d5cf..dd42f19978620 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark @@ -3,6 +3,20 @@ group null_aware_join load sql_benchmarks/null_aware_join/init/load.sql +# Correctness canary: the NOT IN result must match a reference count +# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS. +# As Q04, with NULL outer keys excluded unless the subquery is empty. +assert I +SELECT count(*) = ( + SELECT count(*) FROM small_outer o + WHERE o.z <= (SELECT min(z) FROM small_inner) + OR (o.id_n1 IS NOT NULL AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z)) +) +FROM small_outer o +WHERE o.id_n1 NOT IN (SELECT i.id_n0 FROM small_inner i WHERE i.z < o.z); +---- +true + expect_plan HashJoinExec expect_plan null_aware: true diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark index 668476ded9571..ae87f310222c0 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark @@ -3,6 +3,20 @@ group null_aware_join load sql_benchmarks/null_aware_join/init/load.sql +# Correctness canary: the NOT IN result must match a reference count +# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS. +# As Q04, with NULL outer keys excluded unless the subquery is empty. +assert I +SELECT count(*) = ( + SELECT count(*) FROM small_outer o + WHERE o.z <= (SELECT min(z) FROM small_inner) + OR (o.id_n50 IS NOT NULL AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z)) +) +FROM small_outer o +WHERE o.id_n50 NOT IN (SELECT i.id_n0 FROM small_inner i WHERE i.z < o.z); +---- +true + expect_plan HashJoinExec expect_plan null_aware: true diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark index 38bf7e9048bc9..e210ca89c58f6 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark @@ -3,6 +3,22 @@ group null_aware_join load sql_benchmarks/null_aware_join/init/load.sql +# Correctness canary: the NOT IN result must match a reference count +# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS. +# A row is TRUE when the subquery is empty, or when the subquery holds no NULL +# and the key is not in it. A NULL is in scope when its z is below o.z. +assert I +SELECT count(*) = ( + SELECT count(*) FROM small_outer o + WHERE o.z <= (SELECT min(z) FROM small_inner) + OR (o.z <= (SELECT min(z) FROM small_inner WHERE id_n50 IS NULL) + AND NOT (o.id % 2 = 0 AND (o.id / 2) % 1000 < o.z)) +) +FROM small_outer o +WHERE o.id_n0 NOT IN (SELECT i.id_n50 FROM small_inner i WHERE i.z < o.z); +---- +true + expect_plan HashJoinExec expect_plan null_aware: true diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q08.benchmark index 39d4fbd4ced22..3a770762b6a08 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q08.benchmark +++ b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q08.benchmark @@ -3,7 +3,30 @@ group null_aware_join load sql_benchmarks/null_aware_join/init/load.sql +# Correctness canary: the NOT IN result must match a reference count +# that does not use NOT IN. It holds for every NAJ_ROWS / NAJ_LARGE_ROWS. +# A row is TRUE when o.z > 900, when the subquery for its k is empty, or when +# its key is not NULL and not in that subquery. +assert I +SELECT count(*) = ( + SELECT count(*) + FROM small_outer o + JOIN (SELECT k, min(z) AS min_z FROM small_inner GROUP BY k) m ON m.k = o.k + WHERE o.z > 900 + OR o.z <= m.min_z + OR (o.id_n50 IS NOT NULL + AND NOT (o.id % 2 = 0 AND (o.id / 2) % 16 = o.k AND (o.id / 2) % 1000 < o.z)) +) +FROM small_outer o +WHERE o.z > 900 + OR o.id_n50 NOT IN ( + SELECT i.id_n0 FROM small_inner i WHERE i.k = o.k AND i.z < o.z + ); +---- +true + expect_plan HashJoinExec +expect_plan null_aware: true run -- Q8: NOT IN correlated by both an equality and a non-equality, 50% NULL on From 3b38f6be6c082fb437d0cf15cd8306e6344d3ae4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:06:45 -0500 Subject: [PATCH 5/5] perf: skip build rows already marked UNKNOWN in correlated null-aware joins A build row stays UNKNOWN once it is marked. The candidate pairing now removes pairs whose build row is already marked before it evaluates the join filter. Without correlation scope keys, the pairing also drops the marked build rows after each chunk of pairs, and it stops when no unmarked build row is left. For `NAJ_ROWS=10000` of the `null_aware_join` benchmark, this takes Q06 from 20.6M candidate pairs to 240K, and Q07 from 20.6M to 169K. A consequence is that the join filter is evaluated for fewer pairs, so a filter that gives an error only for a skipped pair no longer gives that error. This is the same as the short-circuit behavior of `AND`. Co-Authored-By: Claude Opus 5 --- .../src/joins/hash_join/stream.rs | 101 +++++++++++++----- 1 file changed, 75 insertions(+), 26 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index de6de146cddd1..e5e8e8da5b8f3 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -1376,6 +1376,11 @@ fn null_aware_left_mark_column( /// all build rows with the NULL-valued probe rows. Scope keys narrow these /// pairs through a hash lookup; without scope keys every pair is a candidate. /// The join filter, if any, then decides 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. #[expect(clippy::too_many_arguments)] fn mark_null_candidates_for_probe_batch( build_side: &BuildSideReadyState, @@ -1407,6 +1412,11 @@ fn mark_null_candidates_for_probe_batch( // Keeps the candidate pairs that pass the join filter and marks their // build 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); + if build_indices.is_empty() { + return Ok(()); + } let build_indices = match filter { Some(filter) => { apply_join_filter_to_indices( @@ -1463,11 +1473,10 @@ fn mark_null_candidates_for_probe_batch( )?; } None => { - let probe_rows = - UInt32Array::from_iter_values(0..state.batch.num_rows() as u32); - for_each_cross_product( - &null_rows.build_indices, - &probe_rows, + for_each_unmarked_cross_product( + left_data, + null_rows.build_indices.values().iter().copied(), + 0..state.batch.num_rows() as u32, batch_size, &mut mark, )?; @@ -1518,11 +1527,10 @@ fn mark_null_candidates_for_probe_batch( )?; } None => { - let build_rows = - UInt64Array::from_iter_values(0..left_data.batch().num_rows() as u64); - for_each_cross_product( - &build_rows, - &null_probe_rows, + for_each_unmarked_cross_product( + left_data, + 0..left_data.batch().num_rows() as u64, + null_probe_rows.values().iter().copied(), batch_size, &mut mark, )?; @@ -1575,30 +1583,71 @@ fn for_each_scope_match( Ok(()) } -/// Calls `f` with every pair of `build_rows` x `probe_rows`, as chunks of at -/// most `batch_size` pairs. -fn for_each_cross_product( - build_rows: &UInt64Array, - probe_rows: &UInt32Array, +/// Removes the candidate pairs whose build row is already marked UNKNOWN. +fn retain_unmarked( + left_data: &JoinLeftData, + 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) { + 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)) + .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. +fn for_each_unmarked_cross_product( + left_data: &JoinLeftData, + build_rows: impl Iterator, + probe_rows: impl Iterator, batch_size: usize, mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, ) -> Result<()> { - let chunk_size = batch_size - .max(1) - .min(build_rows.len().saturating_mul(probe_rows.len())); - let mut build_chunk = Vec::with_capacity(chunk_size); - let mut probe_chunk = Vec::with_capacity(chunk_size); - for probe_row in probe_rows.values() { - for build_row in build_rows.values() { + let retain_unmarked_rows = |rows: &mut Vec| { + let bitmap = left_data.null_indices_bitmap().lock(); + rows.retain(|idx| !bitmap.get_bit(*idx as usize)); + }; + + let mut build_rows: Vec = build_rows.collect(); + retain_unmarked_rows(&mut build_rows); + + let batch_size = batch_size.max(1); + let mut build_chunk = Vec::with_capacity(batch_size); + let mut probe_chunk = Vec::with_capacity(batch_size); + let mut marks_since_refresh = false; + for probe_row in probe_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); + marks_since_refresh = false; + } + if build_rows.is_empty() { + break; + } + for build_row in &build_rows { build_chunk.push(*build_row); - probe_chunk.push(*probe_row); - if build_chunk.len() == chunk_size { + probe_chunk.push(probe_row); + if build_chunk.len() == batch_size { f( - std::mem::replace(&mut build_chunk, Vec::with_capacity(chunk_size)) + std::mem::replace(&mut build_chunk, Vec::with_capacity(batch_size)) .into(), - std::mem::replace(&mut probe_chunk, Vec::with_capacity(chunk_size)) + std::mem::replace(&mut probe_chunk, Vec::with_capacity(batch_size)) .into(), )?; + marks_since_refresh = true; } } }