Make YieldStream public to allow static dispatch

Author: pepijnve

datafusion/physical-plan/src/aggregates/no_grouping.rs

@@ -22,7 +22,7 @@ use crate::aggregates::{
     AggregateMode,
 };
 use crate::metrics::{BaselineMetrics, RecordOutput};
-use crate::{RecordBatchStream, SendableRecordBatchStream};
+use crate::RecordBatchStream;
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
@@ -35,7 +35,7 @@ use std::task::{Context, Poll};
 
 use super::AggregateExec;
 use crate::filter::batch_filter;
-use crate::poll_budget::PollBudget;
+use crate::poll_budget::{PollBudget, YieldStream};
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
 use futures::stream::{Stream, StreamExt};
 
@@ -55,7 +55,7 @@ pub(crate) struct AggregateStream {
 struct AggregateStreamInner {
     schema: SchemaRef,
     mode: AggregateMode,
-    input: SendableRecordBatchStream,
+    input: YieldStream,
     baseline_metrics: BaselineMetrics,
     aggregate_expressions: Vec<Vec<Arc<dyn PhysicalExpr>>>,
     filter_expressions: Vec<Option<Arc<dyn PhysicalExpr>>>,

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> {

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,

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,

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
+                }
+            },
         }
     }
 }

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 {