Skip to content

fix: cast outer projection expr to declared field type in SingleDistinctToGroupBy #26211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jdymitarai wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: cast outer projection expr to declared field type in SingleDistinctToGroupBy #26211

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 51 additions & 11 deletions datafusion/optimizer/src/single_distinct_to_groupby.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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
Expand All @@ -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)),
Expand Down Expand Up @@ -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::<Expr>::new(),
vec![sum(col("d")), count_distinct(col("x"))],
)?
.build()?;

let rule: Arc<dyn OptimizerRule + Send + Sync> =
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(())
}
}