TreeNode refactor code deduplication: Part 3

[ozankabak]

Which issue does this PR close?

Closes #.

Rationale for this change

In the current implementation, maintainability and usability of the apply_children() and map_children() methods require extra effort.

The first step in addressing this was taken by @viirya, with the removal of apply_children(). He introduced the children_nodes() method as a clearer and more understandable implementation. However, a similar approach for map_children() was not feasible due to the diverse nature of dynamic objects (like ExecutionPlan and PhysicalExpr, ) and other structs that combine a plan with a state variable (such as DistributionContext).

This diversity hinders the unification of methods for accessing and updating their children. To overcome this challenge, we implemented a new trait ConcreteTreeNode, akin to DynTreeNode. The unification of map_children and apply_children seemed impossible with the current TreeNode implementation. Thus, we defined this trait to encapsulate the TreeNode. Instead of directly implementing TreeNode, one can now implement either ConcreteTreeNode for node-based state-preserving traversals, or DynTreeNode for stateless traversals, applicable to any tree-based data structure. While DynTreeNode requires two methods for implementation, ConcreteTreeNode demands three. This difference stems from the necessity to distinguish between read-only and read-write accesses in concrete types to avoid unnecessary cloning, which can be costly for heavily loaded nodes.

The ConcreteTreeNode methods are children() for read-only access, take_children() for separating the self node from its children, and with_new_children() for updating the self node with those separated and then updated children.

For rule objects, we have moved away from defining new structs with Arc<dyn ExecutionPlan>. Instead, we've implemented a new type of PlanContext with a state argument. Plans and expressions continue to implement DynTreeNode, but now they also implement ConcreteTreeNode under the wrap of PlanContext and ExprContext to store a payload. This approach is similarly applied to PhysicalExpr and ExprOrdering examples.

We have removed children_nodes from TreeNode, which was initially introduced to unify apply_children.
Additionally, unnecessary tree constructions have been eliminated in the rules to reset the tree.
Inconsistencies in the root node plan and tree traversal plan have been addressed, and tests have been enriched to ensure these issues are resolved.
Some further optimizations and code style improvements have been implemented in the rules.
With these refinements, there is no longer a need to .clone (which may be costly) treenodes.

With the changes in this PR, we make sure that TreeNode states are consistent during traversal. Hence they can be composed across different rules, sub-rules. With this capability we can write rules with more than 1 stages, where subsequent stages use payload from the previous rules. As an example consider the following use case:
Assume we want to satisfy ordering requirements of different executor in the following plan

RequiringExec(requirement=[a ASC]),
  RequiringExec(requirement=[a ASC, b ASC]),
    RequiringExec(requirement=[a ASC, b ASC, c ASC]),
      RequiringExec(requirement=[c ASC]),
        SourceExec   

Assume we have a top-down rule that collects minimum sort requirement that satisfy multiple ancestors. We would store this information similar to following struct

    pub plan: Arc<dyn ExecutionPlan>,
    pub common_sort_exprs: Vec<PhysicalSortExpr>,
    pub children: Vec<Self>,
}

This top-down pass would produce, following tree

RequiringExec(requirement=[a ASC]), , common_sort_exprs=[]
  RequiringExec(requirement=[a ASC, b ASC]), common_sort_exprs=[]
    RequiringExec(requirement=[a ASC, b ASC, c ASC]), common_sort_exprs=[a ASC, b ASC, c ASC]
      RequiringExec(requirement=[c ASC]), , common_sort_exprs=[c ASC]
        SourceExec  

Then, using this tree with a bottom-up pass we can satisfy requirements by inserting SortExecs to non-empty common_sort_exprs.

Bottom-up pass would produce following plan tree

RequiringExec(requirement=[a ASC]),
  RequiringExec(requirement=[a ASC, b ASC])
    RequiringExec(requirement=[a ASC, b ASC, c ASC])
      SortExec(expr=[a ASC, b ASC, c ASC])
        RequiringExec(requirement=[c ASC]),
          SortExec(expr=[c ASC])
            SourceExec   

This composable structure makes it easy to use previous tree results, in subsequent stages.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[ozankabak]

TODO

• Calls to add_sort_above should change PlanContext tree accordingly.
• pushdown_sorts: Has an unnecessary SortPushdown::new_default, should go away when add_sort_above is fixed.
• reorder_partitioned_join_keys: Changes plan, should take in PlanWithKeyRequirements and modify it (and thus remove new_default calls).
• reorder_aggregate_keys -> Changes plan, should take in PlanWithKeyRequirements and modify it (and thus remove new_default calls).
• remove_corresponding_coalesce_in_sub_plan: Changes plan, should modify the PlanWithCorrespondingCoalescePartitions tree it takes in.
  • This way we will get rid of the new_default calls in parallelize_sorts.

TEST TODO

• Tests should validate derived trees match with original trees
• Tests should validate data/payloads are set correctly after transform completes (to ensure composability)

[berkaysynnada]

I will work on TODO's.

[peter-toth]

It seems like this PR reverts @viirya's #8672. Why is that?

[berkaysynnada]

@viirya's PR was a good step towards deduplicating the apply_children() implementations. This PR further deduplicates not only apply_children(), but also map_children(), bringing them together in one place.

