Weijun-H commented on code in PR #23419:
URL: https://github.com/apache/datafusion/pull/23419#discussion_r3552793657
##########
datafusion/physical-plan/src/sorts/merge.rs:
##########
@@ -327,6 +298,59 @@ 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<()>> {
+ // 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 < self.uninitiated_partitions.len() {
+ let partition_idx = self.uninitiated_partitions[idx];
+ match self.maybe_poll_stream(cx, partition_idx) {
+ Poll::Ready(Err(e)) => {
+ return Poll::Ready(Err(e));
Review Comment:
This drops the `self.done = true` from the original error path, so an init
error no longer fuses the stream — a later poll re-enters init and re-polls
already-errored partitions. Add `self.done = true`; before returning the error.
Existing tests miss this because none error during init and poll again.
--
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]
