Cleanup TreeNode implementations #8672

viirya · 2023-12-29T02:43:17Z

Which issue does this PR close?

Closes #.

Rationale for this change

This is motivated by recent several refactoring on TreeNode. Currently we have many duplicate code (e.g., apply_children) around several TreeNode implementations. This patch cleans them up by moving duplicate function implementations back to TreeNode.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

viirya · 2023-12-29T02:47:12Z

datafusion/common/src/tree_node.rs

@@ -33,6 +33,9 @@ use crate::Result;
 /// [`LogicalPlan`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/enum.LogicalPlan.html
 /// [`Expr`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/expr/enum.Expr.html
 pub trait TreeNode: Sized {
+    /// Returns all children of the TreeNode
+    fn children_nodes(&self) -> Vec<Self>;


First tried with Vec<&Self>: https://github.com/apache/arrow-datafusion/compare/main...viirya:arrow-datafusion:refactor_treenode?expand=1.

Most cases it works well. But one trouble is the TreeNode implementation for Arc<T>:

impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> { fn children_nodes(&self) -> Vec<&Arc<T>> { // DynTreeNode.arc_children returns Vec<Arc<Self>> and this function cannot return reference of temporary object. The implementations of `DynTreeNode.arc_children` have same issue. It cannot be changed to return Vec<&Arc<Self>> too unimplemented!("Call arc_children instead") }

So changed to Vec<Self>.

ozankabak · 2023-12-29T14:49:20Z

I like this cleanup. @berkaysynnada is working on an approach that builds on this idea and #8664 to further simplify things. I am hopeful we will arrive at a good refactor very soon

peter-toth · 2023-12-29T18:19:12Z

datafusion/expr/src/tree_node/expr.rs

-    {
-        let children = match self {
-            Expr::Alias(Alias{expr,..})
+    fn children_nodes(&self) -> Vec<Self> {


This PR is a nice refactor but I'm not sure it is needed at all (1.) and I'm also not sure cloning a node's children into a Vec is good direction (2.). Let me share my thoughts on these 2.

I think in DataFusion there are only 4 real tree types: LogicalPlan, Expr, ExecutionPlan and PhysicalExpr and there 7 special tree types based on the above 4: SortPushDown, PlanWithRequitements, ExprOrdering, OrderPreservationContext, PipelineStatePropagator, PlanWithCorrespondingCoalescePartitions and PlanWithCorrespondingSort.
The only reason why DataFusion currently has these 7 special trees is to be able to propagate down/up some additional information during tree transformations. But I think these special trees are simply not needed and can be removed. If we add some transformation helper functions like TreeNode.transform_down_with_payload() and TreeNode.transform_up_with_payload() we can easily remove these 7 special types completely.
Get rid of special TreeNodes #8663 is about this issue and I removed SortPushDown, PlanWithRequitements and ExprOrdering as exaples in this PR: Transform with payload #8664.

If we just focus on the 4 base tree types we can distingsuish 2 different types: Expr uses Boxes but the other 3 (LogicalPlan, ExecutionPlan and PhysicalExpr) use Arcs for the links between the nodes. I think this is very important in their behaviour. E.g. a big issue for Exprs when the nodes are cloned as it basicaly means cloning the whole subtree under the node. But cloning is not an issue for the other 3 as cloning Arcs is cheap. So I think the problem with the current Expr.apply_children() and the Expr.children_nodes() in this PR that it doesn't move the code into the a direction where cloning is not necessary on trees made of Boxes.
What I would suggest long term is in Refactor TreeNode recursions #7942. Instead of the old TreeNode.apply_children() (or collecting the children into a Vec) we could use TreeNode.visit_children() (https://github.com/apache/arrow-datafusion/pull/7942/files#diff-6515fda3c67b9d487ab491fd21c27e32c411a8cbee24725d737290d19c10c199R31). IMO we need to implement visit_children() for the 4 base tree types only (see 1.).
(Unfortunately I didn't noted the visit related changes in the Refactor TreeNode recursions #7942 PR description, but it is a big PR with many different ideas and focusing on transform seemed more important.)

As I mentioned in above comment, the first try I did is to return Vec<&Self> for children_nodes. I think it makes more sense to return reference of children nodes instead of owned types (although for other implementations it is smart pointers, but I don't want to assume that the TreeNode impls must link with children nodes with Arc etc.).

But there is one trouble when implement TreeNode for Arc<T>: #8672 (comment). I've tried some ideas locally to get rid of it, but in the end they don't work so I changed children_nodes to return Vec<Self>. (still thinking if I can remove it but it may be more change. 🤔 )

This can be seen as an incremental cleanup patch as it doesn't change current API and keep current behavior but clean up the code as its purpose. In the long term I'd like to see more fundamental change to TreeNode and this can be seen as a short term effort to make the code easier to look/work on.

Does it make sense to return Vec<Cow<Self>> that would contain borrowed elements for Expr but owned for the others?

Sometimes there may be cases where one needs to do multiple passes over the same tree (e.g. doing one top down and one bottom up pass consecutively). Or we may want to run different logic on the same tree.

So it may be a little premature to conclude we can do away with "derived" trees and just keep the base four.

I suggest we take small steps and progressively discover the requirements/use cases. I think refactoring to get rid of duplicate apply_children and map_children is a good first step.

Yea, that's what I thought so this restrains to be limited change (mostly apply_children).

Sometimes there may be cases where one needs to do multiple passes over the same tree (e.g. doing one top down and one bottom up pass consecutively). Or we may want to run different logic on the same tree.

This is also supported with transform_with_payload() (https://github.com/apache/arrow-datafusion/pull/8664/files#diff-1148652cb6868dfd5d45d595a1013232c2407f3c89d3850a4ac8e48b8a0884e1R127-R151) that does a top-down and a bottom-up transforms too and capable to propagate different payloads in the 2 directions.

So it may be a little premature to conclude we can do away with "derived" trees and just keep the base four.

I'm happy to try refactoring the remaining 4 (OrderPreservationContext, PipelineStatePropagator, PlanWithCorrespondingCoalescePartitions and PlanWithCorrespondingSort) in scope of #8664 next week and check the above statement.
(I thought a smaller PR that shows 1-2 examples for both up and down directions would be easier to review, but probably we can do it in 1 big PR if that's prefered.)

I investigated the Vec<Cow<Self>> idea (#8672 (comment)) a bit further and it seems possible to avoid cloning at many places. This could be especially useful for Exprs where cloning is very expensive. I've opened a PR to this PR to here: viirya#132

…ssible using `Cow`

peter-toth · 2023-12-30T08:30:16Z

datafusion/common/src/tree_node.rs

    fn apply_children<F>(&self, op: &mut F) -> Result<VisitRecursion>
    where
        F: FnMut(&Self) -> Result<VisitRecursion>,
    {
-        for child in self.arc_children() {
-            match op(&child)? {
+        for child in &self.arc_children() {


Do we need this 2nd apply_children() implementation here? The base implementation in TreeNode might be enough.

Yea, missed this. This can be removed too.

Avoid cloning in `TreeNode.children_nodes()` implementations where possible using `Cow`

peter-toth

I think this PR looks very good now.

peter-toth · 2023-12-31T06:59:32Z

datafusion/common/src/tree_node.rs

@@ -368,7 +371,7 @@ impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
            let arc_self = Arc::clone(&self);
            self.with_new_arc_children(arc_self, new_children?)
        } else {
-            Ok(self)
+            Ok(self.clone())


Just one minor note, that this clone seems to be unnecessary.

Yea, probably I added it during trying different ideas locally. Removed.

alamb

Thank you for this PR @viirya and the reviews @ozankabak and @peter-toth .

This PR seems like an improvement to me, but I do also wonder how much overlap it has with #8664 🤔

peter-toth · 2023-12-31T13:58:42Z

This PR seems like an improvement to me, but I do also wonder how much overlap it has with #8664 🤔

This looks like a nice improvement to me too and now I think the 2 PRs are orthogonal.

ozankabak

I think this is a good step towards simplifying TreeNode-related code. Thanks @viirya

viirya · 2023-12-31T18:21:24Z

Thank you @alamb @ozankabak @peter-toth

viirya added 3 commits December 28, 2023 15:35

Refactor TreeNode and cleanup some implementations

2da06d5

More

f725553

More

e05123a

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions core Core DataFusion crate labels Dec 29, 2023

viirya commented Dec 29, 2023

View reviewed changes

Fix clippy

56c3919

viirya force-pushed the refactor_treenode2 branch from edc825e to 56c3919 Compare December 29, 2023 07:25

alamb mentioned this pull request Dec 29, 2023

DataFusion weekly project plan (Andrew Lamb) - Dec 25, 2023 #8655

Closed

7 tasks

peter-toth reviewed Dec 29, 2023

View reviewed changes

avoid cloning in TreeNode.children_nodes() implementations where po…

ee4c9df

…ssible using `Cow`

peter-toth reviewed Dec 30, 2023

View reviewed changes

viirya added 3 commits December 30, 2023 10:42

Remove more unnecessary apply_children

d4eab3a

Merge pull request #132 from peter-toth/refactor_treenode2

3b8a7ad

Avoid cloning in `TreeNode.children_nodes()` implementations where possible using `Cow`

Fix clippy

2958714

peter-toth approved these changes Dec 31, 2023

View reviewed changes

Remove

26736bb

alamb approved these changes Dec 31, 2023

View reviewed changes

peter-toth mentioned this pull request Dec 31, 2023

Transform with payload #8664

Closed

ozankabak approved these changes Dec 31, 2023

View reviewed changes

viirya merged commit bf3bd92 into apache:main Jan 1, 2024
22 checks passed

peter-toth mentioned this pull request Jan 10, 2024

TreeNode refactor code deduplication: Part 3 #8817

Merged

matthewgapp mentioned this pull request Jan 11, 2024

matt/feat/recursive ctes/config flag matthewgapp/arrow-datafusion#3

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Cleanup TreeNode implementations #8672

Cleanup TreeNode implementations #8672

viirya commented Dec 29, 2023 •

edited

Loading

viirya Dec 29, 2023 •

edited

Loading

ozankabak commented Dec 29, 2023

peter-toth Dec 29, 2023 •

edited

Loading

viirya Dec 29, 2023 •

edited

Loading

peter-toth Dec 29, 2023

ozankabak Dec 29, 2023

viirya Dec 29, 2023

peter-toth Dec 29, 2023 •

edited

Loading

peter-toth Dec 30, 2023

peter-toth Dec 30, 2023

viirya Dec 30, 2023

peter-toth left a comment

peter-toth Dec 31, 2023

viirya Dec 31, 2023

alamb left a comment

peter-toth commented Dec 31, 2023 •

edited

Loading

ozankabak left a comment

viirya commented Dec 31, 2023

Cleanup TreeNode implementations #8672

Cleanup TreeNode implementations #8672

Conversation

viirya commented Dec 29, 2023 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

viirya Dec 29, 2023 • edited Loading

Choose a reason for hiding this comment

ozankabak commented Dec 29, 2023

peter-toth Dec 29, 2023 • edited Loading

Choose a reason for hiding this comment

viirya Dec 29, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peter-toth Dec 29, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peter-toth left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

peter-toth commented Dec 31, 2023 • edited Loading

ozankabak left a comment

Choose a reason for hiding this comment

viirya commented Dec 31, 2023

viirya commented Dec 29, 2023 •

edited

Loading

viirya Dec 29, 2023 •

edited

Loading

peter-toth Dec 29, 2023 •

edited

Loading

viirya Dec 29, 2023 •

edited

Loading

peter-toth Dec 29, 2023 •

edited

Loading

peter-toth commented Dec 31, 2023 •

edited

Loading