andygrove commented on code in PR #26277:
URL: https://github.com/apache/datafusion/pull/26277#discussion_r4319432812
##########
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:
Thanks, that's a clear win, and it keeps the module free of `unsafe`. The
chunks now keep their indices in `Vec<u64>`s, with a `NullBufferBuilder` for
the buffered nulls, and `freeze_streamed` builds the arrays with
`UInt64Array::try_new`, as in your diff.
Against the previous head (287d0a0), as the median of seven interleaved runs
of each binary on an M3 Max:
| Case | `append_value` loop | `Vec` + `NullBufferBuilder` | Change |
| --- | --- | --- | --- |
| `inner_1to1` | 1.402 ms | 1.427 ms | +1.8% |
| `inner_1to10` | 10.27 ms | 8.755 ms | −14.7% |
| `inner_1to100` | 18.96 ms | 15.98 ms | −15.7% |
| `left_1to100_filtered` | 47.66 ms | 44.59 ms | −6.5% |
| `left_1to1_unmatched` | 1.399 ms | 1.428 ms | +2.1% |
| `left_semi_1to10`, which doesn't use this stream | 0.607 ms | 0.607 ms |
+0.1% |
| `inner_group_spans_buffered_batches/512` | 3.262 ms | 2.981 ms | −8.6% |
| `inner_group_spans_buffered_batches/64` | 4.504 ms | 3.929 ms | −12.8% |
| `inner_group_spans_buffered_batches/8` | 14.88 ms | 12.41 ms | −16.6% |
The spanning cases are close to what you saw on the M5 Max, and the 1:10 and
1:100 gains are larger here. The second column also has the `saturating_sub`
from your other comment, and that's where the one-to-one cases lose their 2%.
With the `Vec`s alone, they were within 0.7% of the previous head. Thanks for
filing apache/arrow-rs#11717.
##########
datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs:
##########
@@ -722,18 +791,20 @@ impl MaterializingSortMergeJoinStream {
/// paused. Returns true when the group scan is complete and there is
/// room for more pairs.
fn pair_streamed_row_with_group(&mut self) -> bool {
- 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),
+ let mut room = self.batch_size - self.num_unfrozen_pairs();
Review Comment:
Done, thanks. A call with a full batch now appends nothing and returns
false, as on `main`.
It does cost something in my build. Against the plain subtraction, it made
`inner_1to1` and `left_1to1_unmatched` 1.3–1.9% slower over two sets of seven
interleaved runs, and nothing else moved. `checked_sub` with an early return
made them 3.0–3.6% slower. A predicted branch shouldn't cost more than a
`csel`, so I read this as code layout rather than the cost of the check, and
kept `saturating_sub`. Those cases are still about 4% faster than `main`.
--
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]
