rluvaton commented on code in PR #23407:
URL: https://github.com/apache/datafusion/pull/23407#discussion_r3552457316
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -212,118 +226,137 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
}
+ async fn poll_stream_and_wait(&mut self, idx: usize) -> Result<()> {
+ poll_fn(|cx| self.maybe_poll_stream(cx, idx)).await
+ }
+
fn emit_in_progress_batch(&mut self) -> Result<Option<RecordBatch>> {
let rows_before = self.in_progress.len();
let result = self.in_progress.build_record_batch();
self.produced += rows_before - self.in_progress.len();
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);
- }
- // 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.
+ async fn run(&mut self, co: &Co<Result<RecordBatch>>) -> Result<()> {
+ // This vector contains the indices of the partitions that have not
started emitting yet.
+ let mut uninitiated_partitions =
+ (0..self.streams.partitions()).collect::<Vec<_>>();
- if self.loser_tree.is_empty() {
- // 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];
- match self.maybe_poll_stream(cx, partition_idx) {
- Poll::Ready(Err(e)) => {
- self.done = true;
- return Poll::Ready(Some(Err(e)));
- }
- Poll::Pending => {
- // The polled stream is pending which means we're
already set up to
- // be woken when necessary
- // Try the next stream
- idx += 1;
- }
- _ => {
- // The polled stream is ready
- // Remove it from uninitiated_partitions
- // Don't bump idx here, since a new element will have
taken its
- // 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);
- }
- }
- }
+ poll_fn(|cx| self.initialize_all_partitions(&mut
uninitiated_partitions, cx))
+ .await?;
- if self.uninitiated_partitions.is_empty() {
- // If there are no more uninitiated partitions, set up the
loser tree and continue
- // to the next phase.
+ assert_eq!(uninitiated_partitions.len(), 0);
- // Claim the memory for the uninitiated partitions
- self.uninitiated_partitions.shrink_to_fit();
- self.init_loser_tree();
- } else {
- // There are still uninitiated partitions so return pending.
- // We only get here if we've polled all uninitiated streams
and at least one of them
- // returned pending itself. That means we will be woken as
soon as one of the
- // streams would like to be polled again.
- // There is no need to reschedule ourselves eagerly.
- return Poll::Pending;
- }
- }
+ // 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
+ 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();
+ let mut 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)));
- }
- }
+ 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 {
+ if let Some(batch) = self.emit_in_progress_batch()? {
+ drop(timer);
+ co.yield_(Ok(batch)).await;
+ timer = elapsed_compute.timer();
}
- 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);
-
- // 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;
+ // Adjust the loser tree if necessary, returning control if needed
+ 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);
+ self.poll_stream_and_wait(winner).await?;
+ timer = elapsed_compute.timer();
+ }
+ self.update_loser_tree();
+ }
+
+ drop(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()? {
+ co.yield_(Ok(batch)).await;
+ }
+ Ok(())
+ }
+
+ fn initialize_all_partitions(
+ &mut self,
+ uninitiated_partitions: &mut Vec<usize>,
+ cx: &mut Context,
+ ) -> Poll<Result<()>> {
+ // 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.
+
+ assert_eq!(
+ self.loser_tree.len(),
+ 0,
+ "loser tree must be empty when initializing"
+ );
+
+ // Manual indexing since we're iterating over the vector and shrinking
it in the loop
+ let mut idx = 0;
+ while idx < uninitiated_partitions.len() {
+ let partition_idx = uninitiated_partitions[idx];
+ match self.maybe_poll_stream(cx, partition_idx) {
+ Poll::Ready(Err(e)) => {
+ return Poll::Ready(Err(e));
+ }
+ Poll::Pending => {
+ // The polled stream is pending which means we're already
set up to
+ // be woken when necessary
+ // Try the next stream
+ idx += 1;
+ }
+ _ => {
+ // The polled stream is ready
+ // Remove it from uninitiated_partitions
+ // Don't bump idx here, since a new element will have
taken its
+ // 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.
+ uninitiated_partitions.swap_remove(idx);
}
}
+ }
- return Poll::Ready(self.emit_in_progress_batch().transpose());
+ if uninitiated_partitions.is_empty() {
+ Poll::Ready(Ok(()))
+ } else {
+ // There are still uninitiated partitions so return pending.
+ // We only get here if we've polled all uninitiated streams and at
least one of them
+ // returned pending itself. That means we will be woken as soon as
one of the
+ // streams would like to be polled again.
+ // There is no need to reschedule ourselves eagerly.
+ Poll::Pending
}
}
Review Comment:
Extracted this to:
- #23419
--
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]
