rewrite to be more functional

Author: alamb

datafusion/expr/src/logical_plan/ddl.rs

@@ -110,24 +110,6 @@ impl DdlStatement {
         }
     }
 
-    /// Return a mutable reference to the input `LogicalPlan`, if any
-    pub fn input_mut(&mut self) -> Option<&mut Arc<LogicalPlan>> {
-        match self {
-            DdlStatement::CreateMemoryTable(CreateMemoryTable { input, .. }) => {
-                Some(input)
-            }
-            DdlStatement::CreateExternalTable(_) => None,
-            DdlStatement::CreateView(CreateView { input, .. }) => Some(input),
-            DdlStatement::CreateCatalogSchema(_) => None,
-            DdlStatement::CreateCatalog(_) => None,
-            DdlStatement::DropTable(_) => None,
-            DdlStatement::DropView(_) => None,
-            DdlStatement::DropCatalogSchema(_) => None,
-            DdlStatement::CreateFunction(_) => None,
-            DdlStatement::DropFunction(_) => None,
-        }
-    }
-
     /// Return a `format`able structure with the a human readable
     /// description of this LogicalPlan node per node, not including
     /// children.

datafusion/expr/src/logical_plan/rewrite.rs

@@ -19,39 +19,51 @@
 
 use crate::dml::CopyTo;
 use crate::{
-    Aggregate, Analyze, CrossJoin, Distinct, DistinctOn, DmlStatement, EmptyRelation,
-    Explain, Filter, Join, Limit, LogicalPlan, Prepare, Projection, RecursiveQuery,
-    Repartition, Sort, Subquery, SubqueryAlias, Union, Unnest, UserDefinedLogicalNode,
-    Window,
+    Aggregate, Analyze, CreateMemoryTable, CreateView, CrossJoin, DdlStatement, Distinct,
+    DistinctOn, DmlStatement, Explain, Extension, Filter, Join, Limit, LogicalPlan,
+    Prepare, Projection, RecursiveQuery, Repartition, Sort, Subquery, SubqueryAlias,
+    Union, Unnest, Window,
 };
-use datafusion_common::tree_node::{Transformed, TreeNodeIterator};
-use datafusion_common::{map_until_stop_and_collect, DFSchema, DFSchemaRef, Result};
-use std::sync::{Arc, OnceLock};
+use datafusion_common::tree_node::{Transformed, TreeNodeIterator, TreeNodeRecursion};
+use datafusion_common::{map_until_stop_and_collect, Result};
+use std::sync::Arc;
 
-/// A temporary node that is left in place while rewriting the children of a
-/// [`LogicalPlan`]. This is necessary to ensure that the `LogicalPlan` is
-/// always in a valid state (from the Rust perspective)
-static PLACEHOLDER: OnceLock<Arc<LogicalPlan>> = OnceLock::new();
+/// Converts a `Arc<LogicalPlan>` without copying, if possible. Copies the plan
+/// if there is a shared reference
+fn unwrap_arc(plan: Arc<LogicalPlan>) -> LogicalPlan {
+    Arc::try_unwrap(plan)
+        // if None is returned, there is another reference to this
+        // LogicalPlan, so we can not own it, and must clone instead
+        .unwrap_or_else(|node| node.as_ref().clone())
+}
 
-/// its inputs, so this code would not be needed. However, for now we try and
-/// unwrap the `Arc` which avoids `clone`ing in most cases.
-///
-/// On error, node be left with a placeholder logical plan
+/// Applies `f` to rewrite a `Arc<LogicalPlan>` without copying, if possible
 fn rewrite_arc<F>(
-    node: Arc<LogicalPlan>,
+    plan: Arc<LogicalPlan>,
     mut f: F,
 ) -> Result<Transformed<Arc<LogicalPlan>>>
 where
     F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
 {
-    // try to update existing node, if it isn't shared with others
-    let new_node = Arc::try_unwrap(node)
-        // if None is returned, there is another reference to this
-        // LogicalPlan, so we must clone instead
-        .unwrap_or_else(|node| node.as_ref().clone());
+    f(unwrap_arc(plan))?.map_data(|new_plan| Ok(Arc::new(new_plan)))
+}
 
