feat: Allow cancelling of grouping operations which are CPU bound

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

@@ -77,6 +77,11 @@ impl AggregateStream {
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
         let input = agg.input.execute(partition, Arc::clone(&context))?;
 
+        // Only wrap no‐grouping aggregates in our YieldStream
+        // (grouped aggregates tend to produce small streams
+        //  and can rely on Tokio's own task‐yielding)
+        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 {
             AggregateMode::Partial
@@ -170,6 +175,65 @@ impl AggregateStream {
     }
 }
 
+/// A stream that yields batches of data, yielding control back to the executor
+pub struct YieldStream {
+    inner: SendableRecordBatchStream,
+    batches_processed: usize,
+    buffer: Option<Result<RecordBatch>>,
+}
+
+impl YieldStream {
+    pub fn new(inner: SendableRecordBatchStream) -> Self {
+        Self {
+            inner,
+            batches_processed: 0,
+            buffer: None,
+        }
+    }
+}
+
+// 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;
+
+        if let Some(batch) = this.buffer.take() {
+            return Poll::Ready(Some(batch));
+        }
+
+        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;
+                    // We need to buffer the batch when we return Poll::Pending,
+                    // so that we can return it on the next poll.
+                    // Otherwise, the next poll will miss the batch and return None.
+                    this.buffer = Some(Ok(batch));
+                    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>;