From 73382ebe35c586120c373e18c7e6e099d686989c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:45:04 +0300 Subject: [PATCH 1/6] chore: simplify SortPreservingMergeStream to be as textbook-like as possible --- datafusion/physical-plan/src/sorts/merge.rs | 163 +++++++++++------- .../src/sorts/streaming_merge.rs | 20 ++- 2 files changed, 117 insertions(+), 66 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 986da549f75c8..c9108eff8fde5 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -32,9 +32,9 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; -use datafusion_execution::async_try_stream; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{TryEmitter, async_try_stream}; use futures::Stream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] @@ -146,6 +146,9 @@ impl SortPreservingMergeStream { reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, ) -> Self { + assert_ne!(batch_size, 0, "batch size cannot be 0"); + assert_ne!(fetch, Some(0), "fetch must not be Some(0)"); + let stream_count = streams.partitions(); Self { @@ -173,7 +176,6 @@ impl SortPreservingMergeStream { let schema_clone = Arc::clone(self.in_progress.schema()); let cloned_metrics = self.metrics.clone(); - let stream = Box::pin(RecordBatchStreamAdapter::new( schema_clone, self.create_stream(), @@ -212,79 +214,125 @@ impl SortPreservingMergeStream { result } + async fn flush_in_progress( + &mut self, + mut emitter: TryEmitter, + ) -> Result<()> { + if !self.in_progress.is_empty() { + return Ok(()); + } + + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + // When `build_record_batch()` hits an i32 offset overflow (e.g. + // combined string offsets exceed 2 GB), it emits a partial batch + // and keeps the remaining rows in `self.in_progress.indices`. + // Drain those leftover rows before terminating the stream, + // otherwise they would be silently dropped. + // Repeated overflows are fine — each poll emits another partial + // batch until `in_progress` is fully drained. + while let Some(batch) = self.emit_in_progress_batch()? { + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + Ok(()) + } + fn create_stream(mut self) -> impl Stream> { async_try_stream(|mut emitter| async move { - // This vector contains the indices of the partitions that have not started emitting yet. - let mut uninitiated_partitions = - (0..self.streams.partitions()).collect::>(); - - poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) + assert!( + self.fetch.is_none_or(|fetch| fetch != 0), + "fetch {:?} must not be 0", + self.fetch + ); + + // 1. Make sure we have data from each stream so we can initialize the loser tree + { + // This vector contains the indices of the partitions that have not started emitting yet. + let mut uninitiated_partitions = + (0..self.streams.partitions()).collect::>(); + + poll_fn(|cx| { + self.initialize_all_partitions(&mut uninitiated_partitions, cx) + }) .await?; - assert_eq!(uninitiated_partitions.len(), 0); + assert_eq!(uninitiated_partitions.len(), 0); + } - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + let fetch = self.fetch; - // Claim the memory for the uninitiated partitions - drop(uninitiated_partitions); + // 2. Init loser tree self.init_loser_tree(); - // NB timer records time taken on drop, so there are no - // calls to `timer.done()` below. - let elapsed_compute = self.metrics.elapsed_compute().clone(); - let mut timer = elapsed_compute.timer(); + // 3. loop until all streams have been exhausted + while !self.is_exhausted() { + // 3.1. add loser_tree[0] (minimum) stream to pending record batch + let winner_stream = self.loser_tree[0]; + self.in_progress.push_row(winner_stream); - loop { - let stream_idx = self.loser_tree[0]; - if !self.advance_cursors(stream_idx) { + // 3.2. If the new row reached the limit + if fetch + .is_some_and(|fetch| fetch <= self.produced + self.in_progress.len()) + { break; } - self.in_progress.push_row(stream_idx); - // stop sorting if fetch has been reached - if self.fetch_reached() { - break; - } + // 3.3. if there is enough to emit for a full record batch + if self.in_progress.len() >= self.batch_size { + // 3.3.1 build pending record batch and reset builder + let Some(batch) = self.emit_in_progress_batch()? else { + unreachable!("must have batch in progress to emit") + }; - if self.in_progress.len() >= self.batch_size - && let Some(batch) = self.emit_in_progress_batch()? - { + // 3.3.2 emit pending record batch drop(timer); emitter.emit(batch).await; timer = elapsed_compute.timer(); } - let winner = self.loser_tree[0]; - // Fast path: skip the `maybe_poll_stream` call (and its `Poll` - // plumbing) unless the winner's cursor is exhausted and needs a - // fresh batch — it is live for almost every row. - if self.cursors[winner].is_none() { - drop(timer); - poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; - timer = elapsed_compute.timer(); + // 3.4. advance cursor for the winner stream + { + let should_poll_next_batch_for_stream = + self.advance_cursors(winner_stream); + + // Fast path: skip the `maybe_poll_stream` call (and its `Poll` + // plumbing) unless the winner's cursor is exhausted and needs a + // fresh batch — it is live for almost every row. + if should_poll_next_batch_for_stream { + assert!( + self.cursors[winner_stream].is_none(), + "cursor should be exhausted" + ); + + drop(timer); + poll_fn(|cx| self.maybe_poll_stream(cx, winner_stream)).await?; + timer = elapsed_compute.timer(); + } } - // Adjusting the loser tree if necessary + // 3.5. Adjusting the loser tree if necessary self.update_loser_tree(); } - drop(timer); + // 4. Flush any remaining rows in `self.in_progress` + self.flush_in_progress(emitter).await?; - // When `build_record_batch()` hits an i32 offset overflow (e.g. - // combined string offsets exceed 2 GB), it emits a partial batch - // and keeps the remaining rows in `self.in_progress.indices`. - // Drain those leftover rows before terminating the stream, - // otherwise they would be silently dropped. - // Repeated overflows are fine — each poll emits another partial - // batch until `in_progress` is fully drained. - while let Some(batch) = self.emit_in_progress_batch()? { - emitter.emit(batch).await; - } Ok(()) }) } + fn is_exhausted(&self) -> bool { + let winner = self.loser_tree[0]; + + self.cursors[winner].is_none() + } + /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` /// /// This DOES NOT return `Poll::Pending` as soon as the first uninitiated partition returns `Poll::Pending` @@ -360,12 +408,6 @@ impl SortPreservingMergeStream { } } - fn fetch_reached(&mut self) -> bool { - self.fetch - .map(|fetch| self.produced + self.in_progress.len() >= fetch) - .unwrap_or(false) - } - /// Advances the actual cursor. If it reaches its end, update the /// previous cursor with it. /// @@ -373,14 +415,17 @@ impl SortPreservingMergeStream { fn advance_cursors(&mut self, stream_idx: usize) -> bool { if let Some(cursor) = &mut self.cursors[stream_idx] { let _ = cursor.advance(); - if cursor.is_finished() { + return if !cursor.is_finished() { // Take the current cursor, leaving `None` in its place self.prev_cursors[stream_idx] = self.cursors[stream_idx].take(); - } - true - } else { - false + + true + } else { + false + }; } + + true } /// Returns `true` if the cursor at index `a` is greater than at index `b`. diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index e96138ef1306c..5fef8f05bd3c3 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -24,7 +24,7 @@ use crate::sorts::{ merge::SortPreservingMergeStream, stream::{FieldCursorStream, RowCursorStream}, }; -use crate::{SendableRecordBatchStream, SpillManager}; +use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager}; use arrow::array::*; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::human_readable_size; @@ -195,13 +195,22 @@ impl<'a> StreamingMergeBuilder<'a> { let Some(expressions) = expressions else { return internal_err!("Sort expressions cannot be empty for streaming merge"); }; + let schema = schema.expect("Schema cannot be empty for streaming merge"); + + if fetch.is_some_and(|fetch| fetch == 0) { + return Ok(Box::pin(EmptyRecordBatchStream::new(schema))); + } + + let batch_size = + batch_size.expect("Batch size cannot be empty for streaming merge"); + + if batch_size == 0 { + return internal_err!("Batch size cannot be zero for streaming merge"); + } if !sorted_spill_files.is_empty() { // Unwrapping mandatory fields - let schema = schema.expect("Schema cannot be empty for streaming merge"); let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); - let batch_size = - batch_size.expect("Batch size cannot be empty for streaming merge"); let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); @@ -227,10 +236,7 @@ impl<'a> StreamingMergeBuilder<'a> { ); // Unwrapping mandatory fields - let schema = schema.expect("Schema cannot be empty for streaming merge"); let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); - let batch_size = - batch_size.expect("Batch size cannot be empty for streaming merge"); let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); From 61cadd63db6dab66956bad77156eef714f0264ad Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:57:41 +0300 Subject: [PATCH 2/6] add back function --- datafusion/physical-plan/src/sorts/merge.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index c9108eff8fde5..444999939e3c7 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -265,7 +265,6 @@ impl SortPreservingMergeStream { let elapsed_compute = self.metrics.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); - let fetch = self.fetch; // 2. Init loser tree self.init_loser_tree(); @@ -277,9 +276,7 @@ impl SortPreservingMergeStream { self.in_progress.push_row(winner_stream); // 3.2. If the new row reached the limit - if fetch - .is_some_and(|fetch| fetch <= self.produced + self.in_progress.len()) - { + if self.is_fetch_reached() { break; } @@ -408,6 +405,12 @@ impl SortPreservingMergeStream { } } + fn fetch_reached(&mut self) -> bool { + self.fetch + .map(|fetch| self.produced + self.in_progress.len() >= fetch) + .unwrap_or(false) + } + /// Advances the actual cursor. If it reaches its end, update the /// previous cursor with it. /// From c56656bcf6e68d05d48bba83b9237f7e927614bc Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 19:01:13 +0300 Subject: [PATCH 3/6] fix --- datafusion/physical-plan/src/sorts/merge.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 444999939e3c7..13334805adaea 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -276,7 +276,7 @@ impl SortPreservingMergeStream { self.in_progress.push_row(winner_stream); // 3.2. If the new row reached the limit - if self.is_fetch_reached() { + if self.fetch_reached() { break; } @@ -407,8 +407,8 @@ impl SortPreservingMergeStream { fn fetch_reached(&mut self) -> bool { self.fetch - .map(|fetch| self.produced + self.in_progress.len() >= fetch) - .unwrap_or(false) + .map(|fetch| self.produced + self.in_progress.len() >= fetch) + .unwrap_or(false) } /// Advances the actual cursor. If it reaches its end, update the From c02733b945a0ec3ab346ef6b7c1e50252917d2fe Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 19:06:31 +0300 Subject: [PATCH 4/6] update comment --- datafusion/physical-plan/src/sorts/merge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 13334805adaea..bc4bd3ae0b9e0 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -414,7 +414,7 @@ impl SortPreservingMergeStream { /// Advances the actual cursor. If it reaches its end, update the /// previous cursor with it. /// - /// If the given partition is not exhausted, the function returns `true`. + /// If the given partition batch is exhausted, return `true` to signal a poll is needed fn advance_cursors(&mut self, stream_idx: usize) -> bool { if let Some(cursor) = &mut self.cursors[stream_idx] { let _ = cursor.advance(); From ce00c14669b246b8de8d9f43c4b2df0afd59027e Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 19:08:22 +0300 Subject: [PATCH 5/6] fix ! --- datafusion/physical-plan/src/sorts/merge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bc4bd3ae0b9e0..eaaf251f6825a 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -418,7 +418,7 @@ impl SortPreservingMergeStream { fn advance_cursors(&mut self, stream_idx: usize) -> bool { if let Some(cursor) = &mut self.cursors[stream_idx] { let _ = cursor.advance(); - return if !cursor.is_finished() { + return if cursor.is_finished() { // Take the current cursor, leaving `None` in its place self.prev_cursors[stream_idx] = self.cursors[stream_idx].take(); From 7175283b093ac1f1fed6642a8028b5a0d6b03a14 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 19 Jul 2026 19:54:56 +0300 Subject: [PATCH 6/6] fix --- datafusion/physical-plan/src/sorts/merge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index eaaf251f6825a..8654775b2fa98 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -218,7 +218,7 @@ impl SortPreservingMergeStream { &mut self, mut emitter: TryEmitter, ) -> Result<()> { - if !self.in_progress.is_empty() { + if self.in_progress.is_empty() { return Ok(()); }