perf(sort): passthrough last stream that left in SortPreservingMergeStream

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(())
 }

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

@@ -129,6 +129,42 @@ 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<RecordBatch> {
+        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))
+    }
+
+    /// 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,

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

@@ -153,6 +153,23 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
 
     /// This vector contains the indices of the partitions that have not started emitting yet.
     uninitiated_partitions: Vec<usize>,
+
+    /// Which input streams are fully exhausted (their `poll_next` returned `None`)
+    stream_exhausted: Vec<bool>,
+
+    /// 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<usize>,
+
+    /// The not-yet-merged tail of the surviving stream's current batch,
+    /// emitted once before forwarding the stream itself.
+    passthrough_remainder: Option<RecordBatch>,
 }
 
 impl<C: CursorValues> SortPreservingMergeStream<C> {
@@ -186,6 +203,10 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             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<C: CursorValues> SortPreservingMergeStream<C> {
         }
 
         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<C: CursorValues> SortPreservingMergeStream<C> {
             }
             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<C: CursorValues> SortPreservingMergeStream<C> {
                     }
                 }
                 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,52 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
         }
     }
 
+    /// 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<Option<Result<RecordBatch>>> {
+        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. 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)) => {
+                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) {

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 = [
+            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 = [
+            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::<Vec<_>>(),
+                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::<Vec<_>>(),
+                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<dyn MemoryPool> =
+                        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::<Int32Type>().values().to_vec())
+                .collect::<Vec<_>>();
+
+            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())]);
+        }
+    }
+}