Reuse NamePreserver in SimplifyExpressions (#12238)

datafusion/expr/src/expr_rewriter/mod.rs

@@ -311,6 +311,8 @@ impl NamePreserver {
     /// Create a new NamePreserver for rewriting the `expr` that is part of the specified plan
     pub fn new(plan: &LogicalPlan) -> Self {
         Self {
+            // The schema of Filter and Join nodes comes from their inputs rather than their output expressions,
+            // so there is no need to use aliases to preserve expression names.
             use_alias: !matches!(plan, LogicalPlan::Filter(_) | LogicalPlan::Join(_)),
         }
     }

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -27,6 +27,7 @@ use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::utils::merge_schema;
 
 use crate::optimizer::ApplyOrder;
+use crate::utils::NamePreserver;
 use crate::{OptimizerConfig, OptimizerRule};
 
 use super::ExprSimplifier;
@@ -119,18 +120,13 @@ impl SimplifyExpressions {
             simplifier
         };
 
-        // the output schema of a filter or join is the input schema. Thus they
-        // can't handle aliased expressions
-        let use_alias = !matches!(plan, LogicalPlan::Filter(_) | LogicalPlan::Join(_));
+        // Preserve expression names to avoid changing the schema of the plan.
+        let name_preserver = NamePreserver::new(&plan);
         plan.map_expressions(|e| {
-            let new_e = if use_alias {
-                // TODO: unify with `rewrite_preserving_name`
-                let original_name = e.name_for_alias()?;
-                simplifier.simplify(e)?.alias_if_changed(original_name)
-            } else {
-                simplifier.simplify(e)
-            }?;
-
+            let original_name = name_preserver.save(&e)?;
+            let new_e = simplifier
+                .simplify(e)
+                .and_then(|expr| original_name.restore(expr))?;
             // TODO it would be nice to have a way to know if the expression was simplified
             // or not. For now conservatively return Transformed::yes
             Ok(Transformed::yes(new_e))