diff --git a/datafusion/optimizer/src/eliminate_duplicated_expr.rs b/datafusion/optimizer/src/eliminate_duplicated_expr.rs index 97aa6e1d8480..f2f489ab825c 100644 --- a/datafusion/optimizer/src/eliminate_duplicated_expr.rs +++ b/datafusion/optimizer/src/eliminate_duplicated_expr.rs @@ -112,8 +112,12 @@ impl OptimizerRule for EliminateDuplicatedExpr { fetch: sort.fetch, }))) } - LogicalPlan::Aggregate(agg) => { + LogicalPlan::Aggregate(mut agg) => { let len = agg.group_expr.len(); + // With zero or one group expression, there is nothing to deduplicate. + if len <= 1 { + return Ok(Transformed::no(LogicalPlan::Aggregate(agg))); + } let unique_exprs: Vec = agg .group_expr @@ -122,14 +126,15 @@ impl OptimizerRule for EliminateDuplicatedExpr { .into_iter() .collect(); - let transformed = if len != unique_exprs.len() { - Transformed::yes - } else { - Transformed::no - }; + if len == unique_exprs.len() { + // No duplicate expressions were found, so we can reuse the + // existing schema (including functional dependencies). + agg.group_expr = unique_exprs; + return Ok(Transformed::no(LogicalPlan::Aggregate(agg))); + } Aggregate::try_new(agg.input, unique_exprs, agg.aggr_expr) - .map(|f| transformed(LogicalPlan::Aggregate(f))) + .map(|f| Transformed::yes(LogicalPlan::Aggregate(f))) } _ => Ok(Transformed::no(plan)), } @@ -145,7 +150,10 @@ mod tests { use crate::OptimizerContext; use crate::assert_optimized_plan_eq_snapshot; use crate::test::*; - use datafusion_expr::{col, logical_plan::builder::LogicalPlanBuilder}; + use datafusion_expr::{ + GroupingSet, col, lit, logical_plan::builder::LogicalPlanBuilder, + }; + use datafusion_functions_aggregate::expr_fn::sum; use std::sync::Arc; macro_rules! assert_optimized_plan_equal { @@ -200,4 +208,73 @@ mod tests { TableScan: test ") } + + /// A no-op must preserve both the plan and its already computed schema. + fn assert_unchanged_aggregate(plan: LogicalPlan) -> Result<()> { + let schema = Arc::clone(plan.schema()); + let result = EliminateDuplicatedExpr::new() + .rewrite(plan.clone(), &OptimizerContext::new())?; + assert!(!result.transformed); + assert_eq!(result.data, plan); + assert!(Arc::ptr_eq(result.data.schema(), &schema)); + Ok(()) + } + + #[test] + fn unchanged_aggregates_reuse_schema() -> Result<()> { + for group_expr in [vec![], vec![col("a")], vec![col("b"), col("a") + lit(1u32)]] { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate(group_expr, vec![sum(col("c")).alias("total")])? + .build()?; + assert_unchanged_aggregate(plan)?; + } + Ok(()) + } + + #[test] + fn unchanged_grouping_sets_reuse_schema() -> Result<()> { + // Repeat "a" to check that this rule leaves the GroupingSet's contents unchanged. + let grouping_set = GroupingSet::GroupingSets(vec![ + vec![col("a"), col("b")], + vec![col("a")], + vec![col("a")], + vec![], + ]); + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate(vec![Expr::GroupingSet(grouping_set)], vec![sum(col("c"))])? + .build()?; + assert_unchanged_aggregate(plan) + } + + #[test] + fn eliminate_duplicate_group_exprs() -> Result<()> { + let input = Arc::new(test_table_scan()?); + let a = col("test.a"); + let b = col("test.b"); + let computed = a.clone() + lit(1u32); + let aggr_expr = vec![sum(col("test.c")).alias("total")]; + let plan = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::clone(&input), + vec![ + b.clone(), + computed.clone(), + b.clone(), + a.clone(), + computed.clone(), + ], + aggr_expr.clone(), + )?); + let expected = LogicalPlan::Aggregate(Aggregate::try_new( + input, + // Keep the order of the first occurrence of each expression. + vec![b, computed, a], + aggr_expr, + )?); + + let result = + EliminateDuplicatedExpr::new().rewrite(plan, &OptimizerContext::new())?; + assert!(result.transformed); + assert_eq!(result.data, expected); + Ok(()) + } }