sunchao commented on code in PR #27278:
URL: https://github.com/apache/datafusion/pull/27278#discussion_r4528343458
##########
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:
Each newly encountered buffered batch reserves the entire remaining output
capacity for both index vectors, even when the filter keeps a single row. With
one streamed row, 1,024 same-key buffered batches of 256 rows, one survivor per
batch, and `batch_size = 8192`, a probe using the base and head chunk methods
measured peak live index-vector capacity of 2.06 MiB on the base and 120 MiB
here. The head has just 16 KiB of actual indices. The base freezes every 8,192
candidates, while all 1,024 filtered chunks stay live together here.
`append_filtered_pairs` already has `count` before starting the chunk.
Passing it as the initial capacity reduced the same probe to 16 KiB and still
lets a reused chunk grow.
##########
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:
A valid `Dictionary<Int8, Utf8>` payload can now fail here with
`DictionaryKeyOverflowError`. I reproduced it with `batch_size = 256`, one
streamed row, and two 256-row buffered batches for the same key. Each batch
uses 100 distinct strings, disjoint between batches, with dictionary key `row %
100`. A Boolean filter keeps the first 150 rows of each batch. The base returns
all 300 rows. This path puts the first 150 survivors and 106 from the next
batch into the same output. `interleave` now needs 200 distinct dictionary
keys, which `Int8` cannot encode.
Could we preserve the candidate-batch boundary when filtering during
pairing, or split the output when the combined dictionary cannot be
represented? This two-batch case would make a small regression test.
##########
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:
Counting only survivors against `room` can make the join wait for more
streamed input even when it already has a full batch to emit. I tested with
`batch_size = 512`, one buffered batch of 512 rows for a single key, and two
one-row streamed batches whose predicates keep 200 and 360 rows respectively. I
held the streamed child open after its second batch with `BarrierExec`.
The base emits 512 rows before the barrier is released. This head emits
zero, then returns all 560 after release. The coalescer holds the first 200 and
the remaining 360 are still pair indices when the join polls the child again. A
downstream `LIMIT 512` can therefore wait on input it does not need. Could we
trigger the freeze before polling the child when the accumulated rows can
complete an output batch?
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -360,6 +482,24 @@ pub(super) struct MaterializingSortMergeJoinStream {
/// since NULL-padded rows reach filter evaluation. `None` unless
/// `deferred_filtering`.
pub filter_schema: Option<SchemaRef>,
+ /// Whether this is an Inner join with a filter. Such a join drops every
pair
+ /// the filter rejects, so the filter runs before any output column is
+ /// materialized: on a long run of buffered rows as it is paired with a
+ /// streamed row (see [`Self::pair_filtered_run`]), and at freeze for the
+ /// remaining pairs (see [`Self::freeze_inner_filtered`]). Outer joins keep
+ /// a rejected pair for the deferred pipeline, which may null-join its row.
+ /// Only an Inner join filters early, so its streamed side is the left
input.
+ pub early_filtering: bool,
+ /// The indices of every buffered column, for materializing all of them
+ pub buffered_columns: Vec<usize>,
+ /// The buffered columns that the join filter reads, in the filter's order.
+ /// Empty unless `early_filtering`.
+ pub buffered_filter_columns: Vec<usize>,
+ /// How many pairs the Inner join's filter has evaluated at freeze so far
+ /// (see [`Self::filter_first_pays`])
+ pub freeze_filter_evaluated: usize,
Review Comment:
These two counters total pairs across the life of a partition. On wasm32,
65,536 rows with one key on each side yield 2³² candidate pairs even if the
filter rejects every row. Splitting the buffered input into 128-row batches
takes the freeze path for every pair, so `freeze_filter_evaluated` overflows.
Checked builds panic, and in release the wrapped denominator can make the
estimated pass rate exceed 1. Could both totals use `u64`, widening the
per-freeze increments?
--
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]
