Preserve projected constants in aggregate arguments

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -17,10 +17,13 @@
 
 //! Simplify expressions optimizer rule and implementation
 
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use datafusion_common::tree_node::{Transformed, TreeNode};
-use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result};
+use datafusion_common::{
+    Column, DFSchema, DFSchemaRef, DataFusionError, Result, qualified_name,
+};
 use datafusion_expr::Expr;
 use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection};
 use datafusion_expr::simplify::SimplifyContext;
@@ -30,6 +33,7 @@ use datafusion_expr::utils::{
 
 use super::ExprSimplifier;
 use crate::optimizer::ApplyOrder;
+use crate::push_down_filter::replace_cols_by_name;
 use crate::simplify_expressions::linear_aggregates::rewrite_multiple_linear_aggregates;
 use crate::utils::NamePreserver;
 use crate::{OptimizerConfig, OptimizerRule};
@@ -76,6 +80,14 @@ impl SimplifyExpressions {
     fn optimize_internal(
         plan: LogicalPlan,
         config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
+        inline_projection_constants(plan)?
+            .transform_data(|plan| Self::simplify_plan(plan, config))
+    }
+
+    fn simplify_plan(
+        plan: LogicalPlan,
+        config: &dyn OptimizerConfig,
     ) -> Result<Transformed<LogicalPlan>> {
         let schema = if !plan.inputs().is_empty() {
             DFSchemaRef::new(merge_schema(&plan.inputs()))
@@ -145,6 +157,91 @@ impl SimplifyExpressions {
     }
 }
 
+fn inline_projection_constants(plan: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+    let LogicalPlan::Aggregate(Aggregate {
+        input,
+        group_expr,
+        aggr_expr,
+        schema,
+        ..
+    }) = plan
+    else {
+        return Ok(Transformed::no(plan));
+    };
+    let replacements = projection_constants(input.as_ref())?;
+    if replacements.is_empty() {
+        return Ok(Transformed::no(LogicalPlan::Aggregate(
+            Aggregate::try_new_with_schema(input, group_expr, aggr_expr, schema)?,
+        )));
+    }
+
+    let name_preserver = NamePreserver::new_for_projection();
+    let rewritten_aggr_expr = aggr_expr
+        .iter()
+        .cloned()
+        .map(|expr| {
+            let name = name_preserver.save(&expr);
+            replace_cols_by_name(expr, &replacements).map(|expr| name.restore(expr))
+        })
+        .collect::<Result<Vec<_>>>()?;
+    if rewritten_aggr_expr == aggr_expr {
+        return Ok(Transformed::no(LogicalPlan::Aggregate(
+            Aggregate::try_new_with_schema(input, group_expr, aggr_expr, schema)?,
+        )));
+    }
+
+    Ok(Transformed::yes(LogicalPlan::Aggregate(
+        Aggregate::try_new_with_schema(input, group_expr, rewritten_aggr_expr, schema)?,
+    )))
+}
+
+fn projection_constants(plan: &LogicalPlan) -> Result<HashMap<String, Expr>> {
+    match plan {
+        LogicalPlan::Projection(projection) => {
+            let input_constants = projection_constants(projection.input.as_ref())?;
+            projection.schema.iter().zip(&projection.expr).try_fold(
+                HashMap::new(),
+                |mut constants, ((qualifier, field), expr)| {
+                    let expr = replace_cols_by_name(
+                        expr.clone().unalias_nested().data,
+                        &input_constants,
+                    )?;
+                    if !expr.any_column_refs()
+                        && !expr.is_volatile()
+                        && !expr_contains_subquery(&expr)
+                    {
+                        constants.insert(qualified_name(qualifier, field.name()), expr);
+                    }
+                    Ok(constants)
+                },
+            )
+        }
+        LogicalPlan::SubqueryAlias(alias) => {
+            let input_constants = projection_constants(alias.input.as_ref())?;
+            Ok(plan
+                .schema()
+                .iter()
+                .zip(alias.input.schema().iter())
+                .filter_map(|((qualifier, field), (input_qualifier, input_field))| {
+                    input_constants
+                        .get(&qualified_name(input_qualifier, input_field.name()))
+                        .cloned()
+                        .map(|expr| (qualified_name(qualifier, field.name()), expr))
+                })
+                .collect())
+        }
+        _ => Ok(HashMap::new()),
+    }
+}
+
+fn expr_contains_subquery(expr: &Expr) -> bool {
+    expr.exists(|expr| match expr {
+        Expr::ScalarSubquery(_) | Expr::Exists(_) | Expr::InSubquery(_) => Ok(true),
+        _ => Ok(false),
+    })
+    .expect("expression walk is infallible")
+}
+
 impl SimplifyExpressions {
     #[expect(missing_docs)]
     pub fn new() -> Self {

datafusion/sqllogictest/test_files/aggregate.slt

@@ -155,6 +155,28 @@ from values (1), (2), (3) t(col0);
 ----
 2.5
 
+query R
+SELECT approx_percentile_cont(p) WITHIN GROUP (ORDER BY v)
+FROM (
+  SELECT v, CAST(0.5 AS DOUBLE) AS p
+  FROM (VALUES (1), (2), (3)) t(v)
+) q;
+----
+2
+
+query R
+WITH q AS (
+  SELECT v, CAST(1 AS DOUBLE) / 2 AS p
+  FROM (VALUES (1), (2), (3)) t(v)
+)
+SELECT approx_percentile_cont(p) WITHIN GROUP (ORDER BY v) FROM q;
+----
+2
+
+statement error DataFusion error: Error during planning: Percentile value for 'APPROX_PERCENTILE_CONT' must be a literal
+SELECT approx_percentile_cont(p) WITHIN GROUP (ORDER BY v)
+FROM (VALUES (1, 0.25), (2, 0.75)) t(v, p);
+
 #######
 # Error tests
 #######