Skip to content

perf: append sort-merge join output pairs in bulk - #26277

Open
andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:andygrove/smj-bulk-output-pairs
Open

andygrove wants to merge 3 commits into
apache:mainfrom
andygrove:andygrove/smj-bulk-output-pairs

Conversation

@andygrove

@andygrove andygrove commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

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_group calls append_output_pair, which looks up the current chunk and appends one value to each of two UInt64Builders, 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 the take that builds the output batches.

What changes are included in this PR?

  • pair_streamed_row_with_group pairs the streamed row with the rest of the scanning buffered batch at once, capped by the room left in the output batch. The new BufferedData::scanning_remaining gives those rows, and BufferedData::scanning_advance_by moves the scan past them.
  • The new StreamedBatch::append_output_pairs appends a run of one row, as in a one-to-one join, directly. Longer runs go to append_output_run, which appends the streamed index with append_value_n and the buffered indices as a range with append_trusted_len_iter. It is kept out of line so that the one-row path doesn't pay for its setup.
  • The chunk lookup is now chunk_for, which keeps only its check inline and starts a new chunk in a cold start_chunk. With that, one-to-one joins are slightly faster than on main.
  • append_output_pair's only remaining caller passed a null buffered index, so it is now append_null_joined_pair. scanning_advance is replaced by scanning_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?

  • The existing sort-merge join tests cover this path for every join type, including those that run with batch sizes of 1, 2 and 50 and those that spill. The join fuzz tests, which compare sort-merge and hash join results across join types, pass (cargo test -p datafusion --features extended_tests --test fuzz -- join_fuzz).
  • New join_inner_group_spans_buffered_and_output_batches, next to join_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.
  • New benchmark case inner_1to100, in which every streamed row matches 100 buffered rows.

sort_merge_join benchmark against main, the median of five interleaved runs of each binary on an Apple M3 Max:

Case main This PR Change
inner_1to1 1.551 ms 1.481 ms −4.5%
inner_1to10 12.38 ms 9.21 ms −25.6%
inner_1to100 (new) 23.79 ms 16.70 ms −29.8%
left_1to1_unmatched 1.561 ms 1.478 ms −5.3%
left_semi_1to10 0.647 ms 0.624 ms −3.7%
left_anti_partial 0.802 ms 0.812 ms +1.2%
inner_group_spans_buffered_batches/512 4.08 ms 3.32 ms −18.7%
inner_group_spans_buffered_batches/64 5.35 ms 4.58 ms −14.5%
inner_group_spans_buffered_batches/8 16.35 ms 15.54 ms −4.9%

The 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.

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.
@andygrove
andygrove requested a review from comphead October 13, 2026 17:50
@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Oct 13, 2026
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.
@comphead

Copy link
Copy Markdown
Contributor

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 join_fuzz for the bulk path, but I think the fuzz inputs rarely produce a run of more than one row. make_staggered_batches_i32 draws (a, b) from 0..100 x 0..100 for 1000 rows, and on_columns() joins on both columns. That is about 10,000 distinct keys for 1000 rows, so most key groups hold one row and the len == 1 branch is taken almost every time. This is my estimate from the generator and I have not measured it.

The new code only differs from main when len > 1. The new unit test covers the batch boundary and output-batch cuts, but only for Inner and one hand-built shape. Could we add a low-cardinality case to the existing fuzz harness that compares sort-merge join with hash join? For example join on a only, or draw keys from 0..5. Using the existing batch_sizes list ([1, 2, 7, 49, 50, 51, 100]), for Inner, Left and Full, filtered and unfiltered, would exercise runs that end at a buffered batch boundary and runs capped by the output batch. Extending JoinFuzzTestCase should be enough, with no new harness.

The fuzz tests are also gated behind extended_tests, so the PR tier does not run them. A low-cardinality case would be worth having for any later change to this loop.

.streamed_indices
.append_value_n(streamed_idx, len);
// SAFETY: a `Range<u64>` reports its exact length in `size_hint`.
unsafe {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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());

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small simplifications, none blocking:

  • buffered_range.len() is computed here (rows.len()), again in append_output_pairs and again in append_output_run, and num_unfrozen_pairs() is read twice per iteration. Computing room and len once and passing (start, len) down would also let append_output_run drop the Range.
  • In release builds len == 0 falls into the one-row branch of append_output_pairs and appends a bogus pair, because the len > 0 check is only a debug_assert!. I checked that the invariant holds (empty-range batches only trail the group and scanning_advance_by skips them), so matching on len (1 =>, _ =>) would just keep it in the structure.
  • append_output_pairs and append_output_run are hard to tell apart by name. Something like append_output_range for 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-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.83%. Comparing base (2c568b4) to head (626ef73).

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants