- Add TreeNode.transform_with_payload(), `TreeNode.transform_down_w…

…ith_payload()` and `TreeNode.transform_up_with_payload()`

- Refactor `SortPushDown` and `PlanWithRequitements` using `TreeNode.transform_down_with_payload()`
- Refactor `ExprOrdering` using `TreeNode.transform_up_with_payload()`

datafusion/common/src/tree_node.rs

@@ -114,6 +114,42 @@ pub trait TreeNode: Sized {
         self.transform_up(op)
     }
 
+    /// Transforms the tree using `f_down` pre-preorder (top-down) and `f_up` post-order
+    /// (bottom-up) traversals. The `f_down` and `f_up` closures take payloads that they
+    /// propagate down and up during the transformation.
+    ///
+    /// The `f_down` closure takes `FD` type payload from its parent and returns `Vec<FD>`
+    /// type payload to propagate down to its children. One `FD` element is propagated
+    /// down to each child.
+    ///
+    /// The `f_up` closure takes `FU` type payload from its children collected into a
+    /// `Vec<FU>` and returns `FU` type payload to propagate up to its parent.
+    fn transform_with_payload<FD, PD, FU, PU>(
+        self,
+        f_down: &mut FD,
+        payload_down: PD,
+        f_up: &mut FU,
+    ) -> Result<(Self, PU)>
+    where
+        FD: FnMut(Self, PD) -> Result<(Transformed<Self>, Vec<PD>)>,
+        FU: FnMut(Self, Vec<PU>) -> Result<(Transformed<Self>, PU)>,
+    {
+        let (new_node, new_payload_down) = f_down(self, payload_down)?;
+        let mut new_payload_down_iter = new_payload_down.into_iter();
+        let mut payload_up = vec![];
+        let node_with_new_children = new_node.into().map_children(|node| {
+            let (new_node, p) = node.transform_with_payload(
+                f_down,
+                new_payload_down_iter.next().unwrap(),
+                f_up,
+            )?;
+            payload_up.push(p);
+            Ok(new_node)
+        })?;
+        let (new_node, new_payload_up) = f_up(node_with_new_children, payload_up)?;
+        Ok((new_node.into(), new_payload_up))
+    }
+
     /// Convenience utils for writing optimizers rule: recursively apply the given 'op' to the node and all of its
     /// children(Preorder Traversal).
     /// When the `op` does not apply to a given node, it is left unchanged.
@@ -136,6 +172,23 @@ pub trait TreeNode: Sized {
         after_op.map_children(|node| node.transform_down_mut(op))
     }
 
+    /// Transforms the tree using `f` pre-preorder (top-down) traversal. The `f_down`
+    /// closure takes payloads that it propagates down during the transformation.
+    ///
+    /// The `f_down` closure takes `FD` type payload from its parent and returns `Vec<FD>`
+    /// type payload to propagate down to its children. One `FD` element is propagated
+    /// down to each child.
+    fn transform_down_with_payload<F, P>(self, f: &mut F, payload: P) -> Result<Self>
+    where
+        F: FnMut(Self, P) -> Result<(Transformed<Self>, Vec<P>)>,
+    {
+        let (new_node, new_payload) = f(self, payload)?;
+        let mut new_payload_iter = new_payload.into_iter();
+        new_node.into().map_children(|node| {
+            node.transform_down_with_payload(f, new_payload_iter.next().unwrap())
+        })
+    }
+
     /// Convenience utils for writing optimizers rule: recursively apply the given 'op' first to all of its
     /// children and then itself(Postorder Traversal).
     /// When the `op` does not apply to a given node, it is left unchanged.
@@ -162,6 +215,25 @@ pub trait TreeNode: Sized {
         Ok(new_node)
     }
 
+    /// Transforms the tree using `f_up` post-order traversal. The `f_up` closure takes
+    /// payloads that it propagates up during the transformation.
+    ///
+    /// The `f_up` closure takes `FU` type payload from its children collected into a
+    /// `Vec<FU>` and returns `FU` type payload to propagate up to its parent.
+    fn transform_up_with_payload<F, P>(self, f: &mut F) -> Result<(Self, P)>
+    where
+        F: FnMut(Self, Vec<P>) -> Result<(Transformed<Self>, P)>,
+    {
+        let mut payload = vec![];
+        let node_with_new_children = self.map_children(|node| {
+            let (new_node, p) = node.transform_up_with_payload(f)?;
+            payload.push(p);
+            Ok(new_node)
+        })?;
+        let (new_node, new_payload) = f(node_with_new_children, payload)?;
+        Ok((new_node.into(), new_payload))
+    }
+
     /// Transform the tree node using the given [TreeNodeRewriter]
     /// It performs a depth first walk of an node and its children.
     ///