feat: support inability to yeild cpu for loop when it's not using Tokio MPSC (RecordBatchReceiverStream)

Author: zhuqi-lucas

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

@@ -47,6 +47,7 @@ use datafusion_physical_expr::{
     PhysicalSortRequirement,
 };
 
+use crate::aggregates::no_grouping::YieldStream;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
 use itertools::Itertools;
 
@@ -983,8 +984,9 @@ impl ExecutionPlan for AggregateExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        self.execute_typed(partition, context)
-            .map(|stream| stream.into())
+        let raw_stream = self.execute_typed(partition, context)?.into();
+        let wrapped = Box::pin(YieldStream::new(raw_stream));
+        Ok(wrapped)
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

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

@@ -76,6 +76,7 @@ impl AggregateStream {
 
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
         let input = agg.input.execute(partition, Arc::clone(&context))?;
+        let input = Box::pin(YieldStream::new(input)) as SendableRecordBatchStream;
 
         let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?;
         let filter_expressions = match agg.mode {
@@ -170,6 +171,55 @@ impl AggregateStream {
     }
 }
 
+/// A stream that yields batches of data, yielding control back to the executor
+pub struct YieldStream {
+    inner: SendableRecordBatchStream,
+    batches_processed: usize,
+}
+
+impl YieldStream {
+    pub fn new(inner: SendableRecordBatchStream) -> Self {
+        Self {
+            inner,
+            batches_processed: 0,
+        }
+    }
+}
+
+// Stream<Item = Result<RecordBatch>> to poll_next_unpin
+impl Stream for YieldStream {
+    type Item = Result<RecordBatch>;
+
+    fn poll_next(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::Item>> {
+        const YIELD_BATCHES: usize = 64;
+        let this = &mut *self;
+
+        match this.inner.poll_next_unpin(cx) {
+            Poll::Ready(Some(Ok(batch))) => {
+                this.batches_processed += 1;
+                if this.batches_processed >= YIELD_BATCHES {
+                    this.batches_processed = 0;
+                    cx.waker().wake_by_ref();
+                    Poll::Pending
+                } else {
+                    Poll::Ready(Some(Ok(batch)))
+                }
+            }
+            other => other,
+        }
+    }
+}
+
+// RecordBatchStream schema()
+impl RecordBatchStream for YieldStream {
+    fn schema(&self) -> Arc<arrow_schema::Schema> {
+        self.inner.schema()
+    }
+}
+
 impl Stream for AggregateStream {
     type Item = Result<RecordBatch>;