Minor: Remove clone in optimizer (#11315)

* rm clone Signed-off-by: jayzhan211 <[email protected]> * outer join + fix test Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]>
apache · Jul 7, 2024 · 5aa7c4a · 5aa7c4a
1 parent 45599ce
commit 5aa7c4a
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 20 deletions.
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -21,7 +21,9 @@ use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::tree_node::Transformed;
 use datafusion_common::Result;
 use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
+use datafusion_expr::logical_plan::tree_node::unwrap_arc;
 use datafusion_expr::{Distinct, LogicalPlan, Union};
+use itertools::Itertools;
 use std::sync::Arc;
 
 #[derive(Default)]
@@ -56,53 +58,55 @@ impl OptimizerRule for EliminateNestedUnion {
         match plan {
             LogicalPlan::Union(Union { inputs, schema }) => {
                 let inputs = inputs
-                    .iter()
+                    .into_iter()
                     .flat_map(extract_plans_from_union)
                     .collect::<Vec<_>>();
 
                 Ok(Transformed::yes(LogicalPlan::Union(Union {
-                    inputs,
+                    inputs: inputs.into_iter().map(Arc::new).collect_vec(),
                     schema,
                 })))
             }
-            LogicalPlan::Distinct(Distinct::All(ref nested_plan)) => {
-                match nested_plan.as_ref() {
+            LogicalPlan::Distinct(Distinct::All(nested_plan)) => {
+                match unwrap_arc(nested_plan) {
                     LogicalPlan::Union(Union { inputs, schema }) => {
                         let inputs = inputs
-                            .iter()
+                            .into_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,
+                                inputs: inputs.into_iter().map(Arc::new).collect_vec(),
                                 schema: schema.clone(),
                             })),
                         ))))
                     }
-                    _ => Ok(Transformed::no(plan)),
+                    nested_plan => Ok(Transformed::no(LogicalPlan::Distinct(
+                        Distinct::All(Arc::new(nested_plan)),
+                    ))),
                 }
             }
             _ => Ok(Transformed::no(plan)),
         }
     }
 }
 
-fn extract_plans_from_union(plan: &Arc<LogicalPlan>) -> Vec<Arc<LogicalPlan>> {
-    match plan.as_ref() {
+fn extract_plans_from_union(plan: Arc<LogicalPlan>) -> Vec<LogicalPlan> {
+    match unwrap_arc(plan) {
         LogicalPlan::Union(Union { inputs, schema }) => inputs
-            .iter()
-            .map(|plan| Arc::new(coerce_plan_expr_for_schema(plan, schema).unwrap()))
+            .into_iter()
+            .map(|plan| coerce_plan_expr_for_schema(&plan, &schema).unwrap())
             .collect::<Vec<_>>(),
-        _ => vec![plan.clone()],
+        plan => vec![plan],
     }
 }
 
-fn extract_plan_from_distinct(plan: &Arc<LogicalPlan>) -> &Arc<LogicalPlan> {
-    match plan.as_ref() {
+fn extract_plan_from_distinct(plan: Arc<LogicalPlan>) -> Arc<LogicalPlan> {
+    match unwrap_arc(plan) {
         LogicalPlan::Distinct(Distinct::All(plan)) => plan,
-        _ => plan,
+        plan => Arc::new(plan),
     }
 }
 

diff --git a/datafusion/optimizer/src/eliminate_outer_join.rs b/datafusion/optimizer/src/eliminate_outer_join.rs
@@ -18,6 +18,7 @@
 //! [`EliminateOuterJoin`] converts `LEFT/RIGHT/FULL` joins to `INNER` joins
 use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::{Column, DFSchema, Result};
+use datafusion_expr::logical_plan::tree_node::unwrap_arc;
 use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan};
 use datafusion_expr::{Expr, Filter, Operator};
 
@@ -78,7 +79,7 @@ impl OptimizerRule for EliminateOuterJoin {
         _config: &dyn OptimizerConfig,
     ) -> Result<Transformed<LogicalPlan>> {
         match plan {
-            LogicalPlan::Filter(filter) => match filter.input.as_ref() {
+            LogicalPlan::Filter(mut filter) => match unwrap_arc(filter.input) {
                 LogicalPlan::Join(join) => {
                     let mut non_nullable_cols: Vec<Column> = vec![];
 
@@ -109,9 +110,10 @@ impl OptimizerRule for EliminateOuterJoin {
                     } else {
                         join.join_type
                     };
+
                     let new_join = Arc::new(LogicalPlan::Join(Join {
-                        left: Arc::new((*join.left).clone()),
-                        right: Arc::new((*join.right).clone()),
+                        left: join.left,
+                        right: join.right,
                         join_type: new_join_type,
                         join_constraint: join.join_constraint,
                         on: join.on.clone(),
@@ -122,7 +124,10 @@ impl OptimizerRule for EliminateOuterJoin {
                     Filter::try_new(filter.predicate, new_join)
                         .map(|f| Transformed::yes(LogicalPlan::Filter(f)))
                 }
-                _ => Ok(Transformed::no(LogicalPlan::Filter(filter))),
+                filter_input => {
+                    filter.input = Arc::new(filter_input);
+                    Ok(Transformed::no(LogicalPlan::Filter(filter)))
+                }
             },
             _ => Ok(Transformed::no(plan)),
         }

diff --git a/datafusion/optimizer/src/propagate_empty_relation.rs b/datafusion/optimizer/src/propagate_empty_relation.rs
@@ -182,7 +182,9 @@ impl OptimizerRule for PropagateEmptyRelation {
                         },
                     )))
                 } else if new_inputs.len() == 1 {
-                    let child = unwrap_arc(new_inputs[0].clone());
+                    let mut new_inputs = new_inputs;
+                    let input_plan = new_inputs.pop().unwrap(); // length checked
+                    let child = unwrap_arc(input_plan);
                     if child.schema().eq(plan.schema()) {
                         Ok(Transformed::yes(child))
                     } else {