diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index f3afcc8fee11..dcefe67b729a 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -30,7 +30,7 @@ use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; use datafusion_expr::expr_schema::ExprSchemable; use datafusion_expr::{ - AggregateUDF, Expr, col, + AggregateUDF, Expr, cast, col, expr::AggregateFunction, lit, logical_plan::{Aggregate, LogicalPlan}, @@ -406,6 +406,12 @@ impl OptimizerRule for SingleDistinctToGroupBy { .map(|(expr, _)| expr.clone()) .collect(); + let outer_aggr = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::new(inner_agg), + outer_group_exprs, + outer_aggr_exprs, + )?); + // so the aggregates are displayed in the same way even after the rewrite // this optimizer has two kinds of alias: // - group_by aggr @@ -418,18 +424,25 @@ impl OptimizerRule for SingleDistinctToGroupBy { } None => group_expr, }) - .chain(outer_proj_exprs.into_iter().enumerate().map(|(idx, expr)| { - let idx = idx + group_size; - let (qualifier, field) = schema.qualified_field(idx); - expr.alias_qualified(qualifier.cloned(), field.name()) - })) + .chain( + outer_proj_exprs + .into_iter() + .enumerate() + .map(|(idx, expr)| { + let idx = idx + group_size; + let (qualifier, field) = schema.qualified_field(idx); + let current_type = + outer_aggr.schema().field(idx).data_type(); + let expr = if current_type != field.data_type() { + cast(expr, field.data_type().clone()) + } else { + expr + }; + expr.alias_qualified(qualifier.cloned(), field.name()) + }), + ) .collect(); - let outer_aggr = LogicalPlan::Aggregate(Aggregate::try_new( - Arc::new(inner_agg), - outer_group_exprs, - outer_aggr_exprs, - )?); Ok(Transformed::yes(project(outer_aggr, alias_expr)?)) } _ => Ok(Transformed::no(plan)), @@ -1314,4 +1327,31 @@ mod tests { ", ) } + + #[test] + fn sum_decimal_with_distinct_preserves_declared_type() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("d", DataType::Decimal128(5, 2), false), + Field::new("x", DataType::Int32, false), + ]); + let table_scan = table_scan(Some("test"), &schema, None)?.build()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![sum(col("d")), count_distinct(col("x"))], + )? + .build()?; + + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + let optimizer = Optimizer::with_rules(vec![rule]); + let optimized_plan = optimizer + .optimize(plan.clone(), &RegistryOptimizerContext::new(), |_, _| {}) + .expect("failed to optimize plan"); + + assert_eq!(plan.schema().as_ref(), optimized_plan.schema().as_ref()); + Ok(()) + } } +