Repository navigation
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.
| }); | ||
| } | ||
|
|
||
| // 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?
| // 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.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #26277 +/- ##
==========================================
- Coverage 82.83% 82.83% -0.01%
==========================================
Files 1149 1149
Lines 452798 452853 +55
Branches 452798 452853 +55
==========================================
+ Hits 375082 375121 +39
- Misses 54893 54902 +9
- Partials 22823 22830 +7 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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_run, which appends the streamed index withappend_value_nand the buffered indices as a range withappend_trusted_len_iter. It is kept out of line so that the one-row path doesn't pay for its setup.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).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.sort_merge_joinbenchmark againstmain, the median of five interleaved runs of each binary on an Apple M3 Max:maininner_1to1inner_1to10inner_1to100(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, so they show the noise level of this run. In a quieter run, in which the semi case moved 0.2%, the one-to-one cases came out 5.4% and 5.8% faster. Keeping the bulk path and the chunk creation inline instead made them about 3% slower than
main.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.