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..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,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::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, 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,9 +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 (build_indices, probe_indices) = - retain_unmarked(marks, build_indices, probe_indices); + let mark = |build_indices: UInt64Array, probe_indices: UInt32Array| { let (build_indices, probe_indices) = if num_value_keys > 1 { retain_value_mismatch_free( build_value_keys, @@ -841,44 +828,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 +892,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,154 +977,1090 @@ 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 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, - 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( + 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, + /// 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<()>, +{ + fn new(marks: UnknownMarks<'a>, batch_size: usize, f: F) -> Self { + Self { + marks, + batch_size, + 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, + } + } + + /// 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)) + } + } + } + + /// 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. + /// + /// 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, + 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(); + // `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, - build_scope_values, - probe_scope_values, - NullEquality::NullEqualsNothing, - hashes_buffer, + 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( + &[hash], + None, + self.batch_size, + offset, + &mut lookup.probe_indices, + &mut lookup.build_indices, + ); + #[cfg(test)] + { + self.looked_up_matches += lookup.build_indices.len(); + } + 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(()); + }; + offset = next_offset; + } + } + + /// 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, - batch_size, + self.batch_size, offset, - probe_indices_buffer, - build_indices_buffer, - &mut key_comparator, - )?; + 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 + } - if !build_indices.is_empty() { - f(build_indices, probe_indices)?; + /// 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(()) + } - let Some(next_offset) = next_offset else { - break; - }; - offset = next_offset; + /// 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(()) } - 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(()) + } -/// 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); + 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(()) } - 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()), - ), +#[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, } -} -/// 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)); - }; + 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) + } - 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; + 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 + } } - if decided_rows.is_empty() { - break; + + /// 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 } - 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)), + + /// 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(), )?; - marks_since_refresh = true; } + 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, + }, + )) } } - if !decided_chunk.is_empty() { - f(decided_chunk, other_chunk)?; + + #[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, + } + + /// 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); + let mut grouped = [0; 2]; + 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:?}"); + 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(()) + } + + /// 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) + } + + #[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: 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_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)?; + 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: 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_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); + let walked = walked_matches_in_one_scope( + mark_build_rows, + num_build, + num_probe, + 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}" + ); + } + } + Ok(()) } - 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;