fix: SortPreservingMerge round-robin tie breaker reads stale poll counts, so tied keys are mostly drained from one partition

[jayzhan211]

Rationale for this change

SortPreservingMergeExec's round-robin tie breaker (enable_round_robin_repartition, on by default) is supposed to draw equal-key rows from the tied input partitions in turn, so that no partition's upstream buffer (e.g. RepartitionExec's) grows unbounded while another is drained.

It mostly didn't. Poll counts are invalidated lazily with an epoch, but only the winner's count was refreshed before the comparison; the challenger's count was read raw, so a count left over from an earlier run of ties leaked into the next one. The partition with the larger stale count then lost every tie until the other one "caught up", i.e. whole runs of equal keys were drained from a single partition.

On data shaped like our own sort_preserving_merge benchmark (3 identical partitions, 5 distinct keys), 260k of 300k rows were emitted in single-partition runs of 20,000 and only 40k rows actually alternated.

What changes are included in this PR?

• merge.rs: a poll_count() helper that applies the epoch check, used for both sides in is_poll_count_gt. No other behaviour change; the non-round-robin path is untouched.
• streaming_merge.rs: regression test test_round_robin_tie_breaker_resets_poll_counts_between_tie_runs, which fails on main.
• benches/sort_preserving_merge.rs: new bench_merge_tied_keys_slow_producers case (see below) covering the scenario the tie breaker is for.

Performance

cargo bench --bench sort_preserving_merge (clean A/B, --sample-size=20):

case	main	this PR	change
single_u64_column	17.5 ms	17.5 ms	−1%
multiple_u64_columns	32.0 ms	32.1 ms	+0.4% (n.s.)
single_large_string_column	77.3 ms	89.4 ms	+16%
multiple_large_string_columns	633 ms	662 ms	+4.5%

The u64 cases have unique keys, so the tie breaker never engages. The string cases are all ties (5 distinct values over 1M rows), and main was fast on them precisely because of the bug: it drained whole key-runs from one buffer sequentially instead of alternating per row. The slowdown is the cost of the fairness the feature exists to provide, not of the two extra loads in poll_count (which are noise next to the ~300-byte string compare per row).

Those cases have their inputs fully materialised in memory, so nothing upstream benefits from balanced consumption. The situation the tie breaker exists for is inputs that are produced concurrently: SortPreservingMergeExec runs each input in its own task buffered one batch ahead (spawn_buffered(_, 1)), so when the merge drains a single partition through a run of equal keys, that partition's producer is the bottleneck while the others idle. A new bench case, bench_merge_tied_keys_slow_producers (long runs of equal keys, fixed CPU cost per produced batch), shows the fix letting the producers overlap:

partitions	main	this PR	change
2	132.5 ms	82.1 ms	−38%
4	274.3 ms	223.3 ms	−19%

(With 4 inputs the tie breaker only balances the two sub-tree winners at the root, since ties below the root are still broken by index — a pre-existing limit.)

Follow-up worth considering: alternating at batch or N-row granularity instead of per row would keep this producer overlap while recovering most of the sequential-drain speed on the all-ties in-memory cases.

Are these changes tested?

Yes. The new unit test test_round_robin_tie_breaker_resets_poll_counts_between_tie_runs merges two streams of (key, tag) rows:

stream 0:  key=1 ×6,   key=2 ×8     (tag 0)
stream 1:  key=1 ×12,  key=2 ×8     (tag 1)

Stream 0 runs out of 1s first, so the first run of ties ends with stream 0 holding a large poll count while stream 1 drains its remaining 1s alone; then both reach key 2 and a second run of ties starts. Expected tag sequence (derived by hand from the algorithm, not a snapshot): [0,1]×6, [1]×6, [0,1]×8.

On main the test fails — once both streams are on key 2, stream 1 wins five times in a row, then stream 0 six times, before alternation starts:

main:     [0,1,0,1,0,1,0,1,0,1,0,1, 1,1,1,1,1,1, 1,1,1,1,1, 0,0,0,0,0,0, 1,0,1,0,1]
expected: [0,1,0,1,0,1,0,1,0,1,0,1, 1,1,1,1,1,1, 0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1]
                                                  ^ stale count from run one decides run two

With this PR the output matches exactly. The existing memory-limit based test_round_robin_tie_breaker_success / _fail tests still pass (they only bound memory, which is why they did not catch this), as do the full datafusion-physical-plan unit tests, core sort tests and sqllogictest (504 files).

Are there any user-facing changes?

Row order among rows with equal sort keys may differ from before when the round-robin tie breaker is enabled (it was never guaranteed in that mode; with_round_robin_repartition(false) remains stable by partition index).

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.38%. Comparing base (5134a1a) to head (d39a2fd).

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #24585   +/-   ##
=======================================
  Coverage   81.38%   81.38%           
=======================================
  Files        1116     1116           
  Lines      397960   398021   +61     
  Branches   397960   398021   +61     
=======================================
+ Hits       323880   323949   +69     
+ Misses      55120    55111    -9     
- Partials    18960    18961    +1     

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

[kumarUjjawal]

@ariel-miculas you might be interested in this.