From f2e2c3b3b5f12de788f4963a0bb2aa1e1fb0aa16 Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Tue, 23 Apr 2024 14:46:15 -0400
Subject: [PATCH] Avoid copies in TypeCoercion

---
 datafusion/expr/src/logical_plan/plan.rs      |  1 +
 .../optimizer/src/analyzer/type_coercion.rs   | 34 +++++++++++++------
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 6df5516b1bba..bc79ef9e07e3 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -815,6 +815,7 @@ impl LogicalPlan {
             }
         }
     }
+
     /// Replaces placeholder param values (like `$1`, `$2`) in [`LogicalPlan`]
     /// with the specified `param_values`.
     ///
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index 1ef87cd64e16..f8d890bffb58 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -106,16 +106,30 @@ fn analyze_internal(
     };
 
     let preserver = NamePreserver::new(&plan);
-    Ok(plan
-        .map_expressions(|expr| {
-            // ensure aggregate names don't change:
-            // https://github.com/apache/datafusion/issues/3555
-            let original_name = preserver.save(&expr)?;
-            expr.rewrite(&mut expr_rewrite)?
-                .map_data(|expr| original_name.restore(expr))
-        })?
-        // propagate the the transformation information from children
-        .update_transformed(children_transformed))
+    let transformed_plan = plan.map_expressions(|expr| {
+        // ensure aggregate names don't change:
+        // https://github.com/apache/datafusion/issues/3555
+        let original_name = preserver.save(&expr)?;
+        expr.rewrite(&mut expr_rewrite)?
+            .map_data(|expr| original_name.restore(expr))
+    })?;
+
+    // if any of the expressions were rewritten, we need to recreate the plan to
+    // recalculate the schema. At the moment this requires a copy
+    if transformed_plan.transformed || children_transformed {
+        // TODO avoid this copy
+        let plan = transformed_plan.data;
+        let new_inputs = plan
+            .inputs()
+            .into_iter()
+            .map(|input| input.clone())
+            .collect::<Vec<_>>();
+
+        plan.with_new_exprs(plan.expressions(), new_inputs)
+            .map(Transformed::yes)
+    } else {
+        Ok(transformed_plan)
+    }
 }
 
 pub(crate) struct TypeCoercionRewriter {