fix recursion

Author: alamb

datafusion/common/src/tree_node.rs

@@ -215,8 +215,8 @@ pub trait TreeNode: Sized {
         // Traverse children and then call f_up on self if necessary
         match pre_visited.tnr {
             TreeNodeRecursion::Continue => {
-                // recurse to children and then apply f_up to self if needed
-                self.mutate_children(|c| mutator.f_down(c))?
+                // rewrite children recursively with mutator
+                self.mutate_children(|c| c.mutate(mutator))?
                     .try_transform_node_with(
                         |_: ()| mutator.f_up(self),
                         TreeNodeRecursion::Jump,

datafusion/expr/src/logical_plan/mutate.rs

@@ -193,7 +193,10 @@ where
 }
 
 impl LogicalPlan {
-    /// applies the closure `f` to each input of this node, in place
+    /// applies the closure `f` to each input of this node, in place.
+    ///
+    /// Inputs include both direct children as well as any embedded
+    /// `LogicalPlan`s embedded in `Expr::Exists`, etc.
     ///
     /// If Err is returned, the inputs may be in a partially modified state
     pub fn rewrite_inputs<F>(&mut self, mut f: F) -> Result<Transformed<()>>
@@ -225,11 +228,13 @@ impl LogicalPlan {
                 rewrite_arc(input, &mut f)
             }
             LogicalPlan::Extension(extension) => todo!(),
-            LogicalPlan::Union(Union { inputs, .. }) => inputs
-                .iter_mut()
-                .fold(Ok(Transformed::no(())), |acc, input| {
-                    acc.and_then(|acc| rewrite_arc(input, &mut f))
-                }),
+            LogicalPlan::Union(Union { inputs, .. }) => {
+                inputs
+                    .iter_mut()
+                    .fold(Ok(Transformed::no(())), |acc, input| {
+                        acc.and_then(|acc| acc.and_then(|| rewrite_arc(input, &mut f)))
+                    })
+            }
             LogicalPlan::Distinct(
                 Distinct::All(input) | Distinct::On(DistinctOn { input, .. }),
             ) => rewrite_arc(input, &mut f),

datafusion/optimizer/src/optimizer.rs

@@ -185,11 +185,11 @@ pub struct Optimizer {
     pub rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
 }
 
-/// If a rule is with `ApplyOrder`, it means the optimizer will derive to handle children instead of
-/// recursively handling in rule.
-/// We just need handle a subtree pattern itself.
+/// If a rule is with `ApplyOrder`, it means the optimizer will handle
+/// recursion. If it is `None` the rule must handle any required recursion
 ///
-/// Notice: **sometime** result after optimize still can be optimized, we need apply again.
+/// Notice: **sometime** result after optimize still can be optimized, we need
+/// apply again.
 ///
 /// Usage Example: Merge Limit (subtree pattern is: Limit-Limit)
 /// ```rust
@@ -345,7 +345,7 @@ fn optimize_in_place<'a>(
     rule.try_optimize(plan, config).map(|maybe_plan| {
         match maybe_plan {
             Some(new_plan) => {
-                *plan = new_plan; // can avoid this copy with better OptimizerRule::try_optimize
+                *plan = new_plan; // TODO avoid this copy with better OptimizerRule::try_optimize
                 Transformed::yes(())
             }
             None => Transformed::no(()),
@@ -378,7 +378,8 @@ impl Optimizer {
             log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
 
             for rule in &self.rules {
-                // If we need to skip failed rules, must copy plan
+                // If we need to skip failed rules, must copy plan before attempting to rewrite
+                // as rewriting is destructive
                 let prev_plan = options
                     .optimizer
                     .skip_failed_rules