peter-toth commented on code in PR #8891:
URL: https://github.com/apache/arrow-datafusion/pull/8891#discussion_r1490991948
##########
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:
Actually I did. Because does it make any sense to call `f_up(e)` when
`f_down(e)` returned `Jump`? I don't think it does, as you can do everything in
`f_down(e)` (transform the node for example) besides returning `Jump`.
This behavior is aligned with the current `VisitRecursion::Skip` so if we
changed this behaviour and call `f_up(e)` too, then some of the
`TreeNodeVisitor` logics might need to change...
--
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]
