This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 4f63ada240 chore(spm): extract initialize all parititions helper
(#23419)
4f63ada240 is described below
commit 4f63ada24056d9ec9b8decc04f128512bcdd735e
Author: Raz Luvaton <[email protected]>
AuthorDate: Thu Jul 9 19:21:28 2026 +0300
chore(spm): extract initialize all parititions helper (#23419)
## Which issue does this PR close?
N/A
## Rationale for this change
Keeping the core loop tight and readable
## What changes are included in this PR?
Just extracting a function
## Are these changes tested?
existing tests
## Are there any user-facing changes?
Nope
---
datafusion/physical-plan/src/sorts/merge.rs | 107 +++++++++++++++++-----------
1 file changed, 64 insertions(+), 43 deletions(-)
diff --git a/datafusion/physical-plan/src/sorts/merge.rs
b/datafusion/physical-plan/src/sorts/merge.rs
index 4583d19e91..4117789777 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -236,53 +236,24 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
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.
-
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);
- }
- }
- }
-
- if self.uninitiated_partitions.is_empty() {
- // 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
- 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;
- }
+ ready!(self.initialize_all_partitions(cx))?;
+ assert_eq!(
+ self.uninitiated_partitions.len(),
+ 0,
+ "all partitions should be initialized"
+ );
+
+ // 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
+ self.uninitiated_partitions.shrink_to_fit();
+ self.init_loser_tree();
}
// NB timer records time taken on drop, so there are no
@@ -327,6 +298,56 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
}
+ /// 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`
+ /// so we can continue to initialize the remaining partitions
+ fn initialize_all_partitions(&mut self, cx: &mut Context) ->
Poll<Result<()>> {
+ 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 < 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(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);
+ }
+ }
+ }
+
+ if self.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
+ }
+ }
+
/// 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) {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
