Skip to content

Commit 06cdc9f

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 06cdc9f

File tree

1 file changed

+37
-6
lines changed

1 file changed

+37
-6
lines changed

datafusion/optimizer/src/common_subexpr_eliminate.rs

+37-6
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,19 @@ 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
538+
.clone()
539+
.into_iter()
540+
.map(|_| None)
541+
.collect::<Vec<_>>()
542+
};
530543

531544
let mut agg_exprs = common_exprs
532545
.into_iter()
@@ -542,8 +555,11 @@ impl CommonSubexprEliminate {
542555
&mut proj_exprs,
543556
)?
544557
}
545-
for (expr_rewritten, expr_orig) in
546-
rewritten_aggr_expr.into_iter().zip(new_aggr_expr)
558+
for ((expr_rewritten, expr_orig), saved_name) in
559+
rewritten_aggr_expr
560+
.into_iter()
561+
.zip(new_aggr_expr)
562+
.zip(saved_names)
547563
{
548564
if expr_rewritten == expr_orig {
549565
if let Expr::Alias(Alias { expr, name, .. }) =
@@ -557,10 +573,25 @@ impl CommonSubexprEliminate {
557573
config.alias_generator().next(CSE_PREFIX);
558574
let (qualifier, field) =
559575
expr_rewritten.to_field(&new_input_schema)?;
560-
let out_name = qualified_name(
561-
qualifier.as_ref(),
562-
field.name(),
563-
);
576+
let out_name = if let Some(expr) = saved_name {
577+
let name = if let Expr::Alias(Alias {
578+
expr: _,
579+
name,
580+
..
581+
}) =
582+
expr.restore(expr_rewritten.clone())
583+
{
584+
name
585+
} else {
586+
field.name().to_string()
587+
};
588+
qualified_name(qualifier.as_ref(), &name)
589+
} else {
590+
qualified_name(
591+
qualifier.as_ref(),
592+
field.name(),
593+
)
594+
};
564595

565596
agg_exprs.push(expr_rewritten.alias(&expr_alias));
566597
proj_exprs.push(

0 commit comments

Comments
 (0)