mbutrovich commented on code in PR #27278:
URL: https://github.com/apache/datafusion/pull/27278#discussion_r4526296480
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -1771,33 +2008,215 @@ impl MaterializingSortMergeJoinStream {
Ok(())
}
- /// Materializes right-side columns across all matched chunks.
+ /// Freezes the matched pairs of an Inner join with a filter. The filter
runs
+ /// on just the columns it reads, for the pairs of the chunks that
+ /// [`Self::pair_filtered_run`] did not filter already, and only the pairs
that
+ /// pass have their output columns materialized. A selective filter over a
+ /// many-to-many join rejects most pairs, so materializing every column of
+ /// every pair before filtering would spend most of the freeze on rows that
+ /// are then dropped. When no chunk was filtered already and earlier
freezes
+ /// kept enough of their pairs, [`Self::filter_first_pays`] says the
opposite,
+ /// and the batch is materialized whole and then filtered.
+ fn freeze_inner_filtered(
+ &mut self,
+ matched_chunks: &[MatchedChunk],
+ total_matched_rows: usize,
+ ) -> Result<()> {
+ let Some(filter) = &self.filter else {
+ return internal_err!("early filtering requires a join filter");
+ };
+ debug_assert_eq!(
+ self.join_type,
+ JoinType::Inner,
+ "only an Inner join filters early, so the streamed side is the
left one"
+ );
+
+ // Evaluate the filter over the pairs of the chunks not filtered yet.
+ let unfiltered: Vec<&MatchedChunk> =
+ matched_chunks.iter().filter(|c| !c.filtered).collect();
+ if unfiltered.len() == matched_chunks.len() &&
!self.filter_first_pays(filter) {
+ return self.freeze_inner_filtered_whole(matched_chunks,
total_matched_rows);
+ }
+ let passed = if unfiltered.is_empty() {
+ None
+ } else {
+ let num_rows = unfiltered.iter().map(|c| c.len()).sum();
+ let streamed_indices =
+ concat_indices(unfiltered.iter().map(|c|
&c.streamed_indices))?;
+ let mut buffered_arrays = self
+ .materialize_right_columns(
+ &unfiltered,
+ num_rows,
+ &self.buffered_filter_columns,
+ )?
+ .into_iter();
+ let columns = filter
+ .column_indices()
+ .iter()
+ .map(|column_index| match column_index.side {
+ JoinSide::Left => take_column(
+ self.streamed_batch.batch.column(column_index.index),
+ &streamed_indices,
+ ),
+ JoinSide::Right => buffered_arrays.next().ok_or_else(|| {
+ internal_datafusion_err!("missing buffered filter
column")
+ }),
+ JoinSide::None => {
+ internal_err!("Unexpected JoinSide::None in filter")
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+ self.freeze_filter_evaluated += num_rows;
+ Some(evaluate_join_filter(
+ filter,
+ filter.schema(),
+ columns,
+ num_rows,
+ )?)
+ };
+
+ // Keep the pairs that passed, in order.
+ let mut kept: Vec<Cow<MatchedChunk>> =
Vec::with_capacity(matched_chunks.len());
+ let mut kept_rows = 0;
+ let mut offset = 0;
+ for chunk in matched_chunks {
+ let len = chunk.len();
+ match (&passed, chunk.filtered) {
+ (Some(passed), false) => {
+ let mask = passed.slice(offset, len);
+ offset += len;
+ let count = mask.count_set_bits();
+ if count == len {
+ kept.push(Cow::Borrowed(chunk));
+ } else if count > 0 {
+ kept.push(Cow::Owned(chunk.filter(&mask)?));
+ }
+ self.freeze_filter_passed += count;
+ kept_rows += count;
+ }
+ _ => {
+ kept.push(Cow::Borrowed(chunk));
+ kept_rows += len;
+ }
+ }
+ }
+ if kept_rows == 0 {
+ return Ok(());
+ }
+
+ let streamed_indices = concat_indices(kept.iter().map(|c|
&c.streamed_indices))?;
+ let streamed_columns =
+ materialize_left_columns(&self.streamed_batch.batch,
&streamed_indices)?;
+ let buffered_columns =
+ self.materialize_right_columns(&kept, kept_rows,
&self.buffered_columns)?;
+ let output_batch = RecordBatch::try_new(
+ Arc::clone(&self.schema),
+ [streamed_columns, buffered_columns].concat(),
+ )?;
+ self.joined_record_batches
+ .push_batch_without_metadata(output_batch);
+ Ok(())
+ }
+
+ /// Freezes the matched pairs of an Inner join with a filter by
materializing
+ /// every column of every pair, evaluating the filter on the columns it
reads,
+ /// and filtering the batch. None of the pairs was filtered already.
+ fn freeze_inner_filtered_whole(
+ &mut self,
+ matched_chunks: &[MatchedChunk],
+ total_matched_rows: usize,
+ ) -> Result<()> {
+ let streamed_indices =
+ concat_indices(matched_chunks.iter().map(|c|
&c.streamed_indices))?;
+ let streamed_columns =
+ materialize_left_columns(&self.streamed_batch.batch,
&streamed_indices)?;
+ let buffered_columns = self.materialize_right_columns(
+ matched_chunks,
+ total_matched_rows,
+ &self.buffered_columns,
+ )?;
+ let Some(filter) = &self.filter else {
+ return internal_err!("early filtering requires a join filter");
+ };
+ let filter_columns =
+ get_filter_columns(Some(filter), &streamed_columns,
&buffered_columns)?;
+ let passed = evaluate_join_filter(
+ filter,
+ filter.schema(),
+ filter_columns,
+ total_matched_rows,
+ )?;
+ let kept_rows = passed.count_set_bits();
+ self.freeze_filter_evaluated += total_matched_rows;
+ self.freeze_filter_passed += kept_rows;
+ if kept_rows == 0 {
+ return Ok(());
+ }
+
+ let output_batch = RecordBatch::try_new(
+ Arc::clone(&self.schema),
+ [streamed_columns, buffered_columns].concat(),
+ )?;
+ let output_batch = if kept_rows == total_matched_rows {
+ output_batch
+ } else {
+ filter_record_batch(&output_batch, &BooleanArray::new(passed,
None))?
+ };
+ self.joined_record_batches
+ .push_batch_without_metadata(output_batch);
+ Ok(())
+ }
+
+ /// Whether filtering the pairs of a freeze before materializing them is
cheaper
+ /// than materializing every column of every pair and filtering the batch.
+ ///
+ /// Filtering first gathers the filter's `f` columns for every pair,
filters the
+ /// pairs' two index arrays, and gathers all `k` output columns for the
share
+ /// `p` of pairs that pass. The other way gathers the `k` columns for every
+ /// pair and copies the share `p` of them to filter the batch. Counting a
+ /// gathered value as 1 and a copied one as 1/2, filtering first is cheaper
+ /// while `f + 1 + k·p < k + k·p/2`, that is while `p < 2·(k - f - 1)/k`.
`p` is
+ /// the share that passed in the earlier freezes, or 0 before the first.
+ fn filter_first_pays(&self, filter: &JoinFilter) -> bool {
+ let k = self.schema.fields().len() as f64;
+ let f = filter.column_indices().len() as f64;
+ let p = if self.freeze_filter_evaluated == 0 {
+ 0.0
+ } else {
+ self.freeze_filter_passed as f64 / self.freeze_filter_evaluated as
f64
+ };
+ p < 2.0 * (k - f - 1.0) / k
Review Comment:
This model assumes that the whole-batch path gathers all `k` columns for
every pair. When a freeze's indices form a contiguous range, `main` slices the
columns instead
([`materialize_left_columns`](https://github.com/apache/datafusion/blob/9e0a08fd8627e5423f2b09b0ca13a5015b7cadee/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs#L2044-L2053)
and [the buffered
side](https://github.com/apache/datafusion/blob/9e0a08fd8627e5423f2b09b0ca13a5015b7cadee/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs#L2136-L2140)).
That is the usual case in a 1:1 join where every row matches, and there the
whole-batch path costs little more than `filter_record_batch`. Filtering first
keeps the slices only when every pair of a chunk passes. Otherwise
`chunk.filter` leaves gaps in the indices, and every output column goes through
`take`. With the bench schema (`k = 6`, `f = 2`), the threshold is `2(6 - 2 -
1)/6 = 1.0`, so every freeze below 100% pass filt
ers first.
I measured 1:1 Inner joins of 100K rows per side, with the bench's
three-column schema and `left.data < right.data`, against 9d916ab47. These are
medians of 3 alternating rounds on Apple silicon with thin LTO:
| Case | `main` | This PR | Speedup |
|---|---:|---:|---:|
| 1:1, no filter (unchanged path) | 1.403 ms | 1.433 ms | 0.98x |
| 1:1, 100% pass | 1.424 ms | 1.530 ms | 0.93x |
| 1:1, 90% pass, every 10th row fails | 1.692 ms | 2.018 ms | 0.84x |
| 1:1, 90% pass, random | 1.788 ms | 2.037 ms | 0.88x |
| 1:1, 50% pass, random | 1.745 ms | 1.876 ms | 0.93x |
| 1:1, 5% pass, random | 1.483 ms | 1.590 ms | 0.93x |
| N:1, 10 streamed rows per key, 90% pass | 1.531 ms | 1.522 ms | 1.01x |
| N:1, 10 streamed rows per key, 50% pass | 1.452 ms | 1.349 ms | 1.08x |
The 100% case keeps the slices and is still 7% slower, against 2% for the
unfiltered control. Every streamed row of these joins goes through
`pair_streamed_row_with_group_filtered`, so dropping filtering while pairing
may remove that part. The 90% cases, which lose the slices, drop further to
0.84x and 0.88x.
Could the freeze take the whole-batch path when its streamed and buffered
indices are both contiguous? That checks the freeze at hand, so it adds no
state. A 1:1 high-pass case in `benches/sort_merge_join.rs` would cover it.
<details><summary>Bench cases (appended to
<code>benches/sort_merge_join.rs</code>)</summary>
```rust
fn mix(i: u64) -> u64 {
let mut z = i.wrapping_add(0x9E37_79B9_7F4A_7C15);
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
z ^ (z >> 31)
}
/// Batches of 8192 rows with the given (already key-sorted) keys and data.
fn degen_batches(keys: Vec<i64>, data: Vec<i64>, narrow: bool) ->
(SchemaRef, Vec<RecordBatch>) {
let s = if narrow { narrow_schema() } else { schema() };
let mut cols: Vec<arrow::array::ArrayRef> = vec![
Arc::new(Int64Array::from(keys)),
Arc::new(Int64Array::from(data.clone())),
];
if !narrow {
let payload: Vec<String> = (0..data.len()).map(|i|
format!("val_{i}")).collect();
cols.push(Arc::new(StringArray::from(payload)));
}
let batch = RecordBatch::try_new(Arc::clone(&s), cols).unwrap();
let mut out = vec![];
let mut off = 0;
while off < batch.num_rows() {
let len = (batch.num_rows() - off).min(8192);
out.push(batch.slice(off, len));
off += len;
}
(s, out)
}
fn bench_degenerate(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
let mut group = c.benchmark_group("smj_degen");
let n: usize = 100_000;
// 1:1, unique keys on both sides, every streamed row matches.
// left.data < right.data where right.data = t, left.data in 0..100.
for (name, filtered, t, random) in [
("1to1_unfiltered", false, 0, true),
("1to1_pass100", true, 100, true),
("1to1_pass90_regular", true, 90, false),
("1to1_pass90_random", true, 90, true),
("1to1_pass50_random", true, 50, true),
("1to1_pass5_random", true, 5, true),
] {
let keys: Vec<i64> = (0..n as i64).collect();
let ldata: Vec<i64> = (0..n as u64)
.map(|i| if random { (mix(i) % 100) as i64 } else { (i % 10 *
10) as i64 })
.collect();
let (s, left) = degen_batches(keys.clone(), ldata, false);
let (_, right) = degen_batches(keys, vec![t; n], false);
group.bench_function(BenchmarkId::new(name, n), |b| {
b.iter(|| {
do_join_with_filter(
make_exec(&left, &s),
make_exec(&right, &s),
datafusion_common::JoinType::Inner,
filtered.then(|| data_filter(Operator::Lt)),
&rt,
)
})
});
}
// N:1 (FK to PK): 10 streamed rows per key, one buffered row per key.
for (name, t) in [("nto1_pass90_random", 90), ("nto1_pass50_random",
50)] {
let lkeys: Vec<i64> = (0..n as i64).map(|i| i / 10).collect();
let ldata: Vec<i64> = (0..n as u64).map(|i| (mix(i) % 100) as
i64).collect();
let (s, left) = degen_batches(lkeys, ldata, false);
let (_, right) = degen_batches((0..(n / 10) as i64).collect(),
vec![t; n / 10], false);
group.bench_function(BenchmarkId::new(name, n), |b| {
b.iter(|| {
do_join_with_filter(
make_exec(&left, &s),
make_exec(&right, &s),
datafusion_common::JoinType::Inner,
Some(data_filter(Operator::Lt)),
&rt,
)
})
});
}
// 1:1000, 20 keys, 100 streamed and 1000 buffered rows per key (2M
pairs).
// left.data = 90 > right.data in 0..100 keeps 90% of pairs, at random.
for (name, narrow) in [("1to1000_pass90", false),
("1to1000_pass90_narrow", true)] {
let keys = 20i64;
let lkeys: Vec<i64> = (0..keys * 100).map(|i| i / 100).collect();
let rkeys: Vec<i64> = (0..keys * 1000).map(|i| i / 1000).collect();
let rdata: Vec<i64> = (0..(keys * 1000) as u64).map(|i| (mix(i) %
100) as i64).collect();
let (s, left) = degen_batches(lkeys, vec![90; (keys * 100) as
usize], narrow);
let (_, right) = degen_batches(rkeys, rdata, narrow);
group.bench_function(BenchmarkId::new(name, keys * 1000), |b| {
b.iter(|| {
do_join_with_filter(
make_exec(&left, &s),
make_exec(&right, &s),
datafusion_common::JoinType::Inner,
Some(data_filter(Operator::Gt)),
&rt,
)
})
});
}
group.finish();
}
criterion_group!(benches, bench_smj, bench_degenerate);
```
</details>
--
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]
