remove batch wise round robin and rely on row wise round robin

Author: rluvaton

datafusion/physical-plan/src/sorts/merge.rs

@@ -346,9 +346,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                 // The batch is fully emitted: retire the cursor like `advance_cursors` does for a finished one
                 self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
                 self.loser_tree_adjusted = false;
-                if self.enable_round_robin_tie_breaker {
-                    self.update_poll_count_after_skip(stream_idx);
-                }
 
                 if self.in_progress.is_empty() {
                     self.produced += batch.num_rows();
@@ -398,15 +395,13 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             if let Some(other) = &self.cursors[challenger] {
                 let beats = match winner_cursor.compare_last_to(other) {
                     Ordering::Less => true,
-                    // Mirror the loser tree's tie-breaking: round-robin by
-                    // poll count when enabled (see [`Self::handle_tie`]) so
-                    // batch-level ties alternate between streams instead of
-                    // starving one, otherwise by stream index like
-                    // [`Self::is_gt`]
                     Ordering::Equal => {
                         if self.enable_round_robin_tie_breaker {
-                            !self.is_skip_poll_count_gt(winner, challenger)
+                            // true if in round-robin because the winner already decided in a
+                            // round-robin fashion
+                            true
                         } else {
+                            // Keep sort stable
                             winner < challenger
                         }
                     }
@@ -485,38 +480,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         poll_a.cmp(&poll_b).then_with(|| a.cmp(&b)).is_gt()
     }
 
-    /// Poll count of `idx`, treating counts from a stale reset epoch as zero
-    /// (the same lazy reset [`Self::update_poll_count_on_the_same_value`] does)
-    #[inline]
-    fn effective_poll_count(&self, idx: usize) -> usize {
-        if self.poll_reset_epochs[idx] == self.current_reset_epoch {
-            self.num_of_polled_with_same_value[idx]
-        } else {
-            0
-        }
-    }
-
-    /// Like [`Self::is_poll_count_gt`], but epoch-aware since full-batch skips
-    /// happen outside the per-row tie-breaker flow that lazily resets counts
-    #[inline]
-    fn is_skip_poll_count_gt(&self, a: usize, b: usize) -> bool {
-        self.effective_poll_count(a)
-            .cmp(&self.effective_poll_count(b))
-            .then_with(|| a.cmp(&b))
-            .is_gt()
-    }
-
-    /// Count a full-batch skip as a poll so consecutive batch-level ties
-    /// round-robin between the streams instead of starving one
-    #[inline]
-    fn update_poll_count_after_skip(&mut self, idx: usize) {
-        if self.poll_reset_epochs[idx] != self.current_reset_epoch {
-            self.poll_reset_epochs[idx] = self.current_reset_epoch;
-            self.num_of_polled_with_same_value[idx] = 0;
-        }
-        self.num_of_polled_with_same_value[idx] += 1;
-    }
-
     #[inline]
     fn update_winner(&mut self, cmp_node: usize, winner: &mut usize, challenger: usize) {
         self.loser_tree[cmp_node] = *winner;

datafusion/physical-plan/src/sorts/streaming_merge.rs

@@ -437,82 +437,6 @@ mod tests {
         assert_eq!(flatten_vec(output), expected);
     }
 
-    #[tokio::test]
-    async fn test_merge_full_batch_skip_round_robins_between_tied_streams() {
-        // all sort keys are equal, so every batch-level comparison is a tie:
-        // the full-batch skip must round-robin between the streams instead of
-        // draining one stream completely while the other's input backs up
-        // (which is what the round-robin tie breaker exists to prevent)
-        const ROWS: usize = 4;
-        const BATCHES_PER_STREAM: usize = 4;
-
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("sort_key", DataType::Int32, false),
-            Field::new("stream_id", DataType::Int32, false),
-        ]));
-
-        let streams = (0..2_i32)
-            .map(|stream_id| {
-                let batches = (0..BATCHES_PER_STREAM)
-                    .map(|_| {
-                        RecordBatch::try_new(
-                            Arc::clone(&schema),
-                            vec![
-                                Arc::new(Int32Array::from(vec![7; ROWS])),
-                                Arc::new(Int32Array::from(vec![stream_id; ROWS])),
-                            ],
-                        )
-                        .unwrap()
-                    })
-                    .collect();
-                Box::pin(
-                    MemoryStream::try_new(batches, Arc::clone(&schema), None).unwrap(),
-                ) as SendableRecordBatchStream
-            })
-            .collect();
-
-        let merged = StreamingMergeBuilder::new()
-            .with_streams(streams)
-            .with_batch_size(ROWS)
-            .with_schema(Arc::clone(&schema))
-            .with_reservation(
-                MemoryConsumer::new("test")
-                    .register(&(Arc::new(UnboundedMemoryPool::default()) as _)),
-            )
-            .with_expressions(
-                &([
-                    PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap())
-                        .asc(),
-                ]
-                .into()),
-            )
-            .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0))
-            .with_fetch(None)
-            .build()
-            .unwrap();
-
-        let output = collect(merged).await.unwrap();
-
-        let stream_ids = output
-            .iter()
-            .flat_map(|b| b.column(1).as_primitive::<Int32Type>().values().iter())
-            .copied()
-            .collect::<Vec<_>>();
-        assert_eq!(stream_ids.len(), 2 * BATCHES_PER_STREAM * ROWS);
-
-        // neither stream may emit long runs while the other is starved:
-        // batch-level round-robin caps a run at one batch per stream
-        let longest_run = stream_ids
-            .chunk_by(|a, b| a == b)
-            .map(|run| run.len())
-            .max()
-            .unwrap();
-        assert!(
-            longest_run <= 2 * ROWS,
-            "stream emitted {longest_run} consecutive rows, tie breaker starved the other stream: {stream_ids:?}"
-        );
-    }
-
     #[tokio::test]
     async fn test_merge_interleaved_rows() {
         // every batch overlaps batches of the other streams