fix: SortPreservingMerge round-robin tie breaker reads stale poll counts, so tied keys are mostly drained from one partition - #24585
Conversation
Poll counts are invalidated lazily via an epoch, but only the winner's count was refreshed before `is_poll_count_gt`; the challenger's raw count could belong to an earlier run of ties. The partition with the larger stale count then lost every tie until the other caught up, so whole runs of equal keys were drained from a single partition instead of alternating. Read both counts through an epoch-aware `poll_count()` and add a regression test that fails on main.
Covers the scenario the round-robin tie breaker exists for: inputs with long runs of equal keys whose batches cost CPU to produce. Each input runs in its own task buffered one batch ahead, so draining a single partition through a tie run serialises on that producer.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24585 +/- ##
==========================================
+ Coverage 81.38% 81.46% +0.07%
==========================================
Files 1116 1120 +4
Lines 397960 400691 +2731
Branches 397960 400691 +2731
==========================================
+ Hits 323880 326407 +2527
- Misses 55120 55207 +87
- Partials 18960 19077 +117 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@ariel-miculas you might be interested in this. |
|
I took a look, is this a regression or was it always broken? |
The streams are not being pulled as evenly as expected. |
But was this always the case, since the feature was first implemented? |
I think so, based on the change history |
|
Maybe we can invest in some "mid level" tests -- aka make the testing the poll order easier For example, we could have some sort of insta / snapshot test that set up some streams and then polled them and output a trace of what stream got polled and in what order With such a testing framework, I think we could more easily reason about / understand the impacts of PRs like this My personal concern is also that this might accindentally regress something else (or maybe be regressed itself in the future) |
| .chain(std::iter::repeat_n((2, 1), 8)) | ||
| .collect(); | ||
|
|
||
| let tags = merge_tags(vec![stream0, stream1]).await; |
There was a problem hiding this comment.
this tests correctness, but does it validate the order in which the streams are polled? I feel like the change in behavior is not easy to understand from this test
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
@jayzhan211, thanks for working on this. The epoch-aware poll count fix looks good to me, and the regression test covers the stale-count behavior across consecutive equal-key runs nicely. I left one non-blocking suggestion about making the new benchmark more useful for comparing the round-robin tie breaker.
| // balances all producers; with 4 it only balances the two sub-tree winners. | ||
| for partitions in [2, 4] { | ||
| c.bench_function( | ||
| &format!("bench_merge_tied_keys_slow_producers/{partitions}_partitions"), |
There was a problem hiding this comment.
Could we add paired cases with the round-robin tie breaker enabled and disabled here, using with_round_robin_tie_breaker explicitly? I think that would make this benchmark more useful, since we could directly measure the benefit of the tie breaker and catch regressions where the enabled behavior becomes similar to the disabled baseline. Not blocking for this PR.
| row S1 key=1 | ||
| row S1 key=1 | ||
| poll S1 -> done | ||
| "); |
There was a problem hiding this comment.
This makes what’s going on much more straightforward.
| // Each case is run with the tie breaker both enabled and disabled so the | ||
| // pair measures what the tie breaker actually buys on this workload: if | ||
| // the two ever converge, the balancing has regressed into the | ||
| // lowest-index-wins baseline. |
There was a problem hiding this comment.
benches/sort_preserving_merge.rs now runs each partition count with
.with_round_robin_tie_breaker(true) and (false), so the pair directly
measures what the tie breaker buys on this workload.
cargo bench -p datafusion-physical-plan --bench sort_preserving_merge -- \
bench_merge_tied_keys_slow_producers
| case | tie breaker on | tie breaker off | speedup |
|---|---|---|---|
| 2 partitions | 82.8 ms | 150.8 ms | 1.82x |
| 4 partitions | 228.7 ms | 303.3 ms | 1.33x |
The gap is the thing to watch: if the enabled path ever regresses toward the
lowest-index-wins baseline, these two columns converge.
|
Comments addressed! |
6 participants
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_mergebenchmark (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: apoll_count()helper that applies the epoch check, used for both sides inis_poll_count_gt. No other behaviour change; the non-round-robin path is untouched.streaming_merge.rs: regression testtest_round_robin_tie_breaker_resets_poll_counts_between_tie_runs, which fails onmain.benches/sort_preserving_merge.rs: newbench_merge_tied_keys_slow_producerscase (see below) covering the scenario the tie breaker is for.Performance
cargo bench --bench sort_preserving_merge(clean A/B,--sample-size=20):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
mainwas 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 inpoll_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:
SortPreservingMergeExecruns 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:(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_runsmerges two streams of(key, tag)rows: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 remaining1s alone; then both reach key2and a second run of ties starts. Expectedtagsequence (derived by hand from the algorithm, not a snapshot):[0,1]×6, [1]×6, [0,1]×8.On
mainthe test fails — once both streams are on key2, stream 1 wins five times in a row, then stream 0 six times, before alternation starts:With this PR the output matches exactly. The existing memory-limit based
test_round_robin_tie_breaker_success/_failtests still pass (they only bound memory, which is why they did not catch this), as do the fulldatafusion-physical-planunit 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).