martin-g commented on code in PR #19002:
URL: https://github.com/apache/datafusion/pull/19002#discussion_r2585124001
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1540,7 +1550,46 @@ impl Stream for PerPartitionStream {
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
- let poll = self.poll_next_inner(cx);
+ let poll = match self.batch_coalescer.take() {
+ Some(mut coalescer) => {
+ let cloned_time =
self.baseline_metrics.elapsed_compute().clone();
+ let mut completed = false;
+ let poll;
+ loop {
+ if let Some(batch) = coalescer.next_completed_batch() {
+ poll = Poll::Ready(Some(Ok(batch)));
+ break;
+ }
+ if completed {
+ poll = Poll::Ready(None);
+ break;
+ }
+ let inner_poll = self.poll_next_inner(cx);
+ let _timer = cloned_time.timer();
+
+ match inner_poll {
+ Poll::Pending => {
+ poll = Poll::Pending;
+ break;
+ }
+ Poll::Ready(None) => {
+ completed = true;
+ coalescer.finish()?;
Review Comment:
> the loop would continue
The '?' symbol at the end means: `call finish() and if it returns Ok
continue with the execution of the next line, but if it returns Err then return
it, i.e. exit the current method`.
I.e. it won't stay in the loop
--
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]
