From f7fb326538058b3d15633069ca30a87b07496c37 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 15 Jul 2026 20:09:52 +0300 Subject: [PATCH 1/5] fix: store only the last row in prev_cursors instead of the entire cursor --- datafusion/physical-plan/src/sorts/cursor.rs | 115 ++++++++++++++++--- datafusion/physical-plan/src/sorts/merge.rs | 23 +++- 2 files changed, 118 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index d71eaad663410..003de2375ad3f 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -26,7 +26,7 @@ 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`] @@ -34,6 +34,10 @@ use datafusion_execution::memory_pool::MemoryReservation; /// 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 Cursor { t } - pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> 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 PartialEq for Cursor { @@ -135,16 +154,6 @@ impl Cursor { 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 Eq for Cursor {} @@ -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 CursorArray for PrimitiveArray { +impl CursorArray for PrimitiveArray +where + T::Native: Unpin, +{ type Values = PrimitiveValues; fn values(&self) -> Self::Values { @@ -261,7 +286,10 @@ impl PrimitiveValues { } } -impl CursorValues for PrimitiveValues { +impl CursorValues for PrimitiveValues { + // 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 CursorValues for PrimitiveValues { 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 ByteArrayValues { } impl CursorValues for ByteArrayValues { + type SingleRowValue = Box<[u8]>; + #[inline] fn len(&self) -> usize { self.offsets.len() - 1 @@ -339,6 +379,14 @@ impl CursorValues for ByteArrayValues { 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 CursorArray for GenericByteArray { @@ -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 ArrayValues { } impl CursorValues for ArrayValues { + // `None` represents a null value. + type SingleRowValue = Option; + #[inline(always)] fn len(&self) -> usize { self.values.len() @@ -521,6 +582,28 @@ impl CursorValues for ArrayValues { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). self.values.set_offset(offset); } + + #[inline(always)] + fn get_value(&self, idx: usize) -> Option { + 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, + ) -> 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)] diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 647649038766d..aa85a07b0adcd 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -103,9 +103,10 @@ pub(crate) struct SortPreservingMergeStream { /// 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>>>, + /// 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>>, /// Optional number of rows to fetch fetch: Option, @@ -435,7 +436,7 @@ impl SortPreservingMergeStream { // 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] From 3c84b3cf8ccd6401bf93536cef5796529c673f24 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 15 Jul 2026 20:20:06 +0300 Subject: [PATCH 2/5] feat: update test to reflect the improvements --- .../fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index f82d0165f2fdb..b21d372d44a34 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -524,13 +524,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, From 68f46c373aa07873d0a6d1e11b5d5d7fb78b384b Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Fri, 7 Aug 2026 13:02:55 +0300 Subject: [PATCH 3/5] feat: remove the no_round_robin tests --- .../spilling_fuzz_in_memory_constrained_env.rs | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index b21d372d44a34..24b8410955770 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -307,36 +307,18 @@ async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin() run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, false).await } -#[tokio::test] -async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin() --> Result<()> { - run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, false).await -} - #[tokio::test] async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column() -> Result<()> { run_sort_preserving_merge_peak_memory_with_spilled_input(true, true, false).await } -#[tokio::test] -async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column() --> Result<()> { - run_sort_preserving_merge_peak_memory_with_spilled_input(false, true, false).await -} - #[tokio::test] async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values() -> Result<()> { run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, true).await } -#[tokio::test] -async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values() --> Result<()> { - run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, true).await -} - /// Intended to measure the maximum number of record batches held in memory by /// the SortPreservingMergeStream in a convoluted way by measuring the peak /// memory reservation. Relevant for merging spilled streams, where the produced From f3920c867465582b461d31967cded83e6642ec35 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Fri, 7 Aug 2026 13:08:41 +0300 Subject: [PATCH 4/5] fix: comment regarding the value distribution in streams --- .../spilling_fuzz_in_memory_constrained_env.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index 24b8410955770..f92fdaa28fd06 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -355,16 +355,15 @@ async fn run_sort_preserving_merge_peak_memory_with_spilled_input( let mut partition_batches: Vec> = 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 = (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() From 06bde7b495803620701b38df1ae5c166942069ce Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Fri, 7 Aug 2026 16:31:48 +0300 Subject: [PATCH 5/5] Revert "feat: remove the no_round_robin tests" This reverts commit 68f46c373aa07873d0a6d1e11b5d5d7fb78b384b. --- .../spilling_fuzz_in_memory_constrained_env.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index f92fdaa28fd06..761c492d8dde3 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -307,18 +307,36 @@ async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin() run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, false).await } +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, false).await +} + #[tokio::test] async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column() -> Result<()> { run_sort_preserving_merge_peak_memory_with_spilled_input(true, true, false).await } +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, true, false).await +} + #[tokio::test] async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values() -> Result<()> { run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, true).await } +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, true).await +} + /// Intended to measure the maximum number of record batches held in memory by /// the SortPreservingMergeStream in a convoluted way by measuring the peak /// memory reservation. Relevant for merging spilled streams, where the produced