Copilot commented on code in PR #25679:
URL: https://github.com/apache/datafusion/pull/25679#discussion_r4372083060
##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -991,50 +979,281 @@ 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<u32>,
+ build_indices: Vec<u64>,
+}
+
+/// 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<u32>,
- build_indices_buffer: &mut Vec<u64>,
- 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<u64>,
+ probe_chunk: Vec<u32>,
+ /// Whether `f` was called since the rows that the pairs decide were last
+ /// refreshed.
+ marks_since_refresh: bool,
+ /// The matches that [`Self::scope_matches`] dropped because the row they
+ /// decide was marked.
+ skipped_matches: usize,
+ 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),
+ marks_since_refresh: false,
+ skipped_matches: 0,
+ 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<Item = u64>,
+ probe_rows: impl Iterator<Item = u32>,
+ ) -> Result<()> {
+ // Drop the rows that are already marked before the first pair.
+ self.marks_since_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.
+ /// When the marks are on the build rows, each probe row still finds all
+ /// the build rows in its scope, marked or not. So a marked build row costs
+ /// a little for every later probe row in its scope. When those costs add
+ /// up to more than a sort of the probe rows, the rest of the probe rows
+ /// are grouped by hash. Each hash is then looked up once, and its group
+ /// is paired like a cross product that drops the build rows as they are
+ /// marked.
+ fn scope_matches(
+ &mut self,
+ scope_map: &dyn JoinHashMapType,
+ to_build_row: impl Fn(u64) -> u64,
+ probe_rows: impl Iterator<Item = u32>,
+ probe_hashes: &[u64],
+ lookup: &mut LookupBuffers,
+ ) -> Result<()> {
+ let probe_rows: Vec<u32> = probe_rows.collect();
+ let hashes: Vec<u64> = probe_rows
+ .iter()
+ .map(|row| probe_hashes[*row as usize])
+ .collect();
- let Some(next_offset) = next_offset else {
- break;
+ // A sort costs about as much as this many skipped matches per probe
+ // row. A probe row is looked up once, so the marks on probe rows
+ // never make grouping worth it.
+ const SKIPPED_MATCHES_PER_SORTED_ROW: usize = 8;
+ let max_skipped_matches = match self.marks {
+ UnknownMarks::Build(_) => SKIPPED_MATCHES_PER_SORTED_ROW *
probe_rows.len(),
+ UnknownMarks::Probe(_) => usize::MAX,
};
- offset = next_offset;
+ self.skipped_matches = 0;
+ let mut offset = AtRow { row: 0 };
+ let group_from = loop {
+ let next_offset = self.lookup_unmarked(
+ scope_map,
+ &hashes,
+ |position| probe_rows[position as usize],
+ offset,
+ &to_build_row,
+ lookup,
+ );
+ for (probe_row, build_row) in
+ lookup.probe_indices.iter().zip(lookup.build_indices.iter())
+ {
+ self.push(*build_row, *probe_row)?;
+ }
+ let Some(next_offset) = next_offset else {
+ return Ok(());
+ };
+ offset = next_offset;
+ if self.skipped_matches > max_skipped_matches {
+ break match next_offset {
+ AtRow { row } | MidChain { row, .. } => row,
+ };
+ }
+ };
+
+ // Send the pairs that are left, so the lookups below see their marks.
+ self.send()?;
+
+ // 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 mut rest: Vec<(u64, u32)> = hashes[group_from..]
+ .iter()
+ .copied()
+ .zip(probe_rows[group_from..].iter().copied())
+ .collect();
+ rest.sort_unstable_by_key(|(hash, _)| *hash);
+ let groups: Vec<_> = rest.chunk_by(|a, b| a.0 == b.0).collect();
+ let group_hashes: Vec<u64> = groups.iter().map(|group|
group[0].0).collect();
+
+ 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,
+ 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.marks_since_refresh = false;
+ self.pair_unmarked(
+ &mut build_rows,
+ groups[group[0] as usize].iter().map(|(_, row)| *row),
+ |b, p| (b, p),
+ )?;
Review Comment:
The new regression marks all build rows before grouping starts, so it never
exercises grouped pairing with unmarked candidates; its assertions also pass
with grouping disabled. Add a selective-marking case that crosses the switch
threshold while some rows remain unmarked, spans lookup chunk boundaries, and
compares the final bitmap with exhaustive scope matching. Also count raw lookup
matches to verify that repeated hash-chain traversal stays bounded, since
emitted-pair counts alone do not detect that regression.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