Like the children() method introduced by @viirya, which returns a Cow, map_children() requires two different methods: one to separate the children from self for mutation, and another for reattaching them afterwards. Since obtaining ownership of the children without owning self is not possible, I think Cow's purpose becomes not much relevant in this use case. Now we have two children related methods: one with a reference and one with ownership. The original fn children(&self) -> Vec<Cow<Self>> can still be used, but it has become redundant. If the user requires write access, take_children() should be used.

Rather than implementing functions such as apply_children() and map_children(), which provide too much flexibility and are prone to design an unclear flow, I think it is more appropriate to implement these 3 methods, which express the relationships of self nodes with their children, both in terms of ease of use and simplicity.

[peter-toth]

If you check apply_children(), it seems like you duplicate its logic to 4 places in this PR.
Also, the Cow was there to avoid cloning in Expr but now we have cloning again.
I'm ok with the new children(&self) and take_children(self) methods, but it seems like they are not properly used here and there...

[ozankabak]

If you check apply_children(), it seems like you duplicate its logic to 4 places in this PR.\

Right. The point is that there is a fixed number of apply_children implementations (one doesn't need to implement it per rule or per use case). However, 1 is obviously better than 4 and we will iterate on this.

Also, the Cow was there to avoid cloning in Expr but now we have cloning again.

Yes, and I am looking at ways to eliminate it. This is the only use case remaining forCow, and there may be a good way to remove it and still avoid cloning. The strategy in the WIP stage is to first explicitly implement everything and then progressively deduplicate by refining design. Maybe we will converge to the same Cow design, in that case we will simply revert.

Thanks for taking a look 🚀

[alamb]

Related discussion / epic: #8913

[mustafasrepo]

This PR is LGTM!. It helps us to use existing tree nodes in different rules. Also, common implementations for tree nodes are moved to trait implementation.

[alamb]

Is this PR ready for review? I ask because the title still says WIP but maybe that is just outdated.

[ozankabak]

Is this PR ready for review? I ask because the title still says WIP but maybe that is just outdated.

Yes, I forgot to change it

[peter-toth]

I kept my eye on this PR and IMO it looks good (my concers have been addressed) and actually a nice refactor that unifies many different derived trees into PlanContext and ExprContext.

Also, thanks for the general multi-stage example @ozankabak. Once this PR gets merged I would rebase my #8664 to see if it can improve some of the transformations while also simplify the tansform logic.

[alamb]

I plan to review this shortly

[alamb]

This is a really really nice work @ozankabak and @berkaysynnada 👏 I think it is pretty amazing and shows a real mastery of the code

Thank you @peter-toth for helping to review as well as inspiring some of this work and reimagining.

I ran our internal test suite from influxdb_iox against this branch and it passed 🚀 (and didn't require any code changes to comple )

I left some remarks mostly about improving / increasing comments, but I think they could be done as follow on PRs (or never)

cc @yahoNanJing @liukun4515

I recommend leaving this PR open for another few days (or longer if anyone says they would like longer to review) and then merging it in

[alamb]

Perhaps we can document what this macro does (aka returns from the function if Skip or Stop)

[alamb]

Perhaps this documentation can be updated, something like

Suggested change

	/// Use preorder to iterate the node on the tree so that we can
	/// stop fast for some cases.
	///
	/// The `op` closure can be used to collect some info from the
	/// tree node or do some checking for the tree node.
	/// Applies `op` to the node and its children, with traversal controlled by
	/// [`VisitRecursion`]
	///
	/// The `op` closure can be used to collect some info from the
	/// tree node or do some checking for the tree node.

[alamb]

I think we should document this trait, focusing what it should be used for, perhaps using some of the great description from this PR?

I think documenting what the functions do and their expectations, take_children in particular, would be valuable as well

[alamb]

Could we also document what this is used for? It seems it is used for rewriting execution plans where some additional information needs to be stored in addition to the plan itself

Maybe a more descriptive name would be PlanAndContext? Though perhaps that is too verbose

[alamb]

Likewise documentation here would be awesome

[alamb]

Can you document what this is doing?

[alamb]

Can you instead simply compare the plans directly rather than converting them to strings?

[berkaysynnada]

Since ExecutionPlan's do not implement PartialEq, we lack a straightforward way to compare them. However, I believe there are no obstacles to implementing PartialEq support for ExecutionPlan's.

[alamb]

Makes sense -- maybe that is something that can be embedded in a comment to help anyone else that may have the same question

[alamb]

it is really cool to see this same pattern extracted out

[alamb]

are these still TODO items? Did you intent to complete the as part of this PR? If not, perhaps we can file a ticket to track their completion

[berkaysynnada]

Adding payload checks to each test is time-consuming, and verifying the accuracy of the results necessitates a separate analysis. I suggest we track these tasks in a separate ticket as todo, so we don't delay this PR any further.

[alamb]

A separate ticket is a great idea -- can you please file one and add the link to these comments?

[alamb]

Is this still TODO?

[berkaysynnada]

Same for the other todo items

[berkaysynnada]

I left some remarks mostly about improving / increasing comments, but I think they could be done as follow on PRs (or never)

Thank you @alamb for your valuable feedbacks. I will push a commit addressing your suggestions.

[ozankabak]

Thanks everyone!

[alamb]

Epic work!