Allow AggregateUDFImpl::simplify to return non AggregateExpr

Author: alamb

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -20,11 +20,13 @@
 use std::sync::Arc;
 
 use datafusion_common::tree_node::{Transformed, TreeNode};
-use datafusion_common::{DFSchema, DFSchemaRef, DataFusionError, Result};
+use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result};
 use datafusion_expr::Expr;
-use datafusion_expr::logical_plan::LogicalPlan;
+use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection};
 use datafusion_expr::simplify::SimplifyContext;
-use datafusion_expr::utils::merge_schema;
+use datafusion_expr::utils::{
+    columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, merge_schema,
+};
 
 use crate::optimizer::ApplyOrder;
 use crate::utils::NamePreserver;
@@ -130,13 +132,17 @@ impl SimplifyExpressions {
             ))
         };
 
-        plan.map_expressions(|expr| {
+        let transformed = plan.map_expressions(|expr| {
             // Preserve the aliasing of grouping sets.
             if let Expr::GroupingSet(_) = &expr {
                 expr.map_children(&mut rewrite_expr)
             } else {
                 rewrite_expr(expr)
             }
+        })?;
+
+        transformed.transform_data(|plan| {
+            rewrite_aggregate_non_aggregate_aggr_expr(plan)
         })
     }
 }
@@ -148,6 +154,73 @@ impl SimplifyExpressions {
     }
 }
 
+/// Rewrites
+/// `Aggregate(group_expr, aggr_expr=[non_agg_expr(sum(..), count(..), ..)])`
+/// into:
+/// `Projection(..)` over `Aggregate(group_expr, aggr_expr=[sum(..), count(..), ..])`.
+///
+/// Aggregate planning requires each aggregate slot to be an aggregate function
+/// (possibly aliased). Some UDAF simplifiers can return larger expressions that
+/// reference multiple aggregate functions.
+fn rewrite_aggregate_non_aggregate_aggr_expr(
+    plan: LogicalPlan,
+) -> Result<Transformed<LogicalPlan>> {
+    let LogicalPlan::Aggregate(Aggregate {
+        input,
+        group_expr,
+        aggr_expr,
+        ..
+    }) = plan
+    else {
+        return Ok(Transformed::no(plan));
+    };
+
+    if aggr_expr.iter().all(is_top_level_aggregate_expr) {
+        return Ok(Transformed::no(LogicalPlan::Aggregate(Aggregate::try_new(
+            input, group_expr, aggr_expr,
+        )?)));
+    }
+
+    let inner_aggr_expr = find_aggregate_exprs(aggr_expr.iter());
+    let inner_aggregate = LogicalPlan::Aggregate(Aggregate::try_new(
+        Arc::clone(&input),
+        group_expr.clone(),
+        inner_aggr_expr,
+    )?);
+    let inner_aggregate = Arc::new(inner_aggregate);
+
+    let mut projection_exprs = aggregate_output_exprs(&group_expr)?;
+    projection_exprs.extend(aggr_expr);
+    let projection_exprs = projection_exprs
+        .into_iter()
+        .map(|expr| columnize_expr(expr, inner_aggregate.as_ref()))
+        .collect::<Result<Vec<_>>>()?;
+
+    Ok(Transformed::yes(LogicalPlan::Projection(Projection::try_new(
+        projection_exprs,
+        inner_aggregate,
+    )?)))
+}
+
+fn is_top_level_aggregate_expr(expr: &Expr) -> bool {
+    matches!(expr.clone().unalias_nested().data, Expr::AggregateFunction(_))
+}
+
+fn aggregate_output_exprs(group_expr: &[Expr]) -> Result<Vec<Expr>> {
+    let mut output_exprs = grouping_set_to_exprlist(group_expr)?
+        .into_iter()
+        .cloned()
+        .collect::<Vec<_>>();
+
+    if matches!(group_expr, [Expr::GroupingSet(_)]) {
+        output_exprs.push(Expr::Column(Column::from_name(
+            Aggregate::INTERNAL_GROUPING_ID,
+        )));
+    }
+
+    Ok(output_exprs)
+}
+
 #[cfg(test)]
 mod tests {
     use std::ops::Not;
@@ -158,7 +231,7 @@ mod tests {
     use datafusion_expr::logical_plan::builder::table_scan_with_filters;
     use datafusion_expr::logical_plan::table_scan;
     use datafusion_expr::*;
-    use datafusion_functions_aggregate::expr_fn::{max, min};
+    use datafusion_functions_aggregate::expr_fn::{max, min, sum};
 
     use crate::OptimizerContext;
     use crate::assert_optimized_plan_eq_snapshot;
@@ -258,6 +331,28 @@ mod tests {
         )
     }
 
+    #[test]
+    fn test_simplify_udaf_to_non_aggregate_expr() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
+        let table_scan = table_scan(Some("test"), &schema, None)?
+            .build()
+            .expect("building scan");
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(Vec::<Expr>::new(), vec![sum(col("a") + lit(2i64))])?
+            .build()?;
+
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Projection: sum(test.a) + Int64(2) * count(test.a) AS sum(test.a + Int64(2))
+          Aggregate: groupBy=[[]], aggr=[[sum(test.a), count(test.a)]]
+            TableScan: test
+        "
+        )?;
+        Ok(())
+    }
+
     #[test]
     fn test_simplify_optimized_plan_with_or() -> Result<()> {
         let table_scan = test_table_scan();