Implement TreeNode::map_children in place

Author: alamb

datafusion/expr/src/logical_plan/ddl.rs

@@ -112,6 +112,24 @@ 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/plan.rs

@@ -20,7 +20,7 @@
 use std::collections::{HashMap, HashSet};
 use std::fmt::{self, Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 use super::dml::CopyTo;
 use super::DdlStatement;
@@ -1131,6 +1131,159 @@ impl LogicalPlan {
         })?;
         Ok(())
     }
+}
+
+// TODO put this somewhere better than here
+
+/// 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();
+
+/// 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
+fn rewrite_arc<F>(node: &mut Arc<LogicalPlan>, mut f: F) -> Result<Transformed<()>>
+where
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
+{
+    // We need to leave a valid node in the Arc, while we rewrite the existing
+    // one, so use a single global static placeholder node
+    let mut new_node = PLACEHOLDER
+        .get_or_init(|| {
+            Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
+                produce_one_row: false,
+                schema: DFSchemaRef::new(DFSchema::empty()),
+            }))
+        })
+        .clone();
+
+    // take the old value out of the Arc
+    std::mem::swap(node, &mut new_node);
+
+    // try to update existing node, if it isn't shared with others
+    let mut new_node = Arc::try_unwrap(new_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());
+
+    // apply the actual transform
+    let result = f(new_node)?;
+
+    // put the new value back into the Arc
+    let mut new_node = Arc::new(result.data);
+    std::mem::swap(node, &mut new_node);
+
+    // return the `()` back
+    Ok(Transformed::new((), result.transformed, result.tnr))
+}
+
+/*
+/// Rewrties all inputs for an Extension node "in place"
+/// (it currently has to copy values because there are no APIs for in place modification)
+///
+/// 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>,
+    mut f: F,
+) -> Result<Transformed<()>>
+where
+    F: FnMut(&mut LogicalPlan) -> Result<Transformed<()>>,
+{
+    let mut inputs: Vec<_> = node.inputs().into_iter().cloned().collect();
+
+    let result = inputs
+        .iter_mut()
+        .try_fold(Transformed::no(()), |acc, input| acc.and_then(|| f(input)))?;
+    let exprs = node.expressions();
+    let mut new_node = node.from_template(&exprs, &inputs);
+    std::mem::swap(node, &mut new_node);
+    Ok(result)
+}
+*/
+
+impl LogicalPlan {
+    /// applies `f` to each input of this plan node, rewriting them in place.
+    ///
+    /// # Notes
+    /// Inputs include both direct children as well as any embedded subquery
+    /// `LogicalPlan`s, for example such as are in [`Expr::Exists`].
+    ///
+    /// If `f` returns an `Err`, that Err is returned, and the inputs are left
+    /// in a partially modified state
+    pub(crate) fn rewrite_children<F>(&mut self, mut f: F) -> Result<Transformed<()>>
+    where
+        F: FnMut(Self) -> Result<Transformed<Self>>,
+    {
+        let children_result = match self {
+            LogicalPlan::Projection(Projection { input, .. }) => {
+                rewrite_arc(input, &mut f)
+            }
+            LogicalPlan::Filter(Filter { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Repartition(Repartition { input, .. }) => {
+                rewrite_arc(input, &mut f)
+            }
+            LogicalPlan::Window(Window { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Aggregate(Aggregate { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Sort(Sort { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Join(Join { left, right, .. }) => {
+                rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f))
+            }
+            LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => {
+                rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f))
+            }
+            LogicalPlan::Limit(Limit { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Subquery(Subquery { subquery, .. }) => {
+                rewrite_arc(subquery, &mut f)
+            }
+            LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => {
+                rewrite_arc(input, &mut f)
+            }
+            LogicalPlan::Extension(extension) => {
+                todo!();
+                //rewrite_extension_inputs(&mut extension.node, &mut f)
+            }
+            LogicalPlan::Union(Union { inputs, .. }) => inputs
+                .iter_mut()
+                .try_fold(Transformed::no(()), |acc, input| {
+                    acc.and_then(|| rewrite_arc(input, &mut f))
+                }),
+            LogicalPlan::Distinct(
+                Distinct::All(input) | Distinct::On(DistinctOn { input, .. }),
+            ) => rewrite_arc(input, &mut f),
+            LogicalPlan::Explain(explain) => rewrite_arc(&mut explain.plan, &mut f),
+            LogicalPlan::Analyze(analyze) => rewrite_arc(&mut analyze.input, &mut f),
+            LogicalPlan::Dml(write) => rewrite_arc(&mut write.input, &mut f),
+            LogicalPlan::Copy(copy) => rewrite_arc(&mut copy.input, &mut f),
+            LogicalPlan::Ddl(ddl) => {
+                if let Some(input) = ddl.input_mut() {
+                    rewrite_arc(input, &mut f)
+                } else {
+                    Ok(Transformed::no(()))
+                }
+            }
+            LogicalPlan::Unnest(Unnest { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::Prepare(Prepare { input, .. }) => rewrite_arc(input, &mut f),
+            LogicalPlan::RecursiveQuery(RecursiveQuery {
+                static_term,
+                recursive_term,
+                ..
+            }) => rewrite_arc(static_term, &mut f)?
+                .and_then(|| rewrite_arc(recursive_term, &mut f)),
+            // plans without inputs
+            LogicalPlan::TableScan { .. }
+            | LogicalPlan::Statement { .. }
+            | LogicalPlan::EmptyRelation { .. }
+            | LogicalPlan::Values { .. }
+            | LogicalPlan::DescribeTable(_) => Ok(Transformed::no(())),
+        }?;
+
+        // after visiting the actual children we we need to visit any subqueries
+        // that are inside the expressions
+        children_result.and_then(|| self.rewrite_subqueries(&mut f))
+    }
 
     /// Return a `LogicalPlan` with all placeholders (e.g $1 $2,
     /// ...) replaced with corresponding values provided in

datafusion/expr/src/tree_node/plan.rs

@@ -20,7 +20,7 @@
 use crate::LogicalPlan;
 
 use datafusion_common::tree_node::{
-    Transformed, TransformedIterator, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
+    Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
 };
 use datafusion_common::{handle_visit_recursion, Result};
 
@@ -91,23 +91,14 @@ impl TreeNode for LogicalPlan {
         Ok(tnr)
     }
 
-    fn map_children<F>(self, f: F) -> Result<Transformed<Self>>
+    fn map_children<F>(mut self, f: F) -> Result<Transformed<Self>>
     where
         F: FnMut(Self) -> Result<Transformed<Self>>,
     {
-        let new_children = self
-            .inputs()
-            .iter()
-            .map(|&c| c.clone())
-            .map_until_stop_and_collect(f)?;
-        // Propagate up `new_children.transformed` and `new_children.tnr`
-        // along with the node containing transformed children.
-        if new_children.transformed {
-            new_children.map_data(|new_children| {
-                self.with_new_exprs(self.expressions(), new_children)
-            })
-        } else {
-            Ok(new_children.update_data(|_| self))
-        }
+        // Apply the rewrites in place for each child
+        let result = self.rewrite_children(f)?;
+
+        // return a reference to ourself
+        Ok(result.update_data(|_| self))
     }
 }