Implement TreeNode::map_children in place

Author: alamb

datafusion/common/src/tree_node.rs

@@ -530,6 +530,15 @@ impl<T> Transformed<T> {
             })
         }
     }
+
+    /// Discards the data of this [`Transformed`] object transforming it into Transformed<()>
+    pub fn discard_data(self) -> Transformed<()> {
+        Transformed {
+            data: (),
+            transformed: self.transformed,
+            tnr: self.tnr,
+        }
+    }
 }
 
 /// Transformation helper to process a sequence of iterable tree nodes that are siblings.

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;
@@ -45,7 +45,8 @@ use crate::{
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::tree_node::{
-    Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
+    Transformed, TransformedIterator, TransformedResult, TreeNode, TreeNodeRecursion,
+    TreeNodeVisitor,
 };
 use datafusion_common::{
     aggregate_functional_dependencies, internal_err, plan_err, Column, Constraints,
@@ -1131,6 +1132,202 @@ 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<&mut Arc<LogicalPlan>>>
+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 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 `node` back
+    Ok(Transformed::new(node, result.transformed, result.tnr))
+}
+
+/// Rewrite the arc and discard the contents of Transformed
+fn rewrite_arc_no_data<F>(node: &mut Arc<LogicalPlan>, f: F) -> Result<Transformed<()>>
+where
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
+{
+    rewrite_arc(node, f).map(|res| res.discard_data())
+}
+
+/// Rewrites 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>,
+    f: F,
+) -> Result<Transformed<()>>
+where
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
+{
+    let Transformed {
+        data: new_inputs,
+        transformed,
+        tnr,
+    } = 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,
+    })
+}
+
+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_no_data(input, &mut f)
+            }
+            LogicalPlan::Filter(Filter { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Repartition(Repartition { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Window(Window { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Aggregate(Aggregate { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Sort(Sort { input, .. }) => rewrite_arc_no_data(input, &mut f),
+            LogicalPlan::Join(Join { left, right, .. }) => {
+                let results = [left, right]
+                    .into_iter()
+                    .map_until_stop_and_collect(|input| rewrite_arc(input, &mut f))?;
+                Ok(results.discard_data())
+            }
+            LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => {
+                let results = [left, right]
+                    .into_iter()
+                    .map_until_stop_and_collect(|input| rewrite_arc(input, &mut f))?;
+                Ok(results.discard_data())
+            }
+            LogicalPlan::Limit(Limit { input, .. }) => rewrite_arc_no_data(input, &mut f),
+            LogicalPlan::Subquery(Subquery { subquery, .. }) => {
+                rewrite_arc_no_data(subquery, &mut f)
+            }
+            LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Extension(extension) => {
+                rewrite_extension_inputs(&mut extension.node, &mut f)
+            }
+            LogicalPlan::Union(Union { inputs, .. }) => {
+                let results = inputs
+                    .iter_mut()
+                    .map_until_stop_and_collect(|input| rewrite_arc(input, &mut f))?;
+                Ok(results.discard_data())
+            }
+            LogicalPlan::Distinct(
+                Distinct::All(input) | Distinct::On(DistinctOn { input, .. }),
+            ) => rewrite_arc_no_data(input, &mut f),
+            LogicalPlan::Explain(explain) => {
+                rewrite_arc_no_data(&mut explain.plan, &mut f)
+            }
+            LogicalPlan::Analyze(analyze) => {
+                rewrite_arc_no_data(&mut analyze.input, &mut f)
+            }
+            LogicalPlan::Dml(write) => rewrite_arc_no_data(&mut write.input, &mut f),
+            LogicalPlan::Copy(copy) => rewrite_arc_no_data(&mut copy.input, &mut f),
+            LogicalPlan::Ddl(ddl) => {
+                if let Some(input) = ddl.input_mut() {
+                    rewrite_arc_no_data(input, &mut f)
+                } else {
+                    Ok(Transformed::no(()))
+                }
+            }
+            LogicalPlan::Unnest(Unnest { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Prepare(Prepare { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::RecursiveQuery(RecursiveQuery {
+                static_term,
+                recursive_term,
+                ..
+            }) => {
+                let results = [static_term, recursive_term]
+                    .into_iter()
+                    .map_until_stop_and_collect(|input| rewrite_arc(input, &mut f))?;
+                Ok(results.discard_data())
+            }
+            // 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))
+        Ok(children_result)
+    }
 
     /// 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))
     }
 }