andygrove commented on code in PR #27278:
URL: https://github.com/apache/datafusion/pull/27278#discussion_r4520699769
##########
datafusion/physical-plan/src/joins/sort_merge_join/tests.rs:
##########
@@ -8395,3 +8399,322 @@ async fn
left_join_with_partial_filter_preserves_streamed_order() -> Result<()>
);
Ok(())
}
+
+// ============================================================================
+// Early filtering of Inner joins
+// ============================================================================
+
+/// Rows for the early-filtering tests, sorted by the key column `names[0]`:
key
+/// `k` of `0..keys` has `rows_per_key(k)` rows. `names[1]` is the row number
and
+/// `names[2]` a pseudo-random value in `0..1000` that is NULL in every 37th
row.
+/// The rows are split into batches of `rows_per_batch`.
+fn early_filtering_batches(
+ names: [&str; 3],
+ keys: i32,
+ rows_per_key: impl Fn(i32) -> usize,
+ seed: u64,
+ rows_per_batch: usize,
+) -> (Vec<RecordBatch>, SchemaRef) {
+ let mut key_values = vec![];
+ let mut row_numbers = vec![];
+ let mut values = vec![];
+ let mut state = seed;
+ for key in 0..keys {
+ for _ in 0..rows_per_key(key) {
+ state = state
+ .wrapping_mul(6364136223846793005)
+ .wrapping_add(1442695040888963407);
+ let row = key_values.len() as i32;
+ key_values.push(key);
+ row_numbers.push(row);
+ values.push((row % 37 != 0).then_some(((state >> 33) % 1000) as
i32));
+ }
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new(names[0], DataType::Int32, false),
+ Field::new(names[1], DataType::Int32, false),
+ Field::new(names[2], DataType::Int32, true),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(key_values)),
+ Arc::new(Int32Array::from(row_numbers)),
+ Arc::new(Int32Array::from(values)),
+ ],
+ )
+ .unwrap();
+ let batches = (0..batch.num_rows())
+ .step_by(rows_per_batch)
+ .map(|offset| batch.slice(offset, rows_per_batch.min(batch.num_rows()
- offset)))
+ .collect();
+ (batches, schema)
+}
+
+/// Inputs for the early-filtering tests. Every third key has a group of 600
+/// buffered rows, long enough to be filtered as it is paired, and the 250-row
+/// buffered batches split each such group. The other keys have groups of one
or
+/// two buffered rows, whose pairs are filtered at freeze. Keys repeat up to
four
+/// times on the streamed side.
+fn early_filtering_inputs()
+-> ((Vec<RecordBatch>, SchemaRef), (Vec<RecordBatch>, SchemaRef)) {
+ let left =
+ early_filtering_batches(["a1", "b1", "c1"], 40, |k| 1 + (k % 4) as
usize, 1, 7);
+ let right = early_filtering_batches(
+ ["a2", "b2", "c2"],
+ 40,
+ |k| if k % 3 == 0 { 600 } else { (k % 3) as usize },
+ 2,
+ 250,
Review Comment:
A run never crosses a buffered batch, so with 250-row batches every run is
shorter than `FILTERED_RUN_MIN_ROWS` (256). That means no test reaches
`pair_filtered_run`, its fallback for spilled batches, or a freeze holding
chunks that were filtered while pairing. I added probes to those branches, and
none of them fired in any of the 499 SMJ tests.
These mutations all pass the tests today:
- `pair_filtered_run` ignores the filter result and keeps every pair.
- `start_filtered_chunk` sets the chunk key to `Some(buffered_batch_idx)`
rather than `FILTERED_CHUNK_KEY`, so later unfiltered pairs land in the
filtered chunk.
- `freeze_inner_filtered` drops the chunks that were filtered while pairing.
- `pair_filtered_run` returns `Ok(Some(0))` for a spilled batch, which drops
the run.
With 300 here, all of those branches run, and each mutation makes
`join_inner_early_filtering_matches_hash_join` or the spill test fail. Could
you change it to 300? Deriving it from `FILTERED_RUN_MIN_ROWS` (made
`pub(super)`) would keep the tests on this path if the threshold changes.
##########
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)?));
Review Comment:
When a key group spans many small buffered batches, a freeze holds about one
chunk per buffered batch. This line then builds a `FilterBuilder` and runs two
filter kernels for every chunk that partly passes. I tried the shape of
`inner_group_spans_buffered_batches` with `left.data < right.data` keeping 95%
of the pairs. Against `main`, this PR takes 1.19× as long with 8-row buffered
batches (14.8 to 17.7 ms) and 1.05× with 64-row ones.
Filtering the concatenated indices of the unfiltered chunks once, and giving
each partly passing chunk a slice of the result, brings those to 1.11× and
0.94×. It also helps when the filter keeps 5% of the pairs: 0.89× to 0.82× with
8-row batches, and 0.51× to 0.43× with 64-row ones. The diff is below. It
passes the SMJ tests and the differential test from the review summary.
The rest of the 8-row slowdown might need `filter_first_pays` to count
chunks, so that a freeze of many tiny chunks that keep most pairs takes the
whole-batch path.
<details><summary>Diff against e2f465717</summary>
```diff
diff --git
a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs
b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs
index 8f0b5b1ce..5f222db4b 100644
---
a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs
+++
b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs
@@ -79,25 +79,6 @@ impl MatchedChunk {
fn len(&self) -> usize {
self.streamed_indices.len()
}
-
- /// Keeps the pairs whose bit is set in `mask`.
- fn filter(&self, mask: &BooleanBuffer) -> Result<Self> {
- let predicate = FilterBuilder::new(&BooleanArray::new(mask.clone(),
None))
- .optimize()
- .build();
- Ok(Self {
- buffered_batch_idx: self.buffered_batch_idx,
- streamed_indices: as_uint64_array(
- &predicate.filter(&self.streamed_indices)?,
- )?
- .clone(),
- buffered_indices: as_uint64_array(
- &predicate.filter(&self.buffered_indices)?,
- )?
- .clone(),
- filtered: self.filtered,
- })
- }
}
/// Represents a chunk of joined data from streamed and buffered side
@@ -2067,30 +2048,44 @@ impl MaterializingSortMergeJoinStream {
})
.collect::<Result<Vec<_>>>()?;
self.freeze_filter_evaluated += num_rows;
- Some(evaluate_join_filter(
- filter,
- filter.schema(),
- columns,
- num_rows,
- )?)
+ let passed =
+ evaluate_join_filter(filter, filter.schema(), columns,
num_rows)?;
+ // Filter the index arrays of all unfiltered chunks at once;
each partly
+ // passing chunk below takes a slice of the result.
+ let predicate =
FilterBuilder::new(&BooleanArray::new(passed.clone(), None))
+ .optimize()
+ .build();
+ let buffered_indices =
+ concat_indices(unfiltered.iter().map(|c|
&c.buffered_indices))?;
+ let kept_streamed =
+
as_uint64_array(&predicate.filter(&streamed_indices)?)?.clone();
+ let kept_buffered =
+
as_uint64_array(&predicate.filter(&buffered_indices)?)?.clone();
+ Some((passed, kept_streamed, kept_buffered))
};
// 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;
+ let mut kept_offset = 0;
for chunk in matched_chunks {
let len = chunk.len();
match (&passed, chunk.filtered) {
- (Some(passed), false) => {
- let mask = passed.slice(offset, len);
+ (Some((passed, kept_streamed, kept_buffered)), false) => {
+ let count = passed.slice(offset, len).count_set_bits();
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)?));
+ kept.push(Cow::Owned(MatchedChunk {
+ buffered_batch_idx: chunk.buffered_batch_idx,
+ streamed_indices:
kept_streamed.slice(kept_offset, count),
+ buffered_indices:
kept_buffered.slice(kept_offset, count),
+ filtered: false,
+ }));
}
+ kept_offset += count;
self.freeze_filter_passed += count;
kept_rows += count;
}
```
</details>
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -47,10 +50,63 @@ use
datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
use futures::StreamExt;
+/// The shortest run of buffered rows that an Inner join with a filter filters
as
+/// it pairs the run with a streamed row. Each such evaluation has a fixed cost
+/// that a long run amortizes. Shorter runs are paired unfiltered, and the
+/// freeze filters a whole batch of their pairs at once.
+const FILTERED_RUN_MIN_ROWS: usize = 256;
Review Comment:
Since you asked about this threshold, I measured it with Inner joins of 16
streamed rows and R buffered rows per key, 2M pairs each, using the bench's
`key`/`data`/`payload` schema. The filters keep 5%, 50% or 95% of the pairs.
The table shows this PR's time divided by the time of the same build with the
threshold at `usize::MAX`, which filters everything at freeze (median of 5
rounds):
| R | 5% kept | 50% kept | 95% kept |
|---:|---:|---:|---:|
| 256 | 1.17 | 1.14 | 1.19 |
| 512 | 0.87 | 1.05 | 1.13 |
| 2048 | 0.70 | 1.00 | 1.11 |
So filtering while pairing only wins for selective filters, and only from
about 512 rows. A lower threshold is worse: at 64, R = 128 with 5% kept took
1.75× as long as filtering at freeze. Everything in the table is still faster
than `main` (0.13× to 0.96×), so this is speed left on the table rather than a
regression.
Would you consider raising the threshold to 512, and filtering while pairing
only while the share of pairs that pass stays low? The pairing-time evaluations
could feed the same counts that `filter_first_pays` uses.
`inner_1to1000_selective_filtered` would keep its win.
--
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]
