Skip to content

Commit 440b285

Browse files
committed
perf: stop re-pairing marked build rows in the scope-key null-aware join
A correlated NOT IN with an equality correlation gets a correlation scope key. To decide which build rows are UNKNOWN, the join looked up every probe row in the scope map and removed the build rows that were already marked only after the lookup had produced the pairs. A marked build row stays UNKNOWN, but it was still paired with every later probe row in its scope, so the cost grew with the square of the scope size: SELECT count(*) FROM outer_t o WHERE o.z > 900 OR o.id_n50 NOT IN ( SELECT i.id_n0 FROM inner_t i WHERE i.k = o.k AND i.z < o.z); took about 415 ms with 100k rows in each table, against about 13 ms for a NOT IN with only the `i.z < o.z` correlation (release `datafusion-cli`, median of 10 runs). Now the lookup drops the marked build rows before it makes pairs. That alone still walks each scope once per probe row. So when the skipped matches of a probe batch cost more than a sort, the rest of the probe rows are grouped by scope hash: each hash is looked up once, and the group is paired like the cross product of the path without scope keys, which drops the marked build rows after each chunk and stops when none is left. The switch keeps the plain lookup for scopes that hold about one build row, where a sort would only add cost. The query above now takes about 10 ms. The pairs made for one build row no longer grow with the number of probe rows, so the unit test now checks that.
1 parent fddfabb commit 440b285

1 file changed

Lines changed: 141 additions & 20 deletions

File tree

  • datafusion/physical-plan/src/joins/hash_join

‎datafusion/physical-plan/src/joins/hash_join/not_in.rs‎

