From 65eae3bfa5bb3b8e03d96cbeb3a7d2d85207a053 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:19:54 -0500 Subject: [PATCH 1/3] refactor: pair null-aware candidates through one struct Prepares the fix for re-paired marked build rows in the scope-key null-aware join. No behavior changes. The cross product of the path without scope keys and the scope-key lookup now share UnmarkedPairs, which holds the chunk buffers and the refresh of the marked rows. The two lookup index buffers travel together as LookupBuffers. The scope-key lookup now yields full build and probe row indices, and the join checks the scope-key equality on each chunk of pairs. So NullValueBuildRows no longer keeps a copy of the scope values. Probe rows whose scope keys hold a NULL are left out before the lookup, because they are in no scope. --- .../src/joins/hash_join/not_in.rs | 442 +++++++++--------- 1 file changed, 232 insertions(+), 210 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/not_in.rs b/datafusion/physical-plan/src/joins/hash_join/not_in.rs index 268160b421e5..4880cfdf6498 100644 --- a/datafusion/physical-plan/src/joins/hash_join/not_in.rs +++ b/datafusion/physical-plan/src/joins/hash_join/not_in.rs @@ -54,12 +54,12 @@ 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::stream::lookup_join_hashmap; use crate::joins::join_hash_map::JoinHashMapType; use crate::joins::join_hash_map::ProbeOffset::AtRow; use crate::joins::probe_completion::ProbeSideSummary; use crate::joins::utils::{ - BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices, update_hash, + BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices, + equal_rows_arr_with_normalized_left, matchable_join_keys, update_hash, }; /// Which `NOT IN` a null-aware join computes, derived from @@ -409,8 +409,7 @@ fn emit_uniform( #[derive(Debug, Default)] pub(super) struct NotInProbe { hashes: Vec, - probe_indices: Vec, - build_indices: Vec, + lookup: LookupBuffers, /// 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 @@ -489,8 +488,7 @@ impl NotInProbe { random_state, batch_size, &mut self.hashes, - &mut self.probe_indices, - &mut self.build_indices, + &mut self.lookup, )?; } Ok(false) @@ -528,8 +526,7 @@ impl NotInProbe { random_state, batch_size, &mut self.hashes, - &mut self.probe_indices, - &mut self.build_indices, + &mut self.lookup, )?; let unknown = unknown.into_inner().finish(); self.unknown_probe_rows = @@ -646,23 +643,19 @@ fn correlated_row_state( }; 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::(); + // `build_indices` is retained for the lifetime of the join, so charge it + // like the maps and bitmaps. + let retained_size = build_indices.get_array_memory_size(); reservation.try_grow(retained_size)?; metrics.build_mem_used.add(retained_size); - let null_scope_map = if scope_values.is_empty() { + let null_scope_map = if scope_keys.is_empty() { None } else { + let scope_values = normalized[value_keys..] + .iter() + .map(|values| Ok(filter(values.as_ref(), &null_mask)?)) + .collect::>>()?; let null_rows = build_indices.len(); let mut map = new_join_hashmap(null_rows, reservation, metrics)?; let mut hashes_buffer = vec![0; null_rows]; @@ -674,7 +667,6 @@ fn correlated_row_state( scope_map, Some(NullValueBuildRows { scope_map: null_scope_map, - scope_values, build_indices, }), )) @@ -691,14 +683,12 @@ fn correlated_row_state( /// 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. + /// build rows. Stored positions index into `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. + /// Maps positions in `scope_map` back to row indices in the full build + /// batch. build_indices: UInt64Array, } @@ -767,7 +757,7 @@ impl<'a> UnknownMarks<'a> { /// 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. +/// that can be decided is UNKNOWN. See [`UnmarkedPairs`]. #[expect(clippy::too_many_arguments)] fn mark_null_candidates_for_probe_batch( left_data: &JoinLeftData, @@ -782,8 +772,7 @@ fn mark_null_candidates_for_probe_batch( random_state: &RandomState, batch_size: usize, hashes_buffer: &mut Vec, - probe_indices_buffer: &mut Vec, - build_indices_buffer: &mut Vec, + lookup: &mut LookupBuffers, ) -> 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() { @@ -801,7 +790,7 @@ fn mark_null_candidates_for_probe_batch( // 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 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 { @@ -841,44 +830,63 @@ fn mark_null_candidates_for_probe_batch( Ok(()) }; + // The scope maps pair rows by the hash of their scope keys; keep only the + // pairs whose scope keys are equal. + let mut key_comparator = None; + let mut mark_scope_matches = + |build_indices: UInt64Array, probe_indices: UInt32Array| { + let (build_indices, probe_indices) = equal_rows_arr_with_normalized_left( + &build_indices, + &probe_indices, + build_scope_values, + probe_scope_values, + NullEquality::NullEqualsNothing, + &mut key_comparator, + )?; + mark(build_indices, probe_indices) + }; + + // With scope keys, the probe rows whose scope keys hold a NULL are in no + // scope. The others are looked up by the hash of their scope keys. + let probe_rows_in_scope = if probe_scope_values.is_empty() { + None + } else { + hashes_buffer.clear(); + hashes_buffer.resize(probe_batch.num_rows(), 0); + create_hashes(probe_scope_values, random_state, hashes_buffer)?; + matchable_join_keys(probe_scope_values, NullEquality::NullEqualsNothing) + }; + let in_scope = |probe_row: &u32| match &probe_rows_in_scope { + Some(valid) => valid.is_valid(*probe_row as usize), + None => true, + }; + // Case 1: NULL-valued build rows paired with the probe rows in their // correlation scope. if let Some(null_rows) = null_value_build_rows { + let probe_rows = (0..probe_batch.num_rows() as u32).filter(in_scope); 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( + let mut pairs = + UnmarkedPairs::new(marks, batch_size, &mut mark_scope_matches); + // The map indexes only the NULL-valued build rows; translate + // its positions back to build row indices. + pairs.scope_matches( scope_map.as_ref(), - &null_rows.scope_values, - probe_scope_values, + |position| null_rows.build_indices.value(position as usize), + probe_rows, 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) - }, + lookup, )?; + pairs.finish()?; } None => { - for_each_unmarked_cross_product( - marks, + let mut pairs = UnmarkedPairs::new(marks, batch_size, &mark); + pairs.cross_product( null_rows.build_indices.values().iter().copied(), - 0..probe_batch.num_rows() as u32, - batch_size, - &mut mark, + probe_rows, )?; + pairs.finish()?; } } } @@ -886,52 +894,32 @@ fn mark_null_candidates_for_probe_batch( // 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), - ); + let null_probe_rows = null_mask + .values() + .set_indices() + .map(|i| i as u32) + .filter(in_scope); 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( + let mut pairs = + UnmarkedPairs::new(marks, batch_size, &mut mark_scope_matches); + pairs.scope_matches( scope_map, - build_scope_values, - &probe_null_scope_values, + |build_row| build_row, + null_probe_rows, 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) - }, + lookup, )?; + pairs.finish()?; } None => { - for_each_unmarked_cross_product( - marks, + let mut pairs = UnmarkedPairs::new(marks, batch_size, &mark); + pairs.cross_product( 0..left_data.batch().num_rows() as u64, - null_probe_rows.values().iter().copied(), - batch_size, - &mut mark, + null_probe_rows, )?; + pairs.finish()?; } } } @@ -991,50 +979,163 @@ fn retain_value_mismatch_free( } } -/// 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], +/// Reused buffers for the matches of a scope map lookup. +#[derive(Debug, Default)] +pub(super) struct LookupBuffers { + probe_indices: Vec, + build_indices: Vec, +} + +/// Pairs build rows with probe rows for +/// [`mark_null_candidates_for_probe_batch`], and calls `f` with the pairs, as +/// chunks of at most `batch_size` pairs. +/// +/// `f` marks rows (see [`UnknownMarks`]), and a marked row cannot become +/// unmarked. So a cross product drops the marked rows again after each chunk. +struct UnmarkedPairs<'a, F> { + marks: UnknownMarks<'a>, 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, + build_chunk: Vec, + probe_chunk: Vec, + /// Whether the rows that the pairs decide may hold rows marked since they + /// were last refreshed. + needs_refresh: bool, + f: F, +} + +impl<'a, F> UnmarkedPairs<'a, F> +where + F: FnMut(UInt64Array, UInt32Array) -> Result<()>, +{ + fn new(marks: UnknownMarks<'a>, batch_size: usize, f: F) -> Self { + Self { + marks, batch_size, - offset, - probe_indices_buffer, - build_indices_buffer, - &mut key_comparator, - )?; + build_chunk: Vec::with_capacity(batch_size), + probe_chunk: Vec::with_capacity(batch_size), + needs_refresh: false, + f, + } + } - if !build_indices.is_empty() { - f(build_indices, probe_indices)?; + /// Pairs each of `probe_rows` with each of `build_rows`. The pairing stops + /// when every row that the pairs decide is marked. + fn cross_product( + &mut self, + build_rows: impl Iterator, + probe_rows: impl Iterator, + ) -> Result<()> { + // Drop the rows that are already marked before the first pair. + self.needs_refresh = true; + match self.marks { + UnknownMarks::Build(_) => { + self.pair_unmarked(&mut build_rows.collect(), probe_rows, |b, p| (b, p)) + } + UnknownMarks::Probe(_) => { + self.pair_unmarked(&mut probe_rows.collect(), build_rows, |p, b| (b, p)) + } } + } - let Some(next_offset) = next_offset else { - break; - }; - offset = next_offset; + /// Pairs each of `probe_rows` with the build rows that share the hash of + /// its correlation scope keys in `scope_map`. `probe_hashes` holds the + /// hashes of all rows of the probe batch. `to_build_row` translates the + /// positions stored in `scope_map` to build row indices. Rows with equal + /// hashes can have different scope keys, so `f` must still check that the + /// scope keys are equal. + /// + fn scope_matches( + &mut self, + scope_map: &dyn JoinHashMapType, + to_build_row: impl Fn(u64) -> u64, + probe_rows: impl Iterator, + probe_hashes: &[u64], + lookup: &mut LookupBuffers, + ) -> Result<()> { + let probe_rows: Vec = probe_rows.collect(); + let hashes: Vec = probe_rows + .iter() + .map(|row| probe_hashes[*row as usize]) + .collect(); + let mut offset = AtRow { row: 0 }; + loop { + let next_offset = scope_map.get_matched_indices_with_limit_offset( + &hashes, + None, + self.batch_size, + offset, + &mut lookup.probe_indices, + &mut lookup.build_indices, + ); + for i in 0..lookup.build_indices.len() { + self.push( + to_build_row(lookup.build_indices[i]), + probe_rows[lookup.probe_indices[i] as usize], + )?; + } + let Some(next_offset) = next_offset else { + return Ok(()); + }; + offset = next_offset; + } } - Ok(()) + /// Sends the pairs that are left. + fn finish(mut self) -> Result<()> { + self.send() + } + + /// Pairs each of `decided_rows` with each of `other_rows`. `to_pair` + /// orders a pair as (build row, probe row), and `decided_rows` are the + /// rows that the pairs decide (see [`UnknownMarks::unknown_row`]). + /// `decided_rows` must hold no row that was marked before the last chunk + /// was sent, unless `needs_refresh` is set. + fn pair_unmarked, O: Copy>( + &mut self, + decided_rows: &mut Vec, + other_rows: impl Iterator, + to_pair: impl Fn(D, O) -> (u64, u32), + ) -> Result<()> { + 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 self.needs_refresh { + let bitmap = self.marks.bitmap().lock(); + decided_rows.retain(|row| !bitmap.get_bit((*row).into() as usize)); + self.needs_refresh = false; + } + if decided_rows.is_empty() { + break; + } + for decided_row in decided_rows.iter() { + let (build_row, probe_row) = to_pair(*decided_row, other_row); + self.push(build_row, probe_row)?; + } + } + Ok(()) + } + + fn push(&mut self, build_row: u64, probe_row: u32) -> Result<()> { + self.build_chunk.push(build_row); + self.probe_chunk.push(probe_row); + if self.build_chunk.len() == self.batch_size { + self.send()?; + } + Ok(()) + } + + fn send(&mut self) -> Result<()> { + if self.build_chunk.is_empty() { + return Ok(()); + } + let build_chunk = + std::mem::replace(&mut self.build_chunk, Vec::with_capacity(self.batch_size)); + let probe_chunk = + std::mem::replace(&mut self.probe_chunk, Vec::with_capacity(self.batch_size)); + (self.f)(build_chunk.into(), probe_chunk.into())?; + self.needs_refresh = true; + Ok(()) + } } /// Removes the candidate pairs that decide a row already marked UNKNOWN. @@ -1063,82 +1164,3 @@ fn retain_unmarked( .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(()) -} From 5f550984356cff218f389858018cd144ee8bd0c0 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 22:22:18 -0500 Subject: [PATCH 2/3] test: pin how the scope-key null-aware lookup pairs and walks rows A unit test compares the marks of UnmarkedPairs with those of every pair, on random inputs: hash collisions between scopes, rows marked before the pairing or during it by other partitions, filters that pass none to all of the pairs, marks on build or probe rows, and batch sizes from 1 to 8192. Two tests pin the matches that the scope map lookups walk. A lookup per probe row walks every build row of its scope for every probe row, even after the rows are marked, so the walk grows with the product of the rows. The next commit changes this, and the test diff shows it. null_aware_scope_groups.slt compares LeftMark and RightAnti NOT IN queries with a reference over scopes of many rows and a residual filter that leaves some rows unmarked, for batch sizes 1, 7 and 8192 and 1 and 4 partitions. --- .../src/joins/hash_join/not_in.rs | 307 ++++++++++++++++++ .../test_files/null_aware_scope_groups.slt | 110 +++++++ 2 files changed, 417 insertions(+) create mode 100644 datafusion/sqllogictest/test_files/null_aware_scope_groups.slt diff --git a/datafusion/physical-plan/src/joins/hash_join/not_in.rs b/datafusion/physical-plan/src/joins/hash_join/not_in.rs index 4880cfdf6498..340d35b43c09 100644 --- a/datafusion/physical-plan/src/joins/hash_join/not_in.rs +++ b/datafusion/physical-plan/src/joins/hash_join/not_in.rs @@ -1000,6 +1000,9 @@ struct UnmarkedPairs<'a, F> { /// Whether the rows that the pairs decide may hold rows marked since they /// were last refreshed. needs_refresh: bool, + /// All the matches that the scope map lookups returned. + #[cfg(test)] + looked_up_matches: usize, f: F, } @@ -1014,6 +1017,8 @@ where build_chunk: Vec::with_capacity(batch_size), probe_chunk: Vec::with_capacity(batch_size), needs_refresh: false, + #[cfg(test)] + looked_up_matches: 0, f, } } @@ -1067,6 +1072,10 @@ where &mut lookup.probe_indices, &mut lookup.build_indices, ); + #[cfg(test)] + { + self.looked_up_matches += lookup.build_indices.len(); + } for i in 0..lookup.build_indices.len() { self.push( to_build_row(lookup.build_indices[i]), @@ -1164,3 +1173,301 @@ fn retain_unmarked( .unzip(); (build.into(), probe.into()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::joins::join_hash_map::JoinHashMapU64; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + /// The input of one run of [`UnmarkedPairs`]. Each row has a scope key, + /// and the hash of a key is `key % hash_modulus`, so different scopes can + /// share a hash. + #[derive(Debug)] + struct Case { + mark_build_rows: bool, + /// Whether to pair through a scope map, or as a cross product. + through_scope_map: bool, + batch_size: usize, + hash_modulus: u64, + /// The scope key of each build row. + build_keys: Vec, + /// The scope key of each probe row. + probe_keys: Vec, + /// The probe rows to pair, a subset of all the probe rows. + probe_rows: Vec, + /// The build rows to pair, a subset of all the build rows. The scope + /// map stores positions in this list, as for the NULL-valued build + /// rows. + build_rows: Vec, + /// The rows marked before the pairing starts, as by earlier probe + /// batches. + marked_before: Vec, + /// Seed of the pairs that pass the join filter. + filter_seed: u64, + /// One in this many pairs passes the join filter. `u64::MAX` passes + /// none. + filter_one_in: u64, + /// Whether each chunk of pairs also marks a random build row, as + /// another probe partition can while this one pairs. + foreign_marks: bool, + } + + impl Case { + fn random(rng: &mut StdRng) -> Self { + let mark_build_rows = rng.random_bool(0.5); + let num_build = rng.random_range(0..80); + let num_probe = rng.random_range(0..300); + // Few keys give large scopes, which make the pairing switch to + // groups. + let num_keys = rng.random_range(1..12); + let build_keys: Vec = (0..num_build) + .map(|_| rng.random_range(0..num_keys)) + .collect(); + let probe_keys: Vec = (0..num_probe) + .map(|_| rng.random_range(0..num_keys)) + .collect(); + let probe_rows = (0..num_probe as u32) + .filter(|_| rng.random_bool(0.9)) + .collect(); + let build_rows = (0..num_build as u64) + .filter(|_| rng.random_bool(0.8)) + .collect(); + let num_marked = if mark_build_rows { + num_build + } else { + num_probe + }; + let marked_before = (0..num_marked).map(|_| rng.random_bool(0.1)).collect(); + Self { + mark_build_rows, + through_scope_map: rng.random_bool(0.8), + batch_size: [1, 2, 3, 7, 64, 8192][rng.random_range(0..6)], + hash_modulus: rng.random_range(1..=num_keys), + build_keys, + probe_keys, + probe_rows, + build_rows, + marked_before, + filter_seed: rng.random(), + filter_one_in: [1, 2, 5, 50, u64::MAX][rng.random_range(0..5)], + foreign_marks: mark_build_rows && rng.random_bool(0.3), + } + } + + fn passes_filter(&self, build_row: u64, probe_row: u32) -> bool { + let mixed = (build_row ^ self.filter_seed) + .wrapping_mul(0x9E37_79B9_7F4A_7C15) + .wrapping_add(probe_row as u64) + .wrapping_mul(0xC2B2_AE3D_27D4_EB4F); + self.filter_one_in != u64::MAX + && (mixed >> 7).is_multiple_of(self.filter_one_in) + } + + fn decided_row(&self, build_row: u64, probe_row: u32) -> usize { + if self.mark_build_rows { + build_row as usize + } else { + probe_row as usize + } + } + + /// The marks after every pair of the same scope that passes the + /// filter has marked its row. + fn expected_marks(&self) -> Vec { + let mut marks = self.marked_before.clone(); + for build_row in &self.build_rows { + for probe_row in &self.probe_rows { + let same_scope = self.build_keys[*build_row as usize] + == self.probe_keys[*probe_row as usize]; + let paired = !self.through_scope_map || same_scope; + if paired && self.passes_filter(*build_row, *probe_row) { + marks[self.decided_row(*build_row, *probe_row)] = true; + } + } + } + marks + } + + /// Runs [`UnmarkedPairs`] like [`mark_null_candidates_for_probe_batch`] + /// does, and returns the marks and the pairing. + fn run(&self) -> Result<(Vec, Stats)> { + let mut bitmap = BooleanBufferBuilder::new(self.marked_before.len()); + for marked in &self.marked_before { + bitmap.append(*marked); + } + let bitmap = Mutex::new(bitmap); + let marks = if self.mark_build_rows { + UnknownMarks::Build(&bitmap) + } else { + UnknownMarks::Probe(&bitmap) + }; + let hash = |key: u64| key % self.hash_modulus; + + let mut foreign_marks = vec![]; + let mark = |build: UInt64Array, probe: UInt32Array| { + assert!(build.len() <= self.batch_size); + if self.foreign_marks { + let row = (self.filter_seed as usize) + .wrapping_add(foreign_marks.len().wrapping_mul(0x9E37_79B9)) + % self.marked_before.len(); + bitmap.lock().set_bit(row, true); + foreign_marks.push(row); + } + // Like `mark_scope_matches` and `mark`: drop the pairs of + // different scopes, then apply the filter. + let mut bitmap = bitmap.lock(); + for (build_row, probe_row) in build.values().iter().zip(probe.values()) { + let build_key = self.build_keys[*build_row as usize]; + let probe_key = self.probe_keys[*probe_row as usize]; + if self.through_scope_map { + assert_eq!(hash(build_key), hash(probe_key)); + if build_key != probe_key { + continue; + } + } + if self.passes_filter(*build_row, *probe_row) { + bitmap.set_bit(self.decided_row(*build_row, *probe_row), true); + } + } + Ok(()) + }; + + let mut pairs = UnmarkedPairs::new(marks, self.batch_size, mark); + if self.through_scope_map { + let mut scope_map = JoinHashMapU64::with_capacity(self.build_rows.len()); + let build_hashes: Vec = self + .build_rows + .iter() + .map(|row| hash(self.build_keys[*row as usize])) + .collect(); + scope_map + .update_from_iter(Box::new(build_hashes.iter().enumerate().rev()), 0); + let probe_hashes: Vec = + self.probe_keys.iter().map(|key| hash(*key)).collect(); + pairs.scope_matches( + &scope_map, + |position| self.build_rows[position as usize], + self.probe_rows.iter().copied(), + &probe_hashes, + &mut LookupBuffers::default(), + )?; + } else { + pairs.cross_product( + self.build_rows.iter().copied(), + self.probe_rows.iter().copied(), + )?; + } + let looked_up_matches = pairs.looked_up_matches; + pairs.finish()?; + let bitmap = bitmap.into_inner().finish(); + let marks = (0..bitmap.len()).map(|row| bitmap.value(row)).collect(); + Ok(( + marks, + Stats { + looked_up_matches, + foreign_marks, + }, + )) + } + } + + #[derive(Debug)] + struct Stats { + looked_up_matches: usize, + /// The rows that `foreign_marks` marked. + foreign_marks: Vec, + } + + /// Compares the marks of [`UnmarkedPairs`] with those of every pair, on + /// random inputs: hash collisions between scopes, rows marked before or + /// meanwhile by other partitions, filters that pass none to all of the + /// pairs, marks on build or probe rows, and lookups that switch to groups + /// or not. + #[test] + fn unmarked_pairs_match_brute_force() -> Result<()> { + let mut rng = StdRng::seed_from_u64(0x5C0BE); + for i in 0..3000 { + let case = Case::random(&mut rng); + let (marks, stats) = case.run()?; + let mut expected = case.expected_marks(); + for row in &stats.foreign_marks { + expected[*row] = true; + } + assert_eq!(marks, expected, "case {i}: {case:?}"); + } + Ok(()) + } + + /// One scope holds every row, and every pair passes. Returns the matches + /// that the scope map lookups walked. + fn walked_matches_in_one_scope( + mark_build_rows: bool, + num_build: usize, + num_probe: usize, + batch_size: usize, + ) -> Result { + let num_marked = if mark_build_rows { + num_build + } else { + num_probe + }; + let case = Case { + mark_build_rows, + through_scope_map: true, + batch_size, + hash_modulus: 1, + build_keys: vec![0; num_build], + probe_keys: vec![0; num_probe], + probe_rows: (0..num_probe as u32).collect(), + build_rows: (0..num_build as u64).collect(), + marked_before: vec![false; num_marked], + filter_seed: 0, + filter_one_in: 1, + foreign_marks: false, + }; + let (marks, stats) = case.run()?; + assert!(marks.iter().all(|marked| *marked)); + Ok(stats.looked_up_matches) + } + + /// With marks on the probe rows, a scope with far more build rows than + /// probe rows. A lookup per probe row walks the whole hash chain for + /// every probe row, even after the probe row is marked. + #[test] + fn marked_probe_rows_walk_whole_chains() -> Result<()> { + for batch_size in [16, 8192] { + let (num_build, num_probe) = (50_000, 500); + let walked = + walked_matches_in_one_scope(false, num_build, num_probe, batch_size)?; + assert_eq!(walked, num_build * num_probe, "batch_size {batch_size}"); + } + Ok(()) + } + + /// One scope with many rows, where every pair passes. A lookup per probe + /// row walks every build row of the scope for every probe row, even after + /// the rows are marked. So the scope map lookups walk one match per pair + /// of rows. + #[test] + fn scope_lookups_walk_every_pair() -> Result<()> { + for mark_build_rows in [true, false] { + for batch_size in [1, 16, 8192] { + let (num_build, num_probe) = (2_000, 2_000); + let walked = walked_matches_in_one_scope( + mark_build_rows, + num_build, + num_probe, + batch_size, + )?; + assert_eq!( + walked, + num_build * num_probe, + "marks on build rows: {mark_build_rows}, batch_size {batch_size}" + ); + } + } + Ok(()) + } +} diff --git a/datafusion/sqllogictest/test_files/null_aware_scope_groups.slt b/datafusion/sqllogictest/test_files/null_aware_scope_groups.slt new file mode 100644 index 000000000000..7a207f242ff8 --- /dev/null +++ b/datafusion/sqllogictest/test_files/null_aware_scope_groups.slt @@ -0,0 +1,110 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +# A correlated `NOT IN` with an equality correlation pairs rows through a +# hash map on the scope key. When a scope holds many rows, the join stops +# looking up each probe row and groups the probe rows by scope hash. These +# tests give scopes that are large enough for that switch, with a residual +# filter that leaves some rows unmarked, and compare each `NOT IN` with the +# same result computed by an outer join and an aggregate. Small batches split +# the groups and the hash chains across chunks. +# Matrix rules: no EXPLAIN, no in-file SET of a swept knob, rowsort every +# multi-row query. + +# configMatrix: datafusion.execution.batch_size=1,7,8192 +# configMatrix: datafusion.execution.target_partitions=1,4 + +statement ok +CREATE TABLE ns_outer AS SELECT + value AS id, + CASE WHEN value % 4 = 0 THEN NULL ELSE value % 9 END AS a, + CASE WHEN value % 23 = 0 THEN NULL WHEN value <= 50 THEN 1 ELSE 2 END AS b, + value % 11 AS z +FROM generate_series(1, 60); + +statement ok +CREATE TABLE ns_inner AS SELECT + CASE WHEN value % 3 = 0 THEN NULL ELSE value % 9 END AS x, + CASE WHEN value <= 35 THEN 1 ELSE 2 END AS y, + value % 13 AS w +FROM generate_series(1, 40); + +# The `IN` mark of each outer row: TRUE on an equal pair, else UNKNOWN when a +# pair in scope involves a NULL, else FALSE. `NOT IN` is its negation. +statement ok +CREATE VIEW ns_expected AS SELECT + o.id, + CASE + WHEN bool_or(i.x = o.a) THEN false + WHEN count(i.y) > 0 AND bool_or(i.x IS NULL OR o.a IS NULL) THEN NULL + ELSE true + END AS not_in +FROM ns_outer o LEFT JOIN ns_inner i ON i.y = o.b AND i.w < o.z +GROUP BY o.id, o.a; + +query BI rowsort +SELECT not_in, count(*) FROM ns_expected GROUP BY not_in; +---- +NULL 31 +false 19 +true 10 + +# `LeftMark`: the marks are on the outer (build) rows. +query I +SELECT count(*) FROM ( + SELECT id, a NOT IN ( + SELECT x FROM ns_inner i WHERE i.y = o.b AND i.w < o.z + ) AS not_in FROM ns_outer o + EXCEPT SELECT id, not_in FROM ns_expected +); +---- +0 + +query I +SELECT count(*) FROM ( + SELECT id, not_in FROM ns_expected + EXCEPT SELECT id, a NOT IN ( + SELECT x FROM ns_inner i WHERE i.y = o.b AND i.w < o.z + ) AS not_in FROM ns_outer o +); +---- +0 + +# `RightAnti`: the marks are on the outer (probe) rows. +query I rowsort +SELECT id FROM ns_outer o WHERE a NOT IN ( + SELECT x FROM ns_inner i WHERE i.y = o.b AND i.w < o.z +) +EXCEPT SELECT id FROM ns_expected WHERE not_in; +---- + +query I rowsort +SELECT id FROM ns_expected WHERE not_in +EXCEPT SELECT id FROM ns_outer o WHERE a NOT IN ( + SELECT x FROM ns_inner i WHERE i.y = o.b AND i.w < o.z +); +---- + +statement ok +DROP VIEW ns_expected; + +statement ok +DROP TABLE ns_outer; + +statement ok +DROP TABLE ns_inner; From ac9173e6456ebe722795aec448b6135a4bb23e1e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 22:24:27 -0500 Subject: [PATCH 3/3] perf: stop re-pairing marked rows in the scope-key null-aware join A correlated NOT IN with an equality correlation gets a correlation scope key. To decide which rows are UNKNOWN, the join looked up every probe row in the scope map, and paired it with every row of its scope. A marked row stays UNKNOWN, but it was still paired again, so the cost grew with the square of the scope size: SELECT count(*) FROM small_outer o WHERE o.z > 900 OR o.id_n50 NOT IN ( SELECT i.id_n0 FROM small_inner i WHERE i.k = o.k AND i.z < o.z); took about 538 ms with 100k rows in each table, against about 7 ms for a NOT IN with only the `i.z < o.z` correlation (release `datafusion-cli`, median of 10 runs). It now takes about 7 ms. The lookup drops the matches whose build row is marked before it makes pairs. That alone still walks each scope once per probe row, and a RightAnti, which marks probe rows, still paired each probe row with its whole scope. So the join counts the matches that the lookup per probe row wastes. When they cost more than a sort, it estimates from a sample of the rest of the probe rows whether groups of equal hashes save work: - with marks on build rows, when a group holds more than one probe row, which then share one walk of the hash chain; - with marks on probe rows, when the pairs of a group fill more than a chunk, so that the later chunks can skip the marked probe rows. If they do, the join sorts the rest of the probe rows by hash, looks up each hash once, and pairs each group like a cross product that drops the rows as they are marked. A group of probe rows stops its walk of the hash chain when all of them are marked. Otherwise the lookup per probe row goes on, so scopes of a few rows never pay for a sort. With marks on build rows, a check of the marks that drops nothing makes the next checks less frequent, so a filter that marks nothing costs little. Every pair now comes from a pairing that drops the marked rows, so the marking closure no longer checks them again. The unit tests pin the cost decisions and bound the matches walked, and the test diff shows that the walk no longer grows with the product of the rows. --- .../src/joins/hash_join/not_in.rs | 703 ++++++++++++++++-- 1 file changed, 648 insertions(+), 55 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/not_in.rs b/datafusion/physical-plan/src/joins/hash_join/not_in.rs index 340d35b43c09..199fd14cb0be 100644 --- a/datafusion/physical-plan/src/joins/hash_join/not_in.rs +++ b/datafusion/physical-plan/src/joins/hash_join/not_in.rs @@ -54,8 +54,8 @@ 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::join_hash_map::JoinHashMapType; -use crate::joins::join_hash_map::ProbeOffset::AtRow; +use crate::joins::join_hash_map::ProbeOffset::{AtRow, MidChain}; +use crate::joins::join_hash_map::{JoinHashMapType, ProbeOffset}; use crate::joins::probe_completion::ProbeSideSummary; use crate::joins::utils::{ BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices, @@ -791,8 +791,6 @@ fn mark_null_candidates_for_probe_batch( // 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 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, @@ -991,7 +989,8 @@ pub(super) struct LookupBuffers { /// chunks of at most `batch_size` pairs. /// /// `f` marks rows (see [`UnknownMarks`]), and a marked row cannot become -/// unmarked. So a cross product drops the marked rows again after each chunk. +/// unmarked. So the pairing drops the pairs that decide a marked row before +/// it sends them, and drops the marked rows again after each chunk. struct UnmarkedPairs<'a, F> { marks: UnknownMarks<'a>, batch_size: usize, @@ -1000,12 +999,116 @@ struct UnmarkedPairs<'a, F> { /// Whether the rows that the pairs decide may hold rows marked since they /// were last refreshed. needs_refresh: bool, + /// The lookups to skip the check of the build row marks for, after the + /// last check (see [`Self::lookup_unmarked`]). + unchecked_lookups: usize, + /// The lookups left before the next check of the build row marks. + lookups_until_check: usize, + /// The matches of the current [`Self::scope_matches`] that a lookup per + /// probe row pays for without need (see there). + wasted_matches: usize, /// All the matches that the scope map lookups returned. #[cfg(test)] looked_up_matches: usize, + /// Whether a [`Self::scope_matches`] wasted enough matches to consider + /// groups. + #[cfg(test)] + considered_groups: bool, + /// Whether a [`Self::scope_matches`] switched to groups. + #[cfg(test)] + grouped: bool, + /// The lookups with marks on the build rows. + #[cfg(test)] + build_lookups: usize, + /// The lookups that checked the build row marks. + #[cfg(test)] + mark_checks: usize, f: F, } +/// A sort costs about as much as this many wasted matches per probe row (see +/// [`UnmarkedPairs::scope_matches`]). +const WASTED_MATCHES_PER_SORTED_ROW: usize = 8; + +/// The most lookups that skip the check of the build row marks in a row (see +/// [`UnmarkedPairs::lookup_unmarked`]). +const MAX_UNCHECKED_LOOKUPS: usize = 32; + +/// The probe rows sampled to estimate the groups of equal hashes before a +/// sort (see [`UnmarkedPairs::scope_matches`]). +const GROUP_SAMPLE_ROWS: usize = 256; + +/// Estimates how many different values `hashes` holds, from a sample of at +/// most [`GROUP_SAMPLE_ROWS`] of them. The count is exact when `hashes` fits +/// in the sample. +fn estimate_groups(hashes: &[u64]) -> usize { + let mut sample: Vec = if hashes.len() <= GROUP_SAMPLE_ROWS { + hashes.to_vec() + } else { + // Pseudo-random positions, so that the sample sees a hash twice about + // as often as a random sample would. A fixed stride follows the + // period of keys such as `value % n`. + (0..GROUP_SAMPLE_ROWS as u64) + .map(|i| hashes[(splitmix64(i) % hashes.len() as u64) as usize]) + .collect() + }; + sample.sort_unstable(); + let mut seen = 0; + let mut seen_once: usize = 0; + let mut seen_twice = 0; + for run in sample.chunk_by(|a, b| a == b) { + seen += 1; + match run.len() { + 1 => seen_once += 1, + 2 => seen_twice += 1, + _ => {} + } + } + if hashes.len() <= GROUP_SAMPLE_ROWS { + return seen; + } + // The bias-corrected Chao1 estimate: the hashes that the sample saw once + // suggest how many it did not see. + let unseen = seen_once * seen_once.saturating_sub(1) / (2 * (seen_twice + 1)); + (seen + unseen).min(hashes.len()) +} + +/// The SplitMix64 mix of `x`, a fixed pseudo-random sequence. +fn splitmix64(x: u64) -> u64 { + let x = x.wrapping_add(0x9E37_79B9_7F4A_7C15); + let x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + let x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + x ^ (x >> 31) +} + +/// Whether `rows` probe rows in `groups` groups of equal hashes save work when +/// each group is paired at once (see [`UnmarkedPairs::scope_matches`]). +/// `matches_per_row` estimates the build rows of a hash. +fn groups_save_work( + mark_build_rows: bool, + rows: usize, + groups: usize, + matches_per_row: usize, + batch_size: usize, +) -> bool { + if mark_build_rows { + rows >= 2 * groups + } else { + rows * matches_per_row > 2 * batch_size * groups + } +} + +/// The lookups to skip the check of the build row marks for, after a check +/// that followed `unchecked` skipped lookups, and that dropped a match or not +/// (see [`UnmarkedPairs::lookup_unmarked`]). +fn next_unchecked_lookups(unchecked: usize, dropped: bool) -> usize { + if dropped { + 0 + } else { + (2 * unchecked).clamp(1, MAX_UNCHECKED_LOOKUPS) + } +} + impl<'a, F> UnmarkedPairs<'a, F> where F: FnMut(UInt64Array, UInt32Array) -> Result<()>, @@ -1017,8 +1120,19 @@ where build_chunk: Vec::with_capacity(batch_size), probe_chunk: Vec::with_capacity(batch_size), needs_refresh: false, + unchecked_lookups: 0, + lookups_until_check: 0, + wasted_matches: 0, #[cfg(test)] looked_up_matches: 0, + #[cfg(test)] + considered_groups: false, + #[cfg(test)] + grouped: false, + #[cfg(test)] + build_lookups: 0, + #[cfg(test)] + mark_checks: 0, f, } } @@ -1049,6 +1163,26 @@ where /// hashes can have different scope keys, so `f` must still check that the /// scope keys are equal. /// + /// The pairing starts with a lookup per probe row, like the join itself. + /// But each probe row finds all the build rows in its scope: + /// - When the marks are on the build rows, a marked build row costs a + /// little for every later probe row in its scope. + /// - When the marks are on the probe rows, a probe row is paired with all + /// the build rows in its scope, even when its first pair marks it. + /// + /// When those costs add up to more than a sort of the probe rows, the + /// join considers groups of the rest of the probe rows with equal hashes. + /// A sample of the hashes estimates the groups first. If the groups can + /// save work, the join sorts the rest of the probe rows by hash. Each hash + /// is then looked up once, and its group is paired like a cross product, + /// which drops the rows as they are marked. A group saves work when: + /// - with marks on the build rows, it holds more than one probe row, which + /// then share one walk of the hash chain; + /// - with marks on the probe rows, its pairs fill more than a chunk. + /// Within a chunk every pair is sent, so only the later chunks can skip + /// the marked probe rows. + /// + /// Otherwise the lookup per probe row goes on. fn scope_matches( &mut self, scope_map: &dyn JoinHashMapType, @@ -1058,14 +1192,221 @@ where lookup: &mut LookupBuffers, ) -> Result<()> { let probe_rows: Vec = probe_rows.collect(); - let hashes: Vec = probe_rows - .iter() - .map(|row| probe_hashes[*row as usize]) - .collect(); + // `probe_rows` is increasing, so it holds every row exactly when it is + // as long as the batch. Then the positions are the probe rows. + let all_rows = probe_rows.len() == probe_hashes.len(); + let gathered: Vec; + let hashes = if all_rows { + probe_hashes + } else { + gathered = probe_rows + .iter() + .map(|row| probe_hashes[*row as usize]) + .collect(); + &gathered + }; + + self.wasted_matches = 0; + let max_wasted_matches = WASTED_MATCHES_PER_SORTED_ROW * probe_rows.len(); + let Some((offset, matches)) = self.lookup_each_row( + scope_map, + hashes, + (!all_rows).then_some(probe_rows.as_slice()), + AtRow { row: 0 }, + &to_build_row, + max_wasted_matches, + lookup, + )? + else { + return Ok(()); + }; + // Group the probe rows from the one being looked up. Its build rows + // that were already paired are paired again, which changes no mark. + let group_from = match offset { + AtRow { row } | MidChain { row, .. } => row, + }; + + #[cfg(test)] + { + self.considered_groups = true; + } + let mark_build_rows = matches!(self.marks, UnknownMarks::Build(_)); + // The matches per probe row so far estimate the build rows of a hash. + let matches_per_row = matches / group_from.max(1); + let saves_work = |rows: usize, groups: usize| { + groups_save_work( + mark_build_rows, + rows, + groups, + matches_per_row, + self.batch_size, + ) + }; + let rest_hashes = &hashes[group_from..]; + // Estimate the groups from a sample first, because a sort of all the + // rest that finds no group worth it costs as much as the wasted + // matches did. + let worth_sorting = saves_work(rest_hashes.len(), estimate_groups(rest_hashes)); + let mut rest: Vec<(u64, u32)> = vec![]; + let mut groups = vec![]; + if worth_sorting { + rest = rest_hashes + .iter() + .copied() + .zip(probe_rows[group_from..].iter().copied()) + .collect(); + rest.sort_unstable_by_key(|(hash, _)| *hash); + groups = rest.chunk_by(|a, b| a.0 == b.0).collect(); + } + if !worth_sorting || !saves_work(rest.len(), groups.len()) { + self.lookup_each_row( + scope_map, + hashes, + (!all_rows).then_some(probe_rows.as_slice()), + offset, + &to_build_row, + usize::MAX, + lookup, + )?; + return Ok(()); + } + + #[cfg(test)] + { + self.grouped = true; + } + // Send the pairs that are left, so the lookups below see their marks. + self.send()?; + let group_hashes: Vec = groups.iter().map(|group| group[0].0).collect(); + match self.marks { + UnknownMarks::Build(_) => self.pair_build_groups( + scope_map, + &groups, + &group_hashes, + &to_build_row, + lookup, + ), + UnknownMarks::Probe(_) => { + for (group, hash) in groups.iter().zip(&group_hashes) { + self.pair_probe_group( + scope_map, + group, + *hash, + &to_build_row, + lookup, + )?; + } + Ok(()) + } + } + } + + /// Looks up each of `probe_rows`, whose hashes are `hashes`, from `offset`, + /// and pairs it with the unmarked build rows of its hash. `None` stands + /// for every row of the probe batch. Stops once more + /// than `max_wasted_matches` matches are wasted, and then returns where + /// the lookup stopped and how many matches it walked. + #[expect(clippy::too_many_arguments)] + fn lookup_each_row( + &mut self, + scope_map: &dyn JoinHashMapType, + hashes: &[u64], + probe_rows: Option<&[u32]>, + mut offset: ProbeOffset, + to_build_row: impl Fn(u64) -> u64, + max_wasted_matches: usize, + lookup: &mut LookupBuffers, + ) -> Result> { + let mut matches = 0; + loop { + let next_offset = self.lookup_unmarked( + scope_map, + hashes, + |position| match probe_rows { + Some(rows) => rows[position as usize], + None => position, + }, + offset, + &to_build_row, + &mut matches, + lookup, + ); + self.push_all(&lookup.build_indices, &lookup.probe_indices)?; + let Some(next_offset) = next_offset else { + return Ok(None); + }; + offset = next_offset; + if self.wasted_matches > max_wasted_matches { + return Ok(Some((offset, matches))); + } + } + } + + /// Pairs each group of probe rows with the build rows of its hash, when + /// the marks are on the build rows. The groups are looked up together, + /// and the build rows of a group are dropped as they are marked. + fn pair_build_groups( + &mut self, + scope_map: &dyn JoinHashMapType, + groups: &[&[(u64, u32)]], + group_hashes: &[u64], + to_build_row: impl Fn(u64) -> u64, + lookup: &mut LookupBuffers, + ) -> Result<()> { + let mut build_rows = vec![]; + let mut offset = AtRow { row: 0 }; + loop { + // A lookup chunk holds the build rows of one or more groups, in + // group order. + let next_offset = self.lookup_unmarked( + scope_map, + group_hashes, + |group| group, + offset, + &to_build_row, + &mut 0, + lookup, + ); + let mut matches = lookup.build_indices.as_slice(); + for group in lookup.probe_indices.chunk_by(|a, b| a == b) { + let group_matches; + (group_matches, matches) = matches.split_at(group.len()); + build_rows.clear(); + build_rows.extend_from_slice(group_matches); + // A build row belongs to one group only, so the marks made + // for the earlier groups do not concern these build rows. + self.needs_refresh = false; + self.pair_unmarked( + &mut build_rows, + groups[group[0] as usize].iter().map(|(_, row)| *row), + |b, p| (b, p), + )?; + } + let Some(next_offset) = next_offset else { + return Ok(()); + }; + offset = next_offset; + } + } + + /// Pairs one group of probe rows with the build rows of its hash, when + /// the marks are on the probe rows. The probe rows of the group are + /// dropped as they are marked, and the walk of the hash chain stops when + /// none is left. + fn pair_probe_group( + &mut self, + scope_map: &dyn JoinHashMapType, + group: &[(u64, u32)], + hash: u64, + to_build_row: impl Fn(u64) -> u64, + lookup: &mut LookupBuffers, + ) -> Result<()> { + let mut probe_rows: Vec = group.iter().map(|(_, row)| *row).collect(); + let mut build_rows = vec![]; let mut offset = AtRow { row: 0 }; loop { let next_offset = scope_map.get_matched_indices_with_limit_offset( - &hashes, + &[hash], None, self.batch_size, offset, @@ -1076,11 +1417,15 @@ where { self.looked_up_matches += lookup.build_indices.len(); } - for i in 0..lookup.build_indices.len() { - self.push( - to_build_row(lookup.build_indices[i]), - probe_rows[lookup.probe_indices[i] as usize], - )?; + build_rows.clear(); + build_rows.extend(lookup.build_indices.iter().map(|row| to_build_row(*row))); + // Drop the probe rows that the earlier chunks marked. + self.needs_refresh = true; + self.pair_unmarked(&mut probe_rows, build_rows.iter().copied(), |p, b| { + (b, p) + })?; + if probe_rows.is_empty() { + return Ok(()); } let Some(next_offset) = next_offset else { return Ok(()); @@ -1089,6 +1434,95 @@ where } } + /// Looks up the next chunk of matches of `hashes` in `scope_map`, from + /// `offset`, into the buffers: the position in `hashes` translated with + /// `to_probe_row`, and the build row translated with `to_build_row`. Drops + /// the matches that decide a marked row, counts the wasted matches (see + /// [`Self::scope_matches`]), and adds the matches walked to `matches`. + #[expect(clippy::too_many_arguments)] + fn lookup_unmarked( + &mut self, + scope_map: &dyn JoinHashMapType, + hashes: &[u64], + to_probe_row: impl Fn(u32) -> u32, + offset: ProbeOffset, + to_build_row: impl Fn(u64) -> u64, + matches: &mut usize, + lookup: &mut LookupBuffers, + ) -> Option { + let LookupBuffers { + probe_indices, + build_indices, + } = lookup; + let next_offset = scope_map.get_matched_indices_with_limit_offset( + hashes, + None, + self.batch_size, + offset, + probe_indices, + build_indices, + ); + for build_row in build_indices.iter_mut() { + *build_row = to_build_row(*build_row); + } + for probe_row in probe_indices.iter_mut() { + *probe_row = to_probe_row(*probe_row); + } + let len = build_indices.len(); + // With marks on the probe rows, only the lookup per probe row comes + // here, and the matches of a probe row are all pushed before any of + // them is sent. So the check could drop only the rest of a hash chain + // that spans lookup chunks, and it is skipped. + // + // With marks on the build rows, a check that drops nothing makes the + // next checks less frequent, up to one in `MAX_UNCHECKED_LOOKUPS + 1` + // lookups. The pairs of a marked build row that a skipped check lets + // through only cost a filter evaluation, and change no mark. + if let UnknownMarks::Build(bitmap) = self.marks { + #[cfg(test)] + { + self.build_lookups += 1; + } + if self.lookups_until_check > 0 { + self.lookups_until_check -= 1; + } else { + let bitmap = bitmap.lock(); + let mut kept = 0; + for i in 0..len { + let (build, probe) = (build_indices[i], probe_indices[i]); + if !bitmap.get_bit(build as usize) { + build_indices[kept] = build; + probe_indices[kept] = probe; + kept += 1; + } + } + build_indices.truncate(kept); + probe_indices.truncate(kept); + self.unchecked_lookups = + next_unchecked_lookups(self.unchecked_lookups, kept < len); + #[cfg(test)] + { + self.mark_checks += 1; + } + self.lookups_until_check = self.unchecked_lookups; + } + } + *matches += len; + #[cfg(test)] + { + self.looked_up_matches += len; + } + self.wasted_matches += match self.marks { + // The matches whose build row is marked. + UnknownMarks::Build(_) => len - build_indices.len(), + // Every match pairs its probe row again, unless it is the first + // match of that probe row. Counting all matches overcounts by at + // most one per probe row. + UnknownMarks::Probe(_) => len, + }; + next_offset + } + /// Sends the pairs that are left. fn finish(mut self) -> Result<()> { self.send() @@ -1124,6 +1558,24 @@ where Ok(()) } + /// Pushes the pairs of `build_rows` and `probe_rows`. + fn push_all(&mut self, mut build_rows: &[u64], mut probe_rows: &[u32]) -> Result<()> { + while !build_rows.is_empty() { + let n = build_rows + .len() + .min(self.batch_size - self.build_chunk.len()); + let (build_now, build_later) = build_rows.split_at(n); + let (probe_now, probe_later) = probe_rows.split_at(n); + self.build_chunk.extend_from_slice(build_now); + self.probe_chunk.extend_from_slice(probe_now); + (build_rows, probe_rows) = (build_later, probe_later); + if self.build_chunk.len() == self.batch_size { + self.send()?; + } + } + Ok(()) + } + fn push(&mut self, build_row: u64, probe_row: u32) -> Result<()> { self.build_chunk.push(build_row); self.probe_chunk.push(probe_row); @@ -1147,33 +1599,6 @@ where } } -/// 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()) -} - #[cfg(test)] mod tests { use super::*; @@ -1359,14 +1784,20 @@ mod tests { self.probe_rows.iter().copied(), )?; } + let (considered_groups, grouped) = (pairs.considered_groups, pairs.grouped); let looked_up_matches = pairs.looked_up_matches; + let (build_lookups, mark_checks) = (pairs.build_lookups, pairs.mark_checks); pairs.finish()?; let bitmap = bitmap.into_inner().finish(); let marks = (0..bitmap.len()).map(|row| bitmap.value(row)).collect(); Ok(( marks, Stats { + considered_groups, + grouped, looked_up_matches, + build_lookups, + mark_checks, foreign_marks, }, )) @@ -1375,7 +1806,11 @@ mod tests { #[derive(Debug)] struct Stats { + considered_groups: bool, + grouped: bool, looked_up_matches: usize, + build_lookups: usize, + mark_checks: usize, /// The rows that `foreign_marks` marked. foreign_marks: Vec, } @@ -1388,6 +1823,7 @@ mod tests { #[test] fn unmarked_pairs_match_brute_force() -> Result<()> { let mut rng = StdRng::seed_from_u64(0x5C0BE); + let mut grouped = [0; 2]; for i in 0..3000 { let case = Case::random(&mut rng); let (marks, stats) = case.run()?; @@ -1396,6 +1832,73 @@ mod tests { expected[*row] = true; } assert_eq!(marks, expected, "case {i}: {case:?}"); + if stats.grouped { + grouped[case.mark_build_rows as usize] += 1; + } + } + // Both kinds of marks reach the grouped pairing. + assert!(grouped.iter().all(|n| *n > 100), "{grouped:?}"); + Ok(()) + } + + /// Every scope holds `rows_per_scope` build rows and as many probe rows. + /// Returns the pairing of all of them. + fn equal_scopes( + mark_build_rows: bool, + rows_per_scope: usize, + num_scopes: usize, + filter_one_in: u64, + ) -> Result { + let num_rows = rows_per_scope * num_scopes; + let keys: Vec = (0..num_rows as u64) + .map(|row| row % num_scopes as u64) + .collect(); + let case = Case { + mark_build_rows, + through_scope_map: true, + batch_size: 8192, + hash_modulus: num_scopes as u64, + build_keys: keys.clone(), + probe_keys: keys, + probe_rows: (0..num_rows as u32).collect(), + build_rows: (0..num_rows as u64).collect(), + marked_before: vec![false; num_rows], + filter_seed: 0, + filter_one_in, + foreign_marks: false, + }; + let (marks, stats) = case.run()?; + assert_eq!(marks, case.expected_marks()); + Ok(stats) + } + + /// Scopes of a few rows cost less to look up per probe row than to sort. + /// So the pairing must not switch to groups for them, whether every pair + /// marks its row or none does. + #[test] + fn small_scopes_are_not_grouped() -> Result<()> { + for (mark_build_rows, rows_per_scope, filter_one_in) in [ + (true, 1, 1), + (true, 5, u64::MAX), + (false, 1, 1), + (false, 5, 1), + (false, 5, u64::MAX), + ] { + let stats = + equal_scopes(mark_build_rows, rows_per_scope, 400, filter_one_in)?; + assert!( + !stats.considered_groups, + "marks on build rows: {mark_build_rows}, {rows_per_scope} rows per scope, one in {filter_one_in} pairs passes" + ); + } + // Probe rows marked by their first pair waste the other matches of + // their scope. With 16 build rows per scope that is enough to consider + // groups, but a group of 16 probe rows fills no chunk. + let stats = equal_scopes(false, 16, 400, 1)?; + assert!(stats.considered_groups && !stats.grouped); + // A large scope does switch. + for mark_build_rows in [true, false] { + assert!(equal_scopes(mark_build_rows, 400, 1, 1)?.grouped); } Ok(()) } @@ -1432,26 +1935,110 @@ mod tests { Ok(stats.looked_up_matches) } + #[test] + fn estimate_groups_from_a_sample() { + let estimate = |hashes: Vec| estimate_groups(&hashes); + // Exact when the rows fit in the sample. + assert_eq!(estimate(vec![]), 0); + assert_eq!(estimate((0..100).collect()), 100); + assert_eq!(estimate((0..256).map(|row| row % 7).collect()), 7); + // A sample that sees every hash several times. + assert_eq!(estimate(vec![7; 8192]), 1); + assert_eq!(estimate((0..8192).map(|row| row % 4).collect()), 4); + // The sample is fixed, so the estimates are too. + let periodic = |rows: u64, groups: u64| { + estimate((0..rows).map(|row| row % groups).collect()) + }; + assert_eq!(periodic(1000, 300), 262); + assert_eq!(periodic(8192, 1000), 698); + assert_eq!(periodic(8192, 100), 100); + assert_eq!(periodic(8192, 8192), 6377); + // Between those, and with periodic keys as from `value % n`, the + // estimate stays within a factor of two. + for (rows, groups) in [ + (1000, 300), + (8192, 1000), + (8192, 100), + (8192, 32), + (8192, 8192), + ] { + let estimated = estimate((0..rows).map(|row| row % groups).collect()); + assert!( + (groups / 2..=2 * groups).contains(&(estimated as u64)), + "{rows} rows, {groups} groups: {estimated}" + ); + } + } + + #[test] + fn groups_save_work_when_they_share_walks_or_chunks() { + // Marks on build rows: a group shares one walk of the hash chain. + assert!(groups_save_work(true, 10, 5, 0, 8192)); + assert!(!groups_save_work(true, 9, 5, 1_000, 8192)); + // Marks on probe rows: the pairs of a group must fill more than a + // chunk, here 4 groups of 2048 probe rows. + assert!(!groups_save_work(false, 8192, 4, 8, 8192)); + assert!(groups_save_work(false, 8192, 4, 9, 8192)); + assert!(groups_save_work(false, 8192, 4, 2, 1024)); + assert!(!groups_save_work(false, 8192, 4, 2, 4096)); + } + + #[test] + fn mark_checks_back_off_while_they_drop_nothing() { + assert_eq!(next_unchecked_lookups(0, false), 1); + assert_eq!(next_unchecked_lookups(1, false), 2); + assert_eq!(next_unchecked_lookups(5, false), 10); + assert_eq!(next_unchecked_lookups(20, false), MAX_UNCHECKED_LOOKUPS); + assert_eq!( + next_unchecked_lookups(MAX_UNCHECKED_LOOKUPS, false), + MAX_UNCHECKED_LOOKUPS + ); + assert_eq!(next_unchecked_lookups(MAX_UNCHECKED_LOOKUPS, true), 0); + assert_eq!(next_unchecked_lookups(0, true), 0); + } + + /// With marks on the build rows and a filter that no pair passes, no + /// build row is ever marked. So the lookups must skip most checks of the + /// marks. + #[test] + fn lookups_skip_checks_while_nothing_is_marked() -> Result<()> { + let stats = equal_scopes(true, 2_000, 1, u64::MAX)?; + let (checks, lookups) = (stats.mark_checks, stats.build_lookups); + assert!(lookups > 100, "{lookups} lookups"); + assert!(checks > 0); + assert!( + checks * MAX_UNCHECKED_LOOKUPS / 2 <= lookups, + "{checks} checks in {lookups} lookups" + ); + Ok(()) + } + /// With marks on the probe rows, a scope with far more build rows than - /// probe rows. A lookup per probe row walks the whole hash chain for - /// every probe row, even after the probe row is marked. + /// probe rows: once every probe row is marked, the walk of the hash chain + /// must stop. So the walk stays within the matches before the switch to + /// groups and a few lookup chunks, whatever the number of build rows. #[test] - fn marked_probe_rows_walk_whole_chains() -> Result<()> { + fn marked_probe_groups_stop_walking() -> Result<()> { for batch_size in [16, 8192] { let (num_build, num_probe) = (50_000, 500); let walked = walked_matches_in_one_scope(false, num_build, num_probe, batch_size)?; - assert_eq!(walked, num_build * num_probe, "batch_size {batch_size}"); + let bound = WASTED_MATCHES_PER_SORTED_ROW * num_probe + 4 * batch_size; + assert!( + walked <= bound, + "batch_size {batch_size}: {walked} matches walked, bound {bound}" + ); } Ok(()) } - /// One scope with many rows, where every pair passes. A lookup per probe - /// row walks every build row of the scope for every probe row, even after - /// the rows are marked. So the scope map lookups walk one match per pair - /// of rows. + /// One scope with many rows, where every pair passes: once a row is + /// marked, it must not be paired again. So the scope map lookups walk + /// about as many matches as there are rows, not one per pair of rows. A + /// lookup per probe row walks every build row of the scope for every + /// probe row. #[test] - fn scope_lookups_walk_every_pair() -> Result<()> { + fn scope_lookups_stay_linear_in_rows() -> Result<()> { for mark_build_rows in [true, false] { for batch_size in [1, 16, 8192] { let (num_build, num_probe) = (2_000, 2_000); @@ -1461,10 +2048,16 @@ mod tests { num_probe, batch_size, )?; - assert_eq!( - walked, - num_build * num_probe, - "marks on build rows: {mark_build_rows}, batch_size {batch_size}" + // The checks of the build row marks that are skipped can + // delay the switch to groups by that many lookups. + let bound = WASTED_MATCHES_PER_SORTED_ROW * num_probe + + 2 * num_build + + (MAX_UNCHECKED_LOOKUPS + 2) * batch_size; + // Every build row is found at least once. + assert!(walked >= num_build, "{walked} matches walked"); + assert!( + walked <= bound, + "marks on build rows: {mark_build_rows}, batch_size {batch_size}: {walked} matches walked, bound {bound}" ); } }