refactor: build and rewrite AggregateExec through an AggregateExecBuilder

datafusion/core/tests/execution/coop.rs

@@ -243,25 +243,22 @@ async fn agg_grouped_topk_yields(
     let group = binary(value_col.clone(), Divide, lit(1000000i64), &inf.schema())?;
 
     let aggr = Arc::new(
-        AggregateExec::try_new(
-            AggregateMode::Single,
-            PhysicalGroupBy::new(
+        AggregateExec::builder(AggregateMode::Single, inf.clone())
+            .with_group_by(PhysicalGroupBy::new(
                 vec![(group, "group".to_string())],
                 vec![],
                 vec![vec![false]],
                 false,
-            ),
-            vec![Arc::new(
+            ))
+            .with_aggr_exprs(vec![Arc::new(
                 AggregateExprBuilder::new(min_max::max_udaf(), vec![value_col.clone()])
                     .schema(inf.schema())
                     .alias("max")
                     .build()?,
-            )],
-            vec![None],
-            inf.clone(),
-            inf.schema(),
-        )?
-        .with_limit_options(Some(LimitOptions::new(100))),
+            )])
+            .with_input_schema(inf.schema())
+            .with_limit_options(LimitOptions::new(100))
+            .build()?,
     );
 
     query_yields(aggr, session_ctx.task_ctx()).await

datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs

@@ -251,16 +251,12 @@ fn aggregations_with_limit_combined() -> datafusion_common::Result<()> {
 
     let schema = partial_agg.schema();
     let final_agg = Arc::new(
-        AggregateExec::try_new(
-            AggregateMode::Final,
-            final_group_by,
-            aggr_expr,
-            vec![],
-            partial_agg,
-            schema,
-        )
-        .unwrap()
-        .with_limit_options(Some(LimitOptions::new(5))),
+        AggregateExec::builder(AggregateMode::Final, partial_agg)
+            .with_group_by(final_group_by)
+            .with_aggr_exprs(aggr_expr)
+            .with_input_schema(schema)
+            .with_limit_options(LimitOptions::new(5))
+            .build()?,
     );
     let plan: Arc<dyn ExecutionPlan> = final_agg;
     // should combine the Partial/Final AggregateExecs to a Single AggregateExec

datafusion/physical-optimizer/src/combine_partial_final_agg.rs

@@ -90,19 +90,17 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
                 } else {
                     AggregateMode::SinglePartitioned
                 };
-                AggregateExec::try_new(
-                    mode,
-                    input_agg_exec.group_expr().clone(),
-                    input_agg_exec.aggr_expr().to_vec(),
-                    input_agg_exec.filter_expr().to_vec(),
-                    Arc::clone(input_agg_exec.input()),
-                    input_agg_exec.input_schema(),
-                )
-                .map(|combined_agg| {
-                    combined_agg.with_limit_options(agg_exec.limit_options())
-                })
-                .ok()
-                .map(Arc::new)
+                // the partial aggregate, re-planned in the combined mode and
+                // carrying the limit that was pushed into the final aggregate
+                input_agg_exec
+                    .to_builder()
+                    .with_mode(mode)
+                    .with_limit_options(agg_exec.limit_options())
+                    .build()
+                    // the combined aggregate cannot execute the limit: leave
+                    // the two aggregates alone
+                    .ok()
+                    .map(Arc::new)
             } else {
                 None
             };

datafusion/physical-optimizer/src/limited_distinct_aggregation.rs

@@ -54,7 +54,12 @@ impl LimitedDistinctAggregation {
         }
 
         // We found what we want: clone, copy the limit down, and return modified node
-        let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit)));
+        let new_aggr = aggr
+            .to_builder()
+            .with_limit_options(LimitOptions::new(limit))
+            .build()
+            // the aggregate cannot execute the limit: leave the plan alone
+            .ok()?;
 
         Some(Arc::new(new_aggr))
     }

datafusion/physical-optimizer/src/topk_aggregation.rs

@@ -106,10 +106,12 @@ impl TopKAggregation {
         }
 
         // We found what we want: clone, copy the limit down, and return modified node
-        let new_aggr = AggregateExec::with_new_limit_options(
-            aggr,
-            Some(LimitOptions::new_with_order(limit, order_desc)),
-        );
+        let new_aggr = aggr
+            .to_builder()
+            .with_limit_options(LimitOptions::new_with_order(limit, order_desc))
+            .build()
+            // the aggregate cannot execute the limit: leave the plan alone
+            .ok()?;
         Some(Arc::new(new_aggr))
     }
 

datafusion/physical-optimizer/src/update_aggr_exprs.rs

@@ -125,7 +125,8 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder {
                     input.equivalence_properties(),
                 )?;
 
-                let aggr_exec = aggr_exec.with_new_aggr_exprs(aggr_exprs);
+                let aggr_exec =
+                    aggr_exec.to_builder().with_aggr_exprs(aggr_exprs).build()?;
 
                 Ok(Transformed::yes(Arc::new(aggr_exec) as _))
             } else {