andygrove commented on code in PR #26277:
URL: https://github.com/apache/datafusion/pull/26277#discussion_r4298359773
##########
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:
I measured both, as the median of ten interleaved runs of each binary on an
M3 Max:
| Case | `append_trusted_len_iter` | `append_value` loop | Change |
| --- | --- | --- | --- |
| `inner_1to100` | 19.39 ms | 20.27 ms | +4.6% |
| `inner_1to10` | 10.62 ms | 11.21 ms | +5.5% |
| `left_1to100_filtered` | 50.28 ms | 51.13 ms | +1.7% |
| `inner_group_spans_buffered_batches/512` | 3.42 ms | 3.54 ms | +3.5% |
| `inner_1to1`, which never reaches this code | 1.51 ms | 1.55 ms | +2.2% |
| `left_semi_1to10`, which doesn't use this stream | 0.649 ms | 0.643 ms |
−0.9% |
An earlier run of seven rounds gave +5.6% for `inner_1to100`. So the
`unsafe` is worth about 5% at 100 matches per row, where `main` takes 27 ms.
That's small next to the bulk path's own gain, so I switched to the loop, and
the module has no `unsafe` again. The chunk's builders are created with
capacity for the room left in the output batch, so the loop needs no `reserve`.
The safe `extend_from_iter_option` was no faster than the loop.
##########
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:
Added `left_1to100_filtered`. It joins the `inner_1to100` inputs as a left
join with the filter `left.data < right.data`, which keeps about half of the
pairs. The last streamed row of each key fails it for every buffered row, so
the deferred filtering also produces null-joined rows.
Against `main`, as the median of seven interleaved runs, it goes from 58.57
ms to 52.86 ms (−9.8%). The bulk path saves about as much as in `inner_1to100`
(5.7 ms, against 6.7 ms there), so the smaller relative gain comes from the
cost of the filtering, which this PR doesn't change.
##########
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:
Done, thanks:
- `pair_streamed_row_with_group` computes `room` once and counts it down,
and passes `rows.start` and `len` down, so neither `append_output_pairs` nor
the bulk half takes a `Range` now.
- `append_output_pairs` matches on `len`, with the bulk path as the `_` arm,
so a zero `len` in a release build would append no pair. The `debug_assert!`
stays.
- The bulk half is now `append_output_range`.
- The comment above `scanning_remaining()` is gone, and the "cheaper to
append directly" one is on the `1` arm.
--
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]
