Make YieldStream public to allow static dispatch

pepijnve · pepijnve · commit 9d5ad8fe081a · 2025-06-06T15:18:05.000+02:00
diff --git a/datafusion/physical-plan/src/aggregates/no_grouping.rs b/datafusion/physical-plan/src/aggregates/no_grouping.rs
@@ -17,42 +17,29 @@
 
 //! Aggregate without grouping columns
 
+use super::AggregateExec;
 use crate::aggregates::{
     aggregate_expressions, create_accumulators, finalize_aggregation, AccumulatorItem,
     AggregateMode,
 };
+use crate::filter::batch_filter;
 use crate::metrics::{BaselineMetrics, RecordOutput};
+use crate::poll_budget::PollBudget;
 use crate::{RecordBatchStream, SendableRecordBatchStream};
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
+use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::PhysicalExpr;
-use futures::stream::BoxStream;
+use futures::stream::{Stream, StreamExt};
+use futures::FutureExt;
 use std::borrow::Cow;
 use std::sync::Arc;
-use std::task::{Context, Poll};
-
-use super::AggregateExec;
-use crate::filter::batch_filter;
-use crate::poll_budget::PollBudget;
-use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
-use futures::stream::{Stream, StreamExt};
+use std::task::{ready, Context, Poll};
 
 /// stream struct for aggregation without grouping columns
 pub(crate) struct AggregateStream {
-    stream: BoxStream<'static, Result<RecordBatch>>,
-    schema: SchemaRef,
-}
-
-/// Actual implementation of [`AggregateStream`].
-///
-/// This is wrapped into yet another struct because we need to interact with the async memory management subsystem
-/// during poll. To have as little code "weirdness" as possible, we chose to just use [`BoxStream`] together with
-/// [`futures::stream::unfold`].
-///
-/// The latter requires a state object, which is [`AggregateStreamInner`].
-struct AggregateStreamInner {
     schema: SchemaRef,
     mode: AggregateMode,
     input: SendableRecordBatchStream,
@@ -62,6 +49,7 @@ struct AggregateStreamInner {
     accumulators: Vec<AccumulatorItem>,
     reservation: MemoryReservation,
     finished: bool,
+    poll_budget: PollBudget,
 }
 
 impl AggregateStream {
@@ -71,7 +59,6 @@ impl AggregateStream {
         context: Arc<TaskContext>,
         partition: usize,
     ) -> Result<Self> {
-        let agg_schema = Arc::clone(&agg.schema);
         let agg_filter_expr = agg.filter_expr.clone();
 
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
@@ -91,81 +78,17 @@ impl AggregateStream {
         let reservation = MemoryConsumer::new(format!("AggregateStream[{partition}]"))
             .register(context.memory_pool());
 
-        let inner = AggregateStreamInner {
+        Ok(AggregateStream {
             schema: Arc::clone(&agg.schema),
             mode: agg.mode,
-            input: PollBudget::from(context.as_ref()).wrap_stream(input),
+            input,
             baseline_metrics,
             aggregate_expressions,
             filter_expressions,
             accumulators,
             reservation,
             finished: false,
-        };
-        let stream = futures::stream::unfold(inner, |mut this| async move {
-            if this.finished {
-                return None;
-            }
-
-            let elapsed_compute = this.baseline_metrics.elapsed_compute();
-
-            loop {
-                let result = match this.input.next().await {
-                    Some(Ok(batch)) => {
-                        let timer = elapsed_compute.timer();
-                        let result = aggregate_batch(
-                            &this.mode,
-                            batch,
-                            &mut this.accumulators,
-                            &this.aggregate_expressions,
-                            &this.filter_expressions,
-                        );
-
-                        timer.done();
-
-                        // allocate memory
-                        // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with
-                        // overshooting a bit. Also this means we either store the whole record batch or not.
-                        match result
-                            .and_then(|allocated| this.reservation.try_grow(allocated))
-                        {
-                            Ok(_) => continue,
-                            Err(e) => Err(e),
-                        }
-                    }
-                    Some(Err(e)) => Err(e),
-                    None => {
-                        this.finished = true;
-                        let timer = this.baseline_metrics.elapsed_compute().timer();
-                        let result =
-                            finalize_aggregation(&mut this.accumulators, &this.mode)
-                                .and_then(|columns| {
-                                    RecordBatch::try_new(
-                                        Arc::clone(&this.schema),
-                                        columns,
-                                    )
-                                    .map_err(Into::into)
-                                })
-                                .record_output(&this.baseline_metrics);
-
-                        timer.done();
-
-                        result
-                    }
-                };
-
-                this.finished = true;
-                return Some((result, this));
-            }
-        });
-
-        // seems like some consumers call this stream even after it returned `None`, so let's fuse the stream.
-        let stream = stream.fuse();
-        let stream = Box::pin(stream);
-
-        Ok(Self {
-            schema: agg_schema,
-            stream,
+            poll_budget: PollBudget::from(context.as_ref()),
         })
     }
 }
@@ -178,7 +101,61 @@ impl Stream for AggregateStream {
         cx: &mut Context<'_>,
     ) -> Poll<Option<Self::Item>> {
         let this = &mut *self;
-        this.stream.poll_next_unpin(cx)
+
+        if this.finished {
+            return Poll::Ready(None);
+        }
+
+        let elapsed_compute = this.baseline_metrics.elapsed_compute();
+
+        let mut consume_budget = this.poll_budget.consume_budget();
+
+        loop {
+            let result = match ready!(this.input.poll_next_unpin(cx)) {
+                Some(Ok(batch)) => {
+                    let timer = elapsed_compute.timer();
+                    let result = aggregate_batch(
+                        &this.mode,
+                        batch,
+                        &mut this.accumulators,
+                        &this.aggregate_expressions,
+                        &this.filter_expressions,
+                    );
+
+                    timer.done();
+
+                    // allocate memory
+                    // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with
+                    // overshooting a bit. Also this means we either store the whole record batch or not.
+                    match result
+                        .and_then(|allocated| this.reservation.try_grow(allocated))
+                    {
+                        Ok(_) => {
+                            ready!(consume_budget.poll_unpin(cx));
+                            continue;
+                        }
+                        Err(e) => Err(e),
+                    }
+                }
+                Some(Err(e)) => Err(e),
+                None => {
+                    let timer = this.baseline_metrics.elapsed_compute().timer();
+                    let result = finalize_aggregation(&mut this.accumulators, &this.mode)
+                        .and_then(|columns| {
+                            RecordBatch::try_new(Arc::clone(&this.schema), columns)
+                                .map_err(Into::into)
+                        })
+                        .record_output(&this.baseline_metrics);
+
+                    timer.done();
+
+                    result
+                }
+            };
+
+            this.finished = true;
+            return Poll::Ready(Some(result));
+        }
     }
 }
 
diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs
@@ -46,7 +46,7 @@ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::equivalence::join_equivalence_properties;
 
-use crate::poll_budget::PollBudget;
+use crate::poll_budget::{PollBudget, YieldStream};
 use async_trait::async_trait;
 use futures::{ready, Stream, StreamExt, TryStreamExt};
 
@@ -189,7 +189,7 @@ impl CrossJoinExec {
 
 /// Asynchronously collect the result of the left child
 async fn load_left_input(
-    stream: SendableRecordBatchStream,
+    stream: YieldStream,
     metrics: BuildProbeJoinMetrics,
     reservation: MemoryReservation,
 ) -> Result<JoinLeftData> {
diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -81,7 +81,7 @@ use datafusion_physical_expr::equivalence::{
 use datafusion_physical_expr::PhysicalExprRef;
 use datafusion_physical_expr_common::datum::compare_op_for_nested;
 
-use crate::poll_budget::PollBudget;
+use crate::poll_budget::{PollBudget, YieldStream};
 use ahash::RandomState;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
 use futures::{ready, Stream, StreamExt, TryStreamExt};
@@ -953,7 +953,7 @@ impl ExecutionPlan for HashJoinExec {
 /// hash table (`LeftJoinData`)
 async fn collect_left_input(
     random_state: RandomState,
-    left_stream: SendableRecordBatchStream,
+    left_stream: YieldStream,
     on_left: Vec<PhysicalExprRef>,
     metrics: BuildProbeJoinMetrics,
     reservation: MemoryReservation,
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -61,7 +61,7 @@ use datafusion_physical_expr::equivalence::{
     join_equivalence_properties, ProjectionMapping,
 };
 
-use crate::poll_budget::PollBudget;
+use crate::poll_budget::{PollBudget, YieldStream};
 use futures::{ready, Stream, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 
@@ -626,7 +626,7 @@ impl ExecutionPlan for NestedLoopJoinExec {
 
 /// Asynchronously collect input into a single batch, and creates `JoinLeftData` from it
 async fn collect_left_input(
-    stream: SendableRecordBatchStream,
+    stream: YieldStream,
     join_metrics: BuildProbeJoinMetrics,
     reservation: MemoryReservation,
     with_visited_left_side: bool,
diff --git a/datafusion/physical-plan/src/poll_budget.rs b/datafusion/physical-plan/src/poll_budget.rs
@@ -54,16 +54,8 @@ impl PollBudget {
         }
     }
 
-    pub fn wrap_stream(
-        &self,
-        inner: SendableRecordBatchStream,
-    ) -> SendableRecordBatchStream {
-        match self.budget {
-            None => inner,
-            Some(budget) => {
-                Box::pin(YieldStream::new(inner, budget)) as SendableRecordBatchStream
-            }
-        }
+    pub fn wrap_stream(&self, inner: SendableRecordBatchStream) -> YieldStream {
+        YieldStream::new(inner, self.budget)
     }
 }
 
@@ -102,18 +94,18 @@ impl Future for ConsumeBudget {
     }
 }
 
-struct YieldStream {
+pub struct YieldStream {
     inner: SendableRecordBatchStream,
-    budget: u8,
-    remaining: u8,
+    budget: Option<u8>,
+    remaining: Option<u8>,
 }
 
 impl YieldStream {
-    pub fn new(inner: SendableRecordBatchStream, budget: u8) -> Self {
+    pub fn new(inner: SendableRecordBatchStream, budget: Option<u8>) -> Self {
         Self {
             inner,
             budget,
-            remaining: 0,
+            remaining: budget,
         }
     }
 }
@@ -125,21 +117,23 @@ impl Stream for YieldStream {
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Self::Item>> {
-        if self.remaining == 0 {
-            self.remaining = self.budget;
-            cx.waker().wake_by_ref();
-            return Pending;
-        }
-
-        match self.inner.poll_next_unpin(cx) {
-            ready @ Ready(Some(_)) => {
-                self.remaining -= 1;
-                ready
-            }
-            other => {
+        match self.remaining {
+            None => self.inner.poll_next_unpin(cx),
+            Some(0) => {
                 self.remaining = self.budget;
-                other
+                cx.waker().wake_by_ref();
+                Pending
             }
+            Some(remaining) => match self.inner.poll_next_unpin(cx) {
+                ready @ Ready(Some(_)) => {
+                    self.remaining = Some(remaining - 1);
+                    ready
+                }
+                other => {
+                    self.remaining = self.budget;
+                    other
+                }
+            },
         }
     }
 }
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
@@ -19,11 +19,6 @@
 //! It will do in-memory sorting if it has enough memory budget
 //! but spills to disk if needed.
 
-use std::any::Any;
-use std::fmt;
-use std::fmt::{Debug, Formatter};
-use std::sync::Arc;
-
 use crate::common::spawn_buffered;
 use crate::execution_plan::{Boundedness, CardinalityEffect, EmissionType};
 use crate::expressions::PhysicalSortExpr;
@@ -43,6 +38,10 @@ use crate::{
     ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream,
     Statistics,
 };
+use std::any::Any;
+use std::fmt;
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
 
 use arrow::array::{Array, RecordBatch, RecordBatchOptions, StringViewArray};
 use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays};
@@ -1107,10 +1106,9 @@ impl ExecutionPlan for SortExec {
             .equivalence_properties()
             .ordering_satisfy_requirement(requirement);
 
-        let mut input = PollBudget::from(context.as_ref()).wrap_stream(input);
-
         match (sort_satisfied, self.fetch.as_ref()) {
             (true, Some(fetch)) => Ok(Box::pin(LimitStream::new(
+                // limit is not a pipeline breaking stream, so poll budget is not required
                 input,
                 0,
                 Some(*fetch),
@@ -1128,6 +1126,7 @@ impl ExecutionPlan for SortExec {
                     context.runtime_env(),
                     &self.metrics_set,
                 )?;
+                let mut input = PollBudget::from(context.as_ref()).wrap_stream(input);
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
                     futures::stream::once(async move {
@@ -1154,6 +1153,7 @@ impl ExecutionPlan for SortExec {
                     &self.metrics_set,
                     context.runtime_env(),
                 )?;
+                let mut input = PollBudget::from(context.as_ref()).wrap_stream(input);
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
                     futures::stream::once(async move {

Original file line number	Diff line number	Diff line change
`@@ -54,16 +54,8 @@ impl PollBudget {`
`54`	`54`	`}`
`55`	`55`	`}`
`56`	`56`
`57`		`- pub fn wrap_stream(`
`58`		`- &self,`
`59`		`- inner: SendableRecordBatchStream,`
`60`		`- ) -> SendableRecordBatchStream {`
`61`		`- match self.budget {`
`62`		`- None => inner,`
`63`		`- Some(budget) => {`
`64`		`- Box::pin(YieldStream::new(inner, budget)) as SendableRecordBatchStream`
`65`		`- }`
`66`		`- }`
	`57`	`+ pub fn wrap_stream(&self, inner: SendableRecordBatchStream) -> YieldStream {`
	`58`	`+ YieldStream::new(inner, self.budget)`
`67`	`59`	`}`
`68`	`60`	`}`
`69`	`61`
`@@ -102,18 +94,18 @@ impl Future for ConsumeBudget {`
`102`	`94`	`}`
`103`	`95`	`}`
`104`	`96`
`105`		`-struct YieldStream {`
	`97`	`+pub struct YieldStream {`
`106`	`98`	`inner: SendableRecordBatchStream,`
`107`		`- budget: u8,`
`108`		`- remaining: u8,`
	`99`	`+ budget: Option<u8>,`
	`100`	`+ remaining: Option<u8>,`
`109`	`101`	`}`
`110`	`102`
`111`	`103`	`impl YieldStream {`
`112`		`- pub fn new(inner: SendableRecordBatchStream, budget: u8) -> Self {`
	`104`	`+ pub fn new(inner: SendableRecordBatchStream, budget: Option<u8>) -> Self {`
`113`	`105`	`Self {`
`114`	`106`	`inner,`
`115`	`107`	`budget,`
`116`		`- remaining: 0,`
	`108`	`+ remaining: budget,`
`117`	`109`	`}`
`118`	`110`	`}`
`119`	`111`	`}`
`@@ -125,21 +117,23 @@ impl Stream for YieldStream {`
`125`	`117`	`mut self: Pin<&mut Self>,`
`126`	`118`	`cx: &mut Context<'_>,`
`127`	`119`	`) -> Poll<Option<Self::Item>> {`
`128`		`- if self.remaining == 0 {`
`129`		`- self.remaining = self.budget;`
`130`		`- cx.waker().wake_by_ref();`
`131`		`- return Pending;`
`132`		`- }`
`133`		`-`
`134`		`- match self.inner.poll_next_unpin(cx) {`
`135`		`- ready @ Ready(Some(_)) => {`
`136`		`- self.remaining -= 1;`
`137`		`- ready`
`138`		`- }`
`139`		`- other => {`
	`120`	`+ match self.remaining {`
	`121`	`+ None => self.inner.poll_next_unpin(cx),`
	`122`	`+ Some(0) => {`
`140`	`123`	`self.remaining = self.budget;`
`141`		`- other`
	`124`	`+ cx.waker().wake_by_ref();`
	`125`	`+ Pending`
`142`	`126`	`}`
	`127`	`+ Some(remaining) => match self.inner.poll_next_unpin(cx) {`
	`128`	`+ ready @ Ready(Some(_)) => {`
	`129`	`+ self.remaining = Some(remaining - 1);`
	`130`	`+ ready`
	`131`	`+ }`
	`132`	`+ other => {`
	`133`	`+ self.remaining = self.budget;`
	`134`	`+ other`
	`135`	`+ }`
	`136`	`+ },`
`143`	`137`	`}`
`144`	`138`	`}`
`145`	`139`	`}`