feat: Store only the last row in the previous cursor for round robin tie-breaking purposes

datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs

@@ -373,16 +373,15 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     let mut partition_batches: Vec<Vec<RecordBatch>> = Vec::new();
 
     for stream_idx in 0..2usize {
-        // Each stream covers a non-overlapping key range so both are individually
-        // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000). When
-        // `tied_values` is set, every row of every batch in both streams
-        // instead carries the same sort key, so every comparison between the
-        // two streams is a tie.
+        // When `tied_values` is false, each stream covers a non-overlapping key range so both are
+        // individually sorted:
+        // stream 0 → [[0, 100), [200, 300),...]
+        // stream 1 → [[100, 200), [300, 400), ...]
+        //
+        // When `tied_values` is set, every row of every batch in both streams instead carries the
+        // same sort key, so every comparison between the two streams is a tie.
         let batches: Vec<RecordBatch> = (0..num_batches)
             .map(|b| {
-                // Interleave streams: stream 0 → even slots [0,200,400,...],
-                // stream 1 → odd slots [100,300,500,...] so the merge
-                // alternates between them on every batch.
                 let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32;
                 let sort_col: Int32Array = if tied_values {
                     std::iter::repeat_n(0, num_rows_per_batch).collect()
@@ -524,13 +523,7 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
     // BatchBuilder needs to hold 3 Record batches simultaneously to merge two
     // streams (because a stream can cross a record batch boundary)
     // there is also one cursor needed per stream
-    let mut max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size;
-
-    // with round robin enabled, 2 extra cursors live in memory
-    // see https://github.com/apache/datafusion/issues/23604
-    if round_robin {
-        max_peak += 2 * cursor_unit;
-    };
+    let max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size;
 
     assert!(
         peak_bytes > 0,

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

@@ -26,14 +26,18 @@ use arrow::array::{
 use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::SortOptions;
 use arrow::datatypes::ArrowNativeTypeOp;
-use arrow::row::Rows;
+use arrow::row::{OwnedRow, Rows};
 use datafusion_execution::memory_pool::MemoryReservation;
 
 /// A comparable collection of values for use with [`Cursor`]
 ///
 /// This is a trait as there are several specialized implementations, such as for
 /// single columns or for normalized multi column keys ([`Rows`])
 pub trait CursorValues: Debug + Sync + Send {
+    /// An owned copy of a single value, decoupled from any buffer that this
+    /// `CursorValues` holds (e.g. a shared `Buffer`/`Rows`).
+    type SingleRowValue: Send + Sync + Unpin;
+
     fn len(&self) -> usize;
 
     /// Returns true if `l[l_idx] == r[r_idx]`
@@ -46,6 +50,13 @@ pub trait CursorValues: Debug + Sync + Send {
     /// Returns comparison of `l[l_idx]` and `r[r_idx]`
     fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering;
 
+    /// Extract an owned copy of the value at `idx`.
+    fn get_value(&self, idx: usize) -> Self::SingleRowValue;
+
+    /// Returns true if `l[l_idx] == r`, where `r` was previously extracted
+    /// via [`Self::get_value`].
+    fn eq_to_single_row_value(l: &Self, l_idx: usize, r: &Self::SingleRowValue) -> bool;
+
     /// Notifies the values that the owning [`Cursor`] moved to `offset` (always
     /// `< len()`), so caching implementations can refresh the value(s) read by
     /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit).
@@ -113,15 +124,23 @@ impl<T: CursorValues> Cursor<T> {
         t
     }
 
-    pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor<T>>) -> bool {
+    pub fn is_eq_to_prev_one(&self, prev_value: Option<&T::SingleRowValue>) -> bool {
         if self.offset > 0 {
             self.is_eq_to_prev_row()
-        } else if let Some(prev_cursor) = prev_cursor {
-            self.is_eq_to_prev_row_in_prev_batch(prev_cursor)
+        } else if let Some(prev_value) = prev_value {
+            T::eq_to_single_row_value(&self.values, self.offset, prev_value)
         } else {
             false
         }
     }
+
+    /// Extract an owned copy of the last row in this cursor, decoupled from
+    /// any buffer the cursor's [`CursorValues`] holds. Used to remember a
+    /// partition's last row across a batch boundary without keeping the
+    /// whole exhausted batch's memory alive (see [`Self::is_eq_to_prev_one`]).
+    pub fn last_value(&self) -> T::SingleRowValue {
+        self.values.get_value(self.values.len() - 1)
+    }
 }
 
 impl<T: CursorValues> PartialEq for Cursor<T> {
@@ -135,16 +154,6 @@ impl<T: CursorValues> Cursor<T> {
     fn is_eq_to_prev_row(&self) -> bool {
         T::eq_to_previous(&self.values, self.offset)
     }
-
-    fn is_eq_to_prev_row_in_prev_batch(&self, other: &Self) -> bool {
-        assert_eq!(self.offset, 0);
-        T::eq(
-            &self.values,
-            self.offset,
-            &other.values,
-            other.values.len() - 1,
-        )
-    }
 }
 
 impl<T: CursorValues> Eq for Cursor<T> {}
@@ -195,6 +204,11 @@ impl RowValues {
 }
 
 impl CursorValues for RowValues {
+    // Reuse arrow-row's own owned-row type: `Row::owned()` copies just that
+    // row's bytes out of the shared `Rows` buffer, with no `RowConverter`
+    // needed (unlike building a new single-row `Rows`, which would).
+    type SingleRowValue = OwnedRow;
+
     #[inline]
     fn len(&self) -> usize {
         self.rows.num_rows()
@@ -215,6 +229,14 @@ impl CursorValues for RowValues {
     fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
         l.rows.row(l_idx).cmp(&r.rows.row(r_idx))
     }
+
+    fn get_value(&self, idx: usize) -> OwnedRow {
+        self.rows.row(idx).owned()
+    }
+
+    fn eq_to_single_row_value(l: &Self, l_idx: usize, r: &OwnedRow) -> bool {
+        l.rows.row(l_idx) == r.row()
+    }
 }
 
 /// An [`Array`] that can be converted into [`CursorValues`]
@@ -224,7 +246,10 @@ pub trait CursorArray: Array + 'static {
     fn values(&self) -> Self::Values;
 }
 
-impl<T: ArrowPrimitiveType> CursorArray for PrimitiveArray<T> {
+impl<T: ArrowPrimitiveType> CursorArray for PrimitiveArray<T>
+where
+    T::Native: Unpin,
+{
     type Values = PrimitiveValues<T::Native>;
 
     fn values(&self) -> Self::Values {
@@ -261,7 +286,10 @@ impl<T: ArrowNativeTypeOp> PrimitiveValues<T> {
     }
 }
 
-impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
+impl<T: ArrowNativeTypeOp + Unpin> CursorValues for PrimitiveValues<T> {
+    // Already `Copy`, so no buffer is retained by holding one.
+    type SingleRowValue = T;
+
     #[inline(always)]
     fn len(&self) -> usize {
         self.values.len()
@@ -297,6 +325,16 @@ impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
         self.current = self.values[offset];
         self.offset = offset;
     }
+
+    #[inline(always)]
+    fn get_value(&self, idx: usize) -> T {
+        self.values[idx]
+    }
+
+    #[inline(always)]
+    fn eq_to_single_row_value(l: &Self, l_idx: usize, r: &T) -> bool {
+        l.values[l_idx].is_eq(*r)
+    }
 }
 
 #[derive(Debug)]
@@ -319,6 +357,8 @@ impl<T: OffsetSizeTrait> ByteArrayValues<T> {
 }
 
 impl<T: OffsetSizeTrait> CursorValues for ByteArrayValues<T> {
+    type SingleRowValue = Box<[u8]>;
+
     #[inline]
     fn len(&self) -> usize {
         self.offsets.len() - 1
@@ -339,6 +379,14 @@ impl<T: OffsetSizeTrait> CursorValues for ByteArrayValues<T> {
     fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering {
         l.value(l_idx).cmp(r.value(r_idx))
     }
+
+    fn get_value(&self, idx: usize) -> Box<[u8]> {
+        self.value(idx).into()
+    }
+
+    fn eq_to_single_row_value(l: &Self, l_idx: usize, r: &Box<[u8]>) -> bool {
+        l.value(l_idx) == r.as_ref()
+    }
 }
 
 impl<T: ByteArrayType> CursorArray for GenericByteArray<T> {
@@ -360,6 +408,8 @@ impl CursorArray for StringViewArray {
 }
 
 impl CursorValues for StringViewArray {
+    type SingleRowValue = Box<[u8]>;
+
     fn len(&self) -> usize {
         self.views().len()
     }
@@ -422,6 +472,14 @@ impl CursorValues for StringViewArray {
 
         unsafe { GenericByteViewArray::compare_unchecked(l, l_idx, r, r_idx) }
     }
+
+    fn get_value(&self, idx: usize) -> Box<[u8]> {
+        self.value(idx).as_bytes().into()
+    }
+
+    fn eq_to_single_row_value(l: &Self, l_idx: usize, r: &Box<[u8]>) -> bool {
+        l.value(l_idx).as_bytes() == r.as_ref()
+    }
 }
 
 /// A collection of sorted, nullable [`CursorValues`]
@@ -471,6 +529,9 @@ impl<T: CursorValues> ArrayValues<T> {
 }
 
 impl<T: CursorValues> CursorValues for ArrayValues<T> {
+    // `None` represents a null value.
+    type SingleRowValue = Option<T::SingleRowValue>;
+
     #[inline(always)]
     fn len(&self) -> usize {
         self.values.len()
@@ -521,6 +582,28 @@ impl<T: CursorValues> CursorValues for ArrayValues<T> {
         // Forward to the wrapped values (e.g. caching `PrimitiveValues`).
         self.values.set_offset(offset);
     }
+
+    #[inline(always)]
+    fn get_value(&self, idx: usize) -> Option<T::SingleRowValue> {
+        if self.is_null(idx) {
+            None
+        } else {
+            Some(T::get_value(&self.values, idx))
+        }
+    }
+
+    #[inline(always)]
+    fn eq_to_single_row_value(
+        l: &Self,
+        l_idx: usize,
+        r: &Option<T::SingleRowValue>,
+    ) -> bool {
+        match (l.is_null(l_idx), r) {
+            (true, None) => true,
+            (false, Some(r)) => T::eq_to_single_row_value(&l.values, l_idx, r),
+            _ => false,
+        }
+    }
 }
 
 #[cfg(test)]

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

@@ -103,9 +103,10 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// Current reset count
     current_reset_epoch: usize,
 
-    /// Stores the previous value of each partitions for tracking the poll counts on the same value
-    /// Used if and only if round robin tie breaker is enabled, otherwise None
-    prev_cursors: Option<Vec<Option<Cursor<C>>>>,
+    /// Stores an owned copy of the last row of each partition's most
+    /// recently exhausted cursor, for tracking the poll counts on the same
+    /// value across a batch boundary.
+    prev_cursors: Option<Vec<Option<C::SingleRowValue>>>,
 
     /// Optional number of rows to fetch
     fetch: Option<usize>,
@@ -435,7 +436,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                 // Take the current cursor, leaving `None` in its place
                 let taken = self.cursors[stream_idx].take();
                 if let Some(prev_cursors) = &mut self.prev_cursors {
-                    prev_cursors[stream_idx] = taken;
+                    prev_cursors[stream_idx] = taken.map(|c| c.last_value());
                 }
             }
             return finished;
@@ -676,6 +677,8 @@ mod tests {
     struct DummyValues;
 
     impl CursorValues for DummyValues {
+        type SingleRowValue = ();
+
         fn len(&self) -> usize {
             0
         }
@@ -691,6 +694,18 @@ mod tests {
         fn compare(_l: &Self, _l_idx: usize, _r: &Self, _r_idx: usize) -> Ordering {
             unreachable!("done-path test should not compare cursors")
         }
+
+        fn get_value(&self, _idx: usize) -> Self::SingleRowValue {
+            unreachable!("done-path test should not compare cursors")
+        }
+
+        fn eq_to_single_row_value(
+            _l: &Self,
+            _l_idx: usize,
+            _r: &Self::SingleRowValue,
+        ) -> bool {
+            unreachable!("done-path test should not compare cursors")
+        }
     }
 
     #[tokio::test]