Also add grouping case

Author: zhuqi-lucas

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

@@ -77,9 +77,7 @@ 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)
+        // Wrap no‐grouping aggregates in our YieldStream
         let input = Box::pin(YieldStream::new(input)) as SendableRecordBatchStream;
 
         let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?;

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

@@ -49,6 +49,7 @@ use datafusion_physical_expr::{GroupsAccumulatorAdapter, PhysicalSortExpr};
 
 use super::order::GroupOrdering;
 use super::AggregateExec;
+use crate::aggregates::no_grouping::YieldStream;
 use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use futures::ready;
@@ -448,6 +449,10 @@ impl GroupedHashAggregateStream {
 
         let batch_size = context.session_config().batch_size();
         let input = agg.input.execute(partition, Arc::clone(&context))?;
+
+        // Wrap grouping aggregates in our YieldStream
+        let input = Box::pin(YieldStream::new(input)) as SendableRecordBatchStream;
+
         let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
 
         let timer = baseline_metrics.elapsed_compute().timer();