mbutrovich commented on code in PR #26277:
URL: https://github.com/apache/datafusion/pull/26277#discussion_r4316610383
##########
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:
> 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.
What do you think about keeping the chunk's indices in a `Vec<u64>`, with a
`NullBufferBuilder` for the buffered nulls? Hash join already collects its
indices this way
([`chain.rs`](https://github.com/apache/datafusion/blob/fbf77371a0512765ebf81fdc60e4a6cdceb8101f/datafusion/physical-plan/src/joins/chain.rs#L35-L36),
[`hash_join/stream.rs`](https://github.com/apache/datafusion/blob/fbf77371a0512765ebf81fdc60e4a6cdceb8101f/datafusion/physical-plan/src/joins/hash_join/stream.rs#L1324)).
The bulk path then becomes a `Vec::extend` over the `Range<u64>` plus one
`append_n_non_nulls(len)`. That is safe. arrow-rs now recommends `Vec` over
`PrimitiveBuilder` for this
([docs](https://github.com/apache/arrow-rs/blob/3327fbef1c3381a149e9a413a5bb22813e11639e/arrow-array/src/builder/primitive_builder.rs#L100-L116)).
It closed apache/arrow-rs#7591, the request to optimize
`append_trusted_len_iter`, in favor of that approach.
At the head commit, the
[`append_value`](https://github.com/apache/arrow-rs/blob/3327fbef1c3381a149e9a413a5bb22813e11639e/arrow-array/src/builder/primitive_builder.rs#L246-L249)
loop runs 16 instructions per pair on aarch64. It loads and stores the null
builder's length and compares the values' length with their capacity for every
element. With `Vec::extend`, both move out of the loop, which leaves 5
instructions per pair, the same loop as `append_trusted_len_iter`. x86_64 has
the same shape. The listings are below.
<details>
<summary>Generated code</summary>
This is aarch64 in the release profile, which has LTO and one codegen unit.
The first listing is the head loop on the path where the null bitmap isn't
materialized:
```asm
LBB3366_18:
ldur x8, [x23, #-80] ; null bitmap materialized?
cbz x8, LBB3366_23 ; no: count-only path
... ; yes: set one bit
LBB3366_23:
ldur x8, [x23, #-40] ; null builder len, loaded and
add x8, x8, #1
stur x8, [x23, #-40] ; stored back for every element
ldur x26, [x23, #-88] ; values len
ldur x8, [x23, #-104] ; values capacity
cmp x26, x8
b.ne LBB3366_17 ; else grow_one
LBB3366_17:
ldur x8, [x23, #-96] ; values ptr
str x21, [x8, x26, lsl #3]
add x8, x26, #1
stur x8, [x23, #-88] ; values len stored back
add x21, x21, #1
subs x25, x25, #1
b.eq LBB3366_29 ; else back to LBB3366_18
```
That is 6 loads and 3 stores per pair. This is the `Vec::extend` loop:
```asm
LBB3367_9:
str x2, [x8, x1, lsl #3] ; store index
add x2, x2, #1
add x1, x1, #1
subs x9, x9, #1
b.ne LBB3367_9
```
On x86_64 (`--target x86_64-apple-darwin`), the head loop has the same
per-element `cmp qword ptr [r15 - 80], 0`, `inc qword ptr [r15 - 40]` and `cmp
r14, qword ptr [r15 - 104]`. The `Vec` loop is `mov qword ptr [rax + 8*rsi],
rdx; inc rdx; inc rsi; dec rcx; jne`.
</details>
These are `sort_merge_join` criterion medians on an Apple M5 Max, each
version compared with a saved baseline from the head commit. Rerunning the head
commit against its own baseline moved the cases by up to 3.6%.
| Case | Head | `append_trusted_len_iter` | `Vec` + `NullBufferBuilder` |
| --- | --- | --- | --- |
| `inner_1to1` | 1.148 ms | 1.151 ms (+0.5%) | 1.162 ms (+1.4%) |
| `inner_1to10` | 7.803 ms | 7.199 ms (-7.7%) | 7.032 ms (-9.9%) |
| `inner_1to100` | 13.97 ms | 13.06 ms (-6.5%) | 12.96 ms (-7.3%) |
| `left_1to100_filtered` | 37.20 ms | 36.10 ms (-2.9%) | 35.64 ms (-4.2%) |
| `inner_group_spans_buffered_batches/512` | 2.377 ms | 2.216 ms (-6.8%) |
2.180 ms (-8.3%) |
| `inner_group_spans_buffered_batches/64` | 3.493 ms | 3.238 ms (-7.3%) |
2.983 ms (-14.6%) |
| `inner_group_spans_buffered_batches/8` | 12.56 ms | 12.45 ms (-0.9%) |
10.14 ms (-19.3%) |
The extra gain at `/64` and `/8`, where chunks are small, comes from
`finish`.
[`PrimitiveBuilder::finish`](https://github.com/apache/arrow-rs/blob/3327fbef1c3381a149e9a413a5bb22813e11639e/arrow-array/src/builder/primitive_builder.rs#L365-L376)
builds an `ArrayData` for every call, and it is called twice per chunk
([here](https://github.com/apache/datafusion/blob/287d0a00b83674d15435aa8a1597f7309b30180d/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs#L1590-L1594)).
`UInt64Array::try_new` takes the `Vec` without that step. When I built the
`Vec` version's arrays through `ArrayData` again, `/8` went back to -4.2% and
`/64` to -8.6%. I filed apache/arrow-rs#11717 for the `finish` cost in arrow-rs.
This is the change I measured:
```diff
pub(super) struct StreamedJoinedChunk {
/// Index of batch in buffered_data
buffered_batch_idx: Option<usize>,
- /// Array builder for streamed indices
- streamed_indices: UInt64Builder,
- /// Array builder for buffered indices
- /// This could contain nulls if the join is null-joined
- buffered_indices: UInt64Builder,
+ /// Streamed indices, which are never null
+ streamed_indices: Vec<u64>,
+ /// Buffered indices, null where the streamed row is null-joined. A
null's
+ /// slot holds 0.
+ buffered_indices: Vec<u64>,
+ buffered_nulls: NullBufferBuilder,
}
@@ append_null_joined_pair
- current_chunk.streamed_indices.append_value(streamed_idx);
- current_chunk.buffered_indices.append_null();
+ current_chunk.streamed_indices.push(streamed_idx);
+ current_chunk.buffered_indices.push(0);
+ current_chunk.buffered_nulls.append_null();
@@ append_output_pairs, the `1` arm
- current_chunk.streamed_indices.append_value(streamed_idx);
- current_chunk
- .buffered_indices
- .append_value(buffered_start as u64);
+ current_chunk.streamed_indices.push(streamed_idx);
+ current_chunk.buffered_indices.push(buffered_start as u64);
+ current_chunk.buffered_nulls.append_non_null();
@@ append_output_range
current_chunk
.streamed_indices
- .append_value_n(streamed_idx, len);
+ .extend(repeat_n(streamed_idx, len));
let start = buffered_start as u64;
- for idx in start..start + len as u64 {
- current_chunk.buffered_indices.append_value(idx);
- }
+ current_chunk.buffered_indices.extend(start..start + len as u64);
+ current_chunk.buffered_nulls.append_n_non_nulls(len);
@@ start_chunk
- streamed_indices: UInt64Builder::with_capacity(capacity),
- buffered_indices: UInt64Builder::with_capacity(capacity),
+ streamed_indices: Vec::with_capacity(capacity),
+ buffered_indices: Vec::with_capacity(capacity),
+ buffered_nulls: NullBufferBuilder::new(capacity),
@@ freeze_streamed
- let left_indices = chunk.streamed_indices.finish();
+ let left_indices =
+ UInt64Array::try_new(take(&mut
chunk.streamed_indices).into(), None)?;
if left_indices.is_empty() {
continue;
}
- let right_indices: UInt64Array =
chunk.buffered_indices.finish();
+ let right_indices = UInt64Array::try_new(
+ take(&mut chunk.buffered_indices).into(),
+ chunk.buffered_nulls.finish(),
+ )?;
```
It also needs `use std::iter::repeat_n;` and `take` added to the `std::mem`
import.
--
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]
