[branch-52] perf: Cache num_output_rows in sort merge join to avoid O(n) recount (#20478)

[mbutrovich]

DataFusion Comet is still suffering from performance regressions in SMJ introduced with DF 52. We'd like to backport this fix from main to try to help. Benchmarks with Comet look to return a bit of SMJ performance back:

DF 52.3:

This PR

Which issue does this PR close?

N/A - performance optimization

Rationale for this change

In the SMJ tight loop (join_partial), num_unfrozen_pairs() was called twice per iteration: once in the loop guard and once inside append_output_pair. This method iterates all chunks in output_indices and sums their lengths — O(num_chunks). Over a full batch of batch_size iterations, this makes the inner loop O(batch_size

• num_chunks) instead of O(batch_size).

What changes are included in this PR?

Add a num_output_rows field to StreamedBatch that is incremented on each append and reset on freeze, replacing the O(n) summation with an O(1) field read.

• Added num_output_rows: usize field to StreamedBatch, initialized to 0
• Increment num_output_rows in append_output_pair() after each append
• num_output_rows() now returns the cached field directly
• Reset to 0 in freeze_streamed() when output_indices is cleared
• Removed the num_unfrozen_pairs parameter from append_output_pair() since it can now read self.num_output_rows directly

Are these changes tested?

Yes — all 48 existing sort_merge_join tests pass. This is a pure refactor of an internal counter with no behavioral change.

Performance

Very minor improvement.

Before

sort_merge_join/inner_1to1/100000
                        time:   [3.8146 ms 3.8229 ms 3.8314 ms]
sort_merge_join/inner_1to10/100000
                        time:   [16.094 ms 16.125 ms 16.161 ms]
Found 7 outliers among 100 measurements (7.00%)
  6 (6.00%) high mild
  1 (1.00%) high severe
sort_merge_join/left_1to1_unmatched/100000
                        time:   [3.7823 ms 3.7861 ms 3.7902 ms]
Found 4 outliers among 100 measurements (4.00%)
  4 (4.00%) high mild
sort_merge_join/left_semi_1to10/100000
                        time:   [3.0523 ms 3.0755 ms 3.1023 ms]
Found 14 outliers among 100 measurements (14.00%)
  3 (3.00%) high mild
  11 (11.00%) high severe
sort_merge_join/left_anti_partial/100000
                        time:   [3.3458 ms 3.3498 ms 3.3542 ms]
Found 12 outliers among 100 measurements (12.00%)
  8 (8.00%) high mild
  4 (4.00%) high severe

After

sort_merge_join/inner_1to1/100000
                        time:   [3.7162 ms 3.7207 ms 3.7254 ms]
                        change: [−4.2320% −3.9309% −3.6431%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 4 outliers among 100 measurements (4.00%)
  4 (4.00%) high mild
sort_merge_join/inner_1to10/100000
                        time:   [15.556 ms 15.589 ms 15.626 ms]
                        change: [−5.2786% −4.8329% −4.4351%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 4 outliers among 100 measurements (4.00%)
  1 (1.00%) high mild
  3 (3.00%) high severe
sort_merge_join/left_1to1_unmatched/100000
                        time:   [3.7059 ms 3.7101 ms 3.7146 ms]
                        change: [−4.4526% −4.1565% −3.8660%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 2 outliers among 100 measurements (2.00%)
  2 (2.00%) high mild
sort_merge_join/left_semi_1to10/100000
                        time:   [3.0832 ms 3.0899 ms 3.0981 ms]
                        change: [−4.0965% −3.4158% −2.7657%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 3 outliers among 100 measurements (3.00%)
  1 (1.00%) high mild
  2 (2.00%) high severe
sort_merge_join/left_anti_partial/100000
                        time:   [3.2963 ms 3.3048 ms 3.3153 ms]
                        change: [−3.9413% −3.5316% −3.0884%] (p = 0.00 < 0.05)
                        Performance has improved.
Found 8 outliers among 100 measurements (8.00%)
  3 (3.00%) high mild
  5 (5.00%) high severe

Are there any user-facing changes?

No.

🤖 Generated with Claude Code

---

(cherry picked from commit 33b86fe)

Which issue does this PR close?

• Closes #.

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[comphead]

Thanks @mbutrovich

[mbutrovich]

Thanks @mbutrovich

Thanks @comphead for the review! Note that this is already in branch-53.

[alamb]

Thank you!