Skip to content

Commit a40e07c

Browse files
viiryaIsaac
andcommitted
fix: keep few probe rows across cross product chunks, free rejected groups
- With marks on the build rows, `cross_product` called `probe_rows` again for every chunk of build rows. A few NULL-valued probe rows in a large probe batch then walked the batch's bitmap once per chunk. Up to `batch_size` probe rows are now kept for all chunks. More are produced again for each chunk, which costs less than pairing them with it. - When the real group count rejects the grouping, the rejected groups are dropped before the lookups resume, instead of staying allocated until the new groups replace them. - The upgrade guide says that elements of types that can hide a NULL count as elements that can be NULL, so a run-end encoded one is rejected. Co-authored-by: Isaac <[email protected]>
1 parent 0eb4fda commit a40e07c

2 files changed

Lines changed: 50 additions & 14 deletions

File tree

  • datafusion/physical-plan/src/joins/hash_join
  • docs/source/library-user-guide/upgrading

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

Lines changed: 44 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1244,22 +1244,31 @@ where
12441244
/// is marked.
12451245
///
12461246
/// When the pairs decide the build rows, the unmarked build rows are
1247-
/// paired in chunks of `batch_size` rows, and `probe_rows` is called
1248-
/// again for each chunk. Neither side is then held whole: a large build
1249-
/// side, as when this pairing stands in for a scope map that the memory
1250-
/// pool declined, or a probe batch larger than the batch size. A chunk is
1251-
/// filled with unmarked rows only, so a build side with few unmarked rows
1252-
/// left pairs them all at once, as a single chunk would.
1247+
/// paired in chunks of `batch_size` rows. Neither side is then held
1248+
/// whole: a large build side, as when this pairing stands in for a scope
1249+
/// map that the memory pool declined, or a probe batch larger than the
1250+
/// batch size. Up to `batch_size` probe rows are kept for every chunk.
1251+
/// More are produced again for each chunk by calling `probe_rows`, which
1252+
/// costs less than pairing them with the chunk. A chunk is filled with
1253+
/// unmarked rows only, so a build side with few unmarked rows left pairs
1254+
/// them all at once, as a single chunk would.
12531255
fn cross_product<P: Iterator<Item = u32>>(
12541256
&mut self,
12551257
build_rows: impl Iterator<Item = u64>,
12561258
probe_rows: impl Fn() -> P,
12571259
) -> Result<()> {
12581260
match self.marks {
12591261
UnknownMarks::Build(bitmap) => {
1260-
if probe_rows().next().is_none() {
1262+
// A few probe rows can be spread over a large probe batch, so
1263+
// keep them rather than walk the batch for every chunk.
1264+
let mut kept: Vec<u32> = probe_rows().take(self.batch_size + 1).collect();
1265+
if kept.is_empty() {
12611266
return Ok(());
12621267
}
1268+
let keep_probe_rows = kept.len() <= self.batch_size;
1269+
if !keep_probe_rows {
1270+
kept = vec![];
1271+
}
12631272
let mut build_rows = build_rows.peekable();
12641273
let mut chunk = Vec::with_capacity(self.batch_size);
12651274
while build_rows.peek().is_some() {
@@ -1275,7 +1284,13 @@ where
12751284
}
12761285
// The chunk holds no row marked so far.
12771286
self.needs_refresh = false;
1278-
self.pair_unmarked(&mut chunk, probe_rows(), |b, p| (b, p))?;
1287+
if keep_probe_rows {
1288+
self.pair_unmarked(&mut chunk, kept.iter().copied(), |b, p| {
1289+
(b, p)
1290+
})?;
1291+
} else {
1292+
self.pair_unmarked(&mut chunk, probe_rows(), |b, p| (b, p))?;
1293+
}
12791294
}
12801295
Ok(())
12811296
}
@@ -1391,6 +1406,8 @@ where
13911406
}
13921407
let mut groups = hash_groups(&rest);
13931408
if rest.is_empty() || !saves_work(rest.len(), groups.len()) {
1409+
// The groups are found again from where the lookups stop.
1410+
drop(groups);
13941411
self.wasted_matches = 0;
13951412
let Some((offset, _)) = self.lookup_each_row(
13961413
scope_map,
@@ -1916,6 +1933,7 @@ mod tests {
19161933
Ok(())
19171934
};
19181935

1936+
let probe_walks = std::cell::Cell::new(0);
19191937
let mut pairs = UnmarkedPairs::new(marks, self.batch_size, mark);
19201938
if self.through_scope_map {
19211939
let mut scope_map = JoinHashMapU64::with_capacity(self.build_rows.len());
@@ -1937,6 +1955,7 @@ mod tests {
19371955
)?;
19381956
} else {
19391957
pairs.cross_product(self.build_rows.iter().copied(), || {
1958+
probe_walks.set(probe_walks.get() + 1);
19401959
self.probe_rows.iter().copied()
19411960
})?;
19421961
}
@@ -1958,6 +1977,7 @@ mod tests {
19581977
max_decided_rows,
19591978
sends,
19601979
sent_pairs,
1980+
probe_walks: probe_walks.get(),
19611981
foreign_marks,
19621982
},
19631983
))
@@ -1975,6 +1995,8 @@ mod tests {
19751995
/// The chunks of pairs sent, and the pairs in them.
19761996
sends: usize,
19771997
sent_pairs: usize,
1998+
/// The calls of the `probe_rows` of a cross product.
1999+
probe_walks: usize,
19782000
/// The rows that `foreign_marks` marked.
19792001
foreign_marks: Vec<usize>,
19802002
}
@@ -2163,8 +2185,9 @@ mod tests {
21632185
}
21642186

21652187
/// A cross product with marks on the build rows holds at most a batch of
2166-
/// build rows at once, however many build rows there are. No pair passes,
2167-
/// so no row is marked and every build row is paired.
2188+
/// build rows at once, however many build rows there are, and at most a
2189+
/// batch of probe rows. No pair passes, so no row is marked and every
2190+
/// build row is paired.
21682191
#[test]
21692192
fn cross_product_chunks_the_build_rows() -> Result<()> {
21702193
let (num_build, batch_size) = (10_000, 64);
@@ -2185,6 +2208,17 @@ mod tests {
21852208
let (marks, stats) = case.run()?;
21862209
assert!(!marks.contains(&true));
21872210
assert_eq!(stats.max_decided_rows, batch_size);
2211+
// The three probe rows are kept for all 157 chunks.
2212+
assert_eq!(stats.probe_walks, 1);
2213+
// More probe rows than a batch are produced again for each chunk.
2214+
let (marks, stats) = Case {
2215+
probe_keys: vec![0; batch_size + 1],
2216+
probe_rows: (0..=batch_size as u32).collect(),
2217+
..case
2218+
}
2219+
.run()?;
2220+
assert!(!marks.contains(&true));
2221+
assert_eq!(stats.probe_walks, 1 + num_build.div_ceil(batch_size));
21882222
Ok(())
21892223
}
21902224

‎docs/source/library-user-guide/upgrading/56.0.0.md‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -939,10 +939,12 @@ nullable element cannot be NULL here when a correlated filter of the subquery
939939
rejects NULL for it, such as `y > o.c` for `y`. The elements that cannot be
940940
NULL and any equality correlations follow as further equi-join keys when their
941941
types can be hashed, and otherwise stay in the join filter. `Null`,
942-
dictionary, run-end encoded and union elements stay among the first keys,
943-
because they can hold a NULL that their field's nullability does not show. A
944-
null-aware join therefore records how many of its leading keys are `NOT IN`
945-
value keys:
942+
dictionary, run-end encoded and union elements, and elements that read a
943+
column of these types, count as elements that can be NULL, because they can
944+
hold a NULL that their field's nullability does not show. Like any element
945+
that can be NULL, such an element must be a hashable equi-join key, so the
946+
planner rejects a run-end encoded one. A null-aware join therefore records how
947+
many of its leading keys are `NOT IN` value keys:
946948

947949
- `datafusion_expr::logical_plan::Join` has a new `null_aware_value_keys: usize`
948950
field.

0 commit comments

Comments
 (0)