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

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<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.
+    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<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),
+                            )
+                            .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);

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]

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<(i32, i32)>>) -> Vec<i32> {
+        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<i32>, Vec<i32>) = 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::<Int32Type>().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<i32> = [0, 1]
+            .repeat(6)
+            .into_iter()
+            .chain(std::iter::repeat_n(1, 6))
+            .chain([0, 1].repeat(8))
+            .collect();
+        assert_eq!(tags, expected);
+    }
 }