From 2947e23cb774ae3ee55b91690670113a70c40277 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 18:35:11 +0300 Subject: [PATCH 01/10] perf(sort): fast path for full batch when it sorts before all other streams --- datafusion/physical-plan/src/sorts/builder.rs | 9 +++ datafusion/physical-plan/src/sorts/cursor.rs | 12 +++ datafusion/physical-plan/src/sorts/merge.rs | 77 +++++++++++++++++++ 3 files changed, 98 insertions(+) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..f1a545dfe4ae9 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -129,6 +129,15 @@ impl BatchBuilder { &self.schema } + /// The rows of `stream_idx`'s current batch that have not been pushed + /// via [`Self::push_row`] yet, as a zero-copy slice. + pub fn remaining_rows_of(&self, stream_idx: usize) -> Option { + let cursor = &self.cursors[stream_idx]; + let (_, batch) = &self.batches[cursor.batch_idx]; + let remaining = batch.num_rows() - cursor.row_idx; + (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)) + } + /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 8991922779d4a..63348469c6add 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -116,6 +116,18 @@ impl Cursor { t } + /// Returns true if no rows of this cursor were consumed yet + pub fn is_at_start(&self) -> bool { + self.offset == 0 + } + + /// Compare this cursor's *last* row to `other`'s current row. Since the + /// values are sorted, `Less` means every remaining row of this cursor + /// sorts before all of `other`'s remaining rows. + pub fn compare_last_to(&self, other: &Self) -> Ordering { + T::compare(&self.values, self.values.len() - 1, &other.values, other.offset) + } + pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { if self.offset > 0 { self.is_eq_to_prev_row() diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..07bbd1f62f8f5 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -153,6 +153,10 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, + + /// A whole winner batch that was proven to sort before every other + /// stream's head, waiting to be emitted after the in-progress rows are flushed. + pending_skipped_batch: Option, } impl SortPreservingMergeStream { @@ -186,6 +190,7 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + pending_skipped_batch: None, } } @@ -236,6 +241,12 @@ impl SortPreservingMergeStream { } return Poll::Ready(None); } + + if let Some(batch) = self.pending_skipped_batch.take() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Once all partitions have set their corresponding cursors for the loser tree, // we skip the following block. Until then, this function may be called multiple // times and can return Poll::Pending if any partition returns Poll::Pending. @@ -310,6 +321,42 @@ impl SortPreservingMergeStream { } let stream_idx = self.loser_tree[0]; + + // Full-batch skip: + // when the winner sits at the start of a fresh batch whose *last* row sorts before every other stream's + // current row, we skip per-row tree updates and just emit it. + // Checked at most once per input batch so we don't add a lot of overhead when + // the data is not sorted at all + // + // This might produce batches that do not equal to the configured batch_size, + // this is something we are willing to sacrifice + // + // TODO - support fetch + if self.fetch.is_none() + && self + .cursors[stream_idx] + .as_ref() + .is_some_and(|cursor| cursor.is_at_start()) + && self.winner_batch_beats_all(stream_idx) + { + let batch = self + .in_progress + .remaining_rows_of(stream_idx) + .expect("winner cursor is at the start, so rows remain"); + // 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.in_progress.is_empty() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Flush already-merged rows first + // they all sort before the skipped batch + self.pending_skipped_batch = Some(batch); + return Poll::Ready(self.emit_in_progress_batch().transpose()); + } + if self.advance_cursors(stream_idx) { self.loser_tree_adjusted = false; self.in_progress.push_row(stream_idx); @@ -327,6 +374,36 @@ impl SortPreservingMergeStream { } } + /// Returns `true` if every remaining row of `winner`'s current batch + /// sorts before every other stream's current row, using the same index + /// tie-breaking as [`Self::is_gt`]. Exhausted streams cannot compete. + fn winner_batch_beats_all(&self, winner: usize) -> bool { + let Some(winner_cursor) = &self.cursors[winner] else { + return false; + }; + + // Only the streams that lost *directly* to the winner can hold the + // overall runner-up, and those are exactly the losers stored on the + // winner's leaf-to-root path (every other stream lost to one of them + // transitively), so `O(log k)` comparisons suffice instead of `O(k)`. + let mut node = self.lt_leaf_node_index(winner); + while node != 0 { + let challenger = self.loser_tree[node]; + if let Some(other) = &self.cursors[challenger] { + if !winner_cursor + .compare_last_to(other) + .then_with(|| winner.cmp(&challenger)) + .is_lt() + { + return false; + } + } + node = self.lt_parent_node_index(node); + } + + true + } + /// 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) { From c2bb46983c4ff1f5a81fe7b4c4104e181eed202b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:16:00 +0300 Subject: [PATCH 02/10] added tests --- datafusion/physical-plan/src/sorts/cursor.rs | 42 +++- datafusion/physical-plan/src/sorts/merge.rs | 34 ++- .../src/sorts/streaming_merge.rs | 229 ++++++++++++++++++ 3 files changed, 286 insertions(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 63348469c6add..58a4dff2bb384 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -43,8 +43,18 @@ pub trait CursorValues { fn eq_to_previous(cursor: &Self, idx: usize) -> bool; /// Returns comparison of `l[l_idx]` and `r[r_idx]` + /// + /// Only called with the cursors' *current* offsets, so caching + /// implementations may serve it from their cache. fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary* + /// indices, unlike [`Self::compare`]. Cold path: caching implementations + /// must override this to index directly. + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + Self::compare(l, l_idx, r, r_idx) + } + /// 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). @@ -125,7 +135,12 @@ impl Cursor { /// values are sorted, `Less` means every remaining row of this cursor /// sorts before all of `other`'s remaining rows. pub fn compare_last_to(&self, other: &Self) -> Ordering { - T::compare(&self.values, self.values.len() - 1, &other.values, other.offset) + T::compare_at( + &self.values, + self.values.len() - 1, + &other.values, + other.offset, + ) } pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { @@ -302,6 +317,11 @@ impl CursorValues for PrimitiveValues { l.current.compare(r.current) } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Arbitrary indices (cold path), so index directly instead of using the cache. + l.values[l_idx].compare(r.values[r_idx]) + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that @@ -530,6 +550,26 @@ impl CursorValues for ArrayValues { } } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Same null handling as `compare`, but delegates to the inner + // `compare_at` so arbitrary indices work with caching values. + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => Ordering::Equal, + (true, false) => match l.options.nulls_first { + true => Ordering::Less, + false => Ordering::Greater, + }, + (false, true) => match l.options.nulls_first { + true => Ordering::Greater, + false => Ordering::Less, + }, + (false, false) => match l.options.descending { + true => T::compare_at(&r.values, r_idx, &l.values, l_idx), + false => T::compare_at(&l.values, l_idx, &r.values, r_idx), + }, + } + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 07bbd1f62f8f5..8c464a3cf1413 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -246,7 +246,7 @@ impl SortPreservingMergeStream { self.produced += batch.num_rows(); return Poll::Ready(Some(Ok(batch))); } - + // Once all partitions have set their corresponding cursors for the loser tree, // we skip the following block. Until then, this function may be called multiple // times and can return Poll::Pending if any partition returns Poll::Pending. @@ -333,16 +333,15 @@ impl SortPreservingMergeStream { // // TODO - support fetch if self.fetch.is_none() - && self - .cursors[stream_idx] - .as_ref() - .is_some_and(|cursor| cursor.is_at_start()) - && self.winner_batch_beats_all(stream_idx) + && self.cursors[stream_idx] + .as_ref() + .is_some_and(|cursor| cursor.is_at_start()) + && self.winner_batch_beats_all(stream_idx) { let batch = self - .in_progress - .remaining_rows_of(stream_idx) - .expect("winner cursor is at the start, so rows remain"); + .in_progress + .remaining_rows_of(stream_idx) + .expect("winner cursor is at the start, so rows remain"); // 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; @@ -356,7 +355,7 @@ impl SortPreservingMergeStream { self.pending_skipped_batch = Some(batch); return Poll::Ready(self.emit_in_progress_batch().transpose()); } - + if self.advance_cursors(stream_idx) { self.loser_tree_adjusted = false; self.in_progress.push_row(stream_idx); @@ -389,14 +388,13 @@ impl SortPreservingMergeStream { let mut node = self.lt_leaf_node_index(winner); while node != 0 { let challenger = self.loser_tree[node]; - if let Some(other) = &self.cursors[challenger] { - if !winner_cursor - .compare_last_to(other) - .then_with(|| winner.cmp(&challenger)) - .is_lt() - { - return false; - } + if let Some(other) = &self.cursors[challenger] + && !winner_cursor + .compare_last_to(other) + .then_with(|| winner.cmp(&challenger)) + .is_lt() + { + return false; } node = self.lt_parent_node_index(node); } diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..53b704f7aedb1 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -265,3 +265,232 @@ impl<'a> StreamingMergeBuilder<'a> { ))) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::collect; + use crate::memory::MemoryStream; + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::Int32Type; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use std::sync::Arc; + + /// Merge the given streams (each a list of batches of `sort_key` values) + /// and return the merged values. + async fn merge_i32_streams( + inputs: Vec>>, + batch_size: usize, + fetch: Option, + ) -> Vec> { + let schema = Arc::new(Schema::new(vec![Field::new( + "sort_key", + DataType::Int32, + false, + )])); + + let streams = inputs + .into_iter() + .map(|batches| { + let batches = batches + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .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(batch_size) + .with_schema(Arc::clone(&schema)) + .with_reservation({ + let mem_pool: Arc = + Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("merge stream mock memory").register(&mem_pool) + }) + .with_expressions( + &([ + PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) + .asc(), + ] + .into()), + ) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_fetch(fetch) + .build() + .unwrap(); + + let output = collect(merged).await.unwrap(); + + output + .iter() + .map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect() + } + + fn flatten_vec(input: Vec>) -> Vec { + input.into_iter().flatten().collect() + } + + #[tokio::test] + async fn test_merge_non_overlapping_batches() { + // streams take turns holding the smallest batch + let streams = vec![ + vec![ + vec![1, 2, 3, 4, 5], + vec![16, 17, 18, 19, 20], + vec![31, 32, 33], + ], + vec![ + vec![6, 7, 8, 9, 10], + vec![21, 22, 23, 24, 25], + vec![34, 35, 36], + ], + vec![ + vec![11, 12, 13, 14, 15], + vec![26, 27, 28, 29, 30], + vec![37, 38], + ], + ]; + + let output = merge_i32_streams(streams, 5, None).await; + + assert_eq!(flatten_vec(output), (1..=38).collect::>()); + } + + #[tokio::test] + async fn test_merge_interleaved_rows() { + // every batch overlaps batches of the other streams + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 13, 16]], + vec![vec![2, 5, 8], vec![11, 14, 17]], + vec![vec![3, 6, 9], vec![12, 15, 18]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=18).collect::>()); + } + + #[tokio::test] + async fn test_merge_some_rows_from_other_streams_then_a_whole_batch_from_one() { + // one row from stream 0, two from stream 1 (batch_size 4 is still + // not filled), then stream 2's entire batch [4, 5, 6, 7] sorts before + // everything the other streams have left (8 and 11) + let streams = vec![ + vec![vec![1, 8, 9, 10], vec![17, 18, 19, 20]], + vec![vec![2, 3, 11, 12], vec![21, 22, 23, 24]], + vec![vec![4, 5, 6, 7], vec![13, 14, 15, 16]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=24).collect::>()); + } + + #[tokio::test] + async fn test_merge_equal_values_across_batch_and_stream_boundaries() { + // the value 3 ends stream 0's first batch, starts its second + // one, and appears in every other stream as well + let streams = vec![ + vec![vec![1, 2, 3], vec![3, 3, 4]], + vec![vec![3, 4, 5], vec![5, 6, 7]], + vec![vec![3, 3, 3], vec![7, 7, 7]], + ]; + + let output = merge_i32_streams(streams, 3, None).await; + + assert_eq!( + flatten_vec(output), + vec![1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 5, 5, 6, 7, 7, 7, 7] + ); + } + + #[tokio::test] + async fn test_merge_streams_of_different_lengths() { + // stream 1 ends first, then stream 2, stream 0 finishes alone + let streams = vec![ + vec![ + vec![1, 2], + vec![3, 4], + vec![5, 6], + vec![7, 8], + vec![20, 21, 22], + ], + vec![vec![0, 9]], + vec![vec![4, 5], vec![10, 11, 12]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![0, 1, 2, 3, 4, 4, 5, 5, 6, 7, 8, 9, 10, 11, 12, 20, 21, 22] + ); + } + + #[tokio::test] + async fn test_merge_batch_size_smaller_than_input_batches() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 2, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_batch_size_larger_than_entire_input() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 8192, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_fetch_smaller_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(5)).await; + + assert_eq!(flatten_vec(output), vec![1, 2, 3, 4, 5]); + } + + #[tokio::test] + async fn test_merge_fetch_larger_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(100)).await; + + assert_eq!(flatten_vec(output), (1..=12).collect::>()); + } + +} From f3cffebe5de66489918ded717ae7822744f91c83 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 20:23:34 +0300 Subject: [PATCH 03/10] added test --- .../src/sorts/streaming_merge.rs | 144 ++++++++++++------ 1 file changed, 99 insertions(+), 45 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 53b704f7aedb1..d7de13572693e 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -279,65 +279,88 @@ mod tests { use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use std::sync::Arc; - /// Merge the given streams (each a list of batches of `sort_key` values) - /// and return the merged values. - async fn merge_i32_streams( - inputs: Vec>>, + /// Merge the given streams (each a list of batches of nullable `sort_key` + /// values, nulls sorted first) and return the merged values. + async fn merge_nullable_i32_streams( + inputs: Vec>>>, batch_size: usize, fetch: Option, - ) -> Vec> { + ) -> Vec>> { let schema = Arc::new(Schema::new(vec![Field::new( "sort_key", DataType::Int32, - false, + true, )])); let streams = inputs - .into_iter() - .map(|batches| { - let batches = batches - .into_iter() - .map(|values| { - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from(values))], - ) - .unwrap() - }) - .collect::>(); - Box::pin( - MemoryStream::try_new(batches, Arc::clone(&schema), None).unwrap(), - ) as SendableRecordBatchStream - }) - .collect::>(); + .into_iter() + .map(|batches| { + let batches = batches + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .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(batch_size) - .with_schema(Arc::clone(&schema)) - .with_reservation({ - let mem_pool: Arc = - Arc::new(UnboundedMemoryPool::default()); - MemoryConsumer::new("merge stream mock memory").register(&mem_pool) - }) - .with_expressions( - &([ - PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) - .asc(), - ] + .with_streams(streams) + .with_batch_size(batch_size) + .with_schema(Arc::clone(&schema)) + .with_reservation({ + let mem_pool: Arc = + Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("merge stream mock memory").register(&mem_pool) + }) + .with_expressions( + &([ + PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) + .asc(), + ] .into()), - ) - .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) - .with_fetch(fetch) - .build() - .unwrap(); + ) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_fetch(fetch) + .build() + .unwrap(); let output = collect(merged).await.unwrap(); output - .iter() - .map(|b| b.column(0).as_primitive::().values().to_vec()) - .collect() + .iter() + .map(|b| b.column(0).as_primitive::().iter().collect()) + .collect() + } + + /// Same as [`merge_nullable_i32_streams`] for non-nullable values. + async fn merge_i32_streams( + inputs: Vec>>, + batch_size: usize, + fetch: Option, + ) -> Vec> { + let inputs = inputs + .into_iter() + .map(|stream| { + stream + .into_iter() + .map(|batch| batch.into_iter().map(Some).collect()) + .collect() + }) + .collect(); + + merge_nullable_i32_streams(inputs, batch_size, fetch) + .await + .into_iter() + .map(|batch| batch.into_iter().map(|v| v.unwrap()).collect()) + .collect() } fn flatten_vec(input: Vec>) -> Vec { @@ -418,6 +441,38 @@ mod tests { ); } + #[tokio::test] + async fn test_merge_with_nulls() { + // nulls sort first; stream 0's first batch is all nulls, so its last + // row (null) ties with the null heads of the other streams, and nulls + // also compare against non-null values on every batch boundary + let streams = vec![ + vec![vec![None, None], vec![Some(1), Some(2)]], + vec![vec![None, Some(3)], vec![Some(4), Some(5)]], + vec![vec![None, Some(6)], vec![Some(7), Some(8)]], + ]; + + let output = merge_nullable_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![ + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8) + ] + ); + } + #[tokio::test] async fn test_merge_streams_of_different_lengths() { // stream 1 ends first, then stream 2, stream 0 finishes alone @@ -492,5 +547,4 @@ mod tests { assert_eq!(flatten_vec(output), (1..=12).collect::>()); } - } From 6dc6b2dda779c5817913ba185c5aec3fdab8714d Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:05:17 +0300 Subject: [PATCH 04/10] added benches --- .../benches/sort_preserving_merge.rs | 158 +++++++++++++++++- 1 file changed, 157 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 76ebf230a30e0..f6e4466eaf8a8 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -27,6 +27,8 @@ use datafusion_physical_plan::test::TestMemoryExec; use datafusion_physical_plan::{ collect, sorts::sort_preserving_merge::SortPreservingMergeExec, }; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; use std::sync::Arc; @@ -193,5 +195,159 @@ fn bench_merge_sorted_preserving(c: &mut Criterion) { } } -criterion_group!(benches, bench_merge_sorted_preserving); +// --------------------------------------------------------------------------- +// Benchmarks across data orderings (sorted / nearly sorted / reverse / +// unsorted), sort key types (u64 / string / complex) and payload widths +// (5 / 20 / 100 columns). +// --------------------------------------------------------------------------- + +const NUM_PARTITIONS: usize = 4; +const ROWS_PER_PARTITION: usize = 100_000; +const BATCH_SIZE: usize = 8192; + +const ORDERINGS: [&str; 4] = ["sorted", "nearly_sorted", "reverse", "unsorted"]; +const KEY_TYPES: [&str; 3] = ["u64", "string", "complex"]; +const PAYLOAD_WIDTHS: [usize; 3] = [5, 20, 100]; + +/// Generate the keys in their "arrival" order, before partitioning +fn generate_keys(ordering: &str) -> Vec { + let n = (NUM_PARTITIONS * ROWS_PER_PARTITION) as u64; + let mut rng = StdRng::seed_from_u64(42); + match ordering { + "sorted" => (0..n).collect(), + "reverse" => (0..n).rev().collect(), + // Sorted except for ~1% of items misplaced to random positions + "nearly_sorted" => { + let mut keys: Vec = (0..n).collect(); + for _ in 0..(n / 100) { + let a = rng.random_range(0..n as usize); + let b = rng.random_range(0..n as usize); + keys.swap(a, b); + } + keys + } + "unsorted" => (0..n).map(|_| rng.random_range(0..n)).collect(), + _ => unreachable!(), + } +} + +/// Distribute the arrival sequence round-robin (a batch at a time) over the +/// partitions, then sort each partition's keys, as SortExec would before a +/// sort preserving merge +fn partition_keys(keys: Vec) -> Vec> { + let mut partitions = vec![Vec::with_capacity(ROWS_PER_PARTITION); NUM_PARTITIONS]; + for (i, chunk) in keys.chunks(BATCH_SIZE).enumerate() { + partitions[i % NUM_PARTITIONS].extend_from_slice(chunk); + } + for partition in &mut partitions { + partition.sort_unstable(); + } + partitions +} + +fn key_columns(key_type: &str, keys: &[u64]) -> Vec<(String, ArrayRef)> { + let as_string = || { + Arc::new(StringArray::from_iter_values( + keys.iter().map(|k| format!("{k:012}")), + )) as ArrayRef + }; + match key_type { + "u64" => vec![( + "key0".to_string(), + Arc::new(UInt64Array::from(keys.to_vec())) as _, + )], + "string" => vec![("key0".to_string(), as_string())], + // Two sort columns force the row-based (normalized key) cursor + "complex" => vec![ + ( + "key0".to_string(), + Arc::new(UInt64Array::from_iter_values(keys.iter().map(|k| k / 8))) as _, + ), + ("key1".to_string(), as_string()), + ], + _ => unreachable!(), + } +} + +fn create_case( + ordering: &str, + key_type: &str, + payload_width: usize, +) -> (Vec>, SchemaRef, LexOrdering) { + let partitions = partition_keys(generate_keys(ordering)) + .into_iter() + .map(|keys| { + keys.chunks(BATCH_SIZE) + .map(|chunk| { + let mut columns = key_columns(key_type, chunk); + // All payload columns share the same buffer, so wide + // payloads don't blow up memory + let payload = Arc::new(UInt64Array::from(chunk.to_vec())) as ArrayRef; + for i in 0..payload_width { + columns.push((format!("col{i}"), Arc::clone(&payload))); + } + RecordBatch::try_from_iter(columns).unwrap() + }) + .collect::>() + }) + .collect::>(); + + let schema = partitions[0][0].schema(); + let sort_order = LexOrdering::new( + schema + .fields() + .iter() + .filter(|field| field.name().starts_with("key")) + .map(|field| { + PhysicalSortExpr::new( + col(field.name(), &schema).unwrap(), + SortOptions::default(), + ) + }), + ) + .unwrap(); + + (partitions, schema, sort_order) +} + +fn bench_spm_data_patterns(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let task_ctx = Arc::new(TaskContext::default()); + + for ordering in ORDERINGS { + for key_type in KEY_TYPES { + for payload_width in PAYLOAD_WIDTHS { + let (partitions, schema, sort_order) = + create_case(ordering, key_type, payload_width); + + c.bench_function( + &format!("spm/{ordering}/{key_type}/payload_{payload_width}"), + |b| { + b.iter(|| { + let exec = TestMemoryExec::try_new_exec( + &partitions, + Arc::clone(&schema), + None, + ) + .unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new( + sort_order.clone(), + exec, + )); + rt.block_on(async { + collect(merge, Arc::clone(&task_ctx)).await.unwrap() + }) + }) + }, + ); + } + } + } +} + +criterion_group!( + benches, + bench_merge_sorted_preserving, + bench_spm_data_patterns +); criterion_main!(benches); From 55ece4b516fde7edc9f690fb4e3c2d82555a063b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:10:45 +0300 Subject: [PATCH 05/10] change collect to drain to avoid holding the memory --- .../benches/sort_preserving_merge.rs | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index f6e4466eaf8a8..77467285957c8 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -21,17 +21,26 @@ use arrow::{ }; use arrow_schema::{SchemaRef, SortOptions}; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; use datafusion_physical_plan::test::TestMemoryExec; -use datafusion_physical_plan::{ - collect, sorts::sort_preserving_merge::SortPreservingMergeExec, -}; +use datafusion_physical_plan::{collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec}; +use futures::StreamExt; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use std::hint::black_box; use std::sync::Arc; +/// Consume the stream batch by batch, dropping each batch as it arrives +/// instead of holding the whole result in memory like `collect` would +async fn drain(mut stream: SendableRecordBatchStream) { + while let Some(batch) = stream.next().await { + black_box(batch.unwrap()); + } +} + const BENCH_ROWS: usize = 1_000_000; // 1 million rows fn get_large_string(idx: usize) -> String { @@ -334,9 +343,9 @@ fn bench_spm_data_patterns(c: &mut Criterion) { sort_order.clone(), exec, )); - rt.block_on(async { - collect(merge, Arc::clone(&task_ctx)).await.unwrap() - }) + rt.block_on(drain( + execute_stream(merge, Arc::clone(&task_ctx)).unwrap(), + )) }) }, ); From 5a5e1f084d6afaf3fed960be6f1aa17d863cf7ab Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:49:13 +0300 Subject: [PATCH 06/10] revert change and see perf --- datafusion/physical-plan/src/sorts/builder.rs | 9 - datafusion/physical-plan/src/sorts/cursor.rs | 52 ---- datafusion/physical-plan/src/sorts/merge.rs | 75 ----- .../src/sorts/streaming_merge.rs | 283 ------------------ 4 files changed, 419 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index f1a545dfe4ae9..75eb2ff980325 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -129,15 +129,6 @@ impl BatchBuilder { &self.schema } - /// The rows of `stream_idx`'s current batch that have not been pushed - /// via [`Self::push_row`] yet, as a zero-copy slice. - pub fn remaining_rows_of(&self, stream_idx: usize) -> Option { - let cursor = &self.cursors[stream_idx]; - let (_, batch) = &self.batches[cursor.batch_idx]; - let remaining = batch.num_rows() - cursor.row_idx; - (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)) - } - /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 58a4dff2bb384..8991922779d4a 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -43,18 +43,8 @@ pub trait CursorValues { fn eq_to_previous(cursor: &Self, idx: usize) -> bool; /// Returns comparison of `l[l_idx]` and `r[r_idx]` - /// - /// Only called with the cursors' *current* offsets, so caching - /// implementations may serve it from their cache. fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; - /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary* - /// indices, unlike [`Self::compare`]. Cold path: caching implementations - /// must override this to index directly. - fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { - Self::compare(l, l_idx, r, r_idx) - } - /// 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). @@ -126,23 +116,6 @@ impl Cursor { t } - /// Returns true if no rows of this cursor were consumed yet - pub fn is_at_start(&self) -> bool { - self.offset == 0 - } - - /// Compare this cursor's *last* row to `other`'s current row. Since the - /// values are sorted, `Less` means every remaining row of this cursor - /// sorts before all of `other`'s remaining rows. - pub fn compare_last_to(&self, other: &Self) -> Ordering { - T::compare_at( - &self.values, - self.values.len() - 1, - &other.values, - other.offset, - ) - } - pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { if self.offset > 0 { self.is_eq_to_prev_row() @@ -317,11 +290,6 @@ impl CursorValues for PrimitiveValues { l.current.compare(r.current) } - fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { - // Arbitrary indices (cold path), so index directly instead of using the cache. - l.values[l_idx].compare(r.values[r_idx]) - } - #[inline(always)] fn set_offset(&mut self, offset: usize) { // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that @@ -550,26 +518,6 @@ impl CursorValues for ArrayValues { } } - fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { - // Same null handling as `compare`, but delegates to the inner - // `compare_at` so arbitrary indices work with caching values. - match (l.is_null(l_idx), r.is_null(r_idx)) { - (true, true) => Ordering::Equal, - (true, false) => match l.options.nulls_first { - true => Ordering::Less, - false => Ordering::Greater, - }, - (false, true) => match l.options.nulls_first { - true => Ordering::Greater, - false => Ordering::Less, - }, - (false, false) => match l.options.descending { - true => T::compare_at(&r.values, r_idx, &l.values, l_idx), - false => T::compare_at(&l.values, l_idx, &r.values, r_idx), - }, - } - } - #[inline(always)] fn set_offset(&mut self, offset: usize) { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8c464a3cf1413..4583d19e91061 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -153,10 +153,6 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, - - /// A whole winner batch that was proven to sort before every other - /// stream's head, waiting to be emitted after the in-progress rows are flushed. - pending_skipped_batch: Option, } impl SortPreservingMergeStream { @@ -190,7 +186,6 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, - pending_skipped_batch: None, } } @@ -241,12 +236,6 @@ impl SortPreservingMergeStream { } return Poll::Ready(None); } - - if let Some(batch) = self.pending_skipped_batch.take() { - self.produced += batch.num_rows(); - return Poll::Ready(Some(Ok(batch))); - } - // Once all partitions have set their corresponding cursors for the loser tree, // we skip the following block. Until then, this function may be called multiple // times and can return Poll::Pending if any partition returns Poll::Pending. @@ -321,41 +310,6 @@ impl SortPreservingMergeStream { } let stream_idx = self.loser_tree[0]; - - // Full-batch skip: - // when the winner sits at the start of a fresh batch whose *last* row sorts before every other stream's - // current row, we skip per-row tree updates and just emit it. - // Checked at most once per input batch so we don't add a lot of overhead when - // the data is not sorted at all - // - // This might produce batches that do not equal to the configured batch_size, - // this is something we are willing to sacrifice - // - // TODO - support fetch - if self.fetch.is_none() - && self.cursors[stream_idx] - .as_ref() - .is_some_and(|cursor| cursor.is_at_start()) - && self.winner_batch_beats_all(stream_idx) - { - let batch = self - .in_progress - .remaining_rows_of(stream_idx) - .expect("winner cursor is at the start, so rows remain"); - // 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.in_progress.is_empty() { - self.produced += batch.num_rows(); - return Poll::Ready(Some(Ok(batch))); - } - // Flush already-merged rows first - // they all sort before the skipped batch - self.pending_skipped_batch = Some(batch); - return Poll::Ready(self.emit_in_progress_batch().transpose()); - } - if self.advance_cursors(stream_idx) { self.loser_tree_adjusted = false; self.in_progress.push_row(stream_idx); @@ -373,35 +327,6 @@ impl SortPreservingMergeStream { } } - /// Returns `true` if every remaining row of `winner`'s current batch - /// sorts before every other stream's current row, using the same index - /// tie-breaking as [`Self::is_gt`]. Exhausted streams cannot compete. - fn winner_batch_beats_all(&self, winner: usize) -> bool { - let Some(winner_cursor) = &self.cursors[winner] else { - return false; - }; - - // Only the streams that lost *directly* to the winner can hold the - // overall runner-up, and those are exactly the losers stored on the - // winner's leaf-to-root path (every other stream lost to one of them - // transitively), so `O(log k)` comparisons suffice instead of `O(k)`. - let mut node = self.lt_leaf_node_index(winner); - while node != 0 { - let challenger = self.loser_tree[node]; - if let Some(other) = &self.cursors[challenger] - && !winner_cursor - .compare_last_to(other) - .then_with(|| winner.cmp(&challenger)) - .is_lt() - { - return false; - } - node = self.lt_parent_node_index(node); - } - - true - } - /// 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) { diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index d7de13572693e..ade24ff0534ff 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -265,286 +265,3 @@ impl<'a> StreamingMergeBuilder<'a> { ))) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::common::collect; - use crate::memory::MemoryStream; - use arrow::array::{Int32Array, RecordBatch}; - use arrow::datatypes::Int32Type; - use arrow_schema::{DataType, Field, Schema}; - use datafusion_physical_expr::expressions::col; - use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; - use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; - use std::sync::Arc; - - /// Merge the given streams (each a list of batches of nullable `sort_key` - /// values, nulls sorted first) and return the merged values. - async fn merge_nullable_i32_streams( - inputs: Vec>>>, - batch_size: usize, - fetch: Option, - ) -> Vec>> { - let schema = Arc::new(Schema::new(vec![Field::new( - "sort_key", - DataType::Int32, - true, - )])); - - let streams = inputs - .into_iter() - .map(|batches| { - let batches = batches - .into_iter() - .map(|values| { - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from(values))], - ) - .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(batch_size) - .with_schema(Arc::clone(&schema)) - .with_reservation({ - let mem_pool: Arc = - Arc::new(UnboundedMemoryPool::default()); - MemoryConsumer::new("merge stream mock memory").register(&mem_pool) - }) - .with_expressions( - &([ - PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) - .asc(), - ] - .into()), - ) - .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) - .with_fetch(fetch) - .build() - .unwrap(); - - let output = collect(merged).await.unwrap(); - - output - .iter() - .map(|b| b.column(0).as_primitive::().iter().collect()) - .collect() - } - - /// Same as [`merge_nullable_i32_streams`] for non-nullable values. - async fn merge_i32_streams( - inputs: Vec>>, - batch_size: usize, - fetch: Option, - ) -> Vec> { - let inputs = inputs - .into_iter() - .map(|stream| { - stream - .into_iter() - .map(|batch| batch.into_iter().map(Some).collect()) - .collect() - }) - .collect(); - - merge_nullable_i32_streams(inputs, batch_size, fetch) - .await - .into_iter() - .map(|batch| batch.into_iter().map(|v| v.unwrap()).collect()) - .collect() - } - - fn flatten_vec(input: Vec>) -> Vec { - input.into_iter().flatten().collect() - } - - #[tokio::test] - async fn test_merge_non_overlapping_batches() { - // streams take turns holding the smallest batch - let streams = vec![ - vec![ - vec![1, 2, 3, 4, 5], - vec![16, 17, 18, 19, 20], - vec![31, 32, 33], - ], - vec![ - vec![6, 7, 8, 9, 10], - vec![21, 22, 23, 24, 25], - vec![34, 35, 36], - ], - vec![ - vec![11, 12, 13, 14, 15], - vec![26, 27, 28, 29, 30], - vec![37, 38], - ], - ]; - - let output = merge_i32_streams(streams, 5, None).await; - - assert_eq!(flatten_vec(output), (1..=38).collect::>()); - } - - #[tokio::test] - async fn test_merge_interleaved_rows() { - // every batch overlaps batches of the other streams - let streams = vec![ - vec![vec![1, 4, 7], vec![10, 13, 16]], - vec![vec![2, 5, 8], vec![11, 14, 17]], - vec![vec![3, 6, 9], vec![12, 15, 18]], - ]; - - let output = merge_i32_streams(streams, 4, None).await; - - assert_eq!(flatten_vec(output), (1..=18).collect::>()); - } - - #[tokio::test] - async fn test_merge_some_rows_from_other_streams_then_a_whole_batch_from_one() { - // one row from stream 0, two from stream 1 (batch_size 4 is still - // not filled), then stream 2's entire batch [4, 5, 6, 7] sorts before - // everything the other streams have left (8 and 11) - let streams = vec![ - vec![vec![1, 8, 9, 10], vec![17, 18, 19, 20]], - vec![vec![2, 3, 11, 12], vec![21, 22, 23, 24]], - vec![vec![4, 5, 6, 7], vec![13, 14, 15, 16]], - ]; - - let output = merge_i32_streams(streams, 4, None).await; - - assert_eq!(flatten_vec(output), (1..=24).collect::>()); - } - - #[tokio::test] - async fn test_merge_equal_values_across_batch_and_stream_boundaries() { - // the value 3 ends stream 0's first batch, starts its second - // one, and appears in every other stream as well - let streams = vec![ - vec![vec![1, 2, 3], vec![3, 3, 4]], - vec![vec![3, 4, 5], vec![5, 6, 7]], - vec![vec![3, 3, 3], vec![7, 7, 7]], - ]; - - let output = merge_i32_streams(streams, 3, None).await; - - assert_eq!( - flatten_vec(output), - vec![1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 5, 5, 6, 7, 7, 7, 7] - ); - } - - #[tokio::test] - async fn test_merge_with_nulls() { - // nulls sort first; stream 0's first batch is all nulls, so its last - // row (null) ties with the null heads of the other streams, and nulls - // also compare against non-null values on every batch boundary - let streams = vec![ - vec![vec![None, None], vec![Some(1), Some(2)]], - vec![vec![None, Some(3)], vec![Some(4), Some(5)]], - vec![vec![None, Some(6)], vec![Some(7), Some(8)]], - ]; - - let output = merge_nullable_i32_streams(streams, 4, None).await; - - assert_eq!( - flatten_vec(output), - vec![ - None, - None, - None, - None, - Some(1), - Some(2), - Some(3), - Some(4), - Some(5), - Some(6), - Some(7), - Some(8) - ] - ); - } - - #[tokio::test] - async fn test_merge_streams_of_different_lengths() { - // stream 1 ends first, then stream 2, stream 0 finishes alone - let streams = vec![ - vec![ - vec![1, 2], - vec![3, 4], - vec![5, 6], - vec![7, 8], - vec![20, 21, 22], - ], - vec![vec![0, 9]], - vec![vec![4, 5], vec![10, 11, 12]], - ]; - - let output = merge_i32_streams(streams, 4, None).await; - - assert_eq!( - flatten_vec(output), - vec![0, 1, 2, 3, 4, 4, 5, 5, 6, 7, 8, 9, 10, 11, 12, 20, 21, 22] - ); - } - - #[tokio::test] - async fn test_merge_batch_size_smaller_than_input_batches() { - let streams = vec![ - vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], - vec![vec![2, 5, 8, 11]], - vec![vec![3, 6, 9, 12]], - ]; - - let output = merge_i32_streams(streams, 2, None).await; - - assert_eq!(flatten_vec(output), (1..=16).collect::>()); - } - - #[tokio::test] - async fn test_merge_batch_size_larger_than_entire_input() { - let streams = vec![ - vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], - vec![vec![2, 5, 8, 11]], - vec![vec![3, 6, 9, 12]], - ]; - - let output = merge_i32_streams(streams, 8192, None).await; - - assert_eq!(flatten_vec(output), (1..=16).collect::>()); - } - - #[tokio::test] - async fn test_merge_fetch_smaller_than_input() { - let streams = vec![ - vec![vec![1, 4, 7], vec![10, 11, 12]], - vec![vec![2, 5, 8]], - vec![vec![3, 6, 9]], - ]; - - let output = merge_i32_streams(streams, 4, Some(5)).await; - - assert_eq!(flatten_vec(output), vec![1, 2, 3, 4, 5]); - } - - #[tokio::test] - async fn test_merge_fetch_larger_than_input() { - let streams = vec![ - vec![vec![1, 4, 7], vec![10, 11, 12]], - vec![vec![2, 5, 8]], - vec![vec![3, 6, 9]], - ]; - - let output = merge_i32_streams(streams, 4, Some(100)).await; - - assert_eq!(flatten_vec(output), (1..=12).collect::>()); - } -} From 80d85bf390504d6aa259e6db45f97c6ce4fb314d Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:50:34 +0300 Subject: [PATCH 07/10] Revert "revert change and see perf" This reverts commit 5a5e1f084d6afaf3fed960be6f1aa17d863cf7ab. --- datafusion/physical-plan/src/sorts/builder.rs | 9 + datafusion/physical-plan/src/sorts/cursor.rs | 52 ++++ datafusion/physical-plan/src/sorts/merge.rs | 75 +++++ .../src/sorts/streaming_merge.rs | 283 ++++++++++++++++++ 4 files changed, 419 insertions(+) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..f1a545dfe4ae9 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -129,6 +129,15 @@ impl BatchBuilder { &self.schema } + /// The rows of `stream_idx`'s current batch that have not been pushed + /// via [`Self::push_row`] yet, as a zero-copy slice. + pub fn remaining_rows_of(&self, stream_idx: usize) -> Option { + let cursor = &self.cursors[stream_idx]; + let (_, batch) = &self.batches[cursor.batch_idx]; + let remaining = batch.num_rows() - cursor.row_idx; + (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)) + } + /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 8991922779d4a..58a4dff2bb384 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -43,8 +43,18 @@ pub trait CursorValues { fn eq_to_previous(cursor: &Self, idx: usize) -> bool; /// Returns comparison of `l[l_idx]` and `r[r_idx]` + /// + /// Only called with the cursors' *current* offsets, so caching + /// implementations may serve it from their cache. fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary* + /// indices, unlike [`Self::compare`]. Cold path: caching implementations + /// must override this to index directly. + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + Self::compare(l, l_idx, r, r_idx) + } + /// 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). @@ -116,6 +126,23 @@ impl Cursor { t } + /// Returns true if no rows of this cursor were consumed yet + pub fn is_at_start(&self) -> bool { + self.offset == 0 + } + + /// Compare this cursor's *last* row to `other`'s current row. Since the + /// values are sorted, `Less` means every remaining row of this cursor + /// sorts before all of `other`'s remaining rows. + pub fn compare_last_to(&self, other: &Self) -> Ordering { + T::compare_at( + &self.values, + self.values.len() - 1, + &other.values, + other.offset, + ) + } + pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { if self.offset > 0 { self.is_eq_to_prev_row() @@ -290,6 +317,11 @@ impl CursorValues for PrimitiveValues { l.current.compare(r.current) } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Arbitrary indices (cold path), so index directly instead of using the cache. + l.values[l_idx].compare(r.values[r_idx]) + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that @@ -518,6 +550,26 @@ impl CursorValues for ArrayValues { } } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Same null handling as `compare`, but delegates to the inner + // `compare_at` so arbitrary indices work with caching values. + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => Ordering::Equal, + (true, false) => match l.options.nulls_first { + true => Ordering::Less, + false => Ordering::Greater, + }, + (false, true) => match l.options.nulls_first { + true => Ordering::Greater, + false => Ordering::Less, + }, + (false, false) => match l.options.descending { + true => T::compare_at(&r.values, r_idx, &l.values, l_idx), + false => T::compare_at(&l.values, l_idx, &r.values, r_idx), + }, + } + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..8c464a3cf1413 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -153,6 +153,10 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, + + /// A whole winner batch that was proven to sort before every other + /// stream's head, waiting to be emitted after the in-progress rows are flushed. + pending_skipped_batch: Option, } impl SortPreservingMergeStream { @@ -186,6 +190,7 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + pending_skipped_batch: None, } } @@ -236,6 +241,12 @@ impl SortPreservingMergeStream { } return Poll::Ready(None); } + + if let Some(batch) = self.pending_skipped_batch.take() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Once all partitions have set their corresponding cursors for the loser tree, // we skip the following block. Until then, this function may be called multiple // times and can return Poll::Pending if any partition returns Poll::Pending. @@ -310,6 +321,41 @@ impl SortPreservingMergeStream { } let stream_idx = self.loser_tree[0]; + + // Full-batch skip: + // when the winner sits at the start of a fresh batch whose *last* row sorts before every other stream's + // current row, we skip per-row tree updates and just emit it. + // Checked at most once per input batch so we don't add a lot of overhead when + // the data is not sorted at all + // + // This might produce batches that do not equal to the configured batch_size, + // this is something we are willing to sacrifice + // + // TODO - support fetch + if self.fetch.is_none() + && self.cursors[stream_idx] + .as_ref() + .is_some_and(|cursor| cursor.is_at_start()) + && self.winner_batch_beats_all(stream_idx) + { + let batch = self + .in_progress + .remaining_rows_of(stream_idx) + .expect("winner cursor is at the start, so rows remain"); + // 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.in_progress.is_empty() { + self.produced += batch.num_rows(); + return Poll::Ready(Some(Ok(batch))); + } + // Flush already-merged rows first + // they all sort before the skipped batch + self.pending_skipped_batch = Some(batch); + return Poll::Ready(self.emit_in_progress_batch().transpose()); + } + if self.advance_cursors(stream_idx) { self.loser_tree_adjusted = false; self.in_progress.push_row(stream_idx); @@ -327,6 +373,35 @@ impl SortPreservingMergeStream { } } + /// Returns `true` if every remaining row of `winner`'s current batch + /// sorts before every other stream's current row, using the same index + /// tie-breaking as [`Self::is_gt`]. Exhausted streams cannot compete. + fn winner_batch_beats_all(&self, winner: usize) -> bool { + let Some(winner_cursor) = &self.cursors[winner] else { + return false; + }; + + // Only the streams that lost *directly* to the winner can hold the + // overall runner-up, and those are exactly the losers stored on the + // winner's leaf-to-root path (every other stream lost to one of them + // transitively), so `O(log k)` comparisons suffice instead of `O(k)`. + let mut node = self.lt_leaf_node_index(winner); + while node != 0 { + let challenger = self.loser_tree[node]; + if let Some(other) = &self.cursors[challenger] + && !winner_cursor + .compare_last_to(other) + .then_with(|| winner.cmp(&challenger)) + .is_lt() + { + return false; + } + node = self.lt_parent_node_index(node); + } + + true + } + /// 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) { diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..d7de13572693e 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -265,3 +265,286 @@ impl<'a> StreamingMergeBuilder<'a> { ))) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::collect; + use crate::memory::MemoryStream; + use arrow::array::{Int32Array, RecordBatch}; + use arrow::datatypes::Int32Type; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use std::sync::Arc; + + /// Merge the given streams (each a list of batches of nullable `sort_key` + /// values, nulls sorted first) and return the merged values. + async fn merge_nullable_i32_streams( + inputs: Vec>>>, + batch_size: usize, + fetch: Option, + ) -> Vec>> { + let schema = Arc::new(Schema::new(vec![Field::new( + "sort_key", + DataType::Int32, + true, + )])); + + let streams = inputs + .into_iter() + .map(|batches| { + let batches = batches + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + ) + .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(batch_size) + .with_schema(Arc::clone(&schema)) + .with_reservation({ + let mem_pool: Arc = + Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("merge stream mock memory").register(&mem_pool) + }) + .with_expressions( + &([ + PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) + .asc(), + ] + .into()), + ) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_fetch(fetch) + .build() + .unwrap(); + + let output = collect(merged).await.unwrap(); + + output + .iter() + .map(|b| b.column(0).as_primitive::().iter().collect()) + .collect() + } + + /// Same as [`merge_nullable_i32_streams`] for non-nullable values. + async fn merge_i32_streams( + inputs: Vec>>, + batch_size: usize, + fetch: Option, + ) -> Vec> { + let inputs = inputs + .into_iter() + .map(|stream| { + stream + .into_iter() + .map(|batch| batch.into_iter().map(Some).collect()) + .collect() + }) + .collect(); + + merge_nullable_i32_streams(inputs, batch_size, fetch) + .await + .into_iter() + .map(|batch| batch.into_iter().map(|v| v.unwrap()).collect()) + .collect() + } + + fn flatten_vec(input: Vec>) -> Vec { + input.into_iter().flatten().collect() + } + + #[tokio::test] + async fn test_merge_non_overlapping_batches() { + // streams take turns holding the smallest batch + let streams = vec![ + vec![ + vec![1, 2, 3, 4, 5], + vec![16, 17, 18, 19, 20], + vec![31, 32, 33], + ], + vec![ + vec![6, 7, 8, 9, 10], + vec![21, 22, 23, 24, 25], + vec![34, 35, 36], + ], + vec![ + vec![11, 12, 13, 14, 15], + vec![26, 27, 28, 29, 30], + vec![37, 38], + ], + ]; + + let output = merge_i32_streams(streams, 5, None).await; + + assert_eq!(flatten_vec(output), (1..=38).collect::>()); + } + + #[tokio::test] + async fn test_merge_interleaved_rows() { + // every batch overlaps batches of the other streams + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 13, 16]], + vec![vec![2, 5, 8], vec![11, 14, 17]], + vec![vec![3, 6, 9], vec![12, 15, 18]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=18).collect::>()); + } + + #[tokio::test] + async fn test_merge_some_rows_from_other_streams_then_a_whole_batch_from_one() { + // one row from stream 0, two from stream 1 (batch_size 4 is still + // not filled), then stream 2's entire batch [4, 5, 6, 7] sorts before + // everything the other streams have left (8 and 11) + let streams = vec![ + vec![vec![1, 8, 9, 10], vec![17, 18, 19, 20]], + vec![vec![2, 3, 11, 12], vec![21, 22, 23, 24]], + vec![vec![4, 5, 6, 7], vec![13, 14, 15, 16]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!(flatten_vec(output), (1..=24).collect::>()); + } + + #[tokio::test] + async fn test_merge_equal_values_across_batch_and_stream_boundaries() { + // the value 3 ends stream 0's first batch, starts its second + // one, and appears in every other stream as well + let streams = vec![ + vec![vec![1, 2, 3], vec![3, 3, 4]], + vec![vec![3, 4, 5], vec![5, 6, 7]], + vec![vec![3, 3, 3], vec![7, 7, 7]], + ]; + + let output = merge_i32_streams(streams, 3, None).await; + + assert_eq!( + flatten_vec(output), + vec![1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 5, 5, 6, 7, 7, 7, 7] + ); + } + + #[tokio::test] + async fn test_merge_with_nulls() { + // nulls sort first; stream 0's first batch is all nulls, so its last + // row (null) ties with the null heads of the other streams, and nulls + // also compare against non-null values on every batch boundary + let streams = vec![ + vec![vec![None, None], vec![Some(1), Some(2)]], + vec![vec![None, Some(3)], vec![Some(4), Some(5)]], + vec![vec![None, Some(6)], vec![Some(7), Some(8)]], + ]; + + let output = merge_nullable_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![ + None, + None, + None, + None, + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8) + ] + ); + } + + #[tokio::test] + async fn test_merge_streams_of_different_lengths() { + // stream 1 ends first, then stream 2, stream 0 finishes alone + let streams = vec![ + vec![ + vec![1, 2], + vec![3, 4], + vec![5, 6], + vec![7, 8], + vec![20, 21, 22], + ], + vec![vec![0, 9]], + vec![vec![4, 5], vec![10, 11, 12]], + ]; + + let output = merge_i32_streams(streams, 4, None).await; + + assert_eq!( + flatten_vec(output), + vec![0, 1, 2, 3, 4, 4, 5, 5, 6, 7, 8, 9, 10, 11, 12, 20, 21, 22] + ); + } + + #[tokio::test] + async fn test_merge_batch_size_smaller_than_input_batches() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 2, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_batch_size_larger_than_entire_input() { + let streams = vec![ + vec![vec![1, 4, 7, 10], vec![13, 14, 15, 16]], + vec![vec![2, 5, 8, 11]], + vec![vec![3, 6, 9, 12]], + ]; + + let output = merge_i32_streams(streams, 8192, None).await; + + assert_eq!(flatten_vec(output), (1..=16).collect::>()); + } + + #[tokio::test] + async fn test_merge_fetch_smaller_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(5)).await; + + assert_eq!(flatten_vec(output), vec![1, 2, 3, 4, 5]); + } + + #[tokio::test] + async fn test_merge_fetch_larger_than_input() { + let streams = vec![ + vec![vec![1, 4, 7], vec![10, 11, 12]], + vec![vec![2, 5, 8]], + vec![vec![3, 6, 9]], + ]; + + let output = merge_i32_streams(streams, 4, Some(100)).await; + + assert_eq!(flatten_vec(output), (1..=12).collect::>()); + } +} From 17bafed72ba09b4addb65227a52a687dbf075b94 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 23:37:24 +0300 Subject: [PATCH 08/10] fix tests and bugs --- .../physical_optimizer/enforce_sorting.rs | 7 +- datafusion/ffi/tests/ffi_udaf.rs | 15 +- .../benches/sort_preserving_merge.rs | 12 +- datafusion/physical-plan/src/sorts/builder.rs | 73 ++++++---- datafusion/physical-plan/src/sorts/merge.rs | 71 ++++++++-- .../src/sorts/sort_preserving_merge.rs | 26 ++-- .../src/sorts/streaming_merge.rs | 130 +++++++++++++++++- 7 files changed, 272 insertions(+), 62 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index ecff2edbbec16..2cd7bd61b648b 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2839,10 +2839,11 @@ async fn test_sort_with_streaming_table() -> Result<()> { let sql = "SELECT a FROM test_table GROUP BY a ORDER BY a"; let results = ctx.sql(sql).await?.collect().await?; - assert_eq!(results.len(), 1); - assert_eq!(results[0].num_columns(), 1); + // the merge may split the output into several batches, so concat first + let results = arrow::compute::concat_batches(&results[0].schema(), &results)?; + assert_eq!(results.num_columns(), 1); let expected = create_array!(Int32, vec![1, 2, 3]) as ArrayRef; - assert_eq!(results[0].column(0), &expected); + assert_eq!(results.column(0), &expected); Ok(()) } diff --git a/datafusion/ffi/tests/ffi_udaf.rs b/datafusion/ffi/tests/ffi_udaf.rs index 7df3404d7421b..441f8ddc4f88c 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -56,13 +56,16 @@ mod tests { .sort_by(vec![col("a")])?; let result = df.collect().await?; + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); let expected = record_batch!( ("a", Int32, vec![1, 2, 4]), ("sum_b", Float64, vec![1.0, 4.0, 16.0]) )?; - assert_eq!(result[0], expected); + assert_eq!(result, expected); Ok(()) } @@ -105,7 +108,10 @@ mod tests { .sort_by(vec![col("a")])?; let result = df.collect().await?; - let result = result[0].column_by_name("stddev_b").unwrap(); + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); + let result = result.column_by_name("stddev_b").unwrap(); let result = result .as_any() .downcast_ref::() @@ -162,13 +168,16 @@ mod tests { df.clone().show().await?; let result = df.collect().await?; + // the merge may split the output into several batches, so concat first + let result = + arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); let expected = record_batch!( ("a", Int32, vec![1, 2, 4]), ("sum_b", Float64, vec![2.0, 8.0, 32.0]) )?; - assert_eq!(result[0], expected); + assert_eq!(result, expected); Ok(()) } diff --git a/datafusion/physical-plan/benches/sort_preserving_merge.rs b/datafusion/physical-plan/benches/sort_preserving_merge.rs index 77467285957c8..76e39432b1a4a 100644 --- a/datafusion/physical-plan/benches/sort_preserving_merge.rs +++ b/datafusion/physical-plan/benches/sort_preserving_merge.rs @@ -25,7 +25,9 @@ use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; use datafusion_physical_plan::test::TestMemoryExec; -use datafusion_physical_plan::{collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec}; +use datafusion_physical_plan::{ + collect, execute_stream, sorts::sort_preserving_merge::SortPreservingMergeExec, +}; use futures::StreamExt; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -243,8 +245,10 @@ fn generate_keys(ordering: &str) -> Vec { /// Distribute the arrival sequence round-robin (a batch at a time) over the /// partitions, then sort each partition's keys, as SortExec would before a /// sort preserving merge -fn partition_keys(keys: Vec) -> Vec> { - let mut partitions = vec![Vec::with_capacity(ROWS_PER_PARTITION); NUM_PARTITIONS]; +fn partition_keys(keys: &[u64]) -> Vec> { + let mut partitions: Vec> = (0..NUM_PARTITIONS) + .map(|_| Vec::with_capacity(ROWS_PER_PARTITION)) + .collect(); for (i, chunk) in keys.chunks(BATCH_SIZE).enumerate() { partitions[i % NUM_PARTITIONS].extend_from_slice(chunk); } @@ -283,7 +287,7 @@ fn create_case( key_type: &str, payload_width: usize, ) -> (Vec>, SchemaRef, LexOrdering) { - let partitions = partition_keys(generate_keys(ordering)) + let partitions = partition_keys(&generate_keys(ordering)) .into_iter() .map(|keys| { keys.chunks(BATCH_SIZE) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index f1a545dfe4ae9..274f182b83981 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -129,13 +129,21 @@ impl BatchBuilder { &self.schema } - /// The rows of `stream_idx`'s current batch that have not been pushed - /// via [`Self::push_row`] yet, as a zero-copy slice. - pub fn remaining_rows_of(&self, stream_idx: usize) -> Option { - let cursor = &self.cursors[stream_idx]; + /// Consume the rows of `stream_idx`'s current batch that have not been + /// pushed via [`Self::push_row`] yet, as a zero-copy slice. + /// + /// Also releases batches that are no longer referenced: skip-heavy merges + /// never reach [`Self::build_record_batch`], which otherwise does that. + pub fn take_remaining_rows_of(&mut self, stream_idx: usize) -> Option { + let cursor = &mut self.cursors[stream_idx]; let (_, batch) = &self.batches[cursor.batch_idx]; let remaining = batch.num_rows() - cursor.row_idx; - (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)) + let slice = (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)); + cursor.row_idx = batch.num_rows(); + if self.indices.is_empty() { + self.release_unreferenced_batches(); + } + slice } /// Try to interleave all columns using the given index slice. @@ -173,28 +181,41 @@ impl BatchBuilder { // hot path. The retention is bounded and short-lived since leftover // rows are drained over subsequent polls. if self.indices.is_empty() { - // New cursors are only created once the previous cursor for the stream - // is finished. This means all remaining rows from all but the last batch - // for each stream have been yielded to the newly created record batch - // - // We can therefore drop all but the last batch for each stream - let mut batch_idx = 0; - let mut retained = 0; - self.batches.retain(|(stream_idx, batch)| { - let stream_cursor = &mut self.cursors[*stream_idx]; - let retain = stream_cursor.batch_idx == batch_idx; - batch_idx += 1; - - if retain { - stream_cursor.batch_idx = retained; - retained += 1; - } else { - self.batches_mem_used -= get_record_batch_memory_size(batch); - } - retain - }); + self.release_unreferenced_batches(); } + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) + } + + /// Drop buffered batches no stream cursor references anymore and release + /// their memory back to the pool. + /// + /// Must only be called when `indices` is empty, since indices reference + /// batches by position. + fn release_unreferenced_batches(&mut self) { + assert_eq!(self.indices.len(), 0); + + // New cursors are only created once the previous cursor for the stream + // is finished. This means all remaining rows from all but the last batch + // for each stream have been yielded to the newly created record batch + // + // We can therefore drop all but the last batch for each stream + let mut batch_idx = 0; + let mut retained = 0; + self.batches.retain(|(stream_idx, batch)| { + let stream_cursor = &mut self.cursors[*stream_idx]; + let retain = stream_cursor.batch_idx == batch_idx; + batch_idx += 1; + + if retain { + stream_cursor.batch_idx = retained; + retained += 1; + } else { + self.batches_mem_used -= get_record_batch_memory_size(batch); + } + retain + }); + // Release excess memory back to the pool, but never shrink below // initial_reservation to maintain the anti-starvation guarantee // for the merge phase. @@ -202,8 +223,6 @@ impl BatchBuilder { if self.reservation.size() > target { self.reservation.shrink(self.reservation.size() - target); } - - RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) } /// Drains the in_progress row indexes, and builds a new RecordBatch from them diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8c464a3cf1413..bbba884e4225e 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,6 +18,7 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. +use std::cmp::Ordering; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll, ready}; @@ -340,11 +341,14 @@ impl SortPreservingMergeStream { { let batch = self .in_progress - .remaining_rows_of(stream_idx) + .take_remaining_rows_of(stream_idx) .expect("winner cursor is at the start, so rows remain"); // 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(); @@ -388,13 +392,26 @@ impl SortPreservingMergeStream { let mut node = self.lt_leaf_node_index(winner); while node != 0 { let challenger = self.loser_tree[node]; - if let Some(other) = &self.cursors[challenger] - && !winner_cursor - .compare_last_to(other) - .then_with(|| winner.cmp(&challenger)) - .is_lt() - { - return false; + 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) + } else { + winner < challenger + } + } + Ordering::Greater => false, + }; + if !beats { + return false; + } } node = self.lt_parent_node_index(node); } @@ -465,6 +482,38 @@ impl SortPreservingMergeStream { 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; @@ -609,15 +658,15 @@ impl SortPreservingMergeStream { if self.enable_round_robin_tie_breaker { match (&self.cursors[winner], &self.cursors[challenger]) { (Some(ac), Some(bc)) => match ac.cmp(bc) { - std::cmp::Ordering::Equal => { + Ordering::Equal => { self.handle_tie(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Greater => { + Ordering::Greater => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; self.update_winner(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Less => { + Ordering::Less => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 77a7d8f8f2e11..32f63ce6a2d72 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -442,7 +442,7 @@ mod tests { ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray, TimestampNanosecondArray, }; - use arrow::compute::SortOptions; + use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::test_util::batches_to_string; use datafusion_common::{assert_batches_eq, exec_err}; @@ -776,9 +776,9 @@ mod tests { context: Arc, ) -> RecordBatch { let merge = Arc::new(SortPreservingMergeExec::new(sort, input)); - let mut result = collect(merge, context).await.unwrap(); - assert_eq!(result.len(), 1); - result.remove(0) + let result = collect(merge, context).await.unwrap(); + // the merge may split the output into several batches, so concat first + concat_batches(&result[0].schema(), &result).unwrap() } async fn partition_sort( @@ -1139,10 +1139,10 @@ mod tests { .with_reservation(reservation) .build()?; - let mut merged = common::collect(merge_stream).await.unwrap(); + let merged = common::collect(merge_stream).await.unwrap(); - assert_eq!(merged.len(), 1); - let merged = merged.remove(0); + // the merge may split the output into several batches, so concat first + let merged = concat_batches(&merged[0].schema(), &merged).unwrap(); let basic = basic_sort(batches, sort.clone(), Arc::clone(&task_ctx)).await; let basic = arrow::util::pretty::pretty_format_batches(&[basic]) @@ -1557,8 +1557,16 @@ mod tests { .with_reservation(reservation) .build()?; - let first = merge_stream.next().await.unwrap(); - assert!(first.is_err(), "expected merge stream to surface the error"); + // the merge may emit already-proven-sorted batches before polling the + // failing stream again, so consume until the error surfaces + let mut saw_error = false; + while let Some(item) = merge_stream.next().await { + if item.is_err() { + saw_error = true; + break; + } + } + assert!(saw_error, "expected merge stream to surface the error"); assert!( merge_stream.next().await.is_none(), "merge stream yielded data after returning an error" diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index d7de13572693e..7eba1a35db2eb 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -274,6 +274,7 @@ mod tests { use arrow::array::{Int32Array, RecordBatch}; use arrow::datatypes::Int32Type; use arrow_schema::{DataType, Field, Schema}; + use datafusion_execution::memory_pool::GreedyMemoryPool; use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; @@ -285,6 +286,21 @@ mod tests { inputs: Vec>>>, batch_size: usize, fetch: Option, + ) -> Vec>> { + merge_nullable_i32_streams_with_pool( + inputs, + batch_size, + fetch, + Arc::new(UnboundedMemoryPool::default()), + ) + .await + } + + async fn merge_nullable_i32_streams_with_pool( + inputs: Vec>>>, + batch_size: usize, + fetch: Option, + mem_pool: Arc, ) -> Vec>> { let schema = Arc::new(Schema::new(vec![Field::new( "sort_key", @@ -315,11 +331,9 @@ mod tests { .with_streams(streams) .with_batch_size(batch_size) .with_schema(Arc::clone(&schema)) - .with_reservation({ - let mem_pool: Arc = - Arc::new(UnboundedMemoryPool::default()); - MemoryConsumer::new("merge stream mock memory").register(&mem_pool) - }) + .with_reservation( + MemoryConsumer::new("merge stream mock memory").register(&mem_pool), + ) .with_expressions( &([ PhysicalSortExpr::new_default(col("sort_key", &schema).unwrap()) @@ -393,6 +407,112 @@ mod tests { assert_eq!(flatten_vec(output), (1..=38).collect::>()); } + #[tokio::test] + async fn test_merge_full_batch_release_memory() { + // The pool only fits a few batches of the ~400KB total input: + // if skipped batches were not released as the merge progresses, the reservation would exhaust the pool. + const ROWS: i32 = 1024; // ~4KB per batch + const BATCHES_PER_STREAM: i32 = 50; + + let batch_of = |start: i32| (start..start + ROWS).map(Some).collect::>(); + let streams = (0..2) + .map(|stream| { + (0..BATCHES_PER_STREAM) + .map(|batch| batch_of((batch * 2 + stream) * ROWS)) + .collect() + }) + .collect(); + + let output = merge_nullable_i32_streams_with_pool( + streams, + ROWS as usize, + None, + Arc::new(GreedyMemoryPool::new(64 * 1024)), + ) + .await; + + let expected = (0..2 * BATCHES_PER_STREAM * ROWS) + .map(Some) + .collect::>(); + 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::().values().iter()) + .copied() + .collect::>(); + 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 From 92505cfb4bb6317a7a9df973609de3a4d8e4d99e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:00:55 +0300 Subject: [PATCH 09/10] add comment --- datafusion/physical-plan/src/sorts/merge.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bbba884e4225e..4cbde299dc8f3 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -356,6 +356,9 @@ impl SortPreservingMergeStream { } // Flush already-merged rows first // they all sort before the skipped batch + // + // No need to reserve memory for holding this batch: it is a + // zero-copy slice of a batch the builder still buffers which also reserve memory for self.pending_skipped_batch = Some(batch); return Poll::Ready(self.emit_in_progress_batch().transpose()); } From 06d59ac09b1821d16dc0b4bfb0a39cf5db227bca Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 8 Jul 2026 14:27:38 +0300 Subject: [PATCH 10/10] remove batch wise round robin and rely on row wise round robin --- datafusion/physical-plan/src/sorts/merge.rs | 45 +---------- .../src/sorts/streaming_merge.rs | 76 ------------------- 2 files changed, 4 insertions(+), 117 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4cbde299dc8f3..1b9e3ac3cf156 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -346,9 +346,6 @@ impl SortPreservingMergeStream { // 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 SortPreservingMergeStream { 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 SortPreservingMergeStream { 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; diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 7eba1a35db2eb..67f5ff7f0cb3d 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/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::().values().iter()) - .copied() - .collect::>(); - 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