From 0a9b0de974bcd643bd0cc2a73ca9b136cb398c4d Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 23 Aug 2026 11:05:17 +0800 Subject: [PATCH 1/3] fix: SortPreservingMerge round-robin tie breaker reads stale poll counts 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. --- datafusion/physical-plan/src/sorts/merge.rs | 28 ++++++- .../src/sorts/streaming_merge.rs | 77 ++++++++++++++++++- 2 files changed, 101 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bdfdbcf38b839..d08bd55de91ca 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -365,6 +365,21 @@ impl SortPreservingMergeStream { } } + /// 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 SortPreservingMergeStream { } } + /// 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/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 81adad8e9ec84..f5ea8aadd616c 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -281,7 +281,8 @@ mod tests { use super::*; - use arrow::array::{ArrayRef, RecordBatch}; + use arrow::array::{ArrayRef, AsArray, RecordBatch}; + use arrow::datatypes::{Field, Int32Type, Schema}; use arrow_schema::SortOptions; use datafusion_common::Result; use datafusion_execution::TaskContext; @@ -379,4 +380,78 @@ mod tests { Ok(()) } + + /// Merge streams of `(key, tag)` rows sorted on `key` with the round-robin + /// tie breaker enabled, returning the `tag` column of the output in order. + async fn merge_tags(streams: Vec>) -> Vec { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("tag", DataType::Int32, false), + ])); + let streams = streams + .into_iter() + .map(|rows| { + let (keys, tags): (Vec, Vec) = rows.into_iter().unzip(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(tags)), + ], + ) + .unwrap(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )) as SendableRecordBatchStream + }) + .collect(); + let sort: LexOrdering = + [PhysicalSortExpr::new_default(col("key", &schema).unwrap())].into(); + + let merged = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(schema) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(1024) + .with_bypass_mempool() + .with_round_robin_tie_breaker(true) + .build() + .unwrap(); + + collect(merged) + .await + .unwrap() + .iter() + .flat_map(|b| b.column(1).as_primitive::().values().to_vec()) + .collect() + } + + /// The round-robin tie breaker must start every run of equal keys with a + /// clean slate: poll counts left over from an earlier run of ties must not + /// influence which stream wins the next one. + #[tokio::test] + async fn test_round_robin_tie_breaker_resets_poll_counts_between_tie_runs() { + // Stream 0 runs out of `1`s first, so the first tie run ends with + // stream 1 holding several unanswered rows. The second run (key `2`) + // must then alternate from its first row rather than let stream 1 + // "catch up" on the stale count stream 0 accumulated during run one. + let stream0: Vec<_> = std::iter::repeat_n((1, 0), 6) + .chain(std::iter::repeat_n((2, 0), 8)) + .collect(); + let stream1: Vec<_> = std::iter::repeat_n((1, 1), 12) + .chain(std::iter::repeat_n((2, 1), 8)) + .collect(); + + let tags = merge_tags(vec![stream0, stream1]).await; + + let expected: Vec = [0, 1] + .repeat(6) + .into_iter() + .chain(std::iter::repeat_n(1, 6)) + .chain([0, 1].repeat(8)) + .collect(); + assert_eq!(tags, expected); + } } From d39a2fdddddb8c45289a855c4e6ea1f1e1040f32 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Sun, 23 Aug 2026 11:28:51 +0800 Subject: [PATCH 2/3] bench: add SortPreservingMerge case with tied keys and slow producers 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. --- .../benches/sort_preserving_merge.rs | 119 +++++++++++++++++- 1 file changed, 117 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e0..e68fe61a457ef 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,120 @@ 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 = (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::(); + 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. + for partitions in [2, 4] { + c.bench_function( + &format!("bench_merge_tied_keys_slow_producers/{partitions}_partitions"), + |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 = + 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), + ) + .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); From bb97b6d394d59661ee49083195b4d5aba9533cca Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Wed, 26 Aug 2026 21:46:17 +0800 Subject: [PATCH 3/3] feat: add tracing tests for round-robin tie breaker behavior in SortPreservingMerge --- .../benches/sort_preserving_merge.rs | 84 ++-- .../physical-plan/src/sorts/merge_trace.rs | 441 ++++++++++++++++++ datafusion/physical-plan/src/sorts/mod.rs | 2 + .../src/sorts/streaming_merge.rs | 77 +-- 4 files changed, 491 insertions(+), 113 deletions(-) create mode 100644 datafusion/physical-plan/src/sorts/merge_trace.rs diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index e68fe61a457ef..a2858e70c4af6 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -262,45 +262,55 @@ fn bench_merge_tied_keys_slow_producers(c: &mut Criterion) { 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] { - c.bench_function( - &format!("bench_merge_tied_keys_slow_producers/{partitions}_partitions"), - |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 = - 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), - ) - .build() - .unwrap(); - datafusion_physical_plan::common::collect(merged) - .await - .unwrap(); + 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 = + 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(); + }) }) - }) - }, - ); + }, + ); + } } } 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 0000000000000..95f9084e58a64 --- /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` — the merge asked input `n` for another batch +//! * `row S` — 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 }, + /// 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>>; + +/// 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; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + 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::() + .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> = 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 = 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::(); + let sources = batch.column(1).as_primitive::(); + 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::>() + .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::>() + .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 ca8d4a4400c49..4af2a3629fa9c 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; diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index f5ea8aadd616c..81adad8e9ec84 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -281,8 +281,7 @@ mod tests { use super::*; - use arrow::array::{ArrayRef, AsArray, RecordBatch}; - use arrow::datatypes::{Field, Int32Type, Schema}; + use arrow::array::{ArrayRef, RecordBatch}; use arrow_schema::SortOptions; use datafusion_common::Result; use datafusion_execution::TaskContext; @@ -380,78 +379,4 @@ mod tests { Ok(()) } - - /// Merge streams of `(key, tag)` rows sorted on `key` with the round-robin - /// tie breaker enabled, returning the `tag` column of the output in order. - async fn merge_tags(streams: Vec>) -> Vec { - let schema = Arc::new(Schema::new(vec![ - Field::new("key", DataType::Int32, false), - Field::new("tag", DataType::Int32, false), - ])); - let streams = streams - .into_iter() - .map(|rows| { - let (keys, tags): (Vec, Vec) = rows.into_iter().unzip(); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(keys)), - Arc::new(Int32Array::from(tags)), - ], - ) - .unwrap(); - Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&schema), - futures::stream::iter(vec![Ok(batch)]), - )) as SendableRecordBatchStream - }) - .collect(); - let sort: LexOrdering = - [PhysicalSortExpr::new_default(col("key", &schema).unwrap())].into(); - - let merged = StreamingMergeBuilder::new() - .with_streams(streams) - .with_schema(schema) - .with_expressions(&sort) - .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) - .with_batch_size(1024) - .with_bypass_mempool() - .with_round_robin_tie_breaker(true) - .build() - .unwrap(); - - collect(merged) - .await - .unwrap() - .iter() - .flat_map(|b| b.column(1).as_primitive::().values().to_vec()) - .collect() - } - - /// The round-robin tie breaker must start every run of equal keys with a - /// clean slate: poll counts left over from an earlier run of ties must not - /// influence which stream wins the next one. - #[tokio::test] - async fn test_round_robin_tie_breaker_resets_poll_counts_between_tie_runs() { - // Stream 0 runs out of `1`s first, so the first tie run ends with - // stream 1 holding several unanswered rows. The second run (key `2`) - // must then alternate from its first row rather than let stream 1 - // "catch up" on the stale count stream 0 accumulated during run one. - let stream0: Vec<_> = std::iter::repeat_n((1, 0), 6) - .chain(std::iter::repeat_n((2, 0), 8)) - .collect(); - let stream1: Vec<_> = std::iter::repeat_n((1, 1), 12) - .chain(std::iter::repeat_n((2, 1), 8)) - .collect(); - - let tags = merge_tags(vec![stream0, stream1]).await; - - let expected: Vec = [0, 1] - .repeat(6) - .into_iter() - .chain(std::iter::repeat_n(1, 6)) - .chain([0, 1].repeat(8)) - .collect(); - assert_eq!(tags, expected); - } }