-    // apply the actual transform
-    f(new_node).unwrap().map_data(Arc::new)
+/// rewrite a `Vec` of `Arc<LogicalPlan>` without copying, if possible
+fn rewrite_arcs<F>(
+    input_plans: Vec<Arc<LogicalPlan>>,
+    mut f: F,
+) -> Result<Transformed<Vec<Arc<LogicalPlan>>>>
+where
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
+{
+    Ok(input_plans
+        .into_iter()
+        .map(unwrap_arc)
+        .map_until_stop_and_collect(&mut f)?
+        .update_data(|inputs| {
+            // have a Vec<LogicalPlan>, now need to wrap in `Arc`s again
+            inputs.into_iter().map(Arc::new).collect::<Vec<_>>()
+        }))
 }
 
 /// Rewrites all inputs for an Extension node "in place"
@@ -60,30 +72,24 @@ where
 /// Should be removed when we have an API for in place modifications of the
 /// extension to avoid these copies
 fn rewrite_extension_inputs<F>(
-    node: &mut Arc<dyn UserDefinedLogicalNode>,
+    extension: Extension,
     f: F,
-) -> datafusion_common::Result<Transformed<()>>
+) -> Result<Transformed<Extension>>
 where
     F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
 {
-    let Transformed {
-        data: new_inputs,
-        transformed,
-        tnr,
-    } = node
-        .inputs()
+    let Extension { node } = extension;
+
+    node.inputs()
         .into_iter()
         .cloned()
-        .map_until_stop_and_collect(f)?;
-
-    let exprs = node.expressions();
-    let mut new_node = node.from_template(&exprs, &new_inputs);
-    std::mem::swap(node, &mut new_node);
-    Ok(Transformed {
-        data: (),
-        transformed,
-        tnr,
-    })
+        .map_until_stop_and_collect(f)?
+        .map_data(|new_inputs| {
+            let exprs = node.expressions();
+            Ok(Extension {
+                node: node.from_template(&exprs, &new_inputs),
+            })
+        })
 }
 
 impl LogicalPlan {
@@ -192,12 +198,12 @@ impl LogicalPlan {
             }),
             LogicalPlan::Limit(Limit { skip, fetch, input }) => {
                 rewrite_arc(input, &mut f)?
-                    .map_data(|input| LogicalPlan::Limit(Limit { skip, fetch, input }))
+                    .update_data(|input| LogicalPlan::Limit(Limit { skip, fetch, input }))
             }
             LogicalPlan::Subquery(Subquery {
                 subquery,
                 outer_ref_columns,
-            }) => rewrite_arc(subquery, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(subquery, &mut f)?.update_data(|subquery| {
                 LogicalPlan::Subquery(Subquery {
                     subquery,
                     outer_ref_columns,
@@ -214,16 +220,14 @@ impl LogicalPlan {
                     schema,
                 })
             }),
-            LogicalPlan::Extension(extension) => {
-                rewrite_extension_inputs(&mut extension.node, &mut f)
-            }
-            LogicalPlan::Union(Union { inputs, schema }) => inputs
-                .into_iter()
-                .map_until_stop_and_collect(&mut f)?
+            LogicalPlan::Extension(extension) => rewrite_extension_inputs(extension, f)?
+                .update_data(LogicalPlan::Extension),
+            LogicalPlan::Union(Union { inputs, schema }) => rewrite_arcs(inputs, &mut f)?
                 .update_data(|inputs| LogicalPlan::Union(Union { inputs, schema })),
             LogicalPlan::Distinct(distinct) => match distinct {
-                Distinct::All(input) => rewrite_arc(input, &mut f)?
-                    .update_data(|input| Distinct::All(input))?,
+                Distinct::All(input) => {
+                    rewrite_arc(input, &mut f)?.update_data(Distinct::All)
+                }
                 Distinct::On(DistinctOn {
                     on_expr,
                     select_expr,
@@ -238,16 +242,16 @@ impl LogicalPlan {
                         input,
                         schema,
                     })
-                })?,
+                }),
             }
