You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/27 22:25:29 UTC
(datafusion) 01/01: fix: SortPreservingMerge round-robin tie breaker reads stale poll counts, so tied keys are mostly drained from one partition (#24585)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24585-f1a736ec37920d54bd42940277bb9e1db8b4c7f2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 3cabfa9be60255b3e9b05584d00dc92b961b1fed
Author: Jay Zhan <ja...@gmail.com>
AuthorDate: Thu Aug 27 22:25:04 2026 +0000
fix: SortPreservingMerge round-robin tie breaker reads stale poll counts, so tied keys are mostly drained from one partition (#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).
---
.../physical-plan/benches/sort_preserving_merge.rs | 129 +++++-
datafusion/physical-plan/src/sorts/merge.rs | 28 +-
datafusion/physical-plan/src/sorts/merge_trace.rs | 441 +++++++++++++++++++++
datafusion/physical-plan/src/sorts/mod.rs | 2 +
4 files changed, 595 insertions(+), 5 deletions(-)
diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs
index 76ebf230a3..a2858e70c4 100644
--- a/datafusion/physical-plan/benches/sort_preserving_merge.rs
+++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs
@@ -16,7 +16,7 @@
// under the License.
use arrow::{
- array::{ArrayRef, StringArray, UInt64Array},
+ array::{ArrayRef, AsArray, StringArray, UInt64Array},
record_batch::RecordBatch,
};
use arrow_schema::{SchemaRef, SortOptions};
@@ -193,5 +193,130 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) {
}
}
-criterion_group!(benches, bench_merge_sorted_preserving);
+/// Merge inputs whose keys are mostly *tied* and whose producers do real work
+/// per batch.
+///
+/// `SortPreservingMergeExec` runs each input in its own task, buffered one
+/// batch ahead (`spawn_buffered(_, 1)`). If the merge keeps draining a single
+/// partition during a run of equal keys, that partition's producer becomes the
+/// bottleneck while the others idle on their one buffered batch. The
+/// round-robin tie breaker is meant to spread consumption across the tied
+/// partitions so all producers stay busy.
+fn bench_merge_tied_keys_slow_producers(c: &mut Criterion) {
+ use datafusion_execution::memory_pool::{
+ MemoryConsumer, MemoryPool, UnboundedMemoryPool,
+ };
+ use datafusion_physical_plan::common::spawn_buffered;
+ use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
+ use datafusion_physical_plan::sorts::streaming_merge::StreamingMergeBuilder;
+ use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
+ use futures::StreamExt;
+
+ const ROWS: usize = 400_000;
+ const BATCH: usize = 8192;
+ const ROWS_PER_KEY: usize = 100_000;
+
+ let schema: SchemaRef = Arc::new(arrow_schema::Schema::new(vec![
+ arrow_schema::Field::new("key", arrow_schema::DataType::UInt64, false),
+ arrow_schema::Field::new("val", arrow_schema::DataType::UInt64, false),
+ ]));
+ let sort_order = LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("key", &schema).unwrap(),
+ SortOptions::default(),
+ )])
+ .unwrap();
+
+ // Every partition holds the same long runs of equal keys.
+ let batches: Vec<RecordBatch> = (0..ROWS.div_ceil(BATCH))
+ .map(|b| {
+ let start = b * BATCH;
+ let n = BATCH.min(ROWS - start);
+ let keys = UInt64Array::from_iter_values(
+ (start..start + n).map(|i| (i / ROWS_PER_KEY) as u64),
+ );
+ let vals =
+ UInt64Array::from_iter_values((start..start + n).map(|i| i as u64));
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(keys), Arc::new(vals)],
+ )
+ .unwrap()
+ })
+ .collect();
+
+ /// Stand-in for an upstream operator: ~fixed CPU cost per batch.
+ fn produce(batch: RecordBatch) -> RecordBatch {
+ let vals = batch
+ .column(1)
+ .as_primitive::<arrow::datatypes::UInt64Type>();
+ let mut acc = 0u64;
+ for _ in 0..200 {
+ for v in vals.values() {
+ acc = acc.wrapping_mul(6364136223846793005).wrapping_add(*v);
+ }
+ }
+ std::hint::black_box(acc);
+ batch
+ }
+
+ let rt = tokio::runtime::Runtime::new().unwrap();
+ // With 2 inputs the root comparison is the whole tree, so the tie breaker
+ // balances all producers; with 4 it only balances the two sub-tree winners.
+ //
+ // 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.
+ for partitions in [2, 4] {
+ for (label, round_robin) in [("on", true), ("off", false)] {
+ c.bench_function(
+ &format!(
+ "bench_merge_tied_keys_slow_producers/{partitions}_partitions/tie_breaker_{label}"
+ ),
+ |b| {
+ b.iter(|| {
+ rt.block_on(async {
+ let streams = (0..partitions)
+ .map(|_| {
+ let s = RecordBatchStreamAdapter::new(
+ Arc::clone(&schema),
+ futures::stream::iter(batches.clone())
+ .map(|b| Ok(produce(b))),
+ );
+ spawn_buffered(Box::pin(s), 1)
+ })
+ .collect();
+ let pool: Arc<dyn MemoryPool> =
+ Arc::new(UnboundedMemoryPool::default());
+ let merged = StreamingMergeBuilder::new()
+ .with_streams(streams)
+ .with_schema(Arc::clone(&schema))
+ .with_expressions(&sort_order)
+ .with_metrics(BaselineMetrics::new(
+ &ExecutionPlanMetricsSet::new(),
+ 0,
+ ))
+ .with_batch_size(BATCH)
+ .with_reservation(
+ MemoryConsumer::new("bench").register(&pool),
+ )
+ .with_round_robin_tie_breaker(round_robin)
+ .build()
+ .unwrap();
+ datafusion_physical_plan::common::collect(merged)
+ .await
+ .unwrap();
+ })
+ })
+ },
+ );
+ }
+ }
+}
+
+criterion_group!(
+ benches,
+ bench_merge_sorted_preserving,
+ bench_merge_tied_keys_slow_producers
+);
criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index bdfdbcf38b..d08bd55de9 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -365,6 +365,21 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
}
+ /// Returns the poll count of `partition_idx` for the current tie-breaker
+ /// round.
+ ///
+ /// Poll counts are reset lazily by bumping `current_reset_epoch` (see
+ /// [`Self::reset_poll_counts`]), so a count written in an older epoch is
+ /// stale and reads as 0.
+ #[inline]
+ fn poll_count(&self, partition_idx: usize) -> usize {
+ if self.poll_reset_epochs[partition_idx] == self.current_reset_epoch {
+ self.num_of_polled_with_same_value[partition_idx]
+ } else {
+ 0
+ }
+ }
+
/// For the given partition, updates the poll count. If the current value is the same
/// of the previous value, it increases the count by 1; otherwise, it is reset as 0.
fn update_poll_count_on_the_same_value(&mut self, partition_idx: usize) {
@@ -457,11 +472,18 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
}
+ /// Returns `true` if partition `a` has been polled more often than `b` in
+ /// the current tie-breaker round, breaking equal counts by partition index.
+ ///
+ /// Both counts go through [`Self::poll_count`]: only the winner's count is
+ /// refreshed by [`Self::update_poll_count_on_the_same_value`] before this
+ /// is called, so the challenger's raw count may belong to an earlier round.
#[inline]
fn is_poll_count_gt(&self, a: usize, b: usize) -> bool {
- let poll_a = self.num_of_polled_with_same_value[a];
- let poll_b = self.num_of_polled_with_same_value[b];
- poll_a.cmp(&poll_b).then_with(|| a.cmp(&b)).is_gt()
+ self.poll_count(a)
+ .cmp(&self.poll_count(b))
+ .then_with(|| a.cmp(&b))
+ .is_gt()
}
#[inline]
diff --git a/datafusion/physical-plan/src/sorts/merge_trace.rs b/datafusion/physical-plan/src/sorts/merge_trace.rs
new file mode 100644
index 0000000000..95f9084e58
--- /dev/null
+++ b/datafusion/physical-plan/src/sorts/merge_trace.rs
@@ -0,0 +1,441 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Snapshot tests for *which input the merge drains, and in what order*.
+//!
+//! The round-robin tie breaker in [`SortPreservingMergeStream`] never changes
+//! the merged output: rows with equal sort keys are interchangeable, so the
+//! sorted values are identical whether it is on or off. All it changes is
+//! which tied input gets consumed next, and therefore which upstream producers
+//! keep making progress. Ordinary correctness tests are blind to that, which
+//! makes the behavior easy to regress silently.
+//!
+//! This module drives a merge over hand-written inputs and records an
+//! interleaved trace of everything an outside observer can see:
+//!
+//! * `poll S<n>` — the merge asked input `n` for another batch
+//! * `row S<n>` — the merge emitted a row that came from input `n`
+//!
+//! The traces are asserted with inline `insta` snapshots, so a change in poll
+//! order shows up as a snapshot diff instead of passing unnoticed. To accept
+//! an intended change, run the tests with `INSTA_FORCE_UPDATE=1` and apply the
+//! result with `cargo insta accept`, rather than hand-editing the expected
+//! text -- the point is that changing it is a deliberate act.
+//!
+//! [`SortPreservingMergeStream`]: crate::sorts::merge::SortPreservingMergeStream
+
+use std::fmt::Write as _;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll};
+
+use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
+use crate::sorts::streaming_merge::StreamingMergeBuilder;
+use crate::stream::RecordBatchStreamAdapter;
+use crate::{RecordBatchStream, SendableRecordBatchStream};
+
+use arrow::array::{AsArray, Int32Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
+use datafusion_common::Result;
+use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col};
+
+use futures::{Stream, StreamExt};
+use insta::assert_snapshot;
+
+/// One event observed while driving a merge.
+#[derive(Debug)]
+enum TraceEvent {
+ /// The merge polled input `partition`, which returned a batch of `keys`.
+ Batch { partition: usize, keys: Vec<i32> },
+ /// The merge polled input `partition`, which reported end of stream.
+ Eof { partition: usize },
+ /// The merge emitted an output row that originated in `partition`.
+ Row { partition: usize, key: i32 },
+}
+
+/// Shared, append-only log of [`TraceEvent`]s.
+///
+/// The input streams append to it from inside the merge's `poll_next`, and the
+/// driver appends emitted rows in between, which is what makes the two kinds
+/// of event interleave in issue order.
+type TraceLog = Arc<Mutex<Vec<TraceEvent>>>;
+
+/// Wraps one input of the merge and logs every poll that resolves.
+///
+/// `Poll::Pending` is not logged: these inputs are always immediately ready,
+/// so a pending poll never happens.
+struct TracingStream {
+ partition: usize,
+ inner: SendableRecordBatchStream,
+ log: TraceLog,
+}
+
+impl Stream for TracingStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ let partition = self.partition;
+ let polled = self.inner.poll_next_unpin(cx);
+ match &polled {
+ Poll::Ready(Some(Ok(batch))) => {
+ let keys = batch
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .values()
+ .to_vec();
+ self.log
+ .lock()
+ .unwrap()
+ .push(TraceEvent::Batch { partition, keys });
+ }
+ Poll::Ready(None) => {
+ self.log.lock().unwrap().push(TraceEvent::Eof { partition });
+ }
+ Poll::Ready(Some(Err(_))) | Poll::Pending => {}
+ }
+ polled
+ }
+}
+
+impl RecordBatchStream for TracingStream {
+ fn schema(&self) -> SchemaRef {
+ self.inner.schema()
+ }
+}
+
+/// `(key, source partition)` pairs; `key` is the sort key, the partition index
+/// rides along so emitted rows can be attributed back to their input.
+fn trace_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("key", DataType::Int32, false),
+ Field::new("src", DataType::Int32, false),
+ ]))
+}
+
+/// Builds one input stream: `batches` is a list of batches, each a list of
+/// sort keys. All rows are tagged with `partition`.
+fn input_stream(partition: usize, batches: &[&[i32]], log: &TraceLog) -> TracingStream {
+ let schema = trace_schema();
+ let batches: Vec<Result<RecordBatch>> = batches
+ .iter()
+ .map(|keys| {
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(Int32Array::from(keys.to_vec())),
+ Arc::new(Int32Array::from(vec![partition as i32; keys.len()])),
+ ],
+ )
+ .map_err(Into::into)
+ })
+ .collect();
+
+ TracingStream {
+ partition,
+ inner: Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::iter(batches),
+ )),
+ log: Arc::clone(log),
+ }
+}
+
+/// Merges `inputs` (partition -> batch -> sort keys) and returns the rendered
+/// trace of polls and emitted rows.
+///
+/// The merge runs with `batch_size = 1` so every output row is handed back
+/// individually and can be slotted into the log at the point it was produced.
+/// Output batching is independent of the loser tree and the tie breaker, so
+/// this makes the trace finer-grained without changing what is being traced.
+async fn trace_merge(inputs: &[&[&[i32]]], round_robin_tie_breaker: bool) -> String {
+ let schema = trace_schema();
+ let log: TraceLog = Arc::new(Mutex::new(Vec::new()));
+
+ let streams: Vec<SendableRecordBatchStream> = inputs
+ .iter()
+ .enumerate()
+ .map(|(partition, batches)| {
+ Box::pin(input_stream(partition, batches, &log)) as SendableRecordBatchStream
+ })
+ .collect();
+
+ let ordering: LexOrdering =
+ [PhysicalSortExpr::new_default(col("key", &schema).unwrap())].into();
+
+ let mut merged = StreamingMergeBuilder::new()
+ .with_streams(streams)
+ .with_schema(Arc::clone(&schema))
+ .with_expressions(&ordering)
+ .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0))
+ .with_batch_size(1)
+ .with_bypass_mempool()
+ .with_round_robin_tie_breaker(round_robin_tie_breaker)
+ .build()
+ .unwrap();
+
+ while let Some(batch) = merged.next().await {
+ let batch = batch.unwrap();
+ let keys = batch.column(0).as_primitive::<Int32Type>();
+ let sources = batch.column(1).as_primitive::<Int32Type>();
+ let mut log = log.lock().unwrap();
+ for row in 0..batch.num_rows() {
+ log.push(TraceEvent::Row {
+ partition: sources.value(row) as usize,
+ key: keys.value(row),
+ });
+ }
+ }
+
+ render(inputs, round_robin_tie_breaker, &log.lock().unwrap())
+}
+
+/// Renders the inputs and the recorded events as the snapshot text.
+fn render(
+ inputs: &[&[&[i32]]],
+ round_robin_tie_breaker: bool,
+ events: &[TraceEvent],
+) -> String {
+ let mut out = String::new();
+ writeln!(out, "round_robin_tie_breaker = {round_robin_tie_breaker}").unwrap();
+ writeln!(out, "inputs:").unwrap();
+ for (partition, batches) in inputs.iter().enumerate() {
+ let batches = batches
+ .iter()
+ .map(|keys| format!("[{}]", join(keys)))
+ .collect::<Vec<_>>()
+ .join(" ");
+ writeln!(out, " S{partition}: {batches}").unwrap();
+ }
+ writeln!(out, "trace:").unwrap();
+ for event in events {
+ match event {
+ TraceEvent::Batch { partition, keys } => {
+ writeln!(out, " poll S{partition} -> [{}]", join(keys)).unwrap()
+ }
+ TraceEvent::Eof { partition } => {
+ writeln!(out, " poll S{partition} -> done").unwrap()
+ }
+ TraceEvent::Row { partition, key } => {
+ writeln!(out, " row S{partition} key={key}").unwrap()
+ }
+ }
+ }
+ out
+}
+
+fn join(keys: &[i32]) -> String {
+ keys.iter()
+ .map(|k| k.to_string())
+ .collect::<Vec<_>>()
+ .join(",")
+}
+
+/// Traces the same inputs with the tie breaker enabled and disabled, so a
+/// single snapshot shows what the tie breaker changes.
+async fn trace_both(inputs: &[&[&[i32]]]) -> String {
+ format!(
+ "{}\n{}",
+ trace_merge(inputs, true).await,
+ trace_merge(inputs, false).await
+ )
+}
+
+/// With every key tied, the tie breaker must hand rows out in strict
+/// round-robin order; with it off, the lowest-index input is drained first and
+/// the other producer is left idle.
+#[tokio::test]
+async fn test_tied_keys_alternate_between_inputs() {
+ let inputs: &[&[&[i32]]] = &[&[&[1, 1, 1]], &[&[1, 1, 1]]];
+
+ assert_snapshot!(trace_both(inputs).await, @r"
+ round_robin_tie_breaker = true
+ inputs:
+ S0: [1,1,1]
+ S1: [1,1,1]
+ trace:
+ poll S0 -> [1,1,1]
+ poll S1 -> [1,1,1]
+ row S0 key=1
+ row S1 key=1
+ row S0 key=1
+ row S1 key=1
+ row S0 key=1
+ poll S0 -> done
+ row S1 key=1
+ poll S1 -> done
+
+ round_robin_tie_breaker = false
+ inputs:
+ S0: [1,1,1]
+ S1: [1,1,1]
+ trace:
+ poll S0 -> [1,1,1]
+ poll S1 -> [1,1,1]
+ row S0 key=1
+ row S0 key=1
+ row S0 key=1
+ poll S0 -> done
+ row S1 key=1
+ row S1 key=1
+ row S1 key=1
+ poll S1 -> done
+ ");
+}
+
+/// Regression test for stale poll counts leaking across runs of ties.
+///
+/// `S0` runs out of `1`s first, so the first run of ties ends with `S1` still
+/// holding several rows that it drains on its own. The second run (key `2`)
+/// is a fresh tie-breaker round and must alternate from its very first row.
+///
+/// Before the fix, `is_poll_count_gt` compared the challenger's *raw* counter,
+/// which still held the count `S1` accumulated during the key `1` run, and the
+/// key `2` run started `S1 S0 S0 S1` instead of alternating.
+#[tokio::test]
+async fn test_poll_counts_reset_between_runs_of_ties() {
+ let inputs: &[&[&[i32]]] = &[&[&[1, 1], &[2, 2]], &[&[1, 1, 1, 1], &[2, 2]]];
+
+ assert_snapshot!(trace_both(inputs).await, @r"
+ round_robin_tie_breaker = true
+ inputs:
+ S0: [1,1] [2,2]
+ S1: [1,1,1,1] [2,2]
+ trace:
+ poll S0 -> [1,1]
+ poll S1 -> [1,1,1,1]
+ row S0 key=1
+ row S1 key=1
+ row S0 key=1
+ poll S0 -> [2,2]
+ row S1 key=1
+ row S1 key=1
+ row S1 key=1
+ poll S1 -> [2,2]
+ row S0 key=2
+ row S1 key=2
+ row S0 key=2
+ poll S0 -> done
+ row S1 key=2
+ poll S1 -> done
+
+ round_robin_tie_breaker = false
+ inputs:
+ S0: [1,1] [2,2]
+ S1: [1,1,1,1] [2,2]
+ trace:
+ poll S0 -> [1,1]
+ poll S1 -> [1,1,1,1]
+ row S0 key=1
+ row S0 key=1
+ poll S0 -> [2,2]
+ row S1 key=1
+ row S1 key=1
+ row S1 key=1
+ row S1 key=1
+ poll S1 -> [2,2]
+ row S0 key=2
+ row S0 key=2
+ poll S0 -> done
+ row S1 key=2
+ row S1 key=2
+ poll S1 -> done
+ ");
+}
+
+/// Without ties there is nothing to break: both settings must produce the same
+/// poll order. A diff here means the tie breaker started interfering with the
+/// ordinary loser-tree path.
+#[tokio::test]
+async fn test_distinct_keys_are_unaffected_by_the_tie_breaker() {
+ let inputs: &[&[&[i32]]] = &[&[&[1, 3, 5]], &[&[2, 4, 6]]];
+
+ assert_snapshot!(trace_both(inputs).await, @r"
+ round_robin_tie_breaker = true
+ inputs:
+ S0: [1,3,5]
+ S1: [2,4,6]
+ trace:
+ poll S0 -> [1,3,5]
+ poll S1 -> [2,4,6]
+ row S0 key=1
+ row S1 key=2
+ row S0 key=3
+ row S1 key=4
+ row S0 key=5
+ poll S0 -> done
+ row S1 key=6
+ poll S1 -> done
+
+ round_robin_tie_breaker = false
+ inputs:
+ S0: [1,3,5]
+ S1: [2,4,6]
+ trace:
+ poll S0 -> [1,3,5]
+ poll S1 -> [2,4,6]
+ row S0 key=1
+ row S1 key=2
+ row S0 key=3
+ row S1 key=4
+ row S0 key=5
+ poll S0 -> done
+ row S1 key=6
+ poll S1 -> done
+ ");
+}
+
+/// The tie breaker only runs at the root of the loser tree (`cmp_node == 1`),
+/// so with four tied inputs it balances the two sub-tree winners rather than
+/// all four producers: `S1` and `S3` still drain in one go.
+///
+/// This is a known limitation, recorded here so that widening the tie breaker
+/// to inner nodes shows up as a deliberate snapshot change. Note also that the
+/// initial polls are not in partition order — `initialize_all_partitions` uses
+/// `swap_remove`, which reorders the pending list.
+#[tokio::test]
+async fn test_tie_breaker_only_balances_the_root_comparison() {
+ let inputs: &[&[&[i32]]] = &[&[&[1, 1]], &[&[1, 1]], &[&[1, 1]], &[&[1, 1]]];
+
+ assert_snapshot!(trace_merge(inputs, true).await, @r"
+ round_robin_tie_breaker = true
+ inputs:
+ S0: [1,1]
+ S1: [1,1]
+ S2: [1,1]
+ S3: [1,1]
+ trace:
+ poll S0 -> [1,1]
+ poll S3 -> [1,1]
+ poll S2 -> [1,1]
+ poll S1 -> [1,1]
+ row S0 key=1
+ row S2 key=1
+ row S0 key=1
+ poll S0 -> done
+ row S1 key=1
+ row S1 key=1
+ poll S1 -> done
+ row S2 key=1
+ poll S2 -> done
+ row S3 key=1
+ row S3 key=1
+ poll S3 -> done
+ ");
+}
diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs
index ca8d4a4400..4af2a3629f 100644
--- a/datafusion/physical-plan/src/sorts/mod.rs
+++ b/datafusion/physical-plan/src/sorts/mod.rs
@@ -20,6 +20,8 @@
mod builder;
mod cursor;
mod merge;
+#[cfg(test)]
+mod merge_trace;
mod multi_level_merge;
pub mod partial_sort;
pub mod partitioned_topk;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org