From 2d8acc8b9475353a735120ba31a562ee32900ce4 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 00:14:33 +0300 Subject: [PATCH 1/4] perf(sort): passthrough last stream that left in `SortPreservingMergeStream` --- datafusion/physical-plan/src/sorts/builder.rs | 9 ++ datafusion/physical-plan/src/sorts/merge.rs | 98 ++++++++++++++- .../src/sorts/streaming_merge.rs | 119 ++++++++++++++++++ 3 files changed, 225 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..2b6dda9b898cd 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(crate) 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/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..01517a1426e7a 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -153,6 +153,23 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, + + /// Which input streams are fully exhausted (their `poll_next` returned `None`) + stream_exhausted: Vec, + + /// Number of `true` entries in `stream_exhausted` + num_exhausted_streams: usize, + + /// When every stream but one is exhausted there is nothing left to merge + /// against: the surviving stream's index is recorded here and its + /// batches are forwarded as is (skipping the loser tree entirely). + /// + /// Only enabled when `fetch` is `None`. + passthrough_stream: Option, + + /// The not-yet-merged tail of the surviving stream's current batch, + /// emitted once before forwarding the stream itself. + passthrough_remainder: Option, } impl SortPreservingMergeStream { @@ -186,6 +203,10 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + stream_exhausted: vec![false; stream_count], + num_exhausted_streams: 0, + passthrough_stream: None, + passthrough_remainder: None, } } @@ -203,7 +224,13 @@ impl SortPreservingMergeStream { } match futures::ready!(self.streams.poll_next(cx, idx)) { - None => Poll::Ready(Ok(())), + None => { + if !self.stream_exhausted[idx] { + self.stream_exhausted[idx] = true; + self.num_exhausted_streams += 1; + } + Poll::Ready(Ok(())) + } Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { self.cursors[idx] = Some(Cursor::new(cursor)); @@ -236,6 +263,11 @@ impl SortPreservingMergeStream { } return Poll::Ready(None); } + + if self.passthrough_stream.is_some() { + return self.poll_passthrough(cx); + } + // 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. @@ -307,6 +339,29 @@ impl SortPreservingMergeStream { } } self.update_loser_tree(); + + // If every stream but the winner is exhausted there is nothing left to merge against + // switch to forwarding the winner's remaining data as is. + // TODO - support fetch in passthrough + if self.fetch.is_none() + && self.num_exhausted_streams + 1 == self.cursors.len() + { + let winner = self.loser_tree[0]; + if self.cursors[winner].is_some() { + // Rows of the winner's current batch that were not merged yet + // everything before them is already in `in_progress`, + // and everything after them is still in the stream, so ordering is preserved. + self.passthrough_remainder = + self.in_progress.remaining_rows_of(winner); + // No comparisons happen in passthrough mode: + // drop both cursors so the cursor stream's reusable row buffers are free for the batches still converted + // while being forwarded (RowCursorStream keeps only two per stream) + self.prev_cursors[winner] = None; + self.cursors[winner] = None; + self.passthrough_stream = Some(winner); + return self.poll_passthrough(cx); + } + } } let stream_idx = self.loser_tree[0]; @@ -327,6 +382,47 @@ impl SortPreservingMergeStream { } } + /// Drains the stream once only one input stream is left: + /// first the rows already merged into `in_progress`, then the not-yet-merged tail of the + /// surviving stream's current batch, then its batches forwarded as is. + fn poll_passthrough( + &mut self, + cx: &mut Context<'_>, + ) -> Poll>> { + let stream_idx = self + .passthrough_stream + .expect("passthrough stream must be set"); + + // 1. Drain the in-progress batch if we have any + // + // The in-progress might not be full batch size, + // but in order for us to passthrough the stream without doing additional computation + // we are ok with emitting a partial batch now. + if !self.in_progress.is_empty() { + return Poll::Ready(self.emit_in_progress_batch().transpose()); + } + + // 2. Emitting the remainder of the stream's current batch + // Not doing copy or concat to avoid allocating a new batch + if let Some(remainder) = self.passthrough_remainder.take() { + self.produced += remainder.num_rows(); + return Poll::Ready(Some(Ok(remainder))); + } + + // 3. Forwarding the stream's remaining batches as is + match futures::ready!(self.streams.poll_next(cx, stream_idx)) { + None => Poll::Ready(None), + Some(Err(e)) => { + self.done = true; + Poll::Ready(Some(Err(e))) + } + Some(Ok((_, batch))) => { + self.produced += batch.num_rows(); + Poll::Ready(Some(Ok(batch))) + } + } + } + /// 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..08775b07d80a1 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -265,3 +265,122 @@ 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; + + #[tokio::test] + async fn test_sort_preserving_merge_stream_with_one_stream_larger_than_other() { + let schema = Arc::new(Schema::new(vec![Field::new( + "sort_key", + DataType::Int32, + false, + )])); + + let data_left = vec![ + Int32Array::from(vec![1, 3, 5, 7, 9]), + Int32Array::from(vec![11, 13, 15, 17, 19]), + Int32Array::from(vec![21, 23, 25, 27, 29]), + Int32Array::from(vec![31, 33, 35, 37, 39]), + Int32Array::from(vec![41, 43, 45, 47, 49]), + ]; + + let data_right = vec![ + Int32Array::from(vec![0, 2, 4, 6, 8]), + Int32Array::from(vec![9, 10, 11, 12, 13]), + ]; + + for fetch in [ + None, + Some(2), + Some(10), + Some(13), + Some(17), + Some(20), + Some(25), + Some(30), + ] + .iter() + { + let data_left_stream = MemoryStream::try_new( + data_left + .iter() + .map(|a| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(a.clone())], + ) + .unwrap() + }) + .collect::>(), + Arc::clone(&schema), + None, + ) + .unwrap(); + + let data_right_stream = MemoryStream::try_new( + data_right + .iter() + .map(|a| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(a.clone())], + ) + .unwrap() + }) + .collect::>(), + Arc::clone(&schema), + None, + ) + .unwrap(); + + let merged_sort = StreamingMergeBuilder::new() + .with_streams(vec![ + Box::pin(data_left_stream), + Box::pin(data_right_stream), + ]) + .with_batch_size(5) + .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_sort).await.unwrap(); + let sorted_output = output + .into_iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect::>(); + + let expected = vec![ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 10, 11, 11, 12, 13, 13, 15, 17, 19, 21, + 23, 25, 27, 29, 31, 33, 35, 37, 39, 41, 43, 45, 47, 49, + ]; + + assert_eq!(sorted_output, &expected[0..fetch.unwrap_or(expected.len())]); + } + } +} From d7156e6ebc32377d63ebcb7609bba609a2962924 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 6 Jul 2026 00:17:30 +0300 Subject: [PATCH 2/4] lint --- datafusion/physical-plan/src/sorts/streaming_merge.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 08775b07d80a1..483d1763c3916 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -287,7 +287,7 @@ mod tests { false, )])); - let data_left = vec![ + let data_left = [ Int32Array::from(vec![1, 3, 5, 7, 9]), Int32Array::from(vec![11, 13, 15, 17, 19]), Int32Array::from(vec![21, 23, 25, 27, 29]), @@ -295,7 +295,7 @@ mod tests { Int32Array::from(vec![41, 43, 45, 47, 49]), ]; - let data_right = vec![ + let data_right = [ Int32Array::from(vec![0, 2, 4, 6, 8]), Int32Array::from(vec![9, 10, 11, 12, 13]), ]; From ed18da78d14518ced4d4018bac198b0137d4b15c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:19:32 +0300 Subject: [PATCH 3/4] test: don't rely on output batch count in test_sort_with_streaming_table --- .../core/tests/physical_optimizer/enforce_sorting.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index ecff2edbbec16..cfc42f3575692 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 number of output batches is not guaranteed, so concatenate before comparing + 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(()) } From 07411ff75ad605c8a350c96637734218e82b21f9 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:38:51 +0300 Subject: [PATCH 4/4] release memory since we no longer hold on it anymore --- datafusion/physical-plan/src/sorts/builder.rs | 27 +++++++++++++++++++ datafusion/physical-plan/src/sorts/merge.rs | 7 ++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 2b6dda9b898cd..fa09932efdae3 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -138,6 +138,33 @@ impl BatchBuilder { (remaining > 0).then(|| batch.slice(cursor.row_idx, remaining)) } + /// Drop all buffered batches and release their memory back to the pool + /// (never shrinking below the pre-reserved `initial_reservation` floor). + /// + /// Only valid when there are no in-progress rows, and no rows may be + /// pushed or read afterwards: used when the merge switches to forwarding + /// whole batches that bypass this builder entirely. + pub(crate) fn release_buffered_batches(&mut self) { + assert_eq!( + self.indices.len(), + 0, + "in-progress rows still reference the buffered batches" + ); + if self.batches.is_empty() { + return; + } + // Remove data and release capacity + self.batches = vec![]; + self.indices = vec![]; + self.cursors = vec![]; + + self.batches_mem_used = 0; + if self.reservation.size() > self.initial_reservation { + self.reservation + .shrink(self.reservation.size() - self.initial_reservation); + } + } + /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 01517a1426e7a..795dbfbf90970 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -409,7 +409,12 @@ impl SortPreservingMergeStream { return Poll::Ready(Some(Ok(remainder))); } - // 3. Forwarding the stream's remaining batches as is + // 3. Everything buffered was emitted, so release the buffered batches + // and their memory reservation since we now no longer hold on to any memory + // no need to reserve memory + self.in_progress.release_buffered_batches(); + + // 4. Forwarding the stream's remaining batches as is match futures::ready!(self.streams.poll_next(cx, stream_idx)) { None => Poll::Ready(None), Some(Err(e)) => {