You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/26 19:55:49 UTC
(datafusion) branch gh-readonly-queue/main/pr-23702-7576762766d0f8081d7cd487e92be4861cf0c485 created (now 77b172e150)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23702-7576762766d0f8081d7cd487e92be4861cf0c485
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 77b172e150 chore: simplify SortPreservingMergeStream to be as textbook-like as possible (#23702)
This branch includes the following new commits:
new 77b172e150 chore: simplify SortPreservingMergeStream to be as textbook-like as possible (#23702)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: chore: simplify SortPreservingMergeStream to be as textbook-like as possible (#23702)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23702-7576762766d0f8081d7cd487e92be4861cf0c485
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 77b172e1505283f9ac6af050ee8547f910c28262
Author: Raz Luvaton <16...@users.noreply.github.com>
AuthorDate: Sun Jul 26 22:55:27 2026 +0300
chore: simplify SortPreservingMergeStream to be as textbook-like as possible (#23702)
## Which issue does this PR close?
N/A
## Rationale for this change
`SortPreservingMergeStream` is a little complex, so add some guiding
comments and make it as textbook-like as possible
## What changes are included in this PR?
Added comments, reorder code
While this was done this also fixed couple of bugs due to how it work:
1. leftover drain was not counted in the `elapsed_compute`
2. `limit(0)` returns 0 rows and not 1
## Are these changes tested?
Existing tests
## Are there any user-facing changes?
Not API ones.
`limit(0)` now returns 0 rows
---
datafusion/physical-plan/src/sorts/merge.rs | 153 ++++++++++++++-------
.../physical-plan/src/sorts/streaming_merge.rs | 127 ++++++++++++++++-
2 files changed, 220 insertions(+), 60 deletions(-)
diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
index 986da549f7..310416c22d 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, assert_or_internal_err, internal_err};
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<C: CursorValues> SortPreservingMergeStream<C> {
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<C: CursorValues> SortPreservingMergeStream<C> {
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<C: CursorValues> SortPreservingMergeStream<C> {
result
}
- 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<_>>();
+ async fn flush_in_progress(
+ &mut self,
+ mut emitter: TryEmitter<RecordBatch, DataFusionError>,
+ ) -> 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<Item = Result<RecordBatch>> {
+ async_try_stream(|mut emitter| async move {
+ // 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::<Vec<_>>();
+
+ 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 {
+ return internal_err!("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_or_internal_err!(
+ 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(())
})
}
+ /// Returns `true` once every input stream is exhausted.
+ ///
+ /// Should only be called for valid adjusted tree, i.e. the initial tree or after [`Self::update_loser_tree`] call
+ fn is_exhausted(&self) -> bool {
+ let winner = self.loser_tree[0];
+
+ // Checking only the tree root suffices for valid tree
+ // since the winner of the tree cannot be an exhausted stream for a valid tree
+ // as what value is winning over the non exhausted stream?
+ 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,20 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
/// 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() {
+ let finished = cursor.is_finished();
+ if finished {
// Take the current cursor, leaving `None` in its place
self.prev_cursors[stream_idx] = self.cursors[stream_idx].take();
}
- true
- } else {
- false
+ return finished;
}
+
+ // the entire stream is exhausted, so return true (poll won't help here anyway)
+ 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 e96138ef13..81adad8e9e 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");
@@ -267,3 +273,110 @@ impl<'a> StreamingMergeBuilder<'a> {
.into_stream())
}
}
+
+#[cfg(test)]
+mod tests {
+ use crate::{common::collect, stream::RecordBatchStreamAdapter};
+ use std::sync::Arc;
+
+ use super::*;
+
+ use arrow::array::{ArrayRef, RecordBatch};
+ use arrow_schema::SortOptions;
+ use datafusion_common::Result;
+ use datafusion_execution::TaskContext;
+ use datafusion_physical_expr::{PhysicalSortExpr, expressions::col};
+ use datafusion_physical_expr_common::metrics::{
+ ExecutionPlanMetricsSet, SpillMetrics,
+ };
+
+ #[tokio::test]
+ async fn test_sort_merge_fetch_zero_with_only_1_stream() {
+ test_fetch_0_should_output_0_rows(1, 0).await.unwrap();
+ }
+ #[tokio::test]
+ async fn test_sort_merge_fetch_zero_with_2_streams() {
+ test_fetch_0_should_output_0_rows(2, 0).await.unwrap();
+ }
+ #[tokio::test]
+ async fn test_sort_merge_fetch_zero_with_only_1_spill_file() {
+ test_fetch_0_should_output_0_rows(0, 1).await.unwrap();
+ }
+ #[tokio::test]
+ async fn test_sort_merge_fetch_zero_with_2_spill_files() {
+ test_fetch_0_should_output_0_rows(0, 2).await.unwrap();
+ }
+ #[tokio::test]
+ async fn test_sort_merge_fetch_zero_with_1_stream_and_1_spill_file() {
+ test_fetch_0_should_output_0_rows(1, 1).await.unwrap();
+ }
+
+ async fn test_fetch_0_should_output_0_rows(
+ number_of_streams: usize,
+ number_of_spilled_files: usize,
+ ) -> Result<()> {
+ let task_ctx = Arc::new(TaskContext::default());
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3]));
+ let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"]));
+ let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap();
+ let schema = batch.schema();
+
+ let sort: LexOrdering = [PhysicalSortExpr {
+ expr: col("b", &schema).unwrap(),
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ }]
+ .into();
+
+ let streams = (0..number_of_streams)
+ .map(|_| {
+ Box::pin(RecordBatchStreamAdapter::new(
+ Arc::clone(&schema),
+ futures::stream::iter(vec![Ok(batch.clone())]),
+ )) as SendableRecordBatchStream
+ })
+ .collect::<Vec<SendableRecordBatchStream>>();
+
+ let spill_manager = SpillManager::new(
+ task_ctx.runtime_env(),
+ SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ Arc::clone(&schema),
+ );
+
+ let mut sorted_spill_files: Vec<SortedSpillFile> = vec![];
+
+ for _ in 0..number_of_spilled_files {
+ let file = spill_manager
+ .spill_record_batch_and_finish(std::slice::from_ref(&batch), "spill")
+ .unwrap()
+ .unwrap();
+ sorted_spill_files.push(SortedSpillFile {
+ file,
+ max_record_batch_memory: batch.get_array_memory_size(),
+ });
+ }
+
+ let sorted_output_stream = StreamingMergeBuilder::new()
+ .with_batch_size(100)
+ .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0))
+ // Just to avoid having to provide memory pool
+ .with_bypass_mempool()
+ .with_schema(schema)
+ .with_streams(streams)
+ .with_sorted_spill_files(sorted_spill_files)
+ .with_spill_manager(spill_manager)
+ .with_expressions(&sort)
+ // The whole point of the test - fetch is 0
+ .with_fetch(Some(0))
+ .build()
+ .unwrap();
+
+ let collected = collect(sorted_output_stream).await.unwrap();
+ let total: usize = collected.iter().map(|b| b.num_rows()).sum();
+ assert_eq!(total, 0, "fetch=Some(0) must emit zero rows, got {total}");
+
+ Ok(())
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org