compiles

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/plan.rs

@@ -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,
@@ -1144,7 +1145,10 @@ static PLACEHOLDER: OnceLock<Arc<LogicalPlan>> = OnceLock::new();
 /// 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<()>>
+fn rewrite_arc<F>(
+    node: &mut Arc<LogicalPlan>,
+    mut f: F,
+) -> Result<Transformed<&mut Arc<LogicalPlan>>>
 where
     F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
 {
@@ -1163,7 +1167,7 @@ where
     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)
+    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());
@@ -1175,37 +1179,20 @@ where
     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))
+    // return the `node` back
+    Ok(Transformed::new(node, 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<()>>
+/// 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(&mut LogicalPlan) -> Result<Transformed<()>>,
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
 {
-    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)
+    rewrite_arc(node, f).map(|res| res.discard_data())
 }
-*/
 
 impl LogicalPlan {
-    /// applies `f` to each input of this plan node, rewriting them in place.
+    /// 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
@@ -1219,59 +1206,84 @@ impl LogicalPlan {
     {
         let children_result = match self {
             LogicalPlan::Projection(Projection { input, .. }) => {
-                rewrite_arc(input, &mut f)
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Filter(Filter { input, .. }) => {
+                rewrite_arc_no_data(input, &mut f)
             }
-            LogicalPlan::Filter(Filter { input, .. }) => rewrite_arc(input, &mut f),
             LogicalPlan::Repartition(Repartition { input, .. }) => {
-                rewrite_arc(input, &mut f)
+                rewrite_arc_no_data(input, &mut f)
+            }
+            LogicalPlan::Window(Window { input, .. }) => {
+                rewrite_arc_no_data(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::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, .. }) => {
-                rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f))
+                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, .. }) => {
-                rewrite_arc(left, &mut f)?.and_then(|| rewrite_arc(right, &mut f))
+                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(input, &mut f),
+            LogicalPlan::Limit(Limit { input, .. }) => rewrite_arc_no_data(input, &mut f),
             LogicalPlan::Subquery(Subquery { subquery, .. }) => {
-                rewrite_arc(subquery, &mut f)
+                rewrite_arc_no_data(subquery, &mut f)
             }
             LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) => {
-                rewrite_arc(input, &mut f)
+                rewrite_arc_no_data(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::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(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),
+            ) => 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(input, &mut f)
+                    rewrite_arc_no_data(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::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,
                 ..
-            }) => rewrite_arc(static_term, &mut f)?
-                .and_then(|| rewrite_arc(recursive_term, &mut f)),
+            }) => {
+                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 { .. }
@@ -1282,7 +1294,8 @@ impl LogicalPlan {
 
         // 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))
+        // children_result.and_then(|| self.rewrite_subqueries(&mut f))
+        Ok(children_result)
     }
 
     /// Return a `LogicalPlan` with all placeholders (e.g $1 $2,