Lines changed: 141 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ use datafusion_physical_expr_common::datum::apply_cmp;
5454
use crate::joins::PartitionMode;
5555
use crate::joins::SharedBitmapBuilder;
5656
use crate::joins::hash_join::exec::{JoinLeftData, new_join_hashmap};
57-
use crate::joins::join_hash_map::JoinHashMapType;
58-
use crate::joins::join_hash_map::ProbeOffset::AtRow;
57+
use crate::joins::join_hash_map::ProbeOffset::{AtRow, MidChain};
58+
use crate::joins::join_hash_map::{JoinHashMapType, ProbeOffset};
5959
use crate::joins::probe_completion::ProbeSideSummary;
6060
use crate::joins::utils::{
6161
BuildProbeJoinMetrics, JoinFilter, apply_join_filter_to_indices,
@@ -991,7 +991,8 @@ pub(super) struct LookupBuffers {
991991
/// chunks of at most `batch_size` pairs.
992992
///
993993
/// `f` marks rows (see [`UnknownMarks`]), and a marked row cannot become
994-
/// unmarked. So a cross product drops the marked rows again after each chunk.
994+
/// unmarked. So the pairing drops the pairs that decide a marked row before
995+
/// it sends them, and drops the marked rows again after each chunk.
995996
struct UnmarkedPairs<'a, F> {
996997
marks: UnknownMarks<'a>,
997998
batch_size: usize,
@@ -1000,6 +1001,9 @@ struct UnmarkedPairs<'a, F> {
10001001
/// Whether `f` was called since the rows that the pairs decide were last
10011002
/// refreshed.
10021003
marks_since_refresh: bool,
1004+
/// The matches that [`Self::scope_matches`] dropped because the row they
1005+
/// decide was marked.
1006+
skipped_matches: usize,
10031007
f: F,
10041008
}
10051009

@@ -1014,6 +1018,7 @@ where
10141018
build_chunk: Vec::with_capacity(batch_size),
10151019
probe_chunk: Vec::with_capacity(batch_size),
10161020
marks_since_refresh: false,
1021+
skipped_matches: 0,
10171022
f,
10181023
}
10191024
}
@@ -1044,6 +1049,14 @@ where
10441049
/// hashes can have different scope keys, so `f` must still check that the
10451050
/// scope keys are equal.
10461051
///
1052+
/// The pairing starts with a lookup per probe row, like the join itself.
1053+
/// When the marks are on the build rows, each probe row still finds all
1054+
/// the build rows in its scope, marked or not. So a marked build row costs
1055+
/// a little for every later probe row in its scope. When those costs add
1056+
/// up to more than a sort of the probe rows, the rest of the probe rows
1057+
/// are grouped by hash. Each hash is then looked up once, and its group
1058+
/// is paired like a cross product that drops the build rows as they are
1059+
/// marked.
10471060
fn scope_matches(
10481061
&mut self,
10491062
scope_map: &dyn JoinHashMapType,
@@ -1057,20 +1070,82 @@ where
10571070
.iter()
10581071
.map(|row| probe_hashes[*row as usize])
10591072
.collect();
1073+
1074+
// A sort costs about as much as this many skipped matches per probe
1075+
// row. A probe row is looked up once, so the marks on probe rows
1076+
// never make grouping worth it.
1077+
const SKIPPED_MATCHES_PER_SORTED_ROW: usize = 8;
1078+
let max_skipped_matches = match self.marks {
1079+
UnknownMarks::Build(_) => SKIPPED_MATCHES_PER_SORTED_ROW * probe_rows.len(),
1080+
UnknownMarks::Probe(_) => usize::MAX,
1081+
};
1082+
self.skipped_matches = 0;
10601083
let mut offset = AtRow { row: 0 };
1061-
loop {
1062-
let next_offset = scope_map.get_matched_indices_with_limit_offset(
1084+
let group_from = loop {
1085+
let next_offset = self.lookup_unmarked(
1086+
scope_map,
10631087
&hashes,
1064-
None,
1065-
self.batch_size,
1088+
|position| probe_rows[position as usize],
1089+
offset,
1090+
&to_build_row,
1091+
lookup,
1092+
);
1093+
for (probe_row, build_row) in
1094+
lookup.probe_indices.iter().zip(lookup.build_indices.iter())
1095+
{
1096+
self.push(*build_row, *probe_row)?;
1097+
}
1098+
let Some(next_offset) = next_offset else {
1099+
return Ok(());
1100+
};
1101+
offset = next_offset;
1102+
if self.skipped_matches > max_skipped_matches {
1103+
break match next_offset {
1104+
AtRow { row } | MidChain { row, .. } => row,
1105+
};
1106+
}
1107+
};
1108+
1109+
// Send the pairs that are left, so the lookups below see their marks.
1110+
self.send()?;
1111+
1112+
// Group the probe rows from the one being looked up. Its build rows
1113+
// that were already paired are paired again, which changes no mark.
1114+
let mut rest: Vec<(u64, u32)> = hashes[group_from..]
1115+
.iter()
1116+
.copied()
1117+
.zip(probe_rows[group_from..].iter().copied())
1118+
.collect();
1119+
rest.sort_unstable_by_key(|(hash, _)| *hash);
1120+
let groups: Vec<_> = rest.chunk_by(|a, b| a.0 == b.0).collect();
1121+
let group_hashes: Vec<u64> = groups.iter().map(|group| group[0].0).collect();
1122+
1123+
let mut build_rows = vec![];
1124+
let mut offset = AtRow { row: 0 };
1125+
loop {
1126+
// A lookup chunk holds the build rows of one or more groups, in
1127+
// group order.
1128+
let next_offset = self.lookup_unmarked(
1129+
scope_map,
1130+
&group_hashes,
1131+
|group| group,
10661132
offset,
1067-
&mut lookup.probe_indices,
1068-
&mut lookup.build_indices,
1133+
&to_build_row,
1134+
lookup,
10691135
);
1070-
for i in 0..lookup.build_indices.len() {
1071-
self.push(
1072-
to_build_row(lookup.build_indices[i]),
1073-
probe_rows[lookup.probe_indices[i] as usize],
1136+
let mut matches = lookup.build_indices.as_slice();
1137+
for group in lookup.probe_indices.chunk_by(|a, b| a == b) {
1138+
let group_matches;
1139+
(group_matches, matches) = matches.split_at(group.len());
1140+
build_rows.clear();
1141+
build_rows.extend_from_slice(group_matches);
1142+
// A build row belongs to one group only, so the marks made
1143+
// for the earlier groups do not concern these build rows.
1144+
self.marks_since_refresh = false;
1145+
self.pair_unmarked(
1146+
&mut build_rows,
1147+
groups[group[0] as usize].iter().map(|(_, row)| *row),
1148+
|b, p| (b, p),
10741149
)?;
10751150
}
10761151
let Some(next_offset) = next_offset else {
@@ -1080,6 +1155,49 @@ where
10801155
}
10811156
}
10821157

1158+
/// Looks up the next chunk of matches of `hashes` in `scope_map`, from
1159+
/// `offset`, into the buffers: the position in `hashes` translated with
1160+
/// `to_probe_row`, and the build row translated with `to_build_row`. Drops the
1161+
/// matches that decide a marked row, and adds their number to
1162+
/// `skipped_matches`.
1163+
fn lookup_unmarked(
1164+
&mut self,
1165+
scope_map: &dyn JoinHashMapType,
1166+
hashes: &[u64],
1167+
to_probe_row: impl Fn(u32) -> u32,
1168+
offset: ProbeOffset,
1169+
to_build_row: impl Fn(u64) -> u64,
1170+
lookup: &mut LookupBuffers,
1171+
) -> Option<ProbeOffset> {
1172+
let LookupBuffers {
1173+
probe_indices,
1174+
build_indices,
1175+
} = lookup;
1176+
let next_offset = scope_map.get_matched_indices_with_limit_offset(
1177+
hashes,
1178+
None,
1179+
self.batch_size,
1180+
offset,
1181+
probe_indices,
1182+
build_indices,
1183+
);
1184+
let bitmap = self.marks.bitmap().lock();
1185+
let mut kept = 0;
1186+
for i in 0..build_indices.len() {
1187+
let build = to_build_row(build_indices[i]);
1188+
let probe = to_probe_row(probe_indices[i]);
1189+
if !bitmap.get_bit(self.marks.unknown_row(build, probe)) {
1190+
build_indices[kept] = build;
1191+
probe_indices[kept] = probe;
1192+
kept += 1;
1193+
}
1194+
}
1195+
self.skipped_matches += build_indices.len() - kept;
1196+
build_indices.truncate(kept);
1197+
probe_indices.truncate(kept);
1198+
next_offset
1199+
}
1200+
10831201
/// Sends the pairs that are left.
10841202
fn finish(mut self) -> Result<()> {
10851203
self.send()
@@ -1213,12 +1331,13 @@ mod tests {
12131331
Ok(pairs.into_iter().max().unwrap_or(0))
12141332
}
12151333

1216-
/// A correlation scope with many build rows and many probe rows. A lookup
1217-
/// per probe row pairs every build row once per probe row in its scope,
1218-
/// even after the build row is marked. So the pairs made for one build
1219-
/// row grow with the number of probe rows.
1334+
/// A correlation scope with many build rows and many probe rows: once a
1335+
/// build row is marked, the later probe rows in its scope must not pair
1336+
/// with it again. So the pairs made for one build row must not grow with
1337+
/// the number of probe rows. A lookup per probe row pairs every build row
1338+
/// once per probe row in its scope.
12201339
#[test]
1221-
fn scope_lookup_pairs_grow_with_probe_rows() -> Result<()> {
1340+
fn unmarked_pairs_do_not_grow_with_probe_rows() -> Result<()> {
12221341
// Build rows 0..60 are in the scope with hash 1, rows 60..100 in the
12231342
// scope with hash 2. The probe rows alternate between both scopes.
12241343
let build_hashes: Vec<u64> = (0..100).map(|row| 1 + (row >= 60) as u64).collect();
@@ -1229,8 +1348,10 @@ mod tests {
12291348
max_pairs_per_build_row(&build_hashes, &probe_hashes, batch_size)
12301349
});
12311350
let (few, many) = (few?, many?);
1232-
assert_eq!(few, 500, "batch_size {batch_size}");
1233-
assert_eq!(many, 5_000, "batch_size {batch_size}");
1351+
assert!(
1352+
many <= few,
1353+
"batch_size {batch_size}: {few} pairs for 1000 probe rows, {many} for 10000"
1354+
);
12341355
}
12351356
Ok(())
12361357
}

0 commit comments

Comments
 (0)