rluvaton commented on code in PR #23702:
URL: https://github.com/apache/datafusion/pull/23702#discussion_r3610891950
##########
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();
+ }
Review Comment:
I wanted to move the poll stream inside the `advance_cursors` but got
lifetime issues so kept like this
--
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]
