zhuqi-lucas commented on code in PR #23702:
URL: https://github.com/apache/datafusion/pull/23702#discussion_r3645394911
##########
datafusion/physical-plan/src/sorts/streaming_merge.rs:
##########
@@ -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)));
Review Comment:
This is actually a **behavior fix**, not just a chore — I probed both sides
with `SortPreservingMergeExec::with_fetch(Some(0))`:
| input | main | this PR |
|---|---|---|
| 1 partition | 0 rows ✓ (passthrough + `LimitStream` fast path) | 0 rows ✓ |
| **2 partitions** | **1 row emitted** ❌ (merge loop pushes before checking
`fetch_reached`) | 0 rows ✓ |
Two asks:
1. Mention it under "Are there any user-facing changes?" (currently says No).
2. Add a regression test — **it must use ≥ 2 partitions**: a
single-partition test takes the `LimitStream` fast path and passes even without
this fix. Ready-made version (verified failing on main, passing on this
branch), drops into `sort_preserving_merge.rs` tests:
```rust
#[tokio::test]
async fn test_sort_merge_fetch_zero() {
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 = [PhysicalSortExpr {
expr: col("b", &schema).unwrap(),
options: SortOptions {
descending: false,
nulls_first: true,
},
}]
.into();
// Two partitions on purpose: a single-partition input takes the
// passthrough + LimitStream fast path and does not exercise the
// merge loop, so it passes even without the fetch-0 guard.
let batch2 = batch.clone();
let exec =
TestMemoryExec::try_new_exec(&[vec![batch], vec![batch2]], schema,
None)
.unwrap();
let merge =
Arc::new(SortPreservingMergeExec::new(sort,
exec).with_fetch(Some(0)));
let collected = collect(merge, task_ctx).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}");
}
```
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -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 {
+ assert!(
+ self.fetch.is_none_or(|fetch| fetch != 0),
Review Comment:
Third check of the same invariant (builder early-returns at
`streaming_merge.rs:200`, constructor has `assert_ne!` at line 150) — after the
builder guard neither assert can fire. Suggest keeping just the constructor one
and dropping this.
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -369,18 +414,21 @@ 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() {
+ 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
Review Comment:
`return if … { true } else { false }` can collapse to the condition; also
worth noting the trailing `true` is unreachable from the (only) call site now
that the loop is guarded by `is_exhausted`:
```suggestion
fn advance_cursors(&mut self, stream_idx: usize) -> bool {
if let Some(cursor) = &mut self.cursors[stream_idx] {
let _ = cursor.advance();
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();
}
return finished;
}
// Unreachable from the merge loop (guarded by `is_exhausted`),
// kept for defensive completeness: no cursor ⇒ a poll is needed.
true
}
```
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -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<()> {
Review Comment:
Minor: the leftover-drain is now inside `elapsed_compute` (before this PR
the final drain ran after the timer was dropped) — a small metrics behavior
change, arguably more correct, but worth a line in the PR description.
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -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 {
+ 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::<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 {
+ 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()
+ }
Review Comment:
Checking only the tree root is correct, but the reasoning is non-obvious —
worth writing down, especially in a readability PR:
```suggestion
/// Returns `true` once every input stream is exhausted.
///
/// Checking only the tree root suffices: `is_gt` treats a `None`
/// cursor as greater than everything, so a stream with no cursor
/// loses every match — it can only reach the root (the tournament
/// winner, i.e. the minimum) when no stream has a cursor left.
fn is_exhausted(&self) -> bool {
let winner = self.loser_tree[0];
self.cursors[winner].is_none()
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
