peter-toth commented on code in PR #8891:
URL: https://github.com/apache/arrow-datafusion/pull/8891#discussion_r1508780453
##########
datafusion/common/src/tree_node.rs:
##########
@@ -224,100 +319,202 @@ pub trait TreeNode: Sized {
/// tree and makes it easier to add new types of tree node and
/// algorithms.
///
-/// When passed to[`TreeNode::visit`], [`TreeNodeVisitor::pre_visit`]
-/// and [`TreeNodeVisitor::post_visit`] are invoked recursively
+/// When passed to[`TreeNode::visit`], [`TreeNodeVisitor::f_down`]
+/// and [`TreeNodeVisitor::f_up`] are invoked recursively
/// on an node tree.
///
/// If an [`Err`] result is returned, recursion is stopped
/// immediately.
///
-/// If [`VisitRecursion::Stop`] is returned on a call to pre_visit, no
+/// If [`TreeNodeRecursion::Stop`] is returned on a call to pre_visit, no
/// children of that tree node are visited, nor is post_visit
/// called on that tree node
///
-/// If [`VisitRecursion::Stop`] is returned on a call to post_visit, no
+/// If [`TreeNodeRecursion::Stop`] is returned on a call to post_visit, no
/// siblings of that tree node are visited, nor is post_visit
/// called on its parent tree node
///
-/// If [`VisitRecursion::Skip`] is returned on a call to pre_visit, no
+/// If [`TreeNodeRecursion::Jump`] is returned on a call to pre_visit, no
/// children of that tree node are visited.
pub trait TreeNodeVisitor: Sized {
/// The node type which is visitable.
- type N: TreeNode;
+ type Node: TreeNode;
/// Invoked before any children of `node` are visited.
- fn pre_visit(&mut self, node: &Self::N) -> Result<VisitRecursion>;
+ fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion>;
/// Invoked after all children of `node` are visited. Default
/// implementation does nothing.
- fn post_visit(&mut self, _node: &Self::N) -> Result<VisitRecursion> {
- Ok(VisitRecursion::Continue)
+ fn f_up(&mut self, _node: &Self::Node) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
}
}
-/// Trait for potentially recursively transform an [`TreeNode`] node
-/// tree. When passed to `TreeNode::rewrite`, `TreeNodeRewriter::mutate` is
-/// invoked recursively on all nodes of a tree.
+/// Trait for potentially recursively transform a [`TreeNode`] node tree.
pub trait TreeNodeRewriter: Sized {
/// The node type which is rewritable.
- type N: TreeNode;
+ type Node: TreeNode;
- /// Invoked before (Preorder) any children of `node` are rewritten /
- /// visited. Default implementation returns `Ok(Recursion::Continue)`
- fn pre_visit(&mut self, _node: &Self::N) -> Result<RewriteRecursion> {
- Ok(RewriteRecursion::Continue)
+ /// Invoked while traversing down the tree before any children are
rewritten /
+ /// visited.
+ /// Default implementation returns the node unmodified and continues
recursion.
+ fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
+ Ok(Transformed::no(node))
}
- /// Invoked after (Postorder) all children of `node` have been mutated and
- /// returns a potentially modified node.
- fn mutate(&mut self, node: Self::N) -> Result<Self::N>;
+ /// Invoked while traversing up the tree after all children have been
rewritten /
+ /// visited.
+ /// Default implementation returns the node unmodified.
+ fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
+ Ok(Transformed::no(node))
+ }
}
-/// Controls how the [`TreeNode`] recursion should proceed for
[`TreeNode::rewrite`].
-#[derive(Debug)]
-pub enum RewriteRecursion {
- /// Continue rewrite this node tree.
+/// Controls how [`TreeNode`] recursions should proceed.
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum TreeNodeRecursion {
+ /// Continue recursion with the next node.
Continue,
- /// Call 'op' immediately and return.
- Mutate,
- /// Do not rewrite the children of this node.
- Stop,
- /// Keep recursive but skip apply op on this node
- Skip,
-}
-/// Controls how the [`TreeNode`] recursion should proceed for
[`TreeNode::visit`].
-#[derive(Debug)]
-pub enum VisitRecursion {
- /// Continue the visit to this node tree.
- Continue,
- /// Keep recursive but skip applying op on the children
- Skip,
- /// Stop the visit to this node tree.
+ /// In top-down traversals skip recursing into children but continue with
the next
+ /// node, which actually means pruning of the subtree.
+ /// In bottom-up traversals bypass calling bottom-up closures till the
next leaf node.
+ /// In combined traversals bypass calling bottom-up closures till the next
top-down
+ /// closure.
+ Jump,
+
+ /// Stop recursion.
Stop,
}
-pub enum Transformed<T> {
- /// The item was transformed / rewritten somehow
- Yes(T),
- /// The item was not transformed
- No(T),
+#[derive(PartialEq, Debug)]
+pub struct Transformed<T> {
+ pub data: T,
+ pub transformed: bool,
+ pub tnr: TreeNodeRecursion,
}
impl<T> Transformed<T> {
- pub fn into(self) -> T {
- match self {
- Transformed::Yes(t) => t,
- Transformed::No(t) => t,
+ pub fn new(data: T, transformed: bool, tnr: TreeNodeRecursion) -> Self {
+ Self {
+ data,
+ transformed,
+ tnr,
+ }
+ }
+
+ pub fn yes(data: T) -> Self {
+ Self {
+ data,
+ transformed: true,
+ tnr: TreeNodeRecursion::Continue,
+ }
+ }
+
+ pub fn no(data: T) -> Self {
+ Self {
+ data,
+ transformed: false,
+ tnr: TreeNodeRecursion::Continue,
}
}
- pub fn into_pair(self) -> (T, bool) {
- match self {
- Transformed::Yes(t) => (t, true),
- Transformed::No(t) => (t, false),
+ pub fn map_data<U, F: FnOnce(T) -> U>(self, f: F) -> Transformed<U> {
+ Transformed {
+ data: f(self.data),
+ transformed: self.transformed,
+ tnr: self.tnr,
}
}
+
+ pub fn flat_map_data<U, F: FnOnce(T) -> Result<U>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<U>> {
+ Ok(Transformed {
+ data: f(self.data)?,
+ transformed: self.transformed,
+ tnr: self.tnr,
+ })
+ }
+
+ /// This is an important function to decide about recursion continuation
and
+ /// [`TreeNodeRecursion`] state propagation. Handling
[`TreeNodeRecursion::Continue`]
+ /// and [`TreeNodeRecursion::Stop`] is always straightforward, but
+ /// [`TreeNodeRecursion::Jump`] can behave differently when we are
traversing down or
+ /// up on a tree.
+ fn and_then<F: FnOnce(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ return_on_jump: Option<TreeNodeRecursion>,
+ ) -> Result<Transformed<T>> {
+ match self.tnr {
+ TreeNodeRecursion::Continue => {}
+ TreeNodeRecursion::Jump => {
+ if let Some(tnr) = return_on_jump {
+ return Ok(Transformed { tnr, ..self });
+ }
+ }
+ TreeNodeRecursion::Stop => return Ok(self),
+ };
+ let t = f(self.data)?;
+ Ok(Transformed {
+ transformed: t.transformed || self.transformed,
+ ..t
+ })
+ }
+
+ pub fn and_then_transform<F: FnOnce(T) -> Result<Transformed<T>>>(
+ self,
+ f: F,
+ ) -> Result<Transformed<T>> {
+ self.and_then(f, None)
+ }
+}
+
+pub trait TransformedIterator: Iterator {
Review Comment:
Fixed in
https://github.com/apache/arrow-datafusion/pull/8891/commits/08e3c7b8d920ef2a541679db45b5ce8d93a2f51d.
--
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]
