Stop copying LogicalPlan and Exprs in EliminateNestedUnion (#10319)

* First pass of #10296 * Update datafusion/optimizer/src/eliminate_nested_union.rs Co-authored-by: Andrew Lamb <[email protected]> * fix formatting * return original expr instead of clone() --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · May 2, 2024 · d4da80b · d4da80b
1 parent 63e05fa
commit d4da80b
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 30 deletions.
diff --git 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
@@ -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,49 +38,63 @@ 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
                     .iter()
                     .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>> {