From d4da80bfe2e53a98e49e58f0337f2e8cae010d4b Mon Sep 17 00:00:00 2001
From: Matt Green <emgeee@users.noreply.github.com>
Date: Thu, 2 May 2024 06:53:37 -0700
Subject: [PATCH] Stop copying LogicalPlan and Exprs in EliminateNestedUnion
 (#10319)

* First pass of https://github.com/apache/datafusion/issues/10296

* Update datafusion/optimizer/src/eliminate_nested_union.rs

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

* fix formatting

* return original expr instead of clone()

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
---
 datafusion/expr/src/expr_rewriter/mod.rs      |  2 +-
 .../optimizer/src/eliminate_nested_union.rs   | 73 +++++++++++--------
 2 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs
index 14154189a126..700dd560ec0b 100644
--- a/datafusion/expr/src/expr_rewriter/mod.rs
+++ b/datafusion/expr/src/expr_rewriter/mod.rs
@@ -250,7 +250,7 @@ fn coerce_exprs_for_schema(
                     _ => expr.cast_to(new_type, src_schema),
                 }
             } else {
-                Ok(expr.clone())
+                Ok(expr)
             }
         })
         .collect::<Result<_>>()
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs b/datafusion/optimizer/src/eliminate_nested_union.rs
index da2a6a17214e..aa6f2b497531 100644
--- a/datafusion/optimizer/src/eliminate_nested_union.rs
+++ b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -18,7 +18,8 @@
 //! [`EliminateNestedUnion`]: flattens nested `Union` to a single `Union`
 use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
-use datafusion_common::Result;
+use datafusion_common::tree_node::Transformed;
+use datafusion_common::{internal_err, Result};
 use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
 use datafusion_expr::{Distinct, LogicalPlan, Union};
 use std::sync::Arc;
@@ -37,9 +38,29 @@ impl EliminateNestedUnion {
 impl OptimizerRule for EliminateNestedUnion {
     fn try_optimize(
         &self,
-        plan: &LogicalPlan,
+        _plan: &LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
+        internal_err!("Should have called EliminateNestedUnion::rewrite")
+    }
+
+    fn name(&self) -> &str {
+        "eliminate_nested_union"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::BottomUp)
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        true
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
         match plan {
             LogicalPlan::Union(Union { inputs, schema }) => {
                 let inputs = inputs
@@ -47,39 +68,33 @@ impl OptimizerRule for EliminateNestedUnion {
                     .flat_map(extract_plans_from_union)
                     .collect::<Vec<_>>();
 
-                Ok(Some(LogicalPlan::Union(Union {
+                Ok(Transformed::yes(LogicalPlan::Union(Union {
                     inputs,
-                    schema: schema.clone(),
+                    schema,
                 })))
             }
-            LogicalPlan::Distinct(Distinct::All(plan)) => match plan.as_ref() {
-                LogicalPlan::Union(Union { inputs, schema }) => {
-                    let inputs = inputs
-                        .iter()
-                        .map(extract_plan_from_distinct)
-                        .flat_map(extract_plans_from_union)
-                        .collect::<Vec<_>>();
-
-                    Ok(Some(LogicalPlan::Distinct(Distinct::All(Arc::new(
-                        LogicalPlan::Union(Union {
-                            inputs,
-                            schema: schema.clone(),
-                        }),
-                    )))))
+            LogicalPlan::Distinct(Distinct::All(ref nested_plan)) => {
+                match nested_plan.as_ref() {
+                    LogicalPlan::Union(Union { inputs, schema }) => {
+                        let inputs = inputs
+                            .iter()
+                            .map(extract_plan_from_distinct)
+                            .flat_map(extract_plans_from_union)
+                            .collect::<Vec<_>>();
+
+                        Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::All(
+                            Arc::new(LogicalPlan::Union(Union {
+                                inputs,
+                                schema: schema.clone(),
+                            })),
+                        ))))
+                    }
+                    _ => Ok(Transformed::no(plan)),
                 }
-                _ => Ok(None),
-            },
-            _ => Ok(None),
+            }
+            _ => Ok(Transformed::no(plan)),
         }
     }
-
-    fn name(&self) -> &str {
-        "eliminate_nested_union"
-    }
-
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::BottomUp)
-    }
 }
 
 fn extract_plans_from_union(plan: &Arc<LogicalPlan>) -> Vec<Arc<LogicalPlan>> {