berkaysynnada commented on code in PR #8891:
URL: https://github.com/apache/arrow-datafusion/pull/8891#discussion_r1490932349
##########
datafusion/common/src/tree_node.rs:
##########
@@ -377,32 +583,408 @@ 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 {
/// Apply the closure `F` to the node's children
- fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
+ fn apply_children<F>(&self, f: &mut F) -> Result<TreeNodeRecursion>
where
- F: FnMut(&Self) -> Result<VisitRecursion>,
+ F: FnMut(&Self) -> Result<TreeNodeRecursion>,
{
+ let mut tnr = TreeNodeRecursion::Continue;
for child in self.children() {
- handle_tree_recursion!(op(child)?)
+ tnr = f(child)?;
+ handle_visit_recursion!(tnr)
}
- Ok(VisitRecursion::Continue)
+ Ok(tnr)
}
- fn map_children<F>(self, transform: F) -> Result<Self>
+ fn map_children<F>(self, f: F) -> Result<Transformed<Self>>
where
- F: FnMut(Self) -> Result<Self>,
+ F: FnMut(Self) -> Result<Transformed<Self>>,
{
let (new_self, children) = self.take_children();
if !children.is_empty() {
- let new_children =
- children.into_iter().map(transform).collect::<Result<_>>()?;
- new_self.with_new_children(new_children)
+ let t = children.into_iter().map_till_continue_and_collect(f)?;
+ // TODO: Currently `assert_eq!(t.transformed, t2.transformed)`
fails as
+ // `t.transformed` quality comes from if the transformation
closures fill the
+ // field correctly.
+ // Once we trust `t.transformed` we can remove the additional
check in
+ // `with_new_children()`.
+ let t2 = new_self.with_new_children(t.data)?;
+
+ // Propagate up `t.transformed` and `t.tnr` along with the node
containing
+ // transformed children.
+ Ok(Transformed::new(t2.data, t.transformed, t.tnr))
} else {
- Ok(new_self)
+ Ok(Transformed::no(new_self))
}
}
}
+
+#[cfg(test)]
+mod tests {
+ use crate::tree_node::{
+ Transformed, TransformedIterator, TreeNode, TreeNodeRecursion,
TreeNodeRewriter,
+ TreeNodeVisitor,
+ };
+ use crate::Result;
+
+ struct TestTreeNode<T> {
+ children: Vec<TestTreeNode<T>>,
+ data: T,
+ }
+
+ impl<T> TestTreeNode<T> {
+ fn new(children: Vec<TestTreeNode<T>>, data: T) -> Self {
+ Self { children, data }
+ }
+ }
+
+ impl<T> TreeNode for TestTreeNode<T> {
+ fn apply_children<F>(&self, f: &mut F) -> Result<TreeNodeRecursion>
+ where
+ F: FnMut(&Self) -> Result<TreeNodeRecursion>,
+ {
+ let mut tnr = TreeNodeRecursion::Continue;
+ for child in &self.children {
+ tnr = f(child)?;
+ handle_visit_recursion!(tnr);
+ }
+ Ok(tnr)
+ }
+
+ fn map_children<F>(self, f: F) -> Result<Transformed<Self>>
+ where
+ F: FnMut(Self) -> Result<Transformed<Self>>,
+ {
+ Ok(self
+ .children
+ .into_iter()
+ .map_till_continue_and_collect(f)?
+ .map_data(|new_children| Self {
+ children: new_children,
+ ..self
+ }))
+ }
+ }
+
+ fn new_test_tree<'a>() -> TestTreeNode<&'a str> {
+ let node_a = TestTreeNode::new(vec![], "a");
+ let node_b = TestTreeNode::new(vec![], "b");
+ let node_d = TestTreeNode::new(vec![node_a], "d");
+ let node_c = TestTreeNode::new(vec![node_b, node_d], "c");
+ let node_e = TestTreeNode::new(vec![node_c], "e");
+ let node_h = TestTreeNode::new(vec![], "h");
+ let node_g = TestTreeNode::new(vec![node_h], "g");
+ let node_f = TestTreeNode::new(vec![node_e, node_g], "f");
+ let node_i = TestTreeNode::new(vec![node_f], "i");
+ TestTreeNode::new(vec![node_i], "j")
+ }
+
+ fn all_visits<'a>() -> Vec<&'a str> {
+ vec![
+ "f_down(j)",
+ "f_down(i)",
+ "f_down(f)",
+ "f_down(e)",
+ "f_down(c)",
+ "f_down(b)",
+ "f_up(b)",
+ "f_down(d)",
+ "f_down(a)",
+ "f_up(a)",
+ "f_up(d)",
+ "f_up(c)",
+ "f_up(e)",
+ "f_down(g)",
+ "f_down(h)",
+ "f_up(h)",
+ "f_up(g)",
+ "f_up(f)",
+ "f_up(i)",
+ "f_up(j)",
+ ]
+ }
+
+ fn f_down_jump_visits<'a>() -> Vec<&'a str> {
+ vec![
+ "f_down(j)",
+ "f_down(i)",
+ "f_down(f)",
+ "f_down(e)",
+ "f_down(g)",
Review Comment:
Did you intentionally pass to "f_down(g)" after jump from "f_down(e)" ? I
believe it will be more useful and expected to apply "f_up(e)" before jumping
to "f_down(g)". What do you think about that, can you modify it in that way
without much effort?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
