Skip to content

Simplify TreeNode recursions - #9965

Merged
alamb merged 1 commit into
apache:mainfrom
peter-toth:simplify-treenode-recursions
Apr 5, 2024
Merged

Simplify TreeNode recursions#9965
alamb merged 1 commit into
apache:mainfrom
peter-toth:simplify-treenode-recursions

Conversation

@peter-toth

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of #8913.

Rationale for this change

This PR simplifies TreeNode recursion declarations by adding

  • TreeNodeRecursion::visit_children(),
  • TreeNodeRecursion::visit_sibling() and
  • TreeNodeRecursion::visit_parent()

helpers to simplify handling of the current TreeNodeRecursion with continuation closures by specifying the next node position relative to the current one.

Similarly, the transforming recursions are simplified with

  • Transfomed::visit_children(),
  • Transfomed::visit_sibling() and
  • Transfomed::visit_parent()

helpers.

After this PR recursion definition of TreeNode::visit() is as simple as that:

visitor.f_down(self)?
.visit_children(|| self.apply_children(|c| c.visit(visitor)))?
.visit_parent(|| visitor.f_up(self))

and TreeNode::rewrite() became:

rewriter.f_down(self)?
.transform_children(|n| n.map_children(|c| c.rewrite(rewriter)))?
.transform_parent(|n| rewriter.f_up(n))

What changes are included in this PR?

This PR:

  • Adds the above helpers and contains some other minor code cleanup.
  • Renames TransformedIterator to TreeNodeIterator and adds a new TreeNodeIterator::apply_until_stop() helper to simplify TreeNode::apply_children() implementations.

Are these changes tested?

Yes, with existint UTs.

Are there any user-facing changes?

No.

@github-actionsgithub-actionsBot added logical-expr Logical plan and expressions optimizer Optimizer rules labels Apr 5, 2024
fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
&self,
f: &mut F,
f: F,

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not related to handling recursions just a small code cleanup.

@peter-toth

Copy link
Copy Markdown
ContributorAuthor

cc @alamb, @berkaysynnada

@alambalamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @peter-toth -- I think this looks great to me 🙏

I have one small suggestion on naming but otherwise this looks good to me

TreeNodeRecursion::Jump => visitor.f_up(self),
TreeNodeRecursion::Stop => Ok(TreeNodeRecursion::Stop),
}
visitor

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️ -- this is really nice. I would have helped when i was fooling around with the TreeNodeMutator

Stop,
}

impl TreeNodeRecursion {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are much clearer. 👌 very nice


/// Continues visiting nodes with `f` depending on the current [`TreeNodeRecursion`]
/// value and the fact that `f` is visiting the current node's parent.
pub fn visit_parent<F: FnOnce() -> Result<TreeNodeRecursion>>(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was confused at first about the parent term here -- I think it might be clearer to call this function visit_self or visit_node as it is used to visit the node (potentially) after visiting its children, rather than visiting the parent

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm ok with visit_self or visit_node. My rationale for the current visit_children / visit_sibling / visit_parent naming was to specify the relative position to the last tnr.

This is a bit more detailed example that I will implement in #9913 for LogicalPlan::visit_with_subqueries and after we fix #9913 (comment) (so the subquery plans are treated like additional children):

// visit the current node (self)
visitor.f_down(self)?
// `visit_children()` because last tnr is from the current node (self) and the continuation // closure visits the current node (self)'s subqueries (that we treat as children)
.visit_children(|| self.apply_subqueries(|c| c.visit_with_subqueries(visitor)))? // `visit_sibling()` because last tnr is from visiting the last subquery (that we treated
// as a child) and continuation closure visits the current node (self)'s children, that are
// siblings to the subqueries
.visit_sibling(|| self.apply_children(|s| s.visit_with_subqueries(visitor)))?
// `visit_parent()` because last tnr is from transforming the last child and continuation
// cloure visits the current node (self, that is the parent of the last child)
.visit_parent(|| visitor.f_up(self)) 

So I'm ok with renaming visit_parent to visit_self or visit_node, but then visit_sibling sounds weird.

Actually, how about visit_down / visit_next / visit_up?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the current naming is more descriptive than visit_down / visit_next / visit_up -- let's go with what is in this PR for now and we can iterate in the future if needed

/// Maps the [`Transformed`] object to the result of the given `f` depending on the
/// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current
/// node's parent.
pub fn transform_parent<F: FnOnce(T) -> Result<Transformed<T>>>(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

likewise here - maybe it would make more sense to call this transform_self or transform_node

@alamb
alamb merged commit 27bdf3c into apache:mainApr 5, 2024
@alamb

alamb commented Apr 5, 2024

Copy link
Copy Markdown
Contributor

Thanks again @peter-toth

@peter-toth

Copy link
Copy Markdown
ContributorAuthor

Thanks for the review!

@alambalamb mentioned this pull request Apr 8, 2024
9 tasks
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

logical-exprLogical plan and expressionsoptimizerOptimizer rules

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@peter-toth@alamb