diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 3b2241b87bae2..006e94c9b8c71 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -481,7 +481,7 @@ async fn test_join_with_swap_semi() { /// `PartitionMode::Partitioned` is deliberately not a case here: a null-aware /// `LeftAnti` cannot be constructed in that mode at all, so there is no input /// plan to optimize. See -/// `HashJoinExec::null_aware_mode`/`NullAwareMode::try_new`. +/// `HashJoinExec::not_in`/`NotIn::try_new`. #[rstest] #[case::collect_left(PartitionMode::CollectLeft, 1, 1, false)] #[case::auto(PartitionMode::Auto, 1, 1, false)] diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 3559e4497700c..cf4598441ca36 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -320,7 +320,7 @@ fn statistical_join_selection_subrule( && should_swap_join_order(left, right, context)? { // A null-aware join is never in `Partitioned` mode - // (`NullAwareMode::try_new` rejects it), so the swap + // (`NotIn::try_new` rejects it), so the swap // keeps the mode this arm matched on. hash_join .swap_inputs(PartitionMode::Partitioned) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index d3d0c55ee0366..b374764033113 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -32,6 +32,7 @@ use crate::filter_pushdown::{ use crate::joins::Map; use crate::joins::array_map::ArrayMap; use crate::joins::hash_join::inlist_builder::build_struct_inlist_values; +use crate::joins::hash_join::not_in::{NotIn, NotInBuild}; use crate::joins::hash_join::probe_completion::{ProbeCompletion, ProbeSideSummary}; use crate::joins::hash_join::shared_bounds::{ ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator, @@ -70,8 +71,7 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, UInt64Array}; -use arrow::buffer::NullBuffer; +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::compute::concat_batches; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -99,7 +99,7 @@ use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, l use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; -use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::hash_utils::RandomState; use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::TryStreamExt; @@ -202,152 +202,6 @@ fn array_map_key_range( Ok(Some((min_val, max_val))) } -/// The build rows with a NULL `NOT IN` value key (in any element of a -/// multi-column key), used by correlated null-aware joins (see -/// [`NullAwareMode`]). -/// -/// Such rows make `NOT IN` UNKNOWN whenever *any* probe row in their -/// correlation scope passes the join filter (unless another element is a -/// definite mismatch), so every probe row must be tested against them. -/// Restricting this lookup to the NULL-valued build rows keeps it -/// proportional to the number of NULLs instead of enumerating every scope -/// match of every probe row. -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. `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 - /// build batch. - pub(super) build_indices: UInt64Array, -} - -/// Returns the rows where at least one of `value_keys` is NULL, or `None` when -/// there is no such row. -pub(super) fn null_value_key_mask(value_keys: &[ArrayRef]) -> Option { - let logical_nulls: Vec<_> = value_keys.iter().map(|v| v.logical_nulls()).collect(); - let valid = NullBuffer::union_many(logical_nulls.iter().map(Option::as_ref))?; - (valid.null_count() > 0).then(|| BooleanArray::new(!valid.inner(), None)) -} - -/// Null-aware (`NOT IN`) semantics of a hash join, derived from -/// [`HashJoinExec::null_aware`] and the join type. -/// -/// 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 more than one key — correlation scope keys -/// (`on[V..]`, see [`HashJoinExec::null_aware`]), a multi-column value key, or -/// both — or a join filter. A NULL then makes `NOT IN` UNKNOWN only for the -/// rows whose scope, filter and other tuple elements keep that NULL. `LeftAnti` -/// and `LeftMark` record the decision per build row in the null-indices bitmap -/// instead of in shared probe-side flags; `RightAnti` decides it per probe row -/// within each probe batch. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum NullAwareMode { - /// `build.key NOT IN (probe.key)`: emits build rows. When uncorrelated, - /// none of them are emitted once any probe key is NULL. - /// - /// `value_keys` is the number of `NOT IN` value keys, see - /// [`HashJoinExec::null_aware_value_keys`]. - LeftAnti { correlated: bool, value_keys: usize }, - /// `probe.key NOT IN (build.key)`: emits probe rows. When uncorrelated, - /// none of them are emitted once any build key is NULL. When correlated, - /// a probe row is dropped when some build row in its correlation scope - /// matches it exactly or has no definite mismatch with it; `value_keys` - /// as for `LeftAnti`. - RightAnti { correlated: bool, value_keys: usize }, - /// `NOT IN` as a nullable mark column on the build rows; `value_keys` as - /// for `LeftAnti`. - LeftMark { correlated: bool, value_keys: usize }, -} - -impl NullAwareMode { - /// Validates that `null_aware` may be set for this join and returns its mode. - pub(super) fn try_new( - join_type: JoinType, - partition_mode: PartitionMode, - num_keys: usize, - num_value_keys: usize, - has_filter: bool, - ) -> Result { - if num_value_keys == 0 || num_value_keys > num_keys { - return plan_err!( - "null_aware {join_type} join needs between 1 and {num_keys} `NOT IN` value keys, got {num_value_keys}" - ); - } - // `num_keys > 1` also covers a multi-column value key without any - // correlation: a NULL in one tuple element leaves the comparison FALSE - // whenever another element is a definite mismatch, so it too must be - // decided per build row. - let correlated = num_keys > 1 || has_filter; - let mode = match (join_type, partition_mode) { - (JoinType::LeftAnti, _) => Self::LeftAnti { - correlated, - value_keys: num_value_keys, - }, - // `PartitionMode::CollectLeft` is safe because `RightAnti` is - // probe-driven: every probe row sees the whole build side. - (JoinType::RightAnti, PartitionMode::CollectLeft) => Self::RightAnti { - correlated, - value_keys: num_value_keys, - }, - (JoinType::LeftMark, _) => Self::LeftMark { - correlated, - value_keys: num_value_keys, - }, - _ => { - 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}" - ); - } - }; - match mode { - // Build-driven joins decide in the final stage from state shared - // across the probe partitions: the per-build-row null bitmap for - // the correlated forms, and the probe-side NULL flag for the - // uncorrelated ones. `Partitioned` gives each output partition its - // own build side, so neither is global there. - Self::LeftMark { .. } | Self::LeftAnti { .. } - if partition_mode == PartitionMode::Partitioned => - { - plan_err!( - "null_aware joins require PartitionMode::CollectLeft, got PartitionMode::Partitioned" - ) - } - _ => Ok(mode), - } - } - - /// Whether this join decides UNKNOWN per row (see [`NullAwareMode`]). - pub(super) fn is_correlated(self) -> bool { - match self { - Self::LeftAnti { correlated, .. } - | Self::RightAnti { correlated, .. } - | Self::LeftMark { correlated, .. } => correlated, - } - } - - /// Whether this join records UNKNOWN per build row in the shared - /// null-indices bitmap: a correlated `LeftAnti` or `LeftMark`. - pub(super) fn marks_build_rows(self) -> bool { - self.is_correlated() && !matches!(self, Self::RightAnti { .. }) - } - - /// Number of leading keys that are `NOT IN` value keys; the keys after - /// them are correlation scope keys. - pub(super) fn value_keys(self) -> usize { - match self { - Self::LeftAnti { value_keys, .. } - | Self::RightAnti { value_keys, .. } - | Self::LeftMark { value_keys, .. } => value_keys, - } - } -} - /// Immutable build buffers and their durable reservation. struct JoinBuildData { /// Hash table with row indices into `batch`, also shared with dynamic filters. @@ -357,15 +211,9 @@ struct JoinBuildData { /// Evaluated build-side key expressions, with float `-0.0` rewritten to /// `+0.0` so probe batches can compare against them without scanning the /// build side again. The dynamic filters (`bounds`, `membership`) and - /// `build_side_has_null` were computed from the original keys in + /// the `NOT IN` NULL masks were computed from the original keys in /// `collect_left_input`. values: Vec, - /// Logical nulls of the scalar key for uncorrelated null-aware `LeftMark` - /// and `LeftAnti`. - /// - /// Computed once while collecting the build side because dictionary - /// logical-null discovery can scan nested dictionary values. - null_aware_build_key_nulls: Option, /// Bounds computed from the build side; absent for an empty partition. bounds: Option, /// IN-list values or a hash-table reference used for filter pushdown. @@ -376,26 +224,13 @@ struct JoinBuildData { /// A build lease and mutable bookkeeping for one join execution. pub(super) struct JoinLeftData { - /// Hash table over correlated scope keys for correlated null-aware joins. - /// - /// Keys `..V` are the `NOT IN` value keys and keys `V..` the correlated - /// equality scope keys, where `V` is `null_aware_value_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 with a NULL value key (see [`NullValueBuildRows`]). - /// `None` when the build side has no NULL value keys. - null_value_build_rows: Option, + /// The `NOT IN` state of a null-aware join, see [`NotInBuild`]. + not_in: Option, /// Shared bitmap for visited left indices within this consuming join. visited_indices_bitmap: SharedBitmapBuilder, - /// Shared bitmap for NULL marks within this consuming join. - null_indices_bitmap: SharedBitmapBuilder, /// Tracks the final probe partition and what every partition observed. /// See [`ProbeCompletion`] for the completion invariant. probe_completion: ProbeCompletion, - /// Whether the subquery of an uncorrelated null-aware `RightAnti` join - /// contains a NULL key. - pub(super) build_side_has_null: bool, // Drop the reservation after the mutable state it also accounts for. build: Arc, } @@ -406,12 +241,9 @@ impl JoinLeftData { &self.build.map } - pub(super) fn null_aware_scope_map(&self) -> Option<&dyn JoinHashMapType> { - self.null_aware_scope_map.as_deref() - } - - pub(super) fn null_value_build_rows(&self) -> Option<&NullValueBuildRows> { - self.null_value_build_rows.as_ref() + /// This build's `NOT IN` state, for a null-aware join. + pub(super) fn not_in(&self) -> Option<&NotInBuild> { + self.not_in.as_ref() } /// returns a reference to the build side batch @@ -437,25 +269,17 @@ impl JoinLeftData { } /// Returns the build side key values, with float `-0.0` already rewritten - /// to `+0.0` for key comparison. The dynamic filters and - /// `build_side_has_null` were derived from the original keys instead. + /// to `+0.0` for key comparison. The dynamic filters and the `NOT IN` + /// NULL masks were derived from the original keys instead. pub(super) fn values(&self) -> &[ArrayRef] { &self.build.values } - pub(super) fn null_aware_build_key_nulls(&self) -> Option<&NullBuffer> { - self.build.null_aware_build_key_nulls.as_ref() - } - /// returns a reference to the visited indices bitmap pub(super) fn visited_indices_bitmap(&self) -> &SharedBitmapBuilder { &self.visited_indices_bitmap } - pub(super) fn null_indices_bitmap(&self) -> &SharedBitmapBuilder { - &self.null_indices_bitmap - } - /// returns a reference to the InList values for filter pushdown pub(super) fn membership(&self) -> &PushdownStrategy { &self.build.membership @@ -667,7 +491,7 @@ impl HashJoinExecBuilder { } } // Validate null_aware flag - exec.null_aware_mode()?; + exec.not_in()?; if preserve_properties { return Ok(exec); @@ -1375,10 +1199,10 @@ impl HashJoinExec { /// The null-aware semantics of this join, if [`Self::null_aware`] is set. /// /// Errors if `null_aware` is set on a join that does not support it. - pub(super) fn null_aware_mode(&self) -> Result> { + pub(super) fn not_in(&self) -> Result> { self.null_aware .then(|| { - NullAwareMode::try_new( + NotIn::try_new( self.join_type, self.mode, self.on.len(), @@ -1840,7 +1664,7 @@ impl ExecutionPlan for HashJoinExec { let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); - let null_aware = self.null_aware_mode()?; + let not_in = self.not_in()?; let array_map_created_count = MetricBuilder::new(&self.metrics) .with_category(MetricCategory::Rows) @@ -1867,7 +1691,7 @@ impl ExecutionPlan for HashJoinExec { on_right, repartition_random_state, self.null_equality, - null_aware, + not_in, )) }))) }) @@ -1901,7 +1725,7 @@ impl ExecutionPlan for HashJoinExec { enable_dynamic_filter_pushdown, Arc::clone(context.session_config().options()), self.null_equality, - null_aware, + not_in, array_map_created_count, BuildMode::Ordinary, )) @@ -1923,7 +1747,7 @@ impl ExecutionPlan for HashJoinExec { enable_dynamic_filter_pushdown, Arc::clone(context.session_config().options()), self.null_equality, - null_aware, + not_in, array_map_created_count, BuildMode::Ordinary, )) @@ -1975,7 +1799,6 @@ impl ExecutionPlan for HashJoinExec { self.right.output_ordering().is_some(), build_accumulator, self.mode, - null_aware, self.fetch, ))) } @@ -3092,7 +2915,7 @@ fn is_perfect_hash_join_candidate( Ok(ArrayMap::is_supported_type(&data_type)) } -fn new_join_hashmap( +pub(super) fn new_join_hashmap( num_rows: usize, reservation: &mut MemoryReservation, metrics: &BuildProbeJoinMetrics, @@ -3249,7 +3072,7 @@ async fn collect_left_input( should_compute_dynamic_filters: bool, config: Arc, null_equality: NullEquality, - null_aware: Option, + not_in: Option, array_map_created_count: Count, mode: BuildMode, ) -> Result { @@ -3342,7 +3165,7 @@ async fn collect_left_input( should_compute_dynamic_filters, &config, null_equality, - null_aware, + not_in, &array_map_created_count, prepared, ) @@ -3373,20 +3196,10 @@ fn build_in_memory( should_compute_dynamic_filters: bool, config: &ConfigOptions, null_equality: NullEquality, - null_aware: Option, + not_in: Option, array_map_created_count: &Count, prepared: bool, ) -> Result { - // Correlated null-aware joins (see `NullAwareMode`), multi-column value - // keys without scope keys included, pair NULL-valued rows with the other - // side to find UNKNOWN rows: they need the NULL-valued build rows and, with - // scope keys, the scope map. Of those, only the ones that decide build rows - // also need the per-build-row UNKNOWN bitmap. - let with_null_candidate_state = null_aware.is_some_and(NullAwareMode::is_correlated); - let with_build_unknown_bitmap = - null_aware.is_some_and(NullAwareMode::marks_build_rows); - let null_aware_value_keys = null_aware.map_or(1, NullAwareMode::value_keys); - let is_phj_candidate = is_perfect_hash_join_candidate(on_left, schema)?; let (join_hash_map, batch, left_values) = if let Some((min_val, max_val)) = @@ -3512,86 +3325,20 @@ fn build_in_memory( BooleanBufferBuilder::new(0) }; - let null_indices_bitmap = if with_build_unknown_bitmap { - allocate_bitmap()? - } else { - BooleanBufferBuilder::new(0) - }; - - let (null_aware_scope_map, null_value_build_rows) = if with_null_candidate_state { - // Null-aware convention: `on_left[..V]` are the value keys and - // `on_left[V..]` the (possibly empty) correlation scope keys. - let scope_keys = &on_left[null_aware_value_keys..]; - 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, + let not_in = not_in + .map(|not_in| { + NotInBuild::try_new( + not_in, + on_left, &batch, - &mut *scope_map, - 0, + &left_values, + &normalized_values, random_state, - &mut hashes_buffer, - 0, - true, - NullEquality::NullEqualsNothing, - )?; - Some(scope_map) - }; - - // Collect the NULL-valued build rows (see `NullValueBuildRows`). - let null_value_build_rows = if let Some(null_mask) = - null_value_key_mask(&left_values[..null_aware_value_keys]) - { - let build_indices = UInt64Array::from_iter_values( - null_mask.values().set_indices().map(|i| i as u64), - ); - let scope_values = normalized_values[null_aware_value_keys..] - .iter() - .map(|values| Ok(arrow::compute::filter(values.as_ref(), &null_mask)?)) - .collect::>>()?; - - // `scope_values` and `build_indices` are retained for the lifetime - // of the join, so charge them to the reservation like the maps and - // bitmaps above. - let retained_size = build_indices.get_array_memory_size() - + scope_values - .iter() - .map(|values| values.get_array_memory_size()) - .sum::(); - reservation.try_grow(retained_size)?; - metrics.build_mem_used.add(retained_size); - - 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(NullValueBuildRows { - scope_map, - scope_values, - build_indices, - }) - } else { - None - }; - - (scope_map, null_value_build_rows) - } else { - (None, None) - }; + &mut reservation, + metrics, + ) + }) + .transpose()?; let map = Arc::new(join_hash_map); @@ -3654,47 +3401,6 @@ fn build_in_memory( bounds = None; } - let build_has_null = matches!( - null_aware, - Some(NullAwareMode::RightAnti { - correlated: false, - .. - }) - ) && !left_values.is_empty() - && left_values[0].logical_null_count() > 0; - - let null_aware_build_key_nulls = if matches!( - null_aware, - Some( - NullAwareMode::LeftMark { - correlated: false, - .. - } | NullAwareMode::LeftAnti { - correlated: false, - .. - } - ) - ) { - left_values[0] - .logical_nulls() - .filter(|nulls| nulls.null_count() > 0) - } else { - None - }; - if let Some(nulls) = &null_aware_build_key_nulls { - // `logical_nulls` reuses an array's physical null buffer when it can. - // Only charge newly materialized buffers, such as dictionary logical - // nulls, because the build batch's buffers were reserved above. - let reuses_input_buffer = left_values[0] - .nulls() - .is_some_and(|input| input.inner().ptr_eq(nulls.inner())); - if !reuses_input_buffer { - let retained_size = nulls.buffer().capacity(); - reservation.try_grow(retained_size)?; - metrics.build_mem_used.add(retained_size); - } - } - if prepared { batches.clear(); // Prepared keys are direct columns. IN-list arrays share these batch @@ -3713,17 +3419,13 @@ fn build_in_memory( map, batch, values: normalized_values, - null_aware_build_key_nulls, bounds, membership, reservation, }), - null_aware_scope_map, - null_value_build_rows, + not_in, visited_indices_bitmap: Mutex::new(visited_indices_bitmap), - null_indices_bitmap: Mutex::new(null_indices_bitmap), probe_completion: ProbeCompletion::new(probe_threads_count), - build_side_has_null: build_has_null, }; Ok(data) @@ -3732,6 +3434,7 @@ fn build_in_memory( #[cfg(test)] mod tests { use super::*; + use crate::joins::hash_join::not_in::{NotInOutput, NotInScope}; fn assert_phj_used(metrics: &MetricsSet, use_phj: bool) { if use_phj { @@ -9451,7 +9154,7 @@ mod tests { /// /// A dictionary key that points at a NULL dictionary value has /// `null_count() == 0`, so only the `logical_null_count()` check in - /// `null_aware_skip_probe_batch` records it. This pins that check to the + /// `NotInProbe::before_lookup` records it. This pins that check to the /// cross-partition summary: the partition holding the logical NULL is not /// the one that emits the build rows, in either drain order. #[apply(hash_join_exec_configs)] @@ -10452,17 +10155,15 @@ mod tests { NullEquality::NullEqualsNothing, true, )?; + let scope = NotInScope::Correlated { value_keys: 1 }; let expected = match join_type { - JoinType::RightAnti => NullAwareMode::RightAnti { - correlated: true, - value_keys: 1, - }, - _ => NullAwareMode::LeftAnti { - correlated: true, - value_keys: 1, + JoinType::RightAnti => NotIn::ProbeRows { scope }, + _ => NotIn::BuildRows { + scope, + output: NotInOutput::Rows, }, }; - assert_eq!(join.null_aware_mode()?, Some(expected)); + assert_eq!(join.not_in()?, Some(expected)); Ok(()) } @@ -11210,6 +10911,97 @@ mod tests { Ok(Arc::new(swapped.builder().reset_state().build()?)) } + /// A correlated null-aware join whose build value keys are all NULL: under + /// `NullEqualsNothing` its full-key map is empty, but the UNKNOWN marks + /// still depend on each scope. `NULL NOT IN (5)` and `NULL NOT IN (7)` are + /// UNKNOWN; `NULL NOT IN (empty)`, in scope 3, is TRUE. For `RightAnti`, + /// `5 NOT IN (NULL)` and `7 NOT IN (NULL)` are UNKNOWN and `9 NOT IN + /// (empty)`, in scope 4, is TRUE: only `(9, 4)` is emitted, from the same + /// probe batch as the dropped `(7, 2)`. The probe side arrives in two + /// batches. + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_correlated_all_null_build_values( + batch_size: usize, + ) -> Result<()> { + fn table(id: Vec>, g: Vec>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("g", DataType::Int32, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(id)), + Arc::new(Int32Array::from(g)), + ], + ) + .unwrap() + } + let build = table(vec![None, None, None], vec![Some(1), Some(2), Some(3)]); + let probe = [ + table(vec![Some(5)], vec![Some(1)]), + table(vec![Some(7), Some(9)], vec![Some(2), Some(4)]), + ]; + + let mut results = vec![]; + for join_type in [JoinType::LeftAnti, JoinType::LeftMark, JoinType::RightAnti] { + let task_ctx = prepare_task_ctx(batch_size, false); + let left = TestMemoryExec::try_new_exec( + &[vec![build.clone()]], + build.schema(), + None, + )?; + let right = + TestMemoryExec::try_new_exec(&[probe.to_vec()], probe[0].schema(), None)?; + let on = ["id", "g"] + .into_iter() + .map(|name| { + Ok(( + Arc::new(Column::new_with_schema(name, &left.schema())?) as _, + Arc::new(Column::new_with_schema(name, &right.schema())?) as _, + )) + }) + .collect::>()?; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + let batches = common::collect(join.execute(0, task_ctx)?).await?; + results.push(batches_to_sort_string(&batches)); + } + + allow_duplicates! { + assert_snapshot!(results.join("\n"), @r" + +----+---+ + | id | g | + +----+---+ + | | 3 | + +----+---+ + +----+---+-------+ + | id | g | mark | + +----+---+-------+ + | | 1 | | + | | 2 | | + | | 3 | false | + +----+---+-------+ + +----+---+ + | id | g | + +----+---+ + | 9 | 4 | + +----+---+ + "); + } + Ok(()) + } + /// Correlated null-aware `LeftAnti` whose two scope keys are Float64, with /// `-0.0` on the build side and `0.0` on the probe side. Both scopes must /// still match in every direction: the full key lookup, NULL probe values diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs index fb46cedf96cf1..70bba097efcc4 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -70,12 +70,10 @@ impl PreparedHashJoinBuild { pub(super) fn probe_data(&self, probe_threads: usize) -> JoinLeftData { JoinLeftData { build: Arc::clone(&self.build), - null_aware_scope_map: None, - null_value_build_rows: None, + // Prepared builds exclude null-aware joins. + not_in: None, visited_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), - null_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), probe_completion: ProbeCompletion::new(probe_threads), - build_side_has_null: false, } } diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index aef27fa3057fa..9bfd7ecbf60d0 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -29,6 +29,7 @@ pub use partitioned_hash_eval::{HashExpr, SeededRandomState}; mod exec; mod inlist_builder; +mod not_in; mod partitioned_hash_eval; mod probe_completion; mod shared_bounds; diff --git a/datafusion/physical-plan/src/joins/hash_join/not_in.rs b/datafusion/physical-plan/src/joins/hash_join/not_in.rs new file mode 100644 index 0000000000000..eaebd75387b73 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/not_in.rs @@ -0,0 +1,1144 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `NOT IN` (null-aware) semantics of the hash join. +//! +//! `x NOT IN (subquery)` is `NOT (x IN (subquery))` under SQL three-valued +//! logic: `x IN (..)` is TRUE when some subquery value equals `x`, UNKNOWN +//! (NULL) when none does but a NULL on either side might, and FALSE +//! otherwise. `NOT IN` keeps a row only where `IN` is FALSE. +//! +//! The regular hash join finds the equal pairs; what NULLs change in the +//! rows a `NOT IN` emits is here. The join reaches it through: +//! - [`NotIn::try_new`] when the plan is built or executed: is this join +//! legal, and which form is it; +//! - [`NotInBuild::try_new`] once per build, from `collect_left_input`; +//! - [`NotInProbe::before_lookup`] and [`NotInProbe::after_lookup`] around +//! each probe batch's lookup, or [`NotInProbe::after_empty_build`] when +//! the build map is empty and there is no lookup; +//! - [`NotInBuild::final_chunk`] on each chunk of build rows emitted after +//! the probe side is exhausted. + +use std::sync::Arc; + +use parking_lot::Mutex; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, RecordBatch, + UInt32Array, UInt64Array, +}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::compute::{filter, filter_record_batch, take}; +use arrow::util::bit_util; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::{JoinSide, JoinType, NullEquality, Result, plan_err}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_expr::{ColumnarValue, Operator}; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr_common::datum::apply_cmp; + +use crate::joins::PartitionMode; +use crate::joins::SharedBitmapBuilder; +use crate::joins::hash_join::exec::{JoinLeftData, new_join_hashmap}; +use crate::joins::hash_join::probe_completion::ProbeSideSummary; +use crate::joins::hash_join::stream::lookup_join_hashmap; +use crate::joins::join_hash_map::JoinHashMapType; +use crate::joins::join_hash_map::ProbeOffset::AtRow; +use crate::joins::utils::{ + BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices, update_hash, +}; + +/// Which `NOT IN` a null-aware join computes, derived from +/// `HashJoinExec::null_aware`, the join type, the keys and the filter. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum NotIn { + /// `build.key NOT IN (probe.key)`: the final stage gives every build row + /// its `IN` mark, from the visited bitmap and the NULLs seen. + BuildRows { + scope: NotInScope, + output: NotInOutput, + }, + /// `probe.key NOT IN (build.key)`, `CollectLeft` (`RightAnti`): emits the + /// probe rows whose `IN` is FALSE, decided while streaming because every + /// probe row sees the whole build side. + ProbeRows { scope: NotInScope }, +} + +/// Which rows of the other side a NULL can make a row's `IN` UNKNOWN for. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum NotInScope { + /// One value key, no correlation keys, no filter: a NULL anywhere on the + /// other side makes every unmatched row UNKNOWN. [`NotIn::BuildRows`] + /// probe partitions share one summary of what they saw + /// (`ProbeCompletion`); [`NotIn::ProbeRows`] emits nothing once any build + /// key is NULL, and never a probe row whose key is NULL. + Uncorrelated, + /// Correlation scope keys (`on[value_keys..]`), a multi-column value key + /// (`on[..value_keys]`), or a filter: only the rows in the same scope + /// count, so UNKNOWN is decided per row (see + /// [`mark_null_candidates_for_probe_batch`]). + Correlated { value_keys: usize }, +} + +/// What a [`NotIn::BuildRows`] join emits. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum NotInOutput { + /// `LeftAnti`: the build rows whose mark is FALSE. + Rows, + /// `LeftMark`: every build row with its mark, a nullable Boolean column. + Mark, +} + +impl NotIn { + /// Validates that `null_aware` may be set for this join and returns its + /// form. The errors are those of the legality table this replaced. + pub(super) fn try_new( + join_type: JoinType, + partition_mode: PartitionMode, + num_keys: usize, + num_value_keys: usize, + has_filter: bool, + ) -> Result { + if num_value_keys == 0 || num_value_keys > num_keys { + return plan_err!( + "null_aware {join_type} join needs between 1 and {num_keys} `NOT IN` value keys, got {num_value_keys}" + ); + } + // `num_keys > 1` also covers a multi-column value key without any + // correlation: a NULL in one tuple element leaves the comparison FALSE + // whenever another element is a definite mismatch, so it too must be + // decided per row. + let scope = if num_keys > 1 || has_filter { + NotInScope::Correlated { + value_keys: num_value_keys, + } + } else { + NotInScope::Uncorrelated + }; + let output = match (join_type, partition_mode) { + (JoinType::LeftAnti, _) => NotInOutput::Rows, + (JoinType::LeftMark, _) => NotInOutput::Mark, + // `PartitionMode::CollectLeft` is safe because `RightAnti` is + // probe-driven: every probe row sees the whole build side. + (JoinType::RightAnti, PartitionMode::CollectLeft) => { + return Ok(Self::ProbeRows { scope }); + } + _ => { + 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}" + ); + } + }; + // Build-driven joins decide in the final stage from state shared + // across the probe partitions: the per-build-row null bitmap for + // the correlated forms, and the probe-side NULL flag for the + // uncorrelated ones. `Partitioned` gives each output partition its + // own build side, so neither is global there. + if partition_mode == PartitionMode::Partitioned { + return plan_err!( + "null_aware joins require PartitionMode::CollectLeft, got PartitionMode::Partitioned" + ); + } + Ok(Self::BuildRows { scope, output }) + } + + fn scope(self) -> NotInScope { + match self { + Self::BuildRows { scope, .. } | Self::ProbeRows { scope } => scope, + } + } + + /// Number of leading keys that are `NOT IN` value keys; the keys after + /// them are correlation scope keys. + pub(super) fn value_keys(self) -> usize { + match self.scope() { + NotInScope::Correlated { value_keys } => value_keys, + NotInScope::Uncorrelated => 1, + } + } +} + +/// What a build side keeps for its `NOT IN`, one variant per form. +pub(super) enum NotInBuild { + /// Correlated [`NotIn::BuildRows`]: UNKNOWN is recorded per build row. + BuildRowsCorrelated { + value_keys: usize, + output: NotInOutput, + /// One bit per build row, set once its mark is UNKNOWN; shared by + /// every probe partition, and never cleared. + unknown: SharedBitmapBuilder, + /// Hash table over the scope keys (`on[value_keys..]`) of all build + /// rows, probed with NULL-valued probe rows. `None` without scope keys. + scope_map: Option>, + /// The build rows with a NULL value key, `None` when there are none. + null_value_rows: Option, + }, + /// Uncorrelated [`NotIn::BuildRows`]: what every probe partition saw + /// decides every unmatched mark. + BuildRowsUncorrelated { + output: NotInOutput, + /// Logical NULLs of the value key, `None` when it has none. A NULL + /// build key's mark is UNKNOWN once the probe side has a row. + key_nulls: Option, + }, + /// Correlated [`NotIn::ProbeRows`]: UNKNOWN is decided per probe row, + /// within its probe batch, so nothing is shared across probe partitions. + /// The fields are those of [`Self::BuildRowsCorrelated`]. + ProbeRowsCorrelated { + value_keys: usize, + scope_map: Option>, + null_value_rows: Option, + }, + /// Uncorrelated [`NotIn::ProbeRows`]. + ProbeRowsUncorrelated { + /// Whether a build key is NULL: then no probe row is emitted. + build_has_null: bool, + }, +} + +impl NotInBuild { + /// Builds the `NOT IN` state of a collected build side and charges it to + /// `reservation`. `keys` are the build keys as evaluated, `normalized` + /// with `-0.0` rewritten (see `JoinLeftData::values`). + #[expect(clippy::too_many_arguments)] + pub(super) fn try_new( + not_in: NotIn, + on_left: &[PhysicalExprRef], + batch: &RecordBatch, + keys: &[ArrayRef], + normalized: &[ArrayRef], + random_state: &RandomState, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, + ) -> Result { + let num_rows = batch.num_rows(); + let correlated_row_state = |value_keys, reservation: &mut MemoryReservation| { + correlated_row_state( + on_left, + batch, + keys, + normalized, + value_keys, + random_state, + reservation, + metrics, + ) + }; + match not_in { + NotIn::ProbeRows { + scope: NotInScope::Uncorrelated, + } => Ok(Self::ProbeRowsUncorrelated { + build_has_null: keys[0].logical_null_count() > 0, + }), + NotIn::ProbeRows { + scope: NotInScope::Correlated { value_keys }, + } => { + let (scope_map, null_value_rows) = + correlated_row_state(value_keys, reservation)?; + Ok(Self::ProbeRowsCorrelated { + value_keys, + scope_map, + null_value_rows, + }) + } + NotIn::BuildRows { + scope: NotInScope::Uncorrelated, + output, + } => { + let key_nulls = keys[0].logical_nulls().filter(|n| n.null_count() > 0); + if let Some(nulls) = &key_nulls { + // `logical_nulls` reuses an array's physical null buffer + // when it can, which the batch reservation already holds. + // Charge only a newly materialized one (dictionaries). + let reuses_input_buffer = keys[0] + .nulls() + .is_some_and(|input| input.inner().ptr_eq(nulls.inner())); + if !reuses_input_buffer { + let size = nulls.buffer().capacity(); + reservation.try_grow(size)?; + metrics.build_mem_used.add(size); + } + } + Ok(Self::BuildRowsUncorrelated { output, key_nulls }) + } + NotIn::BuildRows { + scope: NotInScope::Correlated { value_keys }, + output, + } => { + let bitmap_size = bit_util::ceil(num_rows, 8); + reservation.try_grow(bitmap_size)?; + metrics.build_mem_used.add(bitmap_size); + let mut unknown = BooleanBufferBuilder::new(num_rows); + unknown.append_n(num_rows, false); + + let (scope_map, null_value_rows) = + correlated_row_state(value_keys, reservation)?; + Ok(Self::BuildRowsCorrelated { + value_keys, + output, + unknown: Mutex::new(unknown), + scope_map, + null_value_rows, + }) + } + } + } + + /// The rows to emit for one chunk of final build rows, once every probe + /// partition has finished (`summary`), with the mark column for + /// [`NotInOutput::Mark`]. + /// + /// A matched build row's `IN` mark is TRUE. An unmatched one's is UNKNOWN + /// (NULL) or FALSE (the "null bucket" and `hadNull` of Neumann, Leis and + /// Kemper, "The Complete Story of Joins (in HyPer)", Section 5.6): + /// - correlated: UNKNOWN where the bitmap recorded it; + /// - uncorrelated, a probe key was NULL: UNKNOWN; + /// - uncorrelated, the probe side was empty: FALSE (`x IN (empty)`); + /// - uncorrelated otherwise: UNKNOWN exactly for a NULL build key. + pub(super) fn final_chunk( + &self, + summary: ProbeSideSummary, + left_side: UInt64Array, + right_side: UInt32Array, + ) -> (UInt64Array, UInt32Array, Option) { + debug_assert!( + !summary.has_null || summary.non_empty, + "a probe batch with a NULL key has rows" + ); + match self { + Self::ProbeRowsCorrelated { .. } | Self::ProbeRowsUncorrelated { .. } => { + (left_side, right_side, None) + } + Self::BuildRowsCorrelated { + output, unknown, .. + } => { + let unknown = unknown.lock(); + emit(*output, left_side, right_side, |row| unknown.get_bit(row)) + } + // A probe batch with a NULL key also has rows: `has_null` implies + // `non_empty`. + Self::BuildRowsUncorrelated { output, .. } if summary.has_null => { + emit_uniform(*output, left_side, right_side, true) + } + Self::BuildRowsUncorrelated { output, .. } if !summary.non_empty => { + emit_uniform(*output, left_side, right_side, false) + } + Self::BuildRowsUncorrelated { output, key_nulls } => match key_nulls { + Some(nulls) => { + emit(*output, left_side, right_side, |row| nulls.is_null(row)) + } + None => emit_uniform(*output, left_side, right_side, false), + }, + } + } +} + +/// Emits a chunk of final build rows given which unmatched rows have an +/// UNKNOWN mark (see [`NotInBuild::final_chunk`]). `LeftAnti` chunks hold +/// only unmatched rows and keep those whose mark is FALSE; `LeftMark` chunks +/// hold every row, a valid probe index marking a matched one. +fn emit( + output: NotInOutput, + left_side: UInt64Array, + right_side: UInt32Array, + is_unknown: impl Fn(usize) -> bool, +) -> (UInt64Array, UInt32Array, Option) { + match output { + NotInOutput::Rows => { + let kept = UInt64Array::from_iter_values( + left_side + .values() + .iter() + .copied() + .filter(|row| !is_unknown(*row as usize)), + ); + let right_side = UInt32Array::new_null(kept.len()); + (kept, right_side, None) + } + NotInOutput::Mark => { + let marks: BooleanArray = left_side + .values() + .iter() + .zip(right_side.iter()) + .map(|(row, probe)| { + if probe.is_some() { + Some(true) + } else { + (!is_unknown(*row as usize)).then_some(false) + } + }) + .collect(); + (left_side, right_side, Some(Arc::new(marks))) + } + } +} + +/// [`emit`] when every unmatched row's mark is the same: an anti join then +/// keeps all of the chunk or none of it. +fn emit_uniform( + output: NotInOutput, + left_side: UInt64Array, + right_side: UInt32Array, + unknown: bool, +) -> (UInt64Array, UInt32Array, Option) { + match (output, unknown) { + (NotInOutput::Rows, true) => { + (UInt64Array::new_null(0), UInt32Array::new_null(0), None) + } + (NotInOutput::Rows, false) => (left_side, right_side, None), + (NotInOutput::Mark, _) => emit(output, left_side, right_side, |_| unknown), + } +} + +/// One probe stream's `NOT IN` work, with scratch space for the correlated +/// marking pass, reused across probe batches. +#[derive(Debug, Default)] +pub(super) struct NotInProbe { + hashes: Vec, + probe_indices: Vec, + build_indices: Vec, + /// For [`NotInBuild::ProbeRowsCorrelated`], the rows of the current probe + /// batch whose `NOT IN` is UNKNOWN, set by [`Self::before_lookup`] on the + /// batch's first lookup; they are never emitted. `None` when there is no + /// such row. + unknown_probe_rows: Option, +} + +impl NotInProbe { + /// Before a lookup of a probe batch. `first_chunk` is whether this is the + /// batch's first lookup: a batch may resume lookup several times, but + /// what it contributes to `NOT IN` is recorded once. + /// + /// - Uncorrelated [`NotIn::ProbeRows`] emits probe rows, so a NULL build + /// key skips every batch. + /// - Uncorrelated [`NotIn::BuildRows`]: record what the batch saw for the + /// final stage. The anti join skips once any partition saw a NULL probe + /// key: every unmatched mark is then UNKNOWN. This is a best-effort + /// early exit, checked on every resumed lookup; the final stage + /// re-checks. + /// - Correlated: mark the rows this batch makes UNKNOWN, build rows for + /// [`NotIn::BuildRows`] and this batch's probe rows for + /// [`NotIn::ProbeRows`]. Must also run when the build map is empty: an + /// all-NULL-key build side still needs its marks. + /// + /// Returns whether the batch cannot contribute any output row, so its + /// lookup is skipped. A no-op returning `false` for other joins. + #[expect(clippy::too_many_arguments)] + pub(super) fn before_lookup( + &mut self, + left_data: &JoinLeftData, + probe_batch: &RecordBatch, + probe_values: &[ArrayRef], + first_chunk: bool, + filter: Option<&JoinFilter>, + random_state: &RandomState, + batch_size: usize, + ) -> Result { + let Some(not_in) = left_data.not_in() else { + return Ok(false); + }; + match not_in { + NotInBuild::ProbeRowsUncorrelated { build_has_null } => Ok(*build_has_null), + NotInBuild::BuildRowsUncorrelated { output, .. } => { + if first_chunk { + // `on[0]` is the `NOT IN` value key. Only batches with rows + // count: `NULL NOT IN (empty)` is TRUE. + let probe_has_null = probe_values[0].logical_null_count() > 0; + left_data + .record_probe_batch(probe_batch.num_rows() > 0, probe_has_null); + } + Ok(*output == NotInOutput::Rows && left_data.probe_side_has_null_hint()) + } + NotInBuild::BuildRowsCorrelated { + value_keys, + output, + unknown, + scope_map, + null_value_rows, + } => { + // The join filter needs the join type the output implies. + let join_type = match output { + NotInOutput::Rows => JoinType::LeftAnti, + NotInOutput::Mark => JoinType::LeftMark, + }; + if first_chunk { + mark_null_candidates_for_probe_batch( + left_data, + UnknownMarks::Build(unknown), + scope_map.as_deref(), + null_value_rows.as_ref(), + probe_batch, + probe_values, + *value_keys, + filter, + join_type, + random_state, + batch_size, + &mut self.hashes, + &mut self.probe_indices, + &mut self.build_indices, + )?; + } + Ok(false) + } + NotInBuild::ProbeRowsCorrelated { + value_keys, + scope_map, + null_value_rows, + } => { + if !first_chunk { + return Ok(false); + } + self.unknown_probe_rows = None; + // Without a NULL-valued row on either side no probe row can be + // UNKNOWN, so the bitmap is only allocated when one exists. + let may_be_unknown = null_value_rows.is_some() + || probe_values[..*value_keys] + .iter() + .any(|values| values.logical_null_count() > 0); + if may_be_unknown { + let num_rows = probe_batch.num_rows(); + let mut unknown = BooleanBufferBuilder::new(num_rows); + unknown.append_n(num_rows, false); + let unknown = Mutex::new(unknown); + mark_null_candidates_for_probe_batch( + left_data, + UnknownMarks::Probe(&unknown), + scope_map.as_deref(), + null_value_rows.as_ref(), + probe_batch, + probe_values, + *value_keys, + filter, + JoinType::RightAnti, + random_state, + batch_size, + &mut self.hashes, + &mut self.probe_indices, + &mut self.build_indices, + )?; + let unknown = unknown.into_inner().finish(); + self.unknown_probe_rows = + (unknown.count_set_bits() > 0).then_some(unknown); + } + Ok(false) + } + } + } + + /// After a lookup, drops the probe rows a [`NotIn::ProbeRows`] join must + /// not emit: when uncorrelated, those whose key is NULL, since + /// `NULL NOT IN (..)` is never TRUE; when correlated, those that + /// [`Self::before_lookup`] marked UNKNOWN. `valid_keys` is the probe keys' + /// validity (`matchable_join_keys`), when the batch has one. + pub(super) fn after_lookup( + &self, + not_in: &NotInBuild, + probe_values: &[ArrayRef], + valid_keys: Option<&NullBuffer>, + right_indices: UInt32Array, + ) -> UInt32Array { + match not_in { + NotInBuild::ProbeRowsUncorrelated { .. } => { + let mask = valid_keys + .cloned() + .or_else(|| probe_values[0].logical_nulls()); + match mask.filter(|m| m.null_count() > 0) { + Some(mask) => { + retain_probe_rows(&right_indices, |idx| !mask.is_null(idx)) + } + None => right_indices, + } + } + NotInBuild::ProbeRowsCorrelated { .. } => match &self.unknown_probe_rows { + Some(unknown) => { + retain_probe_rows(&right_indices, |idx| !unknown.value(idx)) + } + None => right_indices, + }, + NotInBuild::BuildRowsCorrelated { .. } + | NotInBuild::BuildRowsUncorrelated { .. } => right_indices, + } + } + + /// With an empty build map the lookup is skipped and `output` holds the + /// probe batch as the join type emits it. Drops the probe rows that + /// [`Self::before_lookup`] marked UNKNOWN, as [`Self::after_lookup`] does. + pub(super) fn after_empty_build(&self, output: RecordBatch) -> Result { + match &self.unknown_probe_rows { + Some(unknown) => Ok(filter_record_batch( + &output, + &BooleanArray::new(!unknown, None), + )?), + None => Ok(output), + } + } +} + +/// Keeps the probe rows of `right_indices` for which `keep` returns true. +fn retain_probe_rows( + right_indices: &UInt32Array, + keep: impl Fn(usize) -> bool, +) -> UInt32Array { + UInt32Array::from_iter_values( + right_indices + .values() + .iter() + .copied() + .filter(|idx| keep(*idx as usize)), + ) +} + +/// The scope map and NULL-valued build rows of a correlated `NOT IN`. +#[expect(clippy::too_many_arguments, clippy::type_complexity)] +fn correlated_row_state( + on_left: &[PhysicalExprRef], + batch: &RecordBatch, + keys: &[ArrayRef], + normalized: &[ArrayRef], + value_keys: usize, + random_state: &RandomState, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result<(Option>, Option)> { + let num_rows = batch.num_rows(); + // `on_left[..V]` are the value keys and `on_left[V..]` the (possibly + // empty) correlation scope keys. + let scope_keys = &on_left[value_keys..]; + 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, reservation, metrics)?; + let mut hashes_buffer = vec![0; num_rows]; + update_hash( + scope_keys, + batch, + &mut *scope_map, + 0, + random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNothing, + )?; + Some(scope_map) + }; + + let Some(null_mask) = null_value_key_mask(&keys[..value_keys]) else { + return Ok((scope_map, None)); + }; + let build_indices = + UInt64Array::from_iter_values(null_mask.values().set_indices().map(|i| i as u64)); + let scope_values = normalized[value_keys..] + .iter() + .map(|values| Ok(filter(values.as_ref(), &null_mask)?)) + .collect::>>()?; + // `scope_values` and `build_indices` are retained for the lifetime of the + // join, so charge them like the maps and bitmaps. + let retained_size = build_indices.get_array_memory_size() + + scope_values + .iter() + .map(|values| values.get_array_memory_size()) + .sum::(); + reservation.try_grow(retained_size)?; + metrics.build_mem_used.add(retained_size); + + let null_scope_map = if scope_values.is_empty() { + None + } else { + let null_rows = build_indices.len(); + let mut map = new_join_hashmap(null_rows, 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) + }; + Ok(( + scope_map, + Some(NullValueBuildRows { + scope_map: null_scope_map, + scope_values, + build_indices, + }), + )) +} + +/// The build rows with a NULL `NOT IN` value key (in any element of a +/// multi-column key), used by correlated `NOT IN` (see [`NotInScope::Correlated`]). +/// +/// Such rows make `NOT IN` UNKNOWN whenever *any* probe row in their +/// correlation scope passes the join filter (unless another element is a +/// definite mismatch), so every probe row must be tested against them. +/// Restricting this lookup to the NULL-valued build rows keeps it +/// proportional to the number of NULLs instead of enumerating every scope +/// match of every probe row. +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. `None` when the join has no correlation + /// scope keys, so every probe row is in scope. + scope_map: Option>, + /// Correlation scope key values of the NULL-valued build rows. + scope_values: Vec, + /// Maps positions in `map`/`scope_values` back to row indices in the full + /// build batch. + build_indices: UInt64Array, +} + +/// Returns the rows where at least one of `value_keys` is NULL, or `None` when +/// there is no such row. +fn null_value_key_mask(value_keys: &[ArrayRef]) -> Option { + let logical_nulls: Vec<_> = value_keys.iter().map(|v| v.logical_nulls()).collect(); + let valid = NullBuffer::union_many(logical_nulls.iter().map(Option::as_ref))?; + (valid.null_count() > 0).then(|| BooleanArray::new(!valid.inner(), None)) +} + +/// Where a correlated null-aware join records the rows whose `NOT IN` is +/// UNKNOWN, see [`mark_null_candidates_for_probe_batch`]. +#[derive(Clone, Copy)] +enum UnknownMarks<'a> { + /// `LeftAnti` and `LeftMark` mark build rows, in the null-indices bitmap + /// shared by all probe partitions. + Build(&'a SharedBitmapBuilder), + /// `RightAnti` marks the rows of the current probe batch. + Probe(&'a SharedBitmapBuilder), +} + +impl<'a> UnknownMarks<'a> { + fn bitmap(self) -> &'a SharedBitmapBuilder { + match self { + Self::Build(bitmap) | Self::Probe(bitmap) => bitmap, + } + } + + /// The row of a (build, probe) candidate pair whose `NOT IN` the pair + /// decides, and that is marked when the pair makes it UNKNOWN. + fn unknown_row(self, build_idx: u64, probe_idx: u32) -> usize { + match self { + Self::Build(_) => build_idx as usize, + Self::Probe(_) => probe_idx as usize, + } + } +} + +/// Records which rows of a correlated null-aware join are UNKNOWN candidates +/// for this probe batch: the build rows for `LeftAnti` and `LeftMark`, the +/// probe rows for `RightAnti` (see [`UnknownMarks`]). +/// +/// Key layout: `on[..V]` are the `NOT IN` value keys, `on[V..]` the (possibly +/// empty) correlation scope keys (see [`NotInScope::Correlated`]). A row is +/// NULL-valued when any of its value keys is NULL. Each (build, probe) pair +/// decides one of its rows (see [`UnknownMarks::unknown_row`]), and makes that +/// row's `NOT IN` UNKNOWN instead of TRUE (a `LeftMark` mark NULL instead of +/// FALSE) when either: +/// 1. the build row is NULL-valued, or +/// 2. the probe row is NULL-valued, +/// +/// the two rows are in the same correlation scope, the pair passes the join +/// filter, and, for a multi-column value key, the two rows' value tuples are +/// not a definite mismatch: every element pair is equal or involves a NULL. +/// With a single value key a NULL on either side already rules out a mismatch. +/// An exact match, which makes `NOT IN` FALSE, is found by the main hash lookup +/// instead and takes precedence over these marks. +/// +/// Case 1 pairs the NULL-valued build rows with all probe rows; case 2 pairs +/// all build rows with the NULL-valued probe rows. Scope keys narrow these +/// pairs through a hash lookup; without scope keys every pair is a candidate. +/// The value tuples and then the join filter, if any, decide which candidates +/// count. +/// +/// A row stays UNKNOWN once it is marked, so pairs that decide a row already +/// known to be UNKNOWN are skipped, and the join filter is not evaluated for +/// them. Without scope keys this also ends the pairing as soon as every row +/// that can be decided is UNKNOWN. +#[expect(clippy::too_many_arguments)] +fn mark_null_candidates_for_probe_batch( + left_data: &JoinLeftData, + marks: UnknownMarks<'_>, + scope_map: Option<&dyn JoinHashMapType>, + null_value_build_rows: Option<&NullValueBuildRows>, + probe_batch: &RecordBatch, + probe_values: &[ArrayRef], + num_value_keys: usize, + 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 probe_null_value_mask = null_value_key_mask(&probe_values[..num_value_keys]); + if null_value_build_rows.is_none() && probe_null_value_mask.is_none() { + return Ok(()); + } + + debug_assert_eq!( + left_data.values().len(), + probe_values.len(), + "build/probe key counts must match" + ); + let (build_value_keys, build_scope_values) = + left_data.values().split_at(num_value_keys); + let (probe_value_keys, probe_scope_values) = probe_values.split_at(num_value_keys); + + // Keeps the candidate pairs whose value tuples are not a definite mismatch + // and that pass the join filter, and marks their rows as UNKNOWN. + let mut mark = |build_indices: UInt64Array, probe_indices: UInt32Array| { + let (build_indices, probe_indices) = + retain_unmarked(marks, build_indices, probe_indices); + let (build_indices, probe_indices) = if num_value_keys > 1 { + retain_value_mismatch_free( + build_value_keys, + probe_value_keys, + build_indices, + probe_indices, + )? + } else { + (build_indices, probe_indices) + }; + if build_indices.is_empty() { + return Ok(()); + } + let (build_indices, probe_indices) = match filter { + Some(filter) => apply_join_filter_to_indices( + left_data.batch(), + probe_batch, + build_indices, + probe_indices, + filter, + JoinSide::Left, + None, + join_type, + )?, + None => (build_indices, probe_indices), + }; + if !build_indices.is_empty() { + let mut bitmap = marks.bitmap().lock(); + for (build_idx, probe_idx) in + build_indices.values().iter().zip(probe_indices.values()) + { + bitmap.set_bit(marks.unknown_row(*build_idx, *probe_idx), true); + } + } + Ok(()) + }; + + // Case 1: NULL-valued build rows paired with the probe rows in their + // correlation scope. + if let Some(null_rows) = null_value_build_rows { + match &null_rows.scope_map { + Some(scope_map) => { + hashes_buffer.clear(); + hashes_buffer.resize(probe_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 => { + for_each_unmarked_cross_product( + marks, + null_rows.build_indices.values().iter().copied(), + 0..probe_batch.num_rows() as u32, + batch_size, + &mut mark, + )?; + } + } + } + + // Case 2: NULL-valued probe rows paired with the build rows in their + // correlation scope. + if let Some(null_mask) = probe_null_value_mask { + let null_probe_rows = UInt32Array::from_iter_values( + null_mask.values().set_indices().map(|i| i as u32), + ); + + match 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 => { + for_each_unmarked_cross_product( + marks, + 0..left_data.batch().num_rows() as u64, + null_probe_rows.values().iter().copied(), + batch_size, + &mut mark, + )?; + } + } + } + + Ok(()) +} + +/// Keeps the candidate pairs whose multi-column `NOT IN` value tuples are not a +/// definite mismatch, i.e. every element pair is equal or involves a NULL. +/// +/// Such a pair compares UNKNOWN when some element is NULL (TRUE pairs, with no +/// NULL, are found by the hash lookup instead), whereas a pair with a definite +/// mismatch compares FALSE whatever its NULLs are: `(NULL, 1) = (2, 3)` is +/// `UNKNOWN AND FALSE`, which is FALSE. +/// +/// Elements are compared with SQL `=` (see [`apply_cmp`]), which, like the +/// join's own key equality, treats `-0.0` and `+0.0` as equal. +fn retain_value_mismatch_free( + build_value_keys: &[ArrayRef], + probe_value_keys: &[ArrayRef], + build_indices: UInt64Array, + probe_indices: UInt32Array, +) -> Result<(UInt64Array, UInt32Array)> { + if build_indices.is_empty() { + return Ok((build_indices, probe_indices)); + } + let mut keep: Option = None; + for (build, probe) in build_value_keys.iter().zip(probe_value_keys) { + let build = take(build.as_ref(), &build_indices, None)?; + let probe = take(probe.as_ref(), &probe_indices, None)?; + let eq = apply_cmp( + Operator::Eq, + &ColumnarValue::Array(build), + &ColumnarValue::Array(probe), + )? + .into_array(build_indices.len())?; + let eq = eq.as_boolean(); + // A NULL comparison (an element involving NULL) is not a mismatch. + let not_mismatch = match eq.nulls() { + Some(nulls) => eq.values() | &!nulls.inner(), + None => eq.values().clone(), + }; + keep = Some(match keep { + Some(keep) => &keep & ¬_mismatch, + None => not_mismatch, + }); + } + match keep { + Some(keep) if keep.count_set_bits() < keep.len() => { + let keep = BooleanArray::new(keep, None); + Ok(( + filter(&build_indices, &keep)?.as_primitive().clone(), + filter(&probe_indices, &keep)?.as_primitive().clone(), + )) + } + _ => Ok((build_indices, probe_indices)), + } +} + +/// 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`). +/// `build_scope_values` must already be normalized, as for +/// [`lookup_join_hashmap`]. +#[expect(clippy::too_many_arguments)] +fn for_each_scope_match( + scope_map: &dyn JoinHashMapType, + build_scope_values: &[ArrayRef], + probe_scope_values: &[ArrayRef], + hashes_buffer: &[u64], + batch_size: usize, + probe_indices_buffer: &mut Vec, + build_indices_buffer: &mut Vec, + mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, +) -> Result<()> { + let mut offset = AtRow { row: 0 }; + let mut key_comparator = None; + loop { + let (build_indices, probe_indices, next_offset) = lookup_join_hashmap( + scope_map, + build_scope_values, + probe_scope_values, + NullEquality::NullEqualsNothing, + hashes_buffer, + None, + batch_size, + offset, + probe_indices_buffer, + build_indices_buffer, + &mut key_comparator, + )?; + + if !build_indices.is_empty() { + f(build_indices, probe_indices)?; + } + + let Some(next_offset) = next_offset else { + break; + }; + offset = next_offset; + } + + Ok(()) +} + +/// Removes the candidate pairs that decide a row already marked UNKNOWN. +fn retain_unmarked( + marks: UnknownMarks<'_>, + build_indices: UInt64Array, + probe_indices: UInt32Array, +) -> (UInt64Array, UInt32Array) { + let bitmap = marks.bitmap().lock(); + let is_unmarked = |(build_idx, probe_idx): (&u64, &u32)| { + !bitmap.get_bit(marks.unknown_row(*build_idx, *probe_idx)) + }; + if build_indices + .values() + .iter() + .zip(probe_indices.values().iter()) + .all(is_unmarked) + { + return (build_indices, probe_indices); + } + let (build, probe): (Vec, Vec) = build_indices + .values() + .iter() + .zip(probe_indices.values().iter()) + .filter(|pair| is_unmarked(*pair)) + .unzip(); + (build.into(), probe.into()) +} + +/// Calls `f` with the pairs of `build_rows` x `probe_rows` that decide a row +/// (see [`UnknownMarks::unknown_row`]) not yet marked UNKNOWN, as chunks of at +/// most `batch_size` pairs. +fn for_each_unmarked_cross_product( + marks: UnknownMarks<'_>, + build_rows: impl Iterator, + probe_rows: impl Iterator, + batch_size: usize, + mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, +) -> Result<()> { + match marks { + UnknownMarks::Build(bitmap) => for_each_unmarked_pair( + bitmap, + build_rows.collect(), + probe_rows, + batch_size, + |build, probe| f(build.into(), probe.into()), + ), + UnknownMarks::Probe(bitmap) => for_each_unmarked_pair( + bitmap, + probe_rows.collect(), + build_rows, + batch_size, + |probe, build| f(build.into(), probe.into()), + ), + } +} + +/// Calls `f` with the pairs of `decided_rows` x `other_rows` whose row in +/// `decided_rows` is not set in `bitmap`, as chunks of at most `batch_size` +/// pairs. +/// +/// `f` sets bits in `bitmap`, so the unset rows are found again after each +/// chunk. The pairing stops when every row of `decided_rows` is set. +fn for_each_unmarked_pair, O: Copy>( + bitmap: &SharedBitmapBuilder, + mut decided_rows: Vec, + other_rows: impl Iterator, + batch_size: usize, + mut f: impl FnMut(Vec, Vec) -> Result<()>, +) -> Result<()> { + let retain_undecided_rows = |rows: &mut Vec| { + let bitmap = bitmap.lock(); + rows.retain(|idx| !bitmap.get_bit((*idx).into() as usize)); + }; + + retain_undecided_rows(&mut decided_rows); + + let mut decided_chunk = Vec::with_capacity(batch_size); + let mut other_chunk = Vec::with_capacity(batch_size); + let mut marks_since_refresh = false; + for other_row in other_rows { + // Refresh only after a chunk was sent, so the cost of the refresh + // stays proportional to the pairs already evaluated. + if marks_since_refresh { + retain_undecided_rows(&mut decided_rows); + marks_since_refresh = false; + } + if decided_rows.is_empty() { + break; + } + for decided_row in &decided_rows { + decided_chunk.push(*decided_row); + other_chunk.push(other_row); + if decided_chunk.len() == batch_size { + f( + std::mem::replace(&mut decided_chunk, Vec::with_capacity(batch_size)), + std::mem::replace(&mut other_chunk, Vec::with_capacity(batch_size)), + )?; + marks_since_refresh = true; + } + } + } + if !decided_chunk.is_empty() { + f(decided_chunk, other_chunk)?; + } + Ok(()) +} diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index d1a161aa43a6f..1087f31ca2b2d 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -26,8 +26,9 @@ use crate::ExecutionPlanProperties; use crate::Partitioning; use crate::joins::Map; use crate::joins::PartitionMode; -use crate::joins::hash_join::exec::{HASH_JOIN_SEED, NullAwareMode}; +use crate::joins::hash_join::exec::HASH_JOIN_SEED; use crate::joins::hash_join::inlist_builder::build_struct_fields; +use crate::joins::hash_join::not_in::NotIn; use crate::joins::hash_join::partitioned_hash_eval::{ HashExpr, HashTableLookupExpr, SeededRandomState, }; @@ -273,7 +274,7 @@ pub(crate) struct SharedBuildAccumulator { null_equality: NullEquality, /// Null-aware anti join (`NOT IN`). A probe-side NULL must reach the join so its /// three-valued logic can collapse the result, so the pushed filter keeps NULL rows. - null_aware: Option, + not_in: Option, } /// Strategy for filter pushdown (decided at collection time) @@ -385,7 +386,7 @@ impl SharedBuildAccumulator { on_right: Vec, repartition_random_state: SeededRandomState, null_equality: NullEquality, - null_aware: Option, + not_in: Option, ) -> Self { // Troubleshooting: If partition counts are incorrect, verify this logic matches // the actual execution pattern in collect_build_side() @@ -442,7 +443,7 @@ impl SharedBuildAccumulator { probe_schema: right_child.schema(), probe_range_partitioning, null_equality, - null_aware, + not_in, } } @@ -822,13 +823,13 @@ impl SharedBuildAccumulator { // probe NULL makes `NOT IN` unknown for every build row. A null-equal join needs probe // NULLs only to match an actual build-side NULL, so a NULL-free build keeps the filter // at full selectivity. - let needs_probe_nulls = self.null_aware.is_some() + let needs_probe_nulls = self.not_in.is_some() || (self.null_equality == NullEquality::NullEqualsNull && build_keys_have_null); if !needs_probe_nulls { return Ok(filter_expr); } - let keys = match self.null_aware { + let keys = match self.not_in { Some(mode) => { let value_keys = mode.value_keys(); assert_or_internal_err!( @@ -898,7 +899,7 @@ pub(super) fn make_partitioned_accumulator_for_test( probe_schema, probe_range_partitioning: None, null_equality: NullEquality::NullEqualsNothing, - null_aware: None, + not_in: None, } } @@ -918,6 +919,7 @@ pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usi #[cfg(test)] mod tests { use super::*; + use crate::joins::hash_join::not_in::{NotInOutput, NotInScope}; use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array}; use arrow::compute::SortOptions; @@ -963,7 +965,7 @@ mod tests { probe_schema: test_probe_schema(), probe_range_partitioning: None, null_equality: NullEquality::NullEqualsNothing, - null_aware: None, + not_in: None, } } @@ -1451,9 +1453,9 @@ mod tests { probe_schema, probe_range_partitioning: None, null_equality, - null_aware: null_aware.then_some(NullAwareMode::LeftAnti { - correlated: false, - value_keys: 1, + not_in: null_aware.then_some(NotIn::BuildRows { + scope: NotInScope::Uncorrelated, + output: NotInOutput::Rows, }), } } @@ -1666,9 +1668,9 @@ mod tests { probe_schema: Arc::new(Schema::new(fields)), probe_range_partitioning: None, null_equality: NullEquality::NullEqualsNothing, - null_aware: Some(NullAwareMode::LeftMark { - correlated: true, - value_keys, + not_in: Some(NotIn::BuildRows { + scope: NotInScope::Correlated { value_keys }, + output: NotInOutput::Mark, }), }; diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index f7e5aba2cb2a6..2d317fbc6c562 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -26,8 +26,8 @@ use std::task::Poll; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::joins::Map; use crate::joins::PartitionMode; -use crate::joins::SharedBitmapBuilder; -use crate::joins::hash_join::exec::{JoinLeftData, NullAwareMode, null_value_key_mask}; +use crate::joins::hash_join::exec::JoinLeftData; +use crate::joins::hash_join::not_in::NotInProbe; use crate::joins::hash_join::probe_completion::ProbeSideSummary; use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, SharedBuildAccumulator, @@ -45,24 +45,18 @@ use crate::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, build_batch_empty_build_side, build_batch_from_indices, - build_null_aware_left_mark_column, need_produce_result_in_final, + need_produce_result_in_final, }, }; -use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, UInt32Array, - UInt64Array, -}; +use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::{filter, filter_record_batch, take}; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, internal_datafusion_err, internal_err, }; -use datafusion_expr::{ColumnarValue, Operator}; use datafusion_physical_expr::PhysicalExprRef; -use datafusion_physical_expr_common::datum::apply_cmp; use datafusion_common::hash_utils::RandomState; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; @@ -205,11 +199,6 @@ pub(super) struct ProcessProbeBatchState { /// Key comparator for this batch, built on first use and reused by every /// chunk key_comparator: Option, - /// For a correlated null-aware `RightAnti` join, the probe rows whose - /// `NOT IN` is UNKNOWN because of a NULL (see - /// `mark_null_candidates_for_probe_batch`); they are never emitted. `None` - /// when there is no such row; boxed to keep this state small. - unknown_probe_rows: Option>, } impl ProcessProbeBatchState { @@ -414,14 +403,9 @@ pub(super) struct HashJoinStream { /// Scratch space for build indices during hash lookup build_indices_buffer: Vec, - /// Scratch space for scope-key hashes in the null-aware mark pass, reused - /// across probe batches. Separate from `hashes_buffer`, which still holds - /// the probe batch's full-key hashes during the chunked lookup loop. - null_mark_hashes_buffer: Vec, - /// Scratch space for probe indices during null-aware scope lookups - null_mark_probe_indices_buffer: Vec, - /// Scratch space for build indices during null-aware scope lookups - null_mark_build_indices_buffer: Vec, + /// This stream's `NOT IN` work, for a null-aware join (the build side + /// says whether it is one, see [`JoinLeftData::not_in`]). + not_in: NotInProbe, /// Specifies whether the right side has an ordering to potentially preserve right_side_ordered: bool, @@ -432,8 +416,6 @@ pub(super) struct HashJoinStream { /// Output buffer for coalescing small batches into larger ones with optional fetch limit. /// Uses `LimitedBatchCoalescer` to efficiently combine batches and absorb limit with 'fetch' output_buffer: LimitedBatchCoalescer, - /// Null-aware (`NOT IN`) semantics of this join, if any - null_aware: Option, } impl RecordBatchStream for HashJoinStream { @@ -588,7 +570,6 @@ impl HashJoinStream { right_side_ordered: bool, build_accumulator: Option>, mode: PartitionMode, - null_aware: Option, fetch: Option, ) -> Self { // Create output buffer with coalescing and optional fetch limit. @@ -612,16 +593,11 @@ 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 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(), + not_in: NotInProbe::default(), right_side_ordered, build_report: BuildReportHandle::new(partition, mode, build_accumulator), mode, output_buffer, - null_aware, } } @@ -841,7 +817,6 @@ impl HashJoinStream { joined_probe_idx: None, matched_probe_idx: None, key_comparator: None, - unknown_probe_rows: None, }); } Some(Err(err)) => return Poll::Ready(Err(err)), @@ -869,83 +844,33 @@ 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) { - timer.done(); - self.state = HashJoinStreamState::FetchProbeBatch; - return Ok(StatefulStreamResult::Continue); - } - - // For correlated null-aware joins, record this batch's UNKNOWN - // candidates once, before the first chunked lookup - // (offset == AtRow {row: 0}). - // - // Must precede the empty-build-map return below: an all-NULL-key - // build side has an empty full-key map but still needs UNKNOWN marks. - if mode.is_correlated() && state.is_first_chunk() { - let mut mark = |marks| { - mark_null_candidates_for_probe_batch( - build_side, - state, - mode.value_keys(), - self.filter.as_ref(), - self.join_type, - marks, - &self.random_state, - self.batch_size, - &mut self.null_mark_hashes_buffer, - &mut self.null_mark_probe_indices_buffer, - &mut self.null_mark_build_indices_buffer, - ) - }; - match mode { - // `RightAnti` decides the rows of this probe batch. Without - // a NULL-valued row on either side none of them can be - // UNKNOWN, so the bitmap is only allocated when one exists. - NullAwareMode::RightAnti { value_keys, .. } => { - let may_be_unknown = - build_side.left_data.null_value_build_rows().is_some() - || state.values[..value_keys] - .iter() - .any(|values| values.logical_null_count() > 0); - if may_be_unknown { - let num_rows = state.batch.num_rows(); - let mut unknown = BooleanBufferBuilder::new(num_rows); - unknown.append_n(num_rows, false); - let unknown = SharedBitmapBuilder::new(unknown); - mark(UnknownMarks::Probe(&unknown))?; - let unknown = unknown.into_inner().finish(); - state.unknown_probe_rows = - (unknown.count_set_bits() > 0).then(|| Box::new(unknown)); - } - } - // The others decide the build rows shared by all probe - // partitions. - NullAwareMode::LeftAnti { .. } | NullAwareMode::LeftMark { .. } => { - mark(UnknownMarks::Build( - build_side.left_data.null_indices_bitmap(), - ))?; - } - } - } + // Before the empty-build-map return below: a correlated `NOT IN` + // marks UNKNOWN rows even when every build key is NULL. + let skip = self.not_in.before_lookup( + &build_side.left_data, + &state.batch, + &state.values, + state.is_first_chunk(), + self.filter.as_ref(), + &self.random_state, + self.batch_size, + )?; + if skip { + timer.done(); + self.state = HashJoinStreamState::FetchProbeBatch; + return Ok(StatefulStreamResult::Continue); } let is_empty = !build_side.left_data.has_matchable_build_rows(); if is_empty { - let mut result = build_batch_empty_build_side( + let result = self.not_in.after_empty_build(build_batch_empty_build_side( &self.schema, build_side.left_data.batch(), &state.batch, &self.column_indices, self.join_type, - )?; - if let Some(unknown) = &state.unknown_probe_rows { - result = filter_record_batch( - &result, - &BooleanArray::new(!unknown.as_ref(), None), - )?; - } + )?)?; timer.done(); self.output_buffer.push_batch(result)?; self.state = HashJoinStreamState::FetchProbeBatch; @@ -1067,22 +992,13 @@ impl HashJoinStream { self.right_side_ordered, )?; - match self.null_aware { - Some(NullAwareMode::RightAnti { - correlated: false, .. - }) => { - right_indices = drop_null_probe_keys(state, right_indices); - } - Some(NullAwareMode::RightAnti { - correlated: true, .. - }) => { - if let Some(unknown) = &state.unknown_probe_rows { - right_indices = retain_probe_rows(&right_indices, |idx| { - !unknown.value(idx as usize) - }); - } - } - _ => {} + if let Some(not_in) = build_side.left_data.not_in() { + right_indices = self.not_in.after_lookup( + not_in, + &state.values, + state.valid_keys.as_ref(), + right_indices, + ); } // Build output batch and push to coalescer @@ -1222,30 +1138,9 @@ impl HashJoinStream { self.batch_size, ); - // 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 { 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 { correlated, .. }) => { - let mark_column = null_aware_left_mark_column( - &build_side.left_data, - correlated, - probe_summary, - &left_side, - &right_side, - ); - (left_side, right_side, Some(mark_column)) - } - _ => (left_side, right_side, None), + let (left_side, right_side, mark_column) = match build_side.left_data.not_in() { + Some(not_in) => not_in.final_chunk(probe_summary, left_side, right_side), + None => (left_side, right_side, None), }; self.join_metrics.input_rows.add(left_side.len()); @@ -1329,607 +1224,6 @@ fn next_final_indices_chunk( (UInt64Array::from(build_indices), probe_indices) } -/// Applies the pre-lookup bookkeeping of a null-aware join to a probe batch -/// and returns `true` if the batch cannot contribute any output row. -/// -/// Under SQL three-valued logic a NULL on either side of an uncorrelated scalar -/// `NOT IN` makes the predicate UNKNOWN for every row of the other side: -/// - `RightAnti` emits probe rows, so a NULL build key silences every batch. -/// - `LeftAnti` emits build rows in the final stage, so a NULL probe key seen -/// by any partition silences everything; the shared flags carry this to the -/// final stage. -/// - `LeftMark` only records the shared flags; the mark column is built in -/// the final stage from them. -fn null_aware_skip_probe_batch( - mode: NullAwareMode, - state: &ProcessProbeBatchState, - left_data: &JoinLeftData, -) -> bool { - match mode { - NullAwareMode::RightAnti { - correlated: false, .. - } => left_data.build_side_has_null, - // Correlated joins decide UNKNOWN per row instead, in - // `mark_null_candidates_for_probe_batch`. - NullAwareMode::RightAnti { - correlated: true, .. - } - | NullAwareMode::LeftAnti { - correlated: true, .. - } - | NullAwareMode::LeftMark { - correlated: true, .. - } => false, - NullAwareMode::LeftAnti { - correlated: false, .. - } - | NullAwareMode::LeftMark { - correlated: false, .. - } => { - let is_anti = matches!(mode, NullAwareMode::LeftAnti { .. }); - if state.is_first_chunk() { - // `on[0]` is the `NOT IN` value key for both modes. A probe - // batch may resume lookup many times, but its summary is - // invariant and only needs to be computed once. - let probe_has_null = state.values[0].logical_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`. Keep this outside the first- - // chunk guard so a sibling partition can stop an anti join between - // resumed lookups. - is_anti && left_data.probe_side_has_null_hint() - } - } -} - -/// Drops probe rows whose `NOT IN` key is NULL from the output of a null-aware -/// `RightAnti` join: `NULL NOT IN (..)` is never TRUE. -fn drop_null_probe_keys( - state: &ProcessProbeBatchState, - right_indices: UInt32Array, -) -> UInt32Array { - let probe_key = &state.values[0]; - // if the valid_keys mask is available, use that, else use the probe key - let mask = state - .valid_keys - .clone() - .or_else(|| probe_key.logical_nulls()); - // we only need this copy if there are NULLs - match mask.filter(|m| m.null_count() > 0) { - Some(mask) => { - retain_probe_rows(&right_indices, |idx| !mask.is_null(idx as usize)) - } - None => right_indices, - } -} - -/// Keeps the probe rows of `right_indices` for which `keep` returns true. -fn retain_probe_rows( - right_indices: &UInt32Array, - keep: impl Fn(u32) -> bool, -) -> UInt32Array { - UInt32Array::from_iter_values( - right_indices - .values() - .iter() - .copied() - .filter(|idx| keep(*idx)), - ) -} - -/// Final-stage rules of a null-aware `LeftAnti` join, evaluated by the last -/// probe partition. -/// -/// 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 -/// (row dropped) unless the probe side was empty, where it is TRUE (row -/// 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)); - } - if !probe_summary.non_empty { - return (left_side, right_side); - } - // Logical nulls, so that a dictionary key pointing at a NULL value counts - // as a NULL build key. Computed once with the build side. - let Some(build_key_nulls) = left_data.null_aware_build_key_nulls() else { - return (left_side, right_side); - }; - let left_side = UInt64Array::from_iter_values(left_side.iter().filter_map(|idx| { - let idx = - idx.expect("LeftAnti final indices should always contain build-side rows"); - (!build_key_nulls.is_null(idx as usize)).then_some(idx) - })); - let right_side = UInt32Array::new_null(left_side.len()); - (left_side, right_side) -} - -/// Builds the nullable mark column of a null-aware `LeftMark` join from the -/// 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 { - // Correlated joins precomputed the UNKNOWN decision per build row. - let null_indices_bitmap = correlated.then(|| left_data.null_indices_bitmap().lock()); - // Uncorrelated joins use the logical-null bitmap computed once with the - // build side. Correlated joins carry the answer in `null_indices_bitmap`. - let build_key_nulls = if correlated { - None - } else { - left_data.null_aware_build_key_nulls() - }; - build_null_aware_left_mark_column( - left_side, - right_side, - build_key_nulls, - null_indices_bitmap.as_deref(), - probe_summary.has_null, - probe_summary.non_empty, - ) -} - -/// Where a correlated null-aware join records the rows whose `NOT IN` is -/// UNKNOWN, see [`mark_null_candidates_for_probe_batch`]. -#[derive(Clone, Copy)] -enum UnknownMarks<'a> { - /// `LeftAnti` and `LeftMark` mark build rows, in the null-indices bitmap - /// shared by all probe partitions. - Build(&'a SharedBitmapBuilder), - /// `RightAnti` marks the rows of the current probe batch. - Probe(&'a SharedBitmapBuilder), -} - -impl<'a> UnknownMarks<'a> { - fn bitmap(self) -> &'a SharedBitmapBuilder { - match self { - Self::Build(bitmap) | Self::Probe(bitmap) => bitmap, - } - } - - /// The row of a (build, probe) candidate pair whose `NOT IN` the pair - /// decides, and that is marked when the pair makes it UNKNOWN. - fn unknown_row(self, build_idx: u64, probe_idx: u32) -> usize { - match self { - Self::Build(_) => build_idx as usize, - Self::Probe(_) => probe_idx as usize, - } - } -} - -/// Records which rows of a correlated null-aware join are UNKNOWN candidates -/// for this probe batch: the build rows for `LeftAnti` and `LeftMark`, the -/// probe rows for `RightAnti` (see [`UnknownMarks`]). -/// -/// Key layout: `on[..V]` are the `NOT IN` value keys, `on[V..]` the (possibly -/// empty) correlation scope keys (see `HashJoinExec::null_aware`). A row is -/// NULL-valued when any of its value keys is NULL. Each (build, probe) pair -/// decides one of its rows (see [`UnknownMarks::unknown_row`]), and makes that -/// row's `NOT IN` UNKNOWN instead of TRUE (a `LeftMark` mark NULL instead of -/// FALSE) when either: -/// 1. the build row is NULL-valued, or -/// 2. the probe row is NULL-valued, -/// -/// the two rows are in the same correlation scope, the pair passes the join -/// filter, and, for a multi-column value key, the two rows' value tuples are -/// not a definite mismatch: every element pair is equal or involves a NULL. -/// With a single value key a NULL on either side already rules out a mismatch. -/// An exact match, which makes `NOT IN` FALSE, is found by the main hash lookup -/// instead and takes precedence over these marks. -/// -/// Case 1 pairs the NULL-valued build rows with all probe rows; case 2 pairs -/// all build rows with the NULL-valued probe rows. Scope keys narrow these -/// pairs through a hash lookup; without scope keys every pair is a candidate. -/// The value tuples and then the join filter, if any, decide which candidates -/// count. -/// -/// A row stays UNKNOWN once it is marked, so pairs that decide a row already -/// known to be UNKNOWN are skipped, and the join filter is not evaluated for -/// them. Without scope keys this also ends the pairing as soon as every row -/// that can be decided is UNKNOWN. -#[expect(clippy::too_many_arguments)] -fn mark_null_candidates_for_probe_batch( - build_side: &BuildSideReadyState, - state: &ProcessProbeBatchState, - num_value_keys: usize, - filter: Option<&JoinFilter>, - join_type: JoinType, - marks: UnknownMarks<'_>, - random_state: &RandomState, - batch_size: usize, - hashes_buffer: &mut Vec, - probe_indices_buffer: &mut Vec, - build_indices_buffer: &mut Vec, -) -> Result<()> { - let left_data = &build_side.left_data; - let null_value_build_rows = left_data.null_value_build_rows(); - let probe_null_value_mask = null_value_key_mask(&state.values[..num_value_keys]); - if null_value_build_rows.is_none() && probe_null_value_mask.is_none() { - return Ok(()); - } - - debug_assert_eq!( - left_data.values().len(), - state.values.len(), - "build/probe key counts must match" - ); - let (build_value_keys, build_scope_values) = - left_data.values().split_at(num_value_keys); - let (probe_value_keys, probe_scope_values) = state.values.split_at(num_value_keys); - - // Keeps the candidate pairs whose value tuples are not a definite mismatch - // and that pass the join filter, and marks their rows as UNKNOWN. - let mut mark = |build_indices: UInt64Array, probe_indices: UInt32Array| { - let (build_indices, probe_indices) = - retain_unmarked(marks, build_indices, probe_indices); - let (build_indices, probe_indices) = if num_value_keys > 1 { - retain_value_mismatch_free( - build_value_keys, - probe_value_keys, - build_indices, - probe_indices, - )? - } else { - (build_indices, probe_indices) - }; - if build_indices.is_empty() { - return Ok(()); - } - let (build_indices, probe_indices) = match filter { - Some(filter) => apply_join_filter_to_indices( - left_data.batch(), - &state.batch, - build_indices, - probe_indices, - filter, - JoinSide::Left, - None, - join_type, - )?, - None => (build_indices, probe_indices), - }; - if !build_indices.is_empty() { - let mut bitmap = marks.bitmap().lock(); - for (build_idx, probe_idx) in - build_indices.values().iter().zip(probe_indices.values()) - { - bitmap.set_bit(marks.unknown_row(*build_idx, *probe_idx), true); - } - } - Ok(()) - }; - - // Case 1: NULL-valued build rows paired with the probe rows in their - // correlation scope. - if let Some(null_rows) = null_value_build_rows { - match &null_rows.scope_map { - Some(scope_map) => { - 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 => { - for_each_unmarked_cross_product( - marks, - null_rows.build_indices.values().iter().copied(), - 0..state.batch.num_rows() as u32, - batch_size, - &mut mark, - )?; - } - } - } - - // Case 2: NULL-valued probe rows paired with the build rows in their - // correlation scope. - if let Some(null_mask) = probe_null_value_mask { - let null_probe_rows = UInt32Array::from_iter_values( - null_mask.values().set_indices().map(|i| i as u32), - ); - - 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 => { - for_each_unmarked_cross_product( - marks, - 0..left_data.batch().num_rows() as u64, - null_probe_rows.values().iter().copied(), - batch_size, - &mut mark, - )?; - } - } - } - - Ok(()) -} - -/// Keeps the candidate pairs whose multi-column `NOT IN` value tuples are not a -/// definite mismatch, i.e. every element pair is equal or involves a NULL. -/// -/// Such a pair compares UNKNOWN when some element is NULL (TRUE pairs, with no -/// NULL, are found by the hash lookup instead), whereas a pair with a definite -/// mismatch compares FALSE whatever its NULLs are: `(NULL, 1) = (2, 3)` is -/// `UNKNOWN AND FALSE`, which is FALSE. -/// -/// Elements are compared with SQL `=` (see [`apply_cmp`]), which, like the -/// join's own key equality, treats `-0.0` and `+0.0` as equal. -fn retain_value_mismatch_free( - build_value_keys: &[ArrayRef], - probe_value_keys: &[ArrayRef], - build_indices: UInt64Array, - probe_indices: UInt32Array, -) -> Result<(UInt64Array, UInt32Array)> { - if build_indices.is_empty() { - return Ok((build_indices, probe_indices)); - } - let mut keep: Option = None; - for (build, probe) in build_value_keys.iter().zip(probe_value_keys) { - let build = take(build.as_ref(), &build_indices, None)?; - let probe = take(probe.as_ref(), &probe_indices, None)?; - let eq = apply_cmp( - Operator::Eq, - &ColumnarValue::Array(build), - &ColumnarValue::Array(probe), - )? - .into_array(build_indices.len())?; - let eq = eq.as_boolean(); - // A NULL comparison (an element involving NULL) is not a mismatch. - let not_mismatch = match eq.nulls() { - Some(nulls) => eq.values() | &!nulls.inner(), - None => eq.values().clone(), - }; - keep = Some(match keep { - Some(keep) => &keep & ¬_mismatch, - None => not_mismatch, - }); - } - match keep { - Some(keep) if keep.count_set_bits() < keep.len() => { - let keep = BooleanArray::new(keep, None); - Ok(( - filter(&build_indices, &keep)?.as_primitive().clone(), - filter(&probe_indices, &keep)?.as_primitive().clone(), - )) - } - _ => Ok((build_indices, probe_indices)), - } -} - -/// 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`). -/// `build_scope_values` must already be normalized, as for -/// [`lookup_join_hashmap`]. -#[expect(clippy::too_many_arguments)] -fn for_each_scope_match( - scope_map: &dyn JoinHashMapType, - build_scope_values: &[ArrayRef], - probe_scope_values: &[ArrayRef], - hashes_buffer: &[u64], - batch_size: usize, - probe_indices_buffer: &mut Vec, - build_indices_buffer: &mut Vec, - mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, -) -> Result<()> { - let mut offset = AtRow { row: 0 }; - let mut key_comparator = None; - loop { - let (build_indices, probe_indices, next_offset) = lookup_join_hashmap( - scope_map, - build_scope_values, - probe_scope_values, - NullEquality::NullEqualsNothing, - hashes_buffer, - None, - batch_size, - offset, - probe_indices_buffer, - build_indices_buffer, - &mut key_comparator, - )?; - - if !build_indices.is_empty() { - f(build_indices, probe_indices)?; - } - - let Some(next_offset) = next_offset else { - break; - }; - offset = next_offset; - } - - Ok(()) -} - -/// Removes the candidate pairs that decide a row already marked UNKNOWN. -fn retain_unmarked( - marks: UnknownMarks<'_>, - build_indices: UInt64Array, - probe_indices: UInt32Array, -) -> (UInt64Array, UInt32Array) { - let bitmap = marks.bitmap().lock(); - let is_unmarked = |(build_idx, probe_idx): (&u64, &u32)| { - !bitmap.get_bit(marks.unknown_row(*build_idx, *probe_idx)) - }; - if build_indices - .values() - .iter() - .zip(probe_indices.values().iter()) - .all(is_unmarked) - { - return (build_indices, probe_indices); - } - let (build, probe): (Vec, Vec) = build_indices - .values() - .iter() - .zip(probe_indices.values().iter()) - .filter(|pair| is_unmarked(*pair)) - .unzip(); - (build.into(), probe.into()) -} - -/// Calls `f` with the pairs of `build_rows` x `probe_rows` that decide a row -/// (see [`UnknownMarks::unknown_row`]) not yet marked UNKNOWN, as chunks of at -/// most `batch_size` pairs. -fn for_each_unmarked_cross_product( - marks: UnknownMarks<'_>, - build_rows: impl Iterator, - probe_rows: impl Iterator, - batch_size: usize, - mut f: impl FnMut(UInt64Array, UInt32Array) -> Result<()>, -) -> Result<()> { - match marks { - UnknownMarks::Build(bitmap) => for_each_unmarked_pair( - bitmap, - build_rows.collect(), - probe_rows, - batch_size, - |build, probe| f(build.into(), probe.into()), - ), - UnknownMarks::Probe(bitmap) => for_each_unmarked_pair( - bitmap, - probe_rows.collect(), - build_rows, - batch_size, - |probe, build| f(build.into(), probe.into()), - ), - } -} - -/// Calls `f` with the pairs of `decided_rows` x `other_rows` whose row in -/// `decided_rows` is not set in `bitmap`, as chunks of at most `batch_size` -/// pairs. -/// -/// `f` sets bits in `bitmap`, so the unset rows are found again after each -/// chunk. The pairing stops when every row of `decided_rows` is set. -fn for_each_unmarked_pair, O: Copy>( - bitmap: &SharedBitmapBuilder, - mut decided_rows: Vec, - other_rows: impl Iterator, - batch_size: usize, - mut f: impl FnMut(Vec, Vec) -> Result<()>, -) -> Result<()> { - let retain_undecided_rows = |rows: &mut Vec| { - let bitmap = bitmap.lock(); - rows.retain(|idx| !bitmap.get_bit((*idx).into() as usize)); - }; - - retain_undecided_rows(&mut decided_rows); - - let mut decided_chunk = Vec::with_capacity(batch_size); - let mut other_chunk = Vec::with_capacity(batch_size); - let mut marks_since_refresh = false; - for other_row in other_rows { - // Refresh only after a chunk was sent, so the cost of the refresh - // stays proportional to the pairs already evaluated. - if marks_since_refresh { - retain_undecided_rows(&mut decided_rows); - marks_since_refresh = false; - } - if decided_rows.is_empty() { - break; - } - for decided_row in &decided_rows { - decided_chunk.push(*decided_row); - other_chunk.push(other_row); - if decided_chunk.len() == batch_size { - f( - std::mem::replace(&mut decided_chunk, Vec::with_capacity(batch_size)), - std::mem::replace(&mut other_chunk, Vec::with_capacity(batch_size)), - )?; - marks_since_refresh = true; - } - } - } - if !decided_chunk.is_empty() { - f(decided_chunk, other_chunk)?; - } - Ok(()) -} - impl Stream for HashJoinStream { type Item = Result; diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 5136fa9a81372..6bee6c0ebe805 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -1551,70 +1551,6 @@ pub(crate) fn build_batch_from_indices( Ok(RecordBatch::try_new(Arc::clone(schema), columns)?) } -/// Builds the nullable mark column for a null-aware `LeftMark` join. -/// -/// This follows the left mark hash join described in Neumann, Leis, and Kemper, -/// "The Complete Story of Joins (in HyPer)", Section 5.6: -/// -/// -/// `build_indices` and `probe_indices` are the final aligned indices derived from the -/// visited bitmap. At this point: -/// - valid `probe_indices` mean the build row matched at least one probe row, so the mark is `TRUE` -/// - null `probe_indices` mean the build row was unmatched, so the result depends on SQL -/// three-valued logic -/// -/// For the uncorrelated single-key implementation, unmatched rows are classified as follows: -/// 1. if the build key is `NULL` and the probe side is non-empty, the mark is `NULL` -/// 2. if the build key is `NULL` and the probe side is empty, the mark is `FALSE` -/// 3. if the build key is non-null and the probe side contained a `NULL`, the mark is `NULL` -/// 4. otherwise, the mark is `FALSE` -/// -/// For correlated scalar `NOT IN`, `null_indices_bitmap` carries the same UNKNOWN -/// decision per build row, scoped by the correlated equality keys. -/// -/// When present, `build_key_nulls` is indexed by the original build row number. -/// -/// This is the helper equivalent of the paper's "null bucket" and `hadNull` handling. -/// It is intentionally scoped to scalar null-aware mark joins. -pub(crate) fn build_null_aware_left_mark_column( - build_indices: &UInt64Array, - probe_indices: &UInt32Array, - build_key_nulls: Option<&NullBuffer>, - null_indices_bitmap: Option<&BooleanBufferBuilder>, - probe_side_has_null: bool, - probe_side_non_empty: bool, -) -> ArrayRef { - // Whether an unmatched build row's mark is NULL (UNKNOWN) instead of FALSE: - // correlated joins precomputed this per row in `null_indices_bitmap`; the - // uncorrelated rules are cases 1-4 in the doc above. - let unmatched_mark_is_null = |build_idx: usize| match null_indices_bitmap { - Some(bitmap) => bitmap.get_bit(build_idx), - None if build_key_nulls.is_some_and(|nulls| nulls.is_null(build_idx)) => { - probe_side_non_empty - } - None => probe_side_has_null, - }; - - let marks: BooleanArray = build_indices - .iter() - .zip(probe_indices.iter()) - .map(|(build_idx, probe_idx)| { - if probe_idx.is_some() { - return Some(true); - } - let build_idx = build_idx - .expect("LeftMark final indices should always contain build-side rows") - as usize; - if unmatched_mark_is_null(build_idx) { - None - } else { - Some(false) - } - }) - .collect(); - Arc::new(marks) -} - /// Returns a new [RecordBatch] for a probe batch when no probe row can find a /// match: the build-side map is empty, either because the build side has no /// rows or because none of its rows has a matchable (non-NULL) join key.