-            .map_data(|distinct| LogicalPlan::Distinct(distinct)),
+            .update_data(LogicalPlan::Distinct),
             LogicalPlan::Explain(Explain {
                 verbose,
                 plan,
                 stringified_plans,
                 schema,
                 logical_optimization_succeeded,
-            }) => rewrite_arc(plan, &mut f)?.map_data(|plan| {
+            }) => rewrite_arc(plan, &mut f)?.update_data(|plan| {
                 LogicalPlan::Explain(Explain {
                     verbose,
                     plan,
@@ -260,7 +264,7 @@ impl LogicalPlan {
                 verbose,
                 input,
                 schema,
-            }) => rewrite_arc(input, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(input, &mut f)?.update_data(|input| {
                 LogicalPlan::Analyze(Analyze {
                     verbose,
                     input,
@@ -272,7 +276,7 @@ impl LogicalPlan {
                 table_schema,
                 op,
                 input,
-            }) => rewrite_arc(&input, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(input, &mut f)?.update_data(|input| {
                 LogicalPlan::Dml(DmlStatement {
                     table_name,
                     table_schema,
@@ -286,7 +290,7 @@ impl LogicalPlan {
                 partition_by,
                 format_options,
                 options,
-            }) => rewrite_arc(input, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(input, &mut f)?.update_data(|input| {
                 LogicalPlan::Copy(CopyTo {
                     input,
                     output_url,
@@ -296,18 +300,55 @@ impl LogicalPlan {
                 })
             }),
             LogicalPlan::Ddl(ddl) => {
-                if let Some(input) = ddl.input_mut() {
-                    rewrite_arc(input, &mut f)
-                } else {
-                    Ok(Transformed::no(()))
+                match ddl {
+                    DdlStatement::CreateMemoryTable(CreateMemoryTable {
+                        name,
+                        constraints,
+                        input,
+                        if_not_exists,
+                        or_replace,
+                        column_defaults,
+                    }) => rewrite_arc(input, f)?.update_data(|input| {
+                        DdlStatement::CreateMemoryTable(CreateMemoryTable {
+                            name,
+                            constraints,
+                            input,
+                            if_not_exists,
+                            or_replace,
+                            column_defaults,
+                        })
+                    }),
+                    DdlStatement::CreateView(CreateView {
+                        name,
+                        input,
+                        or_replace,
+                        definition,
+                    }) => rewrite_arc(input, f)?.update_data(|input| {
+                        DdlStatement::CreateView(CreateView {
+                            name,
+                            input,
+                            or_replace,
+                            definition,
+                        })
+                    }),
+                    // no inputs in these statements
+                    DdlStatement::CreateExternalTable(_)
+                    | DdlStatement::CreateCatalogSchema(_)
+                    | DdlStatement::CreateCatalog(_)
+                    | DdlStatement::DropTable(_)
+                    | DdlStatement::DropView(_)
+                    | DdlStatement::DropCatalogSchema(_)
+                    | DdlStatement::CreateFunction(_)
+                    | DdlStatement::DropFunction(_) => Transformed::no(ddl),
                 }
+                .update_data(LogicalPlan::Ddl)
             }
             LogicalPlan::Unnest(Unnest {
                 input,
                 column,
                 schema,
                 options,
-            }) => rewrite_arc(input, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(input, &mut f)?.update_data(|input| {
                 LogicalPlan::Unnest(Unnest {
                     input,
                     column,
@@ -319,7 +360,7 @@ impl LogicalPlan {
                 name,
                 data_types,
                 input,
-            }) => rewrite_arc(input, &mut f)?.map_data(|input| {
+            }) => rewrite_arc(input, &mut f)?.update_data(|input| {
                 LogicalPlan::Prepare(Prepare {
                     name,
                     data_types,
@@ -336,7 +377,7 @@ impl LogicalPlan {
                 recursive_term,
                 rewrite_arc(recursive_term, &mut f)
             )?
-            .map_data(|(static_term, recursive_term)| {
+            .update_data(|(static_term, recursive_term)| {
                 LogicalPlan::RecursiveQuery(RecursiveQuery {
                     name,
                     static_term,
@@ -349,7 +390,7 @@ impl LogicalPlan {
             | LogicalPlan::Statement { .. }
             | LogicalPlan::EmptyRelation { .. }
             | LogicalPlan::Values { .. }
-            | LogicalPlan::DescribeTable(_) => Ok(Transformed::no(self)),
+            | LogicalPlan::DescribeTable(_) => Transformed::no(self),
         })
     }
 }

datafusion/expr/src/tree_node/plan.rs

@@ -32,14 +32,10 @@ impl TreeNode for LogicalPlan {
         self.inputs().into_iter().apply_until_stop(f)
     }
 
-    fn map_children<F>(mut self, f: F) -> Result<Transformed<Self>>
+    fn map_children<F>(self, f: F) -> Result<Transformed<Self>>
     where
         F: FnMut(Self) -> Result<Transformed<Self>>,
     {
-        // Apply the rewrite *in place* for each child to avoid cloning
-        let result = self.map_inputs(f)?;
-
-        // return ourself
-        Ok(result.update_data(|_| self))
+        self.map_inputs(f)
     }
 }