Update merge.rs

apache · Sep 2, 2024 · b6aac77 · b6aac77
1 parent b81b99c
commit b6aac77
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs
@@ -171,6 +171,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
             // Ensure all non-exhausted streams have a cursor from which
             // rows can be pulled
             loop {
+                let mut return_pending = false;
                 let uninitiated_cursors = self.uninitiated_cursors.clone();
                 for i in uninitiated_cursors {
                     if self.uninitiated_cursors.contains(&i) {
@@ -179,13 +180,17 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
                                 self.aborted = true;
                                 return Poll::Ready(Some(Err(e)));
                             }
+                            Poll::Pending => {
+                                return_pending = true;
+                            }
                             _ => {}
                         }
                     }
                 }
                 if self.uninitiated_cursors.is_empty() {
                     break;
                 }
+                return Poll::Pending;
             }
             self.init_loser_tree();
         }