diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 986da549f75c8..8654775b2fa98 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,122 @@ impl SortPreservingMergeStream { result } - 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::>(); + async fn flush_in_progress( + &mut self, + mut emitter: TryEmitter, + ) -> Result<()> { + if self.in_progress.is_empty() { + return Ok(()); + } - poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) - .await?; + 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(); + } - assert_eq!(uninitiated_partitions.len(), 0); + Ok(()) + } - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + 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?; - // Claim the memory for the uninitiated partitions - drop(uninitiated_partitions); - self.init_loser_tree(); + assert_eq!(uninitiated_partitions.len(), 0); + } - // 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(); - loop { - let stream_idx = self.loser_tree[0]; - if !self.advance_cursors(stream_idx) { - break; - } - self.in_progress.push_row(stream_idx); + // 2. Init loser tree + self.init_loser_tree(); - // stop sorting if fetch has been reached + // 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); + + // 3.2. If the new row reached the limit if self.fetch_reached() { break; } - if self.in_progress.len() >= self.batch_size - && let Some(batch) = self.emit_in_progress_batch()? - { + // 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") + }; + + // 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` @@ -369,18 +414,21 @@ 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(); - 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");