More Arc to AggregateFunctionExpr

datafusion/core/benches/physical_plan.rs

@@ -24,18 +24,22 @@ extern crate datafusion;
 use std::sync::Arc;
 
 use arrow::{
-    array::{ArrayRef, Int64Array, StringArray},
+    array::{ArrayRef, Int32Array, Int64Array, StringArray},
+    datatypes::{Field, Schema, SchemaRef},
     record_batch::RecordBatch,
 };
 use tokio::runtime::Runtime;
 
-use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::{
     collect,
     expressions::{col, PhysicalSortExpr},
     memory::MemoryExec,
 };
 use datafusion::prelude::SessionContext;
+use datafusion::{
+    physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec,
+    physical_planner::DefaultPhysicalPlanner, physical_planner::PhysicalPlanner,
+};
 
 // Initialise the operator using the provided record batches and the sort key
 // as inputs. All record batches must have the same schema.
@@ -129,7 +133,7 @@ fn batches(
     rbs
 }
 
-fn criterion_benchmark(c: &mut Criterion) {
+fn bench_sort_preserving_merge_operator(c: &mut Criterion) {
     let small_batch = batches(1, 100, 10, 0).remove(0);
     let large_batch = batches(1, 1000, 1, 0).remove(0);
 
@@ -183,5 +187,67 @@ fn criterion_benchmark(c: &mut Criterion) {
     }
 }
 
-criterion_group!(benches, criterion_benchmark);
+/// Aimed at tracking inefficiencies at the stage of creating/optimizing a physical plan.
+fn bench_creation_many_columns(c: &mut Criterion) {
+    const COLUMNS_NUM: usize = 500;
+
+    let ctx = SessionContext::new();
+    let mut fields = vec![];
+    let mut columns = vec![];
+    for i in 0..COLUMNS_NUM {
+        fields.push(Field::new(
+            format!("attribute{}", i),
+            datafusion::arrow::datatypes::DataType::Int32,
+            true,
+        ));
+        columns.push(Int32Array::from(vec![1]));
+    }
+
+    // A fictive batch.
+    let batch = RecordBatch::try_new(
+        SchemaRef::new(Schema::new(fields)),
+        columns
+            .into_iter()
+            .map(|it| Arc::new(it) as Arc<dyn arrow_array::Array>)
+            .collect(),
+    )
+    .unwrap();
+
+    ctx.register_batch("t", batch).unwrap();
+
+    let mut aggregates = String::new();
+    for i in 0..COLUMNS_NUM {
+        if i > 0 {
+            aggregates.push_str(", ");
+        }
+        aggregates.push_str(format!("MAX(attribute{})", i).as_str());
+    }
+
+    // SELECT max(attr0), ..., max(attrN) FROM t.
+    let query = format!("SELECT {} FROM t", aggregates);
+    let rt = Runtime::new().unwrap();
+    let logical_plan = Arc::new(rt.block_on(async {
+        let statement = ctx.state().sql_to_statement(&query, "Generic").unwrap();
+        ctx.state().statement_to_plan(statement).await.unwrap()
+    }));
+
+    c.bench_function("phys_plan_creation_many_columns", |b| {
+        let rt = Runtime::new().unwrap();
+        let planner = DefaultPhysicalPlanner::default();
+
+        b.iter(|| {
+            let plan = Arc::clone(&logical_plan);
+            rt.block_on(async {
+                planner.create_physical_plan(&plan, &ctx.state()).await
+            })
+            .unwrap();
+        });
+    });
+}
+
+criterion_group!(
+    benches,
+    bench_sort_preserving_merge_operator,
+    bench_creation_many_columns
+);
 criterion_main!(benches);

datafusion/core/src/physical_optimizer/update_aggr_exprs.rs

@@ -131,10 +131,10 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder {
 /// successfully. Any errors occurring during the conversion process are
 /// passed through.
 fn try_convert_aggregate_if_better(
-    aggr_exprs: Vec<AggregateFunctionExpr>,
+    aggr_exprs: Vec<Arc<AggregateFunctionExpr>>,
     prefix_requirement: &[PhysicalSortRequirement],
     eq_properties: &EquivalenceProperties,
-) -> Result<Vec<AggregateFunctionExpr>> {
+) -> Result<Vec<Arc<AggregateFunctionExpr>>> {
     aggr_exprs
         .into_iter()
         .map(|aggr_expr| {
@@ -154,7 +154,7 @@ fn try_convert_aggregate_if_better(
                 let reqs = concat_slices(prefix_requirement, &aggr_sort_reqs);
                 if eq_properties.ordering_satisfy_requirement(&reqs) {
                     // Existing ordering satisfies the aggregator requirements:
-                    aggr_expr.with_beneficial_ordering(true)?
+                    aggr_expr.with_beneficial_ordering(true)?.map(Arc::new)
                 } else if eq_properties.ordering_satisfy_requirement(&concat_slices(
                     prefix_requirement,
                     &reverse_aggr_req,
@@ -163,12 +163,14 @@ fn try_convert_aggregate_if_better(
                     // given the existing ordering (if possible):
                     aggr_expr
                         .reverse_expr()
+                        .map(Arc::new)
                         .unwrap_or(aggr_expr)
                         .with_beneficial_ordering(true)?
+                        .map(Arc::new)
                 } else {
                     // There is no beneficial ordering present -- aggregation
                     // will still work albeit in a less efficient mode.
-                    aggr_expr.with_beneficial_ordering(false)?
+                    aggr_expr.with_beneficial_ordering(false)?.map(Arc::new)
                 }
                 .ok_or_else(|| {
                     plan_datafusion_err!(

datafusion/core/src/physical_planner.rs

@@ -1523,7 +1523,7 @@ pub fn create_window_expr(
 }
 
 type AggregateExprWithOptionalArgs = (
-    AggregateFunctionExpr,
+    Arc<AggregateFunctionExpr>,
     // The filter clause, if any
     Option<Arc<dyn PhysicalExpr>>,
     // Ordering requirements, if any
@@ -1587,7 +1587,8 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter(
                         .alias(name)
                         .with_ignore_nulls(ignore_nulls)
                         .with_distinct(*distinct)
-                        .build()?;
+                        .build()
+                        .map(Arc::new)?;
 
                 (agg_expr, filter, physical_sort_exprs)
             };

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -405,6 +405,7 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
                 .schema(Arc::clone(&schema))
                 .alias("sum1")
                 .build()
+                .map(Arc::new)
                 .unwrap(),
         ];
     let expr = group_by_columns

datafusion/core/tests/physical_optimizer/combine_partial_final_agg.rs

@@ -84,7 +84,7 @@ fn parquet_exec(schema: &SchemaRef) -> Arc<ParquetExec> {
 fn partial_aggregate_exec(
     input: Arc<dyn ExecutionPlan>,
     group_by: PhysicalGroupBy,
-    aggr_expr: Vec<AggregateFunctionExpr>,
+    aggr_expr: Vec<Arc<AggregateFunctionExpr>>,
 ) -> Arc<dyn ExecutionPlan> {
     let schema = input.schema();
     let n_aggr = aggr_expr.len();
@@ -104,7 +104,7 @@ fn partial_aggregate_exec(
 fn final_aggregate_exec(
     input: Arc<dyn ExecutionPlan>,
     group_by: PhysicalGroupBy,
-    aggr_expr: Vec<AggregateFunctionExpr>,
+    aggr_expr: Vec<Arc<AggregateFunctionExpr>>,
 ) -> Arc<dyn ExecutionPlan> {
     let schema = input.schema();
     let n_aggr = aggr_expr.len();
@@ -130,11 +130,12 @@ fn count_expr(
     expr: Arc<dyn PhysicalExpr>,
     name: &str,
     schema: &Schema,
-) -> AggregateFunctionExpr {
+) -> Arc<AggregateFunctionExpr> {
     AggregateExprBuilder::new(count_udaf(), vec![expr])
         .schema(Arc::new(schema.clone()))
         .alias(name)
         .build()
+        .map(Arc::new)
         .unwrap()
 }
 
@@ -218,6 +219,7 @@ fn aggregations_with_group_combined() -> datafusion_common::Result<()> {
             .schema(Arc::clone(&schema))
             .alias("Sum(b)")
             .build()
+            .map(Arc::new)
             .unwrap(),
     ];
     let groups: Vec<(Arc<dyn PhysicalExpr>, String)> =

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -347,10 +347,10 @@ fn test_has_aggregate_expression() -> Result<()> {
     let single_agg = AggregateExec::try_new(
         AggregateMode::Single,
         build_group_by(&schema, vec!["a".to_string()]),
-        vec![agg.count_expr(&schema)], /* aggr_expr */
-        vec![None],                    /* filter_expr */
-        source,                        /* input */
-        schema.clone(),                /* input_schema */
+        vec![Arc::new(agg.count_expr(&schema))], /* aggr_expr */
+        vec![None],                              /* filter_expr */
+        source,                                  /* input */
+        schema.clone(),                          /* input_schema */
     )?;
     let limit_exec = LocalLimitExec::new(
         Arc::new(single_agg),
@@ -384,10 +384,10 @@ fn test_has_filter() -> Result<()> {
     let single_agg = AggregateExec::try_new(
         AggregateMode::Single,
         build_group_by(&schema.clone(), vec!["a".to_string()]),
-        vec![agg.count_expr(&schema)], /* aggr_expr */
-        vec![filter_expr],             /* filter_expr */
-        source,                        /* input */
-        schema.clone(),                /* input_schema */
+        vec![Arc::new(agg.count_expr(&schema))], /* aggr_expr */
+        vec![filter_expr],                       /* filter_expr */
+        source,                                  /* input */
+        schema.clone(),                          /* input_schema */
     )?;
     let limit_exec = LocalLimitExec::new(
         Arc::new(single_agg),

datafusion/expr/src/udaf.rs

@@ -225,10 +225,10 @@ impl AggregateUDF {
 
     /// See [`AggregateUDFImpl::with_beneficial_ordering`] for more details.
     pub fn with_beneficial_ordering(
-        self,
+        self: Arc<Self>,
         beneficial_ordering: bool,
     ) -> Result<Option<AggregateUDF>> {
-        self.inner
+        Arc::clone(&self.inner)
             .with_beneficial_ordering(beneficial_ordering)
             .map(|updated_udf| updated_udf.map(|udf| Self { inner: udf }))
     }

datafusion/physical-expr/src/aggregate.rs

@@ -133,11 +133,11 @@ impl AggregateExprBuilder {
         };
 
         Ok(AggregateFunctionExpr {
-            fun: Arc::unwrap_or_clone(fun),
+            fun,
             args,
             data_type,
             name,
-            schema: Arc::unwrap_or_clone(schema),
+            schema,
             ordering_req,
             ignore_nulls,
             ordering_fields,
@@ -197,12 +197,12 @@ impl AggregateExprBuilder {
 /// Physical aggregate expression of a UDAF.
 #[derive(Debug, Clone)]
 pub struct AggregateFunctionExpr {
-    fun: AggregateUDF,
+    fun: Arc<AggregateUDF>,
     args: Vec<Arc<dyn PhysicalExpr>>,
     /// Output / return type of this aggregate
     data_type: DataType,
     name: String,
-    schema: Schema,
+    schema: Arc<Schema>,
     // The physical order by expressions
     ordering_req: LexOrdering,
     // Whether to ignore null values
@@ -328,20 +328,18 @@ impl AggregateFunctionExpr {
     /// not implement the method, returns an error. Order insensitive and hard
     /// requirement aggregators return `Ok(None)`.
     pub fn with_beneficial_ordering(
-        self,
+        self: Arc<Self>,
         beneficial_ordering: bool,
     ) -> Result<Option<AggregateFunctionExpr>> {
-        let Some(updated_fn) = self
-            .fun
-            .clone()
-            .with_beneficial_ordering(beneficial_ordering)?
+        let Some(updated_fn) =
+            Arc::clone(&self.fun).with_beneficial_ordering(beneficial_ordering)?
         else {
             return Ok(None);
         };
 
         AggregateExprBuilder::new(Arc::new(updated_fn), self.args.to_vec())
             .order_by(self.ordering_req.to_vec())
-            .schema(Arc::new(self.schema.clone()))
+            .schema(Arc::clone(&self.schema))
             .alias(self.name().to_string())
             .with_ignore_nulls(self.ignore_nulls)
             .with_distinct(self.is_distinct)
@@ -474,7 +472,7 @@ impl AggregateFunctionExpr {
 
                 AggregateExprBuilder::new(reverse_udf, self.args.to_vec())
                     .order_by(reverse_ordering_req.to_vec())
-                    .schema(Arc::new(self.schema.clone()))
+                    .schema(Arc::clone(&self.schema))
                     .alias(name)
                     .with_ignore_nulls(self.ignore_nulls)
                     .with_distinct(self.is_distinct)

datafusion/physical-expr/src/utils/mod.rs

@@ -86,6 +86,10 @@ pub fn map_columns_before_projection(
     parent_required: &[Arc<dyn PhysicalExpr>],
     proj_exprs: &[(Arc<dyn PhysicalExpr>, String)],
 ) -> Vec<Arc<dyn PhysicalExpr>> {
+    if parent_required.is_empty() {
+        // No need to build mapping.
+        return vec![];
+    }
     let column_mapping = proj_exprs
         .iter()
         .filter_map(|(expr, name)| {

datafusion/physical-expr/src/window/aggregate.rs

@@ -41,7 +41,7 @@ use crate::{expressions::PhysicalSortExpr, reverse_order_bys, PhysicalExpr};
 /// See comments on [`WindowExpr`] for more details.
 #[derive(Debug)]
 pub struct PlainAggregateWindowExpr {
-    aggregate: AggregateFunctionExpr,
+    aggregate: Arc<AggregateFunctionExpr>,
     partition_by: Vec<Arc<dyn PhysicalExpr>>,
     order_by: Vec<PhysicalSortExpr>,
     window_frame: Arc<WindowFrame>,
@@ -50,7 +50,7 @@ pub struct PlainAggregateWindowExpr {
 impl PlainAggregateWindowExpr {
     /// Create a new aggregate window function expression
     pub fn new(
-        aggregate: AggregateFunctionExpr,
+        aggregate: Arc<AggregateFunctionExpr>,
         partition_by: &[Arc<dyn PhysicalExpr>],
         order_by: &[PhysicalSortExpr],
         window_frame: Arc<WindowFrame>,
@@ -137,14 +137,14 @@ impl WindowExpr for PlainAggregateWindowExpr {
             let reverse_window_frame = self.window_frame.reverse();
             if reverse_window_frame.start_bound.is_unbounded() {
                 Arc::new(PlainAggregateWindowExpr::new(
-                    reverse_expr,
+                    Arc::new(reverse_expr),
                     &self.partition_by.clone(),
                     &reverse_order_bys(&self.order_by),
                     Arc::new(self.window_frame.reverse()),
                 )) as _
             } else {
                 Arc::new(SlidingAggregateWindowExpr::new(
-                    reverse_expr,
+                    Arc::new(reverse_expr),
                     &self.partition_by.clone(),
                     &reverse_order_bys(&self.order_by),
                     Arc::new(self.window_frame.reverse()),