comphead commented on code in PR #26277:
URL: https://github.com/apache/datafusion/pull/26277#discussion_r4265221854
##########
datafusion/physical-plan/benches/sort_merge_join.rs:
##########
@@ -168,6 +168,22 @@ fn bench_smj(c: &mut Criterion) {
});
}
+ // 1:100 Inner Join — 20K left, 20K right, 200 distinct keys
Review Comment:
`inner_1to100` is the only high-fanout case, and it is an inner join.
`pair_streamed_row_with_group` is shared by outer joins and by the
deferred-filtering path, so a filtered `Left` or `Full` join with high fanout
takes the same bulk path and then goes through freeze and filter. Could you add
a `left_1to100_filtered` case, or say why the inner case is enough?
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -117,44 +117,112 @@ impl StreamedBatch {
self.num_output_rows
}
- /// Appends new pair consisting of current streamed index and
`buffered_idx`
- /// index of buffered batch with `buffered_batch_idx` index.
- fn append_output_pair(
+ /// Appends a pair of the current streamed index and a null buffered
index, for a
+ /// streamed row that no buffered row matches.
+ fn append_null_joined_pair(
&mut self,
buffered_batch_idx: Option<usize>,
- buffered_idx: Option<usize>,
batch_size: usize,
) {
- // If no current chunk exists or current chunk is not for current
buffered batch,
- // create a new chunk
- if self.output_indices.is_empty() || self.buffered_batch_idx !=
buffered_batch_idx
- {
- // Compute capacity only when creating a new chunk (infrequent
operation).
- // The capacity is the remaining space to reach batch_size.
- // This should always be >= 1 since we only call this when
num_output_rows < batch_size.
- debug_assert!(
- batch_size > self.num_output_rows,
- "batch_size ({batch_size}) must be > num_output_rows ({})",
- self.num_output_rows
- );
- let capacity = batch_size - self.num_output_rows;
- self.output_indices.push(StreamedJoinedChunk {
+ let streamed_idx = self.idx as u64;
+ let current_chunk = self.chunk_for(buffered_batch_idx, batch_size);
+ current_chunk.streamed_indices.append_value(streamed_idx);
+ current_chunk.buffered_indices.append_null();
+ self.num_output_rows += 1;
+ }
+
+ /// Appends one pair for each row in `buffered_range` of the buffered
batch with
+ /// `buffered_batch_idx` index, all with the current streamed index. The
indices
+ /// are appended in bulk, which matters when a streamed row matches many
buffered
+ /// rows.
+ fn append_output_pairs(
+ &mut self,
+ buffered_batch_idx: usize,
+ buffered_range: Range<usize>,
+ batch_size: usize,
+ ) {
+ let len = buffered_range.len();
+ debug_assert!(
+ len > 0 && self.num_output_rows + len <= batch_size,
+ "appending {len} pairs to {} with batch_size {batch_size}",
+ self.num_output_rows
+ );
+ if len > 1 {
+ return self.append_output_run(
buffered_batch_idx,
- streamed_indices: UInt64Builder::with_capacity(capacity),
- buffered_indices: UInt64Builder::with_capacity(capacity),
- });
- self.buffered_batch_idx = buffered_batch_idx;
+ buffered_range,
+ batch_size,
+ );
}
- let current_chunk = self.output_indices.last_mut().unwrap();
+ // A single match, as in a one-to-one join, is cheaper to append
directly.
+ let streamed_idx = self.idx as u64;
+ let current_chunk = self.chunk_for(Some(buffered_batch_idx),
batch_size);
+ current_chunk.streamed_indices.append_value(streamed_idx);
+ current_chunk
+ .buffered_indices
+ .append_value(buffered_range.start as u64);
+ self.num_output_rows += 1;
+ }
- // Append index of streamed batch and index of buffered batch into
current chunk
- current_chunk.streamed_indices.append_value(self.idx as u64);
- if let Some(idx) = buffered_idx {
- current_chunk.buffered_indices.append_value(idx as u64);
- } else {
- current_chunk.buffered_indices.append_null();
+ /// The bulk half of [`Self::append_output_pairs`], kept out of line so
that the
+ /// one-row case, which one-to-one joins take for every row, does not pay
for its
+ /// setup.
+ #[inline(never)]
+ fn append_output_run(
+ &mut self,
+ buffered_batch_idx: usize,
+ buffered_range: Range<usize>,
+ batch_size: usize,
+ ) {
+ let len = buffered_range.len();
+ let streamed_idx = self.idx as u64;
+ let current_chunk = self.chunk_for(Some(buffered_batch_idx),
batch_size);
+ current_chunk
+ .streamed_indices
+ .append_value_n(streamed_idx, len);
+ // SAFETY: a `Range<u64>` reports its exact length in `size_hint`.
+ unsafe {
Review Comment:
This is the first `unsafe` in this module. The `SAFETY` comment is right,
since `Range<u64>` has an exact `size_hint`, and in arrow-rs
`append_trusted_len_iter` is `append_n_non_nulls(len)` followed by
`values_builder.extend(iter)`. The PR table shows the overall gain, but not
what `unsafe` adds over a safe `reserve` plus per-element append. Could you
report that difference for `inner_1to100`? If it is small, I would keep the
safe version.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -725,13 +793,16 @@ impl MaterializingSortMergeJoinStream {
while !self.buffered_data.scanning_finished()
&& self.num_unfrozen_pairs() < self.batch_size
{
- let scanning_idx = self.buffered_data.scanning_idx();
- self.streamed_batch.append_output_pair(
- Some(self.buffered_data.scanning_batch_idx),
- Some(scanning_idx),
+ // Pair the streamed row with the rest of the scanning buffered
batch at
+ // once, up to the room left in the output batch.
+ let rows = self.buffered_data.scanning_remaining();
+ let len = rows.len().min(self.batch_size -
self.num_unfrozen_pairs());
Review Comment:
Small simplifications, none blocking:
- `buffered_range.len()` is computed here (`rows.len()`), again in
`append_output_pairs` and again in `append_output_run`, and
`num_unfrozen_pairs()` is read twice per iteration. Computing `room` and `len`
once and passing `(start, len)` down would also let `append_output_run` drop
the `Range`.
- In release builds `len == 0` falls into the one-row branch of
`append_output_pairs` and appends a bogus pair, because the `len > 0` check is
only a `debug_assert!`. I checked that the invariant holds (empty-range batches
only trail the group and `scanning_advance_by` skips them), so matching on
`len` (`1 =>`, `_ =>`) would just keep it in the structure.
- `append_output_pairs` and `append_output_run` are hard to tell apart by
name. Something like `append_output_range` for the out-of-line half would read
better.
- The comment above `scanning_remaining()` repeats the function docs and can
go. The "cheaper to append directly" comment explains why, so it should stay.
--
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]
