Projection Optimization Rule

datafusion/core/src/physical_optimizer/aggregate_statistics.rs

@@ -35,9 +35,6 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 #[derive(Default)]
 pub struct AggregateStatistics {}
 
-/// The name of the column corresponding to [`COUNT_STAR_EXPANSION`]
-const COUNT_STAR_NAME: &str = "COUNT(*)";
-
 impl AggregateStatistics {
     #[allow(missing_docs)]
     pub fn new() -> Self {
@@ -144,7 +141,7 @@ fn take_optimizable(node: &dyn ExecutionPlan) -> Option<Arc<dyn ExecutionPlan>>
 fn take_optimizable_table_count(
     agg_expr: &dyn AggregateExpr,
     stats: &Statistics,
-) -> Option<(ScalarValue, &'static str)> {
+) -> Option<(ScalarValue, String)> {
     if let (&Precision::Exact(num_rows), Some(casted_expr)) = (
         &stats.num_rows,
         agg_expr.as_any().downcast_ref::<expressions::Count>(),
@@ -158,7 +155,7 @@ fn take_optimizable_table_count(
                 if lit_expr.value() == &COUNT_STAR_EXPANSION {
                     return Some((
                         ScalarValue::Int64(Some(num_rows as i64)),
-                        COUNT_STAR_NAME,
+                        casted_expr.name().to_owned(),
                     ));
                 }
             }
@@ -427,7 +424,7 @@ pub(crate) mod tests {
         /// What name would this aggregate produce in a plan?
         fn column_name(&self) -> &'static str {
             match self {
-                Self::CountStar => COUNT_STAR_NAME,
+                Self::CountStar => "COUNT(*)",
                 Self::ColumnA(_) => "COUNT(a)",
             }
         }

datafusion/core/src/physical_optimizer/combine_partial_final_agg.rs

@@ -70,12 +70,12 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
                                         AggregateMode::Partial
                                     ) && can_combine(
                                         (
-                                            agg_exec.group_by(),
+                                            agg_exec.group_expr(),
                                             agg_exec.aggr_expr(),
                                             agg_exec.filter_expr(),
                                         ),
                                         (
-                                            input_agg_exec.group_by(),
+                                            input_agg_exec.group_expr(),
                                             input_agg_exec.aggr_expr(),
                                             input_agg_exec.filter_expr(),
                                         ),
@@ -88,7 +88,7 @@ impl PhysicalOptimizerRule for CombinePartialFinalAggregate {
                                             };
                                         AggregateExec::try_new(
                                             mode,
-                                            input_agg_exec.group_by().clone(),
+                                            input_agg_exec.group_expr().clone(),
                                             input_agg_exec.aggr_expr().to_vec(),
                                             input_agg_exec.filter_expr().to_vec(),
                                             input_agg_exec.input().clone(),

datafusion/core/src/physical_optimizer/convert_first_last.rs

@@ -79,7 +79,7 @@ fn get_common_requirement_of_aggregate_input(
     if let Some(aggr_exec) = plan.as_any().downcast_ref::<AggregateExec>() {
         let input = aggr_exec.input();
         let mut aggr_expr = try_get_updated_aggr_expr_from_child(aggr_exec);
-        let group_by = aggr_exec.group_by();
+        let group_by = aggr_exec.group_expr();
         let mode = aggr_exec.mode();
 
         let input_eq_properties = input.equivalence_properties();
@@ -113,7 +113,7 @@ fn get_common_requirement_of_aggregate_input(
                 InputOrderMode::Linear
             };
         let projection_mapping =
-            ProjectionMapping::try_new(group_by.expr(), &input.schema())?;
+            ProjectionMapping::try_new(group_by.expr().to_vec(), &input.schema())?;
 
         let cache = AggregateExec::compute_properties(
             input,

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -461,7 +461,7 @@ fn reorder_aggregate_keys(
 ) -> Result<PlanWithKeyRequirements> {
     let parent_required = &agg_node.data;
     let output_columns = agg_exec
-        .group_by()
+        .group_expr()
         .expr()
         .iter()
         .enumerate()
@@ -474,15 +474,15 @@ fn reorder_aggregate_keys(
         .collect::<Vec<_>>();
 
     if parent_required.len() == output_exprs.len()
-        && agg_exec.group_by().null_expr().is_empty()
+        && agg_exec.group_expr().null_expr().is_empty()
         && !physical_exprs_equal(&output_exprs, parent_required)
     {
         if let Some(positions) = expected_expr_positions(&output_exprs, parent_required) {
             if let Some(agg_exec) =
                 agg_exec.input().as_any().downcast_ref::<AggregateExec>()
             {
                 if matches!(agg_exec.mode(), &AggregateMode::Partial) {
-                    let group_exprs = agg_exec.group_by().expr();
+                    let group_exprs = agg_exec.group_expr().expr();
                     let new_group_exprs = positions
                         .into_iter()
                         .map(|idx| group_exprs[idx].clone())

datafusion/core/src/physical_optimizer/limited_distinct_aggregation.rs

@@ -55,7 +55,7 @@ impl LimitedDistinctAggregation {
         // We found what we want: clone, copy the limit down, and return modified node
         let new_aggr = AggregateExec::try_new(
             *aggr.mode(),
-            aggr.group_by().clone(),
+            aggr.group_expr().clone(),
             aggr.aggr_expr().to_vec(),
             aggr.filter_expr().to_vec(),
             aggr.input().clone(),
@@ -116,7 +116,7 @@ impl LimitedDistinctAggregation {
                     if let Some(parent_aggr) =
                         match_aggr.as_any().downcast_ref::<AggregateExec>()
                     {
-                        if !parent_aggr.group_by().eq(aggr.group_by()) {
+                        if !parent_aggr.group_expr().eq(aggr.group_expr()) {
                             // a partial and final aggregation with different groupings disqualifies
                             // rewriting the child aggregation
                             rewrite_applicable = false;

datafusion/core/src/physical_optimizer/mod.rs

@@ -29,10 +29,10 @@ pub mod enforce_distribution;
 pub mod enforce_sorting;
 pub mod join_selection;
 pub mod limited_distinct_aggregation;
+mod optimize_projections;
 pub mod optimizer;
 pub mod output_requirements;
 pub mod pipeline_checker;
-mod projection_pushdown;
 pub mod pruning;
 pub mod replace_with_order_preserving_variants;
 mod sort_pushdown;