jayzhan211 opened a new pull request, #24585:
URL: https://github.com/apache/datafusion/pull/24585
## 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 `1`s first, so the first run of ties ends with stream 0
holding a large poll count while stream 1 drains its remaining `1`s 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).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
