don't trust t.transformed coming from transformation closures, keep…

… the old way of detecting if changes were made

datafusion/common/src/tree_node.rs

@@ -479,9 +479,13 @@ impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
         let children = self.arc_children();
         if !children.is_empty() {
             let t = children.into_iter().map_till_continue_and_collect(f)?;
-            // TODO: once we trust `t.transformed` don't create new node if not necessary
+            // TODO: Currently `t.transformed` quality comes from if the transformation
+            //  closures fill the field correctly. Once we trust `t.transformed` we can
+            //  remove the additional `t2` check.
+            //  Please note that we need to propagate up `t.tnr` though.
             let arc_self = Arc::clone(&self);
-            self.with_new_arc_children(arc_self, t.data)
+            let t2 = self.with_new_arc_children(arc_self, t.data)?;
+            Ok(Transformed::new(t2.data, t2.transformed, t.tnr))
         } else {
             Ok(Transformed::no(self))
         }
@@ -499,7 +503,7 @@ pub trait ConcreteTreeNode: Sized {
     fn take_children(self) -> (Self, Vec<Self>);
 
     /// Reattaches updated child nodes to the node, returning the updated node.
-    fn with_new_children(self, children: Vec<Self>) -> Result<Self>;
+    fn with_new_children(self, children: Vec<Self>) -> Result<Transformed<Self>>;
 }
 
 impl<T: ConcreteTreeNode> TreeNode for T {
@@ -520,11 +524,13 @@ impl<T: ConcreteTreeNode> TreeNode for T {
     {
         let (new_self, children) = self.take_children();
         if !children.is_empty() {
-            children
-                .into_iter()
-                .map_till_continue_and_collect(f)?
-                // TODO: once we trust `transformed` don't create new node if not necessary
-                .flat_map_data(|new_children| new_self.with_new_children(new_children))
+            let t = children.into_iter().map_till_continue_and_collect(f)?;
+            // TODO: Currently `t.transformed` quality comes from if the transformation
+            //  closures fill the field correctly. Once we trust `t.transformed` we can
+            //  remove the additional `t2` check.
+            //  Please note that we need to propagate up `t.tnr` though.
+            let t2 = new_self.with_new_children(t.data)?;
+            Ok(Transformed::new(t2.data, t2.transformed, t.tnr))
         } else {
             Ok(Transformed::no(new_self))
         }

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -1030,7 +1030,7 @@ fn replace_order_preserving_variants(
         }
     }
 
-    context.update_plan_from_children()
+    context.update_plan_from_children().map(|t| t.data)
 }
 
 /// This utility function adds a [`SortExec`] above an operator according to the

datafusion/core/src/physical_optimizer/enforce_sorting.rs

@@ -113,7 +113,7 @@ fn update_sort_ctx_children(
     }
 
     node.data = data;
-    node.update_plan_from_children()
+    node.update_plan_from_children().map(|t| t.data)
 }
 
 /// This object is used within the [`EnforceSorting`] rule to track the closest
@@ -477,7 +477,7 @@ fn remove_corresponding_coalesce_in_sub_plan(
             .collect::<Result<_>>()?;
     }
 
-    requirements.update_plan_from_children()
+    requirements.update_plan_from_children().map(|t| t.data)
 }
 
 /// Updates child to remove the unnecessary sort below it.

datafusion/core/src/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -138,7 +138,7 @@ fn plan_with_order_preserving_variants(
         }
     }
 
-    sort_input.update_plan_from_children()
+    sort_input.update_plan_from_children().map(|t| t.data)
 }
 
 /// Calculates the updated plan by replacing operators that preserve ordering
@@ -184,7 +184,7 @@ fn plan_with_order_breaking_variants(
         let coalesce = CoalescePartitionsExec::new(child);
         sort_input.plan = Arc::new(coalesce) as _;
     } else {
-        return sort_input.update_plan_from_children();
+        return sort_input.update_plan_from_children().map(|t| t.data);
     }
 
     sort_input.children[0].data = false;

datafusion/expr/src/tree_node/plan.rs

@@ -87,12 +87,15 @@ impl TreeNode for LogicalPlan {
             .iter()
             .map(|c| (*c).clone())
             .map_till_continue_and_collect(f)?;
-        // TODO: once we trust `t.transformed` remove additional check
-        if old_children
+        // TODO: Currently `t.transformed` quality comes from if the transformation
+        //  closures fill the field correctly. Once we trust `t.transformed` we can remove
+        //  the additional `t2` check.
+        //  Please note that we need to propagate up `t.tnr` though.
+        let t2 = old_children
             .into_iter()
             .zip(t.data.iter())
-            .any(|(c1, c2)| c1 != c2)
-        {
+            .any(|(c1, c2)| c1 != c2);
+        if t2 {
             Ok(Transformed::new(
                 self.with_new_exprs(self.expressions(), t.data)?,
                 true,

datafusion/physical-expr/src/tree_node.rs

@@ -61,10 +61,11 @@ impl<T> ExprContext<T> {
         }
     }
 
-    pub fn update_expr_from_children(mut self) -> Result<Self> {
+    pub fn update_expr_from_children(mut self) -> Result<Transformed<Self>> {
         let children_expr = self.children.iter().map(|c| c.expr.clone()).collect();
-        self.expr = with_new_children_if_necessary(self.expr, children_expr)?.data;
-        Ok(self)
+        let t = with_new_children_if_necessary(self.expr, children_expr)?;
+        self.expr = t.data;
+        Ok(Transformed::new(self, t.transformed, t.tnr))
     }
 }
 
@@ -93,7 +94,7 @@ impl<T> ConcreteTreeNode for ExprContext<T> {
         (self, children)
     }
 
-    fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
+    fn with_new_children(mut self, children: Vec<Self>) -> Result<Transformed<Self>> {
         self.children = children;
         self.update_expr_from_children()
     }

datafusion/physical-plan/src/tree_node.rs

@@ -61,10 +61,11 @@ impl<T> PlanContext<T> {
         }
     }
 
-    pub fn update_plan_from_children(mut self) -> Result<Self> {
+    pub fn update_plan_from_children(mut self) -> Result<Transformed<Self>> {
         let children_plans = self.children.iter().map(|c| c.plan.clone()).collect();
-        self.plan = with_new_children_if_necessary(self.plan, children_plans)?.data;
-        Ok(self)
+        let t = with_new_children_if_necessary(self.plan, children_plans)?;
+        self.plan = t.data;
+        Ok(Transformed::new(self, t.transformed, t.tnr))
     }
 }
 
@@ -94,7 +95,7 @@ impl<T> ConcreteTreeNode for PlanContext<T> {
         (self, children)
     }
 
-    fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> {
+    fn with_new_children(mut self, children: Vec<Self>) -> Result<Transformed<Self>> {
         self.children = children;
         self.update_plan_from_children()
     }