chore: Simplifying SortPreservingMergeStream to use generators instead of state machine

datafusion/physical-plan/src/sorts/cursor.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::cmp::Ordering;
+use std::fmt::Debug;
 use std::sync::Arc;
 
 use arrow::array::{
@@ -32,7 +33,7 @@ use datafusion_execution::memory_pool::MemoryReservation;
 ///
 /// This is a trait as there are several specialized implementations, such as for
 /// single columns or for normalized multi column keys ([`Rows`])
-pub trait CursorValues {
+pub trait CursorValues: Debug + Sync + Send {
     fn len(&self) -> usize;
 
     /// Returns true if `l[l_idx] == r[r_idx]`
@@ -302,6 +303,7 @@ impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> {
     }
 }
 
+#[derive(Debug)]
 pub struct ByteArrayValues<T: OffsetSizeTrait> {
     offsets: OffsetBuffer<T>,
     values: Buffer,

datafusion/physical-plan/src/sorts/merge.rs

@@ -18,21 +18,23 @@
 //! Merge that deals with an arbitrary size of streaming inputs.
 //! This is an order-preserving merge.
 
-use std::pin::Pin;
+use std::fmt::Debug;
+use std::future::poll_fn;
 use std::sync::Arc;
-use std::task::{Context, Poll, ready};
+use std::task::{Context, Poll};
 
-use crate::RecordBatchStream;
+use crate::SendableRecordBatchStream;
 use crate::metrics::BaselineMetrics;
 use crate::sorts::builder::BatchBuilder;
 use crate::sorts::cursor::{Cursor, CursorValues};
 use crate::sorts::stream::PartitionedStream;
+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_execution::memory_pool::MemoryReservation;
-
 use futures::Stream;
 
 /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`]
@@ -49,18 +51,6 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// used to record execution metrics
     metrics: BaselineMetrics,
 
-    /// If the stream has encountered an error or reaches the
-    /// `fetch` limit.
-    done: bool,
-
-    /// Whether buffered rows should be drained after `done` is set.
-    ///
-    /// This is enabled when we stop because the `fetch` limit has been
-    /// reached, allowing partial batches left over after overflow handling to
-    /// be emitted on subsequent polls. It remains disabled for terminal
-    /// errors so the stream does not yield data after returning `Err`.
-    drain_in_progress_on_done: bool,
-
     /// A loser tree that always produces the minimum cursor
     ///
     /// Node 0 stores the top winner, Nodes 1..num_streams store
@@ -93,12 +83,6 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
     /// reference: <https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree>
     loser_tree: Vec<usize>,
 
-    /// If the most recently yielded overall winner has been replaced
-    /// within the loser tree. A value of `false` indicates that the
-    /// overall winner has been yielded but the loser tree has not
-    /// been updated
-    loser_tree_adjusted: bool,
-
     /// Target batch size
     batch_size: usize,
 
@@ -150,9 +134,6 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
 
     /// number of rows produced
     produced: usize,
-
-    /// This vector contains the indices of the partitions that have not started emitting yet.
-    uninitiated_partitions: Vec<usize>,
 }
 
 impl<C: CursorValues> SortPreservingMergeStream<C> {
@@ -171,24 +152,36 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation),
             streams,
             metrics,
-            done: false,
-            drain_in_progress_on_done: false,
             cursors: (0..stream_count).map(|_| None).collect(),
             prev_cursors: (0..stream_count).map(|_| None).collect(),
             round_robin_tie_breaker_mode: false,
             num_of_polled_with_same_value: vec![0; stream_count],
             current_reset_epoch: 0,
             poll_reset_epochs: vec![0; stream_count],
             loser_tree: vec![],
-            loser_tree_adjusted: false,
             batch_size,
             fetch,
             produced: 0,
-            uninitiated_partitions: (0..stream_count).collect(),
             enable_round_robin_tie_breaker,
         }
     }
 
+    pub(crate) fn into_stream(self) -> SendableRecordBatchStream
+    where
+        C: 'static,
+    {
+        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(),
+        ));
+
+        Box::pin(ObservedStream::new(stream, cloned_metrics, None))
+    }
+
     /// If the stream at the given index is not exhausted, and the last cursor for the
     /// stream is finished, poll the stream for the next RecordBatch and create a new
     /// cursor for the stream from the returned result
@@ -219,90 +212,88 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         result
     }
 
-    fn poll_next_inner(
-        &mut self,
-        cx: &mut Context<'_>,
-    ) -> Poll<Option<Result<RecordBatch>>> {
-        if self.done {
-            // 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.
-            if self.drain_in_progress_on_done && !self.in_progress.is_empty() {
-                return Poll::Ready(self.emit_in_progress_batch().transpose());
-            }
-            return Poll::Ready(None);
-        }
+    fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
+        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::<Vec<_>>();
 
-        // 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.
-        if self.loser_tree.is_empty() {
-            ready!(self.initialize_all_partitions(cx))?;
-            assert_eq!(
-                self.uninitiated_partitions.len(),
-                0,
-                "all partitions should be initialized"
-            );
+            poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx))
+                .await?;
+
+            assert_eq!(uninitiated_partitions.len(), 0);
 
             // If there are no more uninitiated partitions, set up the loser tree and continue
             // to the next phase.
 
             // Claim the memory for the uninitiated partitions
-            self.uninitiated_partitions.shrink_to_fit();
+            drop(uninitiated_partitions);
             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 _timer = elapsed_compute.timer();
+            // 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);
+
+                // stop sorting if fetch has been reached
+                if self.fetch_reached() {
+                    break;
+                }
+
+                if self.in_progress.len() >= self.batch_size
+                    && let Some(batch) = self.emit_in_progress_batch()?
+                {
+                    drop(timer);
+                    emitter.emit(batch).await;
+                    timer = elapsed_compute.timer();
+                }
 
-        loop {
-            // Adjust the loser tree if necessary, returning control if needed
-            if !self.loser_tree_adjusted {
                 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() {
-                    match ready!(self.maybe_poll_stream(cx, winner)) {
-                        Ok(()) => {}
-                        Err(e) => {
-                            self.done = true;
-                            return Poll::Ready(Some(Err(e)));
-                        }
-                    }
+                    drop(timer);
+                    poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?;
+                    timer = elapsed_compute.timer();
                 }
+
+                // Adjusting the loser tree if necessary
                 self.update_loser_tree();
             }
 
-            let stream_idx = self.loser_tree[0];
-            if self.advance_cursors(stream_idx) {
-                self.loser_tree_adjusted = false;
-                self.in_progress.push_row(stream_idx);
+            drop(timer);
 
-                // stop sorting if fetch has been reached
-                if self.fetch_reached() {
-                    self.done = true;
-                    self.drain_in_progress_on_done = true;
-                } else if self.in_progress.len() < self.batch_size {
-                    continue;
-                }
+            // 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;
             }
-
-            return Poll::Ready(self.emit_in_progress_batch().transpose());
-        }
+            Ok(())
+        })
     }
 
     /// 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`
     /// so we can continue to initialize the remaining partitions
-    fn initialize_all_partitions(&mut self, cx: &mut Context) -> Poll<Result<()>> {
+    fn initialize_all_partitions(
+        &mut self,
+        uninitiated_partitions: &mut Vec<usize>,
+        cx: &mut Context,
+    ) -> Poll<Result<()>> {
         assert_eq!(
             self.loser_tree.len(),
             0,
@@ -311,11 +302,10 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
 
         // Manual indexing since we're iterating over the vector and shrinking it in the loop
         let mut idx = 0;
-        while idx < self.uninitiated_partitions.len() {
-            let partition_idx = self.uninitiated_partitions[idx];
+        while idx < uninitiated_partitions.len() {
+            let partition_idx = uninitiated_partitions[idx];
             match self.maybe_poll_stream(cx, partition_idx) {
                 Poll::Ready(Err(e)) => {
-                    self.done = true;
                     return Poll::Ready(Err(e));
                 }
                 Poll::Pending => {
@@ -331,12 +321,12 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                     // place which we'll try in the next loop iteration
                     // swap_remove will change the partition poll order, but that shouldn't
                     // make a difference since we're waiting for all streams to be ready.
-                    self.uninitiated_partitions.swap_remove(idx);
+                    uninitiated_partitions.swap_remove(idx);
                 }
             }
         }
 
-        if self.uninitiated_partitions.is_empty() {
+        if uninitiated_partitions.is_empty() {
             Poll::Ready(Ok(()))
         } else {
             // There are still uninitiated partitions so return pending.
@@ -474,7 +464,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             }
             self.loser_tree[cmp_node] = winner;
         }
-        self.loser_tree_adjusted = true;
     }
 
     /// Resets the poll count by incrementing the reset epoch.
@@ -586,25 +575,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         }
 
         self.loser_tree[0] = winner;
-        self.loser_tree_adjusted = true;
-    }
-}
-
-impl<C: CursorValues + Unpin> Stream for SortPreservingMergeStream<C> {
-    type Item = Result<RecordBatch>;
-
-    fn poll_next(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Option<Self::Item>> {
-        let poll = self.poll_next_inner(cx);
-        self.metrics.record_poll(poll)
-    }
-}
-
-impl<C: CursorValues + Unpin> RecordBatchStream for SortPreservingMergeStream<C> {
-    fn schema(&self) -> SchemaRef {
-        Arc::clone(self.in_progress.schema())
     }
 }
 
@@ -618,7 +588,7 @@ mod tests {
     use datafusion_execution::memory_pool::{
         MemoryConsumer, MemoryPool, UnboundedMemoryPool,
     };
-    use futures::task::noop_waker_ref;
+    use futures::TryStreamExt;
     use std::cmp::Ordering;
 
     #[derive(Debug)]
@@ -661,8 +631,8 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_done_drains_buffered_rows() {
+    #[tokio::test]
+    async fn test_done_drains_buffered_rows() {
         let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
         let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default());
         let reservation = MemoryConsumer::new("test").register(&pool);
@@ -678,24 +648,20 @@ mod tests {
             true,
         );
 
+        // Simulate rows left buffered in `in_progress` (as happens when
+        // `build_record_batch` emits a partial batch on offset overflow). With
+        // an empty input stream the merge loop breaks immediately, so the only
+        // way these rows reach the consumer is the generator's final drain loop.
         let batch =
             RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))])
                 .unwrap();
         stream.in_progress.push_batch(0, batch).unwrap();
         stream.in_progress.push_row(0);
-        stream.done = true;
-        stream.drain_in_progress_on_done = true;
 
-        let waker = noop_waker_ref();
-        let mut cx = Context::from_waker(waker);
+        // Drive the actual stream and confirm the buffered row is drained.
+        let batches: Vec<RecordBatch> = stream.into_stream().try_collect().await.unwrap();
 
-        match stream.poll_next_inner(&mut cx) {
-            Poll::Ready(Some(Ok(batch))) => assert_eq!(batch.num_rows(), 1),
-            other => {
-                panic!("expected buffered rows to be drained after done, got {other:?}")
-            }
-        }
-        assert!(stream.in_progress.is_empty());
-        assert!(matches!(stream.poll_next_inner(&mut cx), Poll::Ready(None)));
+        assert_eq!(batches.len(), 1);
+        assert_eq!(batches[0].num_rows(), 1);
     }
 }