Repository navigation
perf: append sort-merge join output pairs in bulk - #26277
Conversation
When a streamed row matches a buffered key group, pair it with the rest of the scanning buffered batch at once, up to the room left in the output batch, instead of one buffered row at a time. The pairs and their order are unchanged. Adds a test whose key group spans buffered and output batches, and a benchmark case in which every streamed row matches 100 buffered rows.
Match the other sort-merge join tests, which compare the whole output with batches_to_string.
- append_output_pairs hands runs of more than one row to an out-of-line append_output_run, so the one-row case that one-to-one joins take for every row doesn't pay for the bulk path's setup. - chunk_for keeps only its check inline and starts a new chunk in a cold start_chunk. - append_output_pair's only caller passes a null buffered index, so it is now append_null_joined_pair. - BufferedData::scanning_remaining returns the unscanned rows of the scanning batch, and the one-caller scanning_advance is replaced by scanning_advance_by(1). - Move the new test next to the related one-row-batch test.
|
Thanks for the PR. The change reads correct to me, and keeping the one-row path inline while moving the bulk path out of line makes sense. One question on test coverage. The description cites The new code only differs from The fuzz tests are also gated behind |
| .streamed_indices | ||
| .append_value_n(streamed_idx, len); | ||
| // SAFETY: a `Range<u64>` reports its exact length in `size_hint`. | ||
| unsafe { |
There was a problem hiding this 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.
There was a problem hiding this 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.
There was a problem hiding this comment.
So the
unsafeis worth about 5% at 100 matches per row, wheremaintakes 27 ms. That's small next to the bulk path's own gain, so I switched to the loop, and the module has nounsafeagain.
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, hash_join/stream.rs). 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). 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 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.
Generated code
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:
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_18That is 6 loads and 3 stores per pair. This is the Vec::extend loop:
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_9On 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.
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 builds an ArrayData for every call, and it is called twice per chunk (here). 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:
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.
There was a problem hiding this 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 Vecs alone, they were within 0.7% of the previous head. Thanks for filing apache/arrow-rs#11717.
| }); | ||
| } | ||
|
|
||
| // 1:100 Inner Join — 20K left, 20K right, 200 distinct keys |
There was a problem hiding this 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?
There was a problem hiding this 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.
| // 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()); |
There was a problem hiding this comment.
Small simplifications, none blocking:
buffered_range.len()is computed here (rows.len()), again inappend_output_pairsand again inappend_output_run, andnum_unfrozen_pairs()is read twice per iteration. Computingroomandlenonce and passing(start, len)down would also letappend_output_rundrop theRange.- In release builds
len == 0falls into the one-row branch ofappend_output_pairsand appends a bogus pair, because thelen > 0check is only adebug_assert!. I checked that the invariant holds (empty-range batches only trail the group andscanning_advance_byskips them), so matching onlen(1 =>,_ =>) would just keep it in the structure. append_output_pairsandappend_output_runare hard to tell apart by name. Something likeappend_output_rangefor 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.
There was a problem hiding this comment.
Done, thanks:
pair_streamed_row_with_groupcomputesroomonce and counts it down, and passesrows.startandlendown, so neitherappend_output_pairsnor the bulk half takes aRangenow.append_output_pairsmatches onlen, with the bulk path as the_arm, so a zerolenin a release build would append no pair. Thedebug_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 the1arm.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #26277 +/- ##
==========================================
+ Coverage 82.83% 82.92% +0.08%
==========================================
Files 1149 1154 +5
Lines 452798 457757 +4959
Branches 452798 457757 +4959
==========================================
+ Hits 375082 379600 +4518
- Misses 54893 55147 +254
- Partials 22823 23010 +187 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
- pair_streamed_row_with_group computes the room left in the output batch once and counts it down, and passes the start and length of each run instead of a Range. - append_output_pairs matches on the length, with the bulk path as the catch-all arm, so a run of no rows would append no pair in a release build. The loop never produces one, and the debug_assert stays. - Rename append_output_run to append_output_range, and drop a comment that repeats the docs.
Replace append_trusted_len_iter with a loop of append_value. The chunk's builders are created with capacity for the room left in the output batch, so the loop needs no reserve. The loop is about 5% slower at 100 matches per row, but leaves the module without unsafe.
left_1to100_filtered joins the inputs of inner_1to100 as a left join with the filter left.data < right.data, so that the pairs go through deferred filtering. The filter keeps about half of them, and none of the last streamed row of each key, which is null-joined.
Most key groups in the join fuzz inputs hold one row, so a sort-merge join almost never pairs a streamed row with a run of more than one buffered row there, and never cuts such a run at the output batch. Add Inner, Left, Right and Full cases, unfiltered and filtered, whose keys take 25 values, so that each streamed row matches about 40 buffered rows and runs end both at buffered batch boundaries and where the output batch fills.
You were right. I counted the appends with a temporary counter. Across the 21 joins of I added Their joins return about 40,000 rows, so each new case takes about 40 s of CPU in the |
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Thanks for looking at this, @andygrove! I left one minor comment, and a reply to earlier discussion (as well as filing an upstream arrow-rs issue).
| 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(); |
There was a problem hiding this comment.
On main, the loop condition num_unfrozen_pairs() < self.batch_size means a call with a full batch appends nothing and returns false, so the caller freezes and calls again. Here that state underflows the subtraction. Debug builds panic. In release, room wraps to near usize::MAX, and the loop appends the rest of the group past batch_size. As I read the callers, the Ordering::Less arm and the while !self.pair_streamed_row_with_group() loop keep the count below batch_size, so this only matters if a later change breaks that. Should we keep main's behavior here with saturating_sub?
| let mut room = self.batch_size - self.num_unfrozen_pairs(); | |
| let mut room = self.batch_size.saturating_sub(self.num_unfrozen_pairs()); |
There was a problem hiding this 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.
StreamedJoinedChunk keeps its indices in Vec<u64>s, with a NullBufferBuilder for the buffered nulls, instead of in UInt64Builders, as hash join does. A run of buffered rows is then one Vec::extend over the range and one append_n_non_nulls, so the per-row null and capacity checks move out of the loop. freeze_streamed builds the arrays with UInt64Array::try_new, which skips the ArrayData that PrimitiveBuilder::finish builds for every chunk.
On main, pair_streamed_row_with_group appends nothing and returns false when it is called with a full output batch, so the caller freezes and calls again. Computing the room with saturating_sub keeps that behavior where the subtraction would underflow. The callers keep the count below batch_size, so this only matters if a later change breaks that.
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Thank you @andygrove!
Which issue does this PR close?
Rationale for this change
A sort-merge join in which each row matches many rows on the other side spends most of its time recording the matches. For every buffered row of the key group,
pair_streamed_row_with_groupcallsappend_output_pair, which looks up the current chunk and appends one value to each of twoUInt64Builders, and then advances the scan by one row.In DataFusion Comet, which runs DataFusion 55.2.0, a sort-merge join of 1M rows in which every row matches 100 rows (105M joined rows feeding a
sum) spent 46% of its active CPU in these two functions, against 14% in thetakethat builds the output batches.What changes are included in this PR?
pair_streamed_row_with_grouppairs the streamed row with the rest of the scanning buffered batch at once, capped by the room left in the output batch. The newBufferedData::scanning_remaininggives those rows, andBufferedData::scanning_advance_bymoves the scan past them.StreamedBatch::append_output_pairsappends a run of one row, as in a one-to-one join, directly. Longer runs go toappend_output_range, which extends the chunk's streamed indices withrepeat_nand its buffered indices with the range of rows, and appends their validity with oneappend_n_non_nulls. It is kept out of line so that the one-row path doesn't pay for its setup.StreamedJoinedChunkkeeps its indices inVec<u64>s, as hash join does, with aNullBufferBuilderfor the buffered nulls, instead of inUInt64Builders. That takes the per-row null and capacity checks out of the bulk append, andfreeze_streamedbuilds the arrays withUInt64Array::try_new, which skips theArrayDatathatPrimitiveBuilder::finishbuilds for every chunk (Build arrays directly in builderfinishinstead of throughArrayDataarrow-rs#11717).chunk_for, which keeps only its check inline and starts a new chunk in a coldstart_chunk. With that, one-to-one joins are slightly faster than onmain.append_output_pair's only remaining caller passed a null buffered index, so it is nowappend_null_joined_pair.scanning_advanceis replaced byscanning_advance_by(1).The pairs, their order and the chunks they go to are unchanged, so freezing, filtering and spilling see the same indices as before.
What is the testing strategy for this PR?
cargo test -p datafusion --features extended_tests --test fuzz -- join_fuzz). Most of their key groups hold one row, though, so only 3–5% of their appends are runs of more than one row, and none is cut by the output batch. The new low-cardinality cases,test_{inner,left,right,full}_join_1k_low_cardinalityand their_filteredvariants, draw the keys from 25 values. In each of them about 30% of the appends are runs of more than one row, 14,000–19,000 runs are cut by the output batch, and 51,000–56,000 end at a buffered batch boundary with the group continuing in the next one.join_inner_group_spans_buffered_and_output_batches, next tojoin_with_group_spanning_batches_rescanned_per_streamed_row, which uses one-row buffered batches. In the new test the key group spans buffered batches of 3, 2 and 1 rows and the batch size is 4, so runs of several rows end both at the end of a buffered batch and when the output batch fills.inner_1to100, in which every streamed row matches 100 buffered rows, andleft_1to100_filtered, which joins the same inputs with a filter, so that the pairs go through deferred filtering.sort_merge_joinbenchmark againstmain, the median of seven interleaved runs of each binary on an Apple M3 Max:maininner_1to1inner_1to10inner_1to100(new)left_1to100_filtered(new)left_1to1_unmatchedleft_semi_1to10left_anti_partialinner_group_spans_buffered_batches/512inner_group_spans_buffered_batches/64inner_group_spans_buffered_batches/8The semi and anti cases don't use this stream, and they moved by 0.1% at most. Keeping the bulk path and the chunk creation inline instead made the one-to-one cases about 5% slower than
main. Keeping the indices inUInt64Builders, with a loop ofappend_valuefor a run's buffered indices, made the cases with several matches per row 7–20% slower.In DataFusion Comet (apache/datafusion-comet#6907), a prototype of the bulk path on 55.2.0 cut the native sort-merge join's CPU time by 39% at 100 matches per row and by 46% at 1000. At 100 matches per row that took it from twice the time of Spark's whole-stage-codegen sort-merge join to the same time.
Are there any user-facing changes?
No. Sort-merge joins in which rows match many rows run faster.