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/7] 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/7] 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/7] 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/7] 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/7] 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; } } } From 283b8cabcfdef80b7e71182ba01cc4dd258ac81c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:52:51 -0500 Subject: [PATCH 6/7] fix: keep a positive `IN` mark join non-null-aware under `AND`/`OR` The previous commits make a correlated `IN` mark join null-aware whenever a join key can be NULL. That is necessary only when a NULL mark can give a different answer than a FALSE mark. A `Filter` keeps a row only when its predicate is TRUE, and `AND` and `OR` give TRUE only when an operand is TRUE. A non-negated `IN` or `EXISTS` that a `WHERE` conjunct reaches only through `AND`/`OR` thus keeps the same rows with a FALSE mark as with a NULL mark. There the null-aware join only costs more: it pins the outer table as the `CollectLeft` build side, it cannot be swapped, and it evaluates the join filter for every pair of a build row and a NULL probe row. `subqueries_only_positive` finds that shape. The `Filter` path passes the result down to `build_join`, and every other caller asks for the null-aware mark. `NOT IN`, `NOT (x IN ...)`, `(x IN ...) IS NULL`, `CASE` and a mark that goes into a projection thus stay null-aware, and the fix of this pull request is not affected. Three of the four `subquery.slt` plans that the previous commits changed go back to their earlier form. The fourth is a `NOT IN` and keeps `null_aware`. Co-Authored-By: Claude Opus 5 --- .../src/decorrelate_predicate_subquery.rs | 162 ++++++++++++++++-- .../test_files/null_aware_mark_join.slt | 45 +++++ .../sqllogictest/test_files/subquery.slt | 8 +- 3 files changed, 199 insertions(+), 16 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 8ea4dbd0451c6..90a32e8e79d43 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -80,7 +80,7 @@ impl OptimizerRule for DecorrelatePredicateSubquery { for expr in projection.expr { let original_name = expr.schema_name().to_string(); let (new_input, mut rewritten_expr) = - rewrite_inner_subqueries(cur_input, expr, config, true)?; + rewrite_inner_subqueries(cur_input, expr, config, true, true)?; if has_subquery(&rewritten_expr) { return Ok(Transformed::no(LogicalPlan::Projection( original_projection, @@ -134,8 +134,18 @@ impl OptimizerRule for DecorrelatePredicateSubquery { } // The subquery expression is embedded within another expression SubqueryPredicate::Embedded(expr) => { - let (plan, expr_without_subqueries) = - rewrite_inner_subqueries(cur_input, expr, config, false)?; + // A `Filter` keeps a row only when the predicate is TRUE, and + // `AND`/`OR` make TRUE only out of TRUE. A mark that is NULL + // thus acts exactly like a mark that is FALSE here, and the + // subquery does not need the more expensive null-aware join. + let needs_null_aware_mark = !subqueries_only_positive(&expr); + let (plan, expr_without_subqueries) = rewrite_inner_subqueries( + cur_input, + expr, + config, + false, + needs_null_aware_mark, + )?; cur_input = plan; other_exprs.push(expr_without_subqueries); } @@ -166,11 +176,14 @@ impl OptimizerRule for DecorrelatePredicateSubquery { } } +/// `needs_null_aware_mark` is `false` only when the caller can prove that a NULL +/// mark and a FALSE mark give the same answer. See [`subqueries_only_positive`]. fn rewrite_inner_subqueries( outer: LogicalPlan, expr: Expr, config: &dyn OptimizerConfig, materialize_in_value: bool, + needs_null_aware_mark: bool, ) -> Result<(LogicalPlan, Expr)> { let mut cur_input = outer; let alias = config.alias_generator(); @@ -178,7 +191,14 @@ fn rewrite_inner_subqueries( Expr::Exists(Exists { subquery: Subquery { subquery, .. }, negated, - }) => match mark_join(&cur_input, &subquery, None, negated, alias)? { + }) => match mark_join( + &cur_input, + &subquery, + None, + negated, + alias, + needs_null_aware_mark, + )? { Some((plan, exists_expr)) => { cur_input = plan; Ok(Transformed::yes(exists_expr)) @@ -205,7 +225,14 @@ fn rewrite_inner_subqueries( .map_or(plan_err!("single expression required."), |output_expr| { Ok(Expr::eq(*expr.clone(), output_expr)) })?; - mark_join(&cur_input, &subquery, Some(&in_predicate), negated, alias)? + mark_join( + &cur_input, + &subquery, + Some(&in_predicate), + negated, + alias, + needs_null_aware_mark, + )? }; match rewritten { Some((plan, exists_expr)) => { @@ -233,7 +260,7 @@ fn in_subquery_value_mark_join( .map_or(plan_err!("single expression required."), Ok)?; let in_predicate = Expr::eq(expr.clone(), output_expr.clone()); let Some((matched_plan, matched)) = - mark_join(left, subquery, Some(&in_predicate), false, alias)? + mark_join(left, subquery, Some(&in_predicate), false, alias, true)? else { return Ok(None); }; @@ -243,12 +270,12 @@ fn in_subquery_value_mark_join( .filter(output_expr.is_null())? .build()?; let Some((null_plan, subquery_has_null)) = - mark_join(&matched_plan, &null_subquery, None, false, alias)? + mark_join(&matched_plan, &null_subquery, None, false, alias, true)? else { return Ok(None); }; let Some((final_plan, subquery_non_empty)) = - mark_join(&null_plan, subquery, None, false, alias)? + mark_join(&null_plan, subquery, None, false, alias, true)? else { return Ok(None); }; @@ -309,6 +336,26 @@ fn has_subquery(expr: &Expr) -> bool { .unwrap() } +/// True when every subquery in `expr` is a non-negated `IN`/`EXISTS` reached only +/// through `AND`/`OR`. +/// +/// `AND` and `OR` give TRUE only when an operand is TRUE, so a `Filter` on such an +/// expression keeps the same rows whether a mark is NULL or FALSE. The mark join +/// then does not need to be null-aware. `NOT`, `IS NULL`, `CASE` and a mark that +/// goes into a projection can tell NULL from FALSE, so they give `false` here. +fn subqueries_only_positive(expr: &Expr) -> bool { + match expr { + Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::And | Operator::Or, + right, + }) => subqueries_only_positive(left) && subqueries_only_positive(right), + Expr::InSubquery(InSubquery { negated, .. }) => !negated, + Expr::Exists(Exists { negated, .. }) => !negated, + other => !has_subquery(other), + } +} + /// Optimize the subquery to left-anti/left-semi join. /// If the subquery is a correlated subquery, we need extract the join predicate from the subquery. /// @@ -370,6 +417,7 @@ fn build_join_top( in_predicate_opt.as_ref(), join_type, subquery_alias, + true, ) } @@ -394,16 +442,22 @@ fn mark_join( in_predicate_opt: Option<&Expr>, negated: bool, alias_generator: &Arc, + needs_null_aware_mark: bool, ) -> Result> { let alias = alias_generator.next("__correlated_sq"); let exists_col = Expr::Column(Column::new(Some(alias.clone()), "mark")); let exists_expr = if negated { !exists_col } else { exists_col }; - Ok( - build_join(left, subquery, in_predicate_opt, JoinType::LeftMark, alias)? - .map(|plan| (plan, exists_expr)), - ) + Ok(build_join( + left, + subquery, + in_predicate_opt, + JoinType::LeftMark, + alias, + needs_null_aware_mark, + )? + .map(|plan| (plan, exists_expr))) } /// Check if join keys in the join filter may contain NULL values @@ -445,6 +499,7 @@ fn build_join( in_predicate_opt: Option<&Expr>, join_type: JoinType, alias: String, + needs_null_aware_mark: bool, ) -> Result> { let mut pull_up = PullUpCorrelatedExpr::new() .with_in_predicate_opt(in_predicate_opt.cloned()) @@ -589,6 +644,7 @@ fn build_join( // whether a NULL makes the mark UNKNOWN. let null_aware = join_type == JoinType::LeftMark && in_predicate_opt.is_some() + && needs_null_aware_mark && join_keys_may_be_null( &join_filter, left.schema(), @@ -742,6 +798,18 @@ 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(), @@ -1560,6 +1628,76 @@ mod tests { Ok(()) } + /// A `Filter` drops a row whose predicate is NULL and a row whose predicate is + /// FALSE, and `OR` gives TRUE only when an operand is TRUE. A non-negated `IN` + /// under `OR` thus gives the same rows without the null-aware join, which costs + /// more because it pins the build side and cannot be swapped. + #[test] + fn correlated_in_mark_join_under_or_is_not_null_aware() -> Result<()> { + let outer_scan = nullable_scalar_mark_scan("outer_t")?; + let inner_scan = nullable_scalar_mark_scan("inner_t")?; + + let subquery = Arc::new( + LogicalPlanBuilder::from(inner_scan) + .filter( + out_ref_col(DataType::Int32, "outer_t.grp").lt(col("inner_t.grp")), + )? + .project(vec![col("inner_t.id")])? + .build()?, + ); + + let plan = LogicalPlanBuilder::from(outer_scan) + .filter( + col("outer_t.grp") + .gt(lit(0i32)) + .or(in_subquery(col("outer_t.id"), subquery)), + )? + .build()?; + + let optimized = optimize_with_decorrelate(plan)?; + assert!( + has_non_null_aware_left_mark_join(&optimized), + "{}", + optimized.display_indent_schema() + ); + + Ok(()) + } + + /// The negated form of [`correlated_in_mark_join_under_or_is_not_null_aware`]: + /// `NOT mark` tells NULL from FALSE, so this one stays null-aware. + #[test] + fn correlated_not_in_mark_join_under_or_is_null_aware() -> Result<()> { + let outer_scan = nullable_scalar_mark_scan("outer_t")?; + let inner_scan = nullable_scalar_mark_scan("inner_t")?; + + let subquery = Arc::new( + LogicalPlanBuilder::from(inner_scan) + .filter( + out_ref_col(DataType::Int32, "outer_t.grp").lt(col("inner_t.grp")), + )? + .project(vec![col("inner_t.id")])? + .build()?, + ); + + let plan = LogicalPlanBuilder::from(outer_scan) + .filter( + col("outer_t.grp") + .gt(lit(0i32)) + .or(not_in_subquery(col("outer_t.id"), subquery)), + )? + .build()?; + + let optimized = optimize_with_decorrelate(plan)?; + assert!( + has_null_aware_left_mark_join(&optimized), + "{}", + optimized.display_indent_schema() + ); + + Ok(()) + } + #[test] fn in_subquery_both_side_expr() -> Result<()> { let table_scan = test_table_scan()?; diff --git a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt index 27faf4aa9ba34..769369eef862a 100644 --- a/datafusion/sqllogictest/test_files/null_aware_mark_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_mark_join.slt @@ -741,6 +741,51 @@ FROM nam_res_outer; NULL 1 true NULL 30 NULL +# A non-negated `IN` under `OR` does not need the null-aware join. A `Filter` +# drops the row whether the mark is NULL or FALSE, and the null-aware join costs +# more, because it pins the outer table as the build side and cannot be swapped. +query TT +EXPLAIN SELECT id FROM nam_res_outer o +WHERE o.z > 35 OR o.id IN (SELECT i.id FROM nam_res_inner i WHERE i.z < o.z); +---- +logical_plan +01)Projection: o.id +02)--Filter: o.z > Int32(35) OR __correlated_sq_1.mark +03)----LeftMark Join: o.id = __correlated_sq_1.id Filter: __correlated_sq_1.z < o.z +04)------SubqueryAlias: o +05)--------TableScan: nam_res_outer projection=[id, z] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: i +08)----------TableScan: nam_res_inner projection=[id, z] +physical_plan +01)FilterExec: z@1 > 35 OR mark@2, projection=[id@0] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----HashJoinExec: mode=CollectLeft, join_type=RightMark, on=[(id@0, id@0)], filter=z@1 < z@0 +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] + +# The negated form of the same query stays null-aware, because `NOT mark` tells +# NULL from FALSE. +query TT +EXPLAIN SELECT id FROM nam_res_outer o +WHERE o.z > 35 OR o.id NOT IN (SELECT i.id FROM nam_res_inner i WHERE i.z < o.z); +---- +logical_plan +01)Projection: o.id +02)--Filter: o.z > Int32(35) OR NOT __correlated_sq_1.mark +03)----LeftMark Join: o.id = __correlated_sq_1.id Filter: __correlated_sq_1.z < o.z null_aware +04)------SubqueryAlias: o +05)--------TableScan: nam_res_outer projection=[id, z] +06)------SubqueryAlias: __correlated_sq_1 +07)--------SubqueryAlias: i +08)----------TableScan: nam_res_inner projection=[id, z] +physical_plan +01)FilterExec: z@1 > 35 OR NOT mark@2, projection=[id@0] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +03)----HashJoinExec: mode=CollectLeft, join_type=LeftMark, on=[(id@0, id@0)], filter=z@1 < z@0, null_aware +04)------DataSourceExec: partitions=1, partition_sizes=[1] +05)------DataSourceExec: partitions=1, partition_sizes=[1] + # Output batches of one row split the candidate pairs checked against the # residual filter into many chunks. statement ok diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 0dad17ccfb687..c208d7ab1976f 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) null_aware +03)----LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) 04)------TableScan: t1 projection=[t1_id, t1_name, t1_int] 05)------SubqueryAlias: __correlated_sq_1 06)--------TableScan: t2 projection=[t2_id] @@ -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) null_aware +04)------LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) 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=LeftMark, on=[(t1_id@0, t2_id@0)], filter=t1_int@0 > 0, projection=[t1_id@0, t1_name@1, mark@3], null_aware +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] 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) null_aware +03)----LeftMark Join: t1.t1_id = __correlated_sq_2.t2_id Filter: t1.t1_int > Int32(0) 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 78b49ba214af90b329685db1fe9ab1ed42231bb0 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco Date: Sat, 19 Sep 2026 13:39:09 -0500 Subject: [PATCH 7/7] fix: keep the `IN` equality as `on[0]` for a constant correlated `NOT IN` A null-aware hash join reads `on[0]` as the `NOT IN` value key and `on[1..]` as correlation scope keys. Decorrelation projected a constant value expression as an outer column only for an uncorrelated subquery, so a correlated one kept ` = __sq.col` in the join filter and gave `on[0]` to the correlation. The join then applied the value-key NULL rules to the correlation key and returned wrong rows. CREATE TABLE t3(z INT) AS VALUES (10), (NULL), (7); CREATE TABLE t4(id INT, z INT) AS VALUES (NULL, 10), (1, 99), (2, 7); SELECT z FROM t3 WHERE 1 NOT IN (SELECT t4.id FROM t4 WHERE t4.z = t3.z) ORDER BY z; returned `7, 10`; DuckDB and the SQL standard give `7, NULL`. Project the constant for a correlated subquery too, and keep the `IN` equality as the leading conjunct so that it stays `on[0]`. This also removes the planning gap for a constant value with a non-equality correlation, which had no equi-join key at all. The projection is now taken only when the join really ends up null-aware, so a positive `IN` mark join keeps pushing the equality into the subquery. Co-Authored-By: Claude Opus 5 --- .../src/decorrelate_predicate_subquery.rs | 82 +++++++++++-------- .../test_files/null_aware_anti_join.slt | 66 +++++++++++++-- 2 files changed, 105 insertions(+), 43 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 90a32e8e79d43..c288538d802f8 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -525,29 +525,15 @@ fn build_join( replace_qualified_name(filter, &all_correlated_cols, &alias).map(Some) })?; - // The outer value expression of an `IN`/`NOT IN` predicate whose join filter - // is nothing but that predicate, recorded together with the subquery column - // it is compared against and a name for the column it can be projected as. - // Correlated subqueries are excluded on purpose: their correlation predicate - // is a second join key, and null-aware hash joins accept only a single key. + // The outer value expression of an `IN`/`NOT IN` predicate, recorded together + // with the subquery column it is compared against, a name for the column it + // can be projected as, and the correlation predicates that share the join + // filter with it. let mut in_value_expr = None; let mut join_filter = match (join_filter_opt, in_predicate_opt.cloned()) { ( - Some(join_filter), - Some(Expr::BinaryExpr(BinaryExpr { - left, - op: Operator::Eq, - right, - })), - ) => { - let right_col = create_col_from_scalar_expr(&right, alias)?; - let in_predicate = Expr::eq(left.deref().clone(), Expr::Column(right_col)); - in_predicate.and(join_filter) - } - (Some(join_filter), _) => join_filter, - ( - _, + correlation_opt, Some(Expr::BinaryExpr(BinaryExpr { left, op: Operator::Eq, @@ -557,10 +543,19 @@ fn build_join( let value_name = format!("{alias}_value"); let right_col = create_col_from_scalar_expr(&right, alias)?; let value = left.deref().clone(); - in_value_expr = Some((value.clone(), right_col.clone(), value_name)); + let in_predicate = Expr::eq(value.clone(), Expr::Column(right_col.clone())); + // The `IN` equality is the leading conjunct so that it becomes + // `on[0]`, the key position a null-aware hash join reads as the + // `NOT IN` value key (see `HashJoinExec::null_aware`). + let join_filter = match &correlation_opt { + Some(correlation) => in_predicate.and(correlation.clone()), + None => in_predicate, + }; + in_value_expr = Some((value, right_col, value_name, correlation_opt)); - Expr::eq(value, Expr::Column(right_col)) + join_filter } + (Some(join_filter), _) => join_filter, (None, None) => lit(true), _ => return Ok(None), }; @@ -572,12 +567,21 @@ fn build_join( // right-only, so `push_down_filter` moves it into the subquery and drops the // very NULLs that make `NOT IN` UNKNOWN, and a join without equi-join keys // is planned as a nested loop join, which has no null-aware implementation. + // A correlated subquery hits a third problem: its correlation predicate is a + // valid equi-join key, so it takes `on[0]`, the position the null-aware hash + // join reads as the `NOT IN` value key. // Projecting the constant as a column of the outer side turns the predicate - // into a real equi-join key so the null-aware hash join handles it. + // into a real equi-join key, which fixes all three. + // The projection only pays for itself on a join that ends up null-aware, so + // its guard repeats the `null_aware` conditions below. let mut projected_left = None; - if let Some((value, right_col, mut value_name)) = in_value_expr + if let Some((value, right_col, mut value_name, correlation_opt)) = in_value_expr && value.column_refs().is_empty() - && matches!(join_type, JoinType::LeftAnti | JoinType::LeftMark) + && match join_type { + JoinType::LeftAnti => true, + JoinType::LeftMark => needs_null_aware_mark, + _ => false, + } && join_keys_may_be_null(&join_filter, left.schema(), sub_query_alias.schema())? { // The projected column is unqualified, so a left field that already has @@ -598,9 +602,13 @@ fn build_join( .project(projections)? .build()?, ); - // `in_value_expr` is only set when the `IN` equality is the whole join - // filter, so it can simply be rebuilt against the projected column. - join_filter = Expr::eq(Expr::Column(value_col), Expr::Column(right_col)); + // Rebuild the `IN` equality against the projected column, keeping it + // ahead of the correlation predicates so that it stays `on[0]`. + let in_predicate = Expr::eq(Expr::Column(value_col), Expr::Column(right_col)); + join_filter = match correlation_opt { + Some(correlation) => in_predicate.and(correlation), + None => in_predicate, + }; } let left = projected_left.as_ref().unwrap_or(left); @@ -1537,11 +1545,11 @@ mod tests { ) } - /// The same rewrite must not fire for a correlated subquery: the - /// correlation predicate is a second equi-join key, and null-aware hash - /// joins accept only one. + /// The same rewrite fires for a correlated subquery, where it also keeps the + /// `IN` equality ahead of the correlation predicate so that the equality + /// stays `on[0]`, the null-aware hash join's value key. #[test] - fn constant_not_in_correlated_subquery_is_not_rewritten() -> Result<()> { + fn constant_not_in_correlated_subquery_is_rewritten() -> Result<()> { let outer_scan = nullable_scalar_mark_scan("outer_t")?; let inner_scan = nullable_scalar_mark_scan("inner_t")?; @@ -1561,11 +1569,13 @@ mod tests { assert_optimized_plan_equal!( plan, @r" - LeftAnti Join: Filter: Int32(3) = __correlated_sq_1.id AND outer_t.grp = __correlated_sq_1.grp null_aware [id:Int32;N, grp:Int32;N] - TableScan: outer_t [id:Int32;N, grp:Int32;N] - SubqueryAlias: __correlated_sq_1 [id:Int32;N, grp:Int32;N] - Projection: inner_t.id, inner_t.grp [id:Int32;N, grp:Int32;N] - TableScan: inner_t [id:Int32;N, grp:Int32;N] + Projection: outer_t.id, outer_t.grp [id:Int32;N, grp:Int32;N] + LeftAnti Join: Filter: __correlated_sq_1_value = __correlated_sq_1.id AND outer_t.grp = __correlated_sq_1.grp null_aware [id:Int32;N, grp:Int32;N, __correlated_sq_1_value:Int32] + Projection: outer_t.id, outer_t.grp, Int32(3) AS __correlated_sq_1_value [id:Int32;N, grp:Int32;N, __correlated_sq_1_value:Int32] + TableScan: outer_t [id:Int32;N, grp:Int32;N] + SubqueryAlias: __correlated_sq_1 [id:Int32;N, grp:Int32;N] + Projection: inner_t.id, inner_t.grp [id:Int32;N, grp:Int32;N] + TableScan: inner_t [id:Int32;N, grp:Int32;N] " ) } diff --git a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt index 1944c1aa988e0..e781aa0ffe96b 100644 --- a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt @@ -641,19 +641,40 @@ ORDER BY 1; statement ok DROP TABLE naconst_clash; -# A constant value expression with a non-equality correlation leaves the -# null-aware join without any equi-join key. Only `HashJoinExec` implements -# null-aware semantics and it needs a key, so the planner reports the gap -# instead of falling back to a nested loop join that ignores the NULLs and -# silently returns wrong results. +# The rewrite also fires for a correlated subquery. The projected value column +# is the first equi-join key, which is the position the null-aware hash join +# reads as the `NOT IN` value key; the correlation follows it. statement ok CREATE TABLE naconst_corr_t1(id INT, g INT) AS VALUES (1, 1), (2, 2); statement ok CREATE TABLE naconst_corr_t2(id INT, g INT) AS VALUES (1, 1), (NULL, 2); -query error DataFusion error: Error during planning: null_aware LeftAnti join requires equi\-join keys, but the join has none -SELECT id FROM naconst_corr_t1 WHERE 3 NOT IN (SELECT id FROM naconst_corr_t2 WHERE naconst_corr_t2.g > naconst_corr_t1.g); +# A non-equality correlation stays in the join filter. Row (1, 1) sees the NULL +# of (NULL, 2), so its `NOT IN` is UNKNOWN; row (2, 2) sees an empty subquery, +# so its `NOT IN` is TRUE. Expected results verified with DuckDB. +query I +SELECT id FROM naconst_corr_t1 WHERE 3 NOT IN (SELECT id FROM naconst_corr_t2 WHERE naconst_corr_t2.g > naconst_corr_t1.g) ORDER BY id; +---- +2 + +query I +SELECT id FROM naconst_corr_t1 WHERE NOT (3 IN (SELECT id FROM naconst_corr_t2 WHERE naconst_corr_t2.g > naconst_corr_t1.g)) ORDER BY id; +---- +2 + +# An equality correlation becomes a second equi-join key. Without the rewrite +# the correlation would take the first key position and the join would apply +# the value-key NULL rules to it. +query II +SELECT id, g FROM naconst_corr_t1 WHERE 3 NOT IN (SELECT id FROM naconst_corr_t2 WHERE naconst_corr_t2.g = naconst_corr_t1.g) ORDER BY id; +---- +1 1 + +# With `1` as the value, row (1, 1) is FALSE rather than TRUE, so no row survives. +query II +SELECT id, g FROM naconst_corr_t1 WHERE 1 NOT IN (SELECT id FROM naconst_corr_t2 WHERE naconst_corr_t2.g = naconst_corr_t1.g) ORDER BY id; +---- statement ok DROP TABLE naconst_corr_t1; @@ -661,6 +682,37 @@ DROP TABLE naconst_corr_t1; statement ok DROP TABLE naconst_corr_t2; +# https://github.com/apache/datafusion/pull/25339#issuecomment-5738402844 +# The outer row with a NULL correlation value matches no subquery row at all, +# so its `NOT IN` is TRUE. Expected results verified with DuckDB. +statement ok +CREATE TABLE naconst_corr_t3(z INT) AS VALUES (10), (NULL), (7); + +statement ok +CREATE TABLE naconst_corr_t4(id INT, z INT) AS VALUES (NULL, 10), (1, 99), (2, 7); + +query I +SELECT z FROM naconst_corr_t3 WHERE 1 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z = naconst_corr_t3.z) ORDER BY z; +---- +7 +NULL + +query I +SELECT z FROM naconst_corr_t3 WHERE 2 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z = naconst_corr_t3.z) ORDER BY z; +---- +NULL + +query I +SELECT z FROM naconst_corr_t3 WHERE 1 NOT IN (SELECT naconst_corr_t4.id FROM naconst_corr_t4 WHERE naconst_corr_t4.z > naconst_corr_t3.z) ORDER BY z; +---- +NULL + +statement ok +DROP TABLE naconst_corr_t3; + +statement ok +DROP TABLE naconst_corr_t4; + statement ok DROP TABLE naconst_t1;