Skip to content

Commit f823a32

Browse files
committed
Fix: common_sub_expression_eliminate optimizer rule failed
Common_sub_expression_eliminate rule failed with error: `SchemaError(FieldNotFound {field: <name>}, valid_fields: []})` due to the schema being changed by the second application of `find_common_exprs`. Used NamePreserver mechanism to restore original schema name and generate Projection with original names at the end of aggregate optimization.
1 parent 9e17b4f commit f823a32

File tree

1 file changed

+29
-6
lines changed

1 file changed

+29
-6
lines changed

datafusion/optimizer/src/common_subexpr_eliminate.rs

+29-6
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,15 @@ impl CommonSubexprEliminate {
527527
} => {
528528
let rewritten_aggr_expr = new_exprs_list.pop().unwrap();
529529
let new_aggr_expr = original_exprs_list.pop().unwrap();
530+
let saved_names = if let Some(aggr_expr) = aggr_expr {
531+
let name_perserver = NamePreserver::new_for_projection();
532+
aggr_expr
533+
.iter()
534+
.map(|expr| Some(name_perserver.save(expr)))
535+
.collect::<Vec<_>>()
536+
} else {
537+
new_aggr_expr.clone().into_iter().map(|_| None).collect::<Vec<_>>()
538+
};
530539

531540
let mut agg_exprs = common_exprs
532541
.into_iter()
@@ -542,8 +551,8 @@ impl CommonSubexprEliminate {
542551
&mut proj_exprs,
543552
)?
544553
}
545-
for (expr_rewritten, expr_orig) in
546-
rewritten_aggr_expr.into_iter().zip(new_aggr_expr)
554+
for ((expr_rewritten, expr_orig), saved_name) in
555+
rewritten_aggr_expr.into_iter().zip(new_aggr_expr).zip(saved_names)
547556
{
548557
if expr_rewritten == expr_orig {
549558
if let Expr::Alias(Alias { expr, name, .. }) =
@@ -557,10 +566,24 @@ impl CommonSubexprEliminate {
557566
config.alias_generator().next(CSE_PREFIX);
558567
let (qualifier, field) =
559568
expr_rewritten.to_field(&new_input_schema)?;
560-
let out_name = qualified_name(
561-
qualifier.as_ref(),
562-
field.name(),
563-
);
569+
let out_name = if let Some(expr) = saved_name {
570+
let name = if let Expr::Alias(Alias { expr: _, name, .. }) =
571+
expr.restore(expr_rewritten.clone())
572+
{
573+
name
574+
} else {
575+
field.name().to_string()
576+
};
577+
qualified_name(
578+
qualifier.as_ref(),
579+
&name
580+
)
581+
} else {
582+
qualified_name(
583+
qualifier.as_ref(),
584+
field.name(),
585+
)
586+
};
564587

565588
agg_exprs.push(expr_rewritten.alias(&expr_alias));
566589
proj_exprs.push(

0 commit comments

Comments
 (0)