comphead commented on code in PR #27278:
URL: https://github.com/apache/datafusion/pull/27278#discussion_r4540670056
##########
datafusion/physical-plan/src/joins/sort_merge_join/tests.rs:
##########
Review Comment:
Nothing did. In 8012102c4, `assert_streamed_order` asserts that `(b1, b2)`
ascends. These are the row numbers of the two inputs, so the check covers the
order within a key. Reversing the kept chunks now fails it.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
Review Comment:
The fixed threshold belonged to filtering while pairing, which 8012102c4
removes. Its high-pass regression was real: up to 1.5x slower than `main` with
a narrow schema (table in the top-level comment). The freeze path keeps the
cost model and adds a range check for 1:1 joins.
##########
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:
Thanks, reproduced and fixed in 8012102c4. The freeze takes the whole-batch
path when both sides' indices span a range. The check looks only at the first
and last index, because the slicing path checks the full range again anyway. A
full scan here kept 1:1 joins about 5% slower than `main`. I also added
`inner_1to1_filtered` to the bench. Medians as a ratio to `main` (9d916ab47):
| Case | e2f465717 | 8012102c4 |
|---|---:|---:|
| 1:1, no filter | 1.03 | 1.00 |
| 1:1, 100% pass | 1.09 | 0.99 |
| 1:1, 90% pass, every 10th row fails | 1.20 | 0.98 |
| 1:1, 90% pass, random | 1.17 | 0.98 |
| 1:1, 50% pass, random | 1.09 | 0.99 |
| 1:1, 5% pass, random | 1.09 | 0.99 |
| `inner_1to1_filtered` (new) | n/a | 0.99 |
@viirya, I added both caveats to the `filter_first_pays` doc: it counts
columns rather than bytes, and `p` is a lifetime average. Since pairing no
longer evaluates the filter, every evaluation now feeds `p`.
##########
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)?;
Review Comment:
@viirya is right that `main` already fails this way. On 9d916ab47, with this
schema and `batch_size = 256`, both variants fail with `Dictionary key bigger
than the key type`:
- four 128-row buffered batches, keeping every row, which goes through
`interleave`
- two 256-row batches keeping 100 rows each, which goes through the
coalescer's `concat`
With filtering while pairing removed in 8012102c4, your exact input returns
all 300 rows, as on `main`.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -805,6 +967,91 @@ impl MaterializingSortMergeJoinStream {
true
}
+ /// [`Self::pair_streamed_row_with_group`] for an Inner join with a filter,
+ /// which keeps only the pairs that pass it. For a run of at least
+ /// [`FILTERED_RUN_MIN_ROWS`] buffered rows, the filter runs as the run is
+ /// paired, and the pairs it rejects are never recorded.
+ fn pair_streamed_row_with_group_filtered(&mut self) -> Result<bool> {
+ let mut room =
self.batch_size.saturating_sub(self.num_unfrozen_pairs());
+ while !self.buffered_data.scanning_finished() && room > 0 {
+ let rows = self.buffered_data.scanning_remaining();
+ let len = rows.len().min(room);
+ let filtered = if len >= FILTERED_RUN_MIN_ROWS {
+ self.pair_filtered_run(rows.start, len)?
+ } else {
+ None
+ };
+ match filtered {
+ Some(appended) => room -= appended,
Review Comment:
This came with filtering while pairing, which 8012102c4 removes. The join
again freezes once `batch_size` candidate pairs accumulate, as on `main`. As
@viirya noted, freezing before polling the child when the coalesced rows and
the unfrozen pairs can fill a batch would help `main` as well, so I'd keep that
as a separate change.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -211,12 +317,28 @@ impl StreamedBatch {
let capacity = batch_size - self.num_output_rows;
self.output_indices.push(StreamedJoinedChunk {
buffered_batch_idx,
+ filtered: false,
streamed_indices: Vec::with_capacity(capacity),
buffered_indices: Vec::with_capacity(capacity),
buffered_nulls: NullBufferBuilder::new(capacity),
});
self.buffered_batch_idx = buffered_batch_idx;
}
+
+ /// Starts a chunk of filtered pairs for `buffered_batch_idx`. It sets the
key
+ /// that [`Self::chunk_for`] compares to [`FILTERED_CHUNK_KEY`], so that
the next
+ /// unfiltered pair starts a chunk of its own.
+ fn start_filtered_chunk(&mut self, buffered_batch_idx: usize, batch_size:
usize) {
+ let capacity = batch_size - self.num_output_rows;
Review Comment:
This came with filtering while pairing. 8012102c4 removes it along with
`start_filtered_chunk`.
--
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]
