Epic: Unified TreeNode rewrite API

[alamb]

Is your feature request related to a problem or challenge?

(note this is mostly from @peter-toth 's description on #8891 (comment))

As @peter-toth noted, the current TreeNode APIs are somewhat inconsistent which has the following challenges:

1. adds a barrier to entry for newcomers to the DataFusion codebase
2. The APIs sometimes can not (or do no) implement pruning / early tree traversal termination and therefore do more tree walking / transforming than necessary
3. Increases code maintenance costs

Specific Challenges of the Current APIs

Some specific challenges with the current APIs. These can cause a newcommer (like me) some confusion:

• The apply / visit functions are capable to prune (skip) subtrees or return immediately (stop) but the transform functions are not. To make it more confusion rewrite is also capable to to prune, but instead of using a common way to control recursion, visit and rewrite use different enums with conflicting semantics.
  See this PR (Consolidate TreeNode transform and rewrite APIs #8891) for details on VisitRecursion vs RewriteRecursion.

• There is this Transformed enum whose purpose is to explicitely define if any change was made to a node. This enum is used in transform clousres but for some reason it is missing from rewrite. Moreover this explicit information is neither propogatged up to API callee nor it is used for detecting changes in nodes.
  See details in Remove Transformed enum #8835. I believe the current state simply doesn't make sense and just causes confusion.

• rewrite also seems to be inconsistent with transform_up and transform_down. rewrite probably organically ended up in its current state and capable to do its current job, but what I would expect from such a function is that it simply incorporates 2 closures from transform_up and transform_down to define:

  • what transformation should be done before (pre-order) and after (post-order) transformation of the node's children
  • and how the recursion should continue .

  See this PR (Consolidate TreeNode transform and rewrite APIs #8891) for the details. IMO the current TreeNodeRewriter.pre_visit() seems like a mess and its logic is hard to grasp.

Specific missing APIs

• Pruning capability of transform functions:
  Pruning would be very important as many of the transformations wouldn't require traversing on the whole tree and could improve performance a lot.
• Payload propagation:
  I think the the reason why there are so many (4 base + 9 derived) tree implementations (with many code duplications) in DataFusion is that there is no API to describe a transformation while also propagating state. In Transform with payload #8664 with the help of transform_with_payload I not just eleminated 7 derived trees but also improved some of the algorightms to require only one traversal (also a perf improvement).

Describe the solution you'd like

The proposal is a general API like this:

    /// Transforms the tree using `f_down` and `f_up` closures. `f_down` is applied on a
    /// node while traversing the tree top-down (pre-order, before the node's children are
    /// visited) while `f_up` is applied on a node while traversing the tree bottom-up
    /// (post-order, after the the node's children are visited).
    ///
    /// The `f_down` closure takes
    /// - a `PD` type payload from its parent
    /// and returns a tuple made of:
    /// - a possibly modified node,
    /// - a `PC` type payload to pass to `f_up`,
    /// - a `Vec<FD>` type payload to propagate down to its children
    ///    (one `FD` element is propagated down to each child),
    /// - a `TreeNodeRecursion` enum element to control recursion flow.
    ///
    /// The `f_up` closure takes
    /// - a `PC` type payload from `f_down` and
    /// - a `Vec<PU>` type payload collected from its children
    /// and returns a tuple made of:
    /// - a possibly modified node,
    /// - a `FU` type payload to propagate up to its parent,
    /// - a `TreeNodeRecursion` enum element to control recursion flow.
    fn transform_with_payload<FD, PD, PC, FU, PU>(
        self,
        f_down: &mut FD,
        payload_down: PD,
        f_up: &mut FU,
    ) -> Result<(Self, PU)>
    where
        FD: FnMut(Self, PD) -> Result<(Self, Vec<PD>, PC, TreeNodeRecursion)>,
        FU: FnMut(Self, PC, Vec<PU>) -> Result<(Self, PU, TreeNodeRecursion)>,

Obviously the closure return types can be extracted to a type alias or strucs like in the case of f_down could be:

pub struct TransformDownResult<N, PD, PC> {
    pub transformed_node: N,
    pub payload_to_children: Vec<PD>,
    pub payload_to_f_up: PC,
    pub recursion_control: TreeNodeRecursion,
}

All the other functions can then be turned into a specialization of the above:

• apply (visit_down): Only f_down closure that doesn't return a modified node and payload (only retrurns TreeNodeRecursion)
• visit / TreeNodeVisitor: Both f_down and f_up closures needed. The closures can be incorporated into a TreeNodeVisitor object but they don't return a modified node and payload
• transform_down: Only f_down closure that doesn't return payload
• transform_up: Only f_up closure that doesn't return payload
• transform: Both f_down and f_up closures needed, but they don't return payloads
• rewrite / TreeNodeRewriter: Both f_down and f_up are incorporated into a TreeNodeRewriter object, but they don't return a payload
• transform_down_with_payload: Only f_down closure
• transform_up_with_payload: Only f_up closure

Describe alternatives you've considered

As noted above, there are a variety of PRs that test out various ideas in one way or another and the major challenge I think is figuring out what changes to make, in what order, and how much code churn to apply)

Additional context

Related tickets / PRs:

• Consolidate TreeNode transform and rewrite APIs #8891
• TreeNode refactor code deduplication: Part 3 #8817
• Transform with payload #8664
• Remove Transformed enum #8835
• Simplify Expr::map_children #9876
• Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit  #9913
• Simplify TreeNode recursions #9965
• Consolidate LogicalPlan tree node walking / rewriting code in one place #9994
• Deprecate TreeNode transform_xx_mut methods #10097

[ozankabak]

The main concern I have is composability. Downstream, we have many use cases where the flow goes something like this:

1. Visit/transform a given tree, and save per-node data during the process.
2. Make a decision based depending on the final result of the visitation.
3. Visit/transform the tree again, re-use the saved per-node data and only modify when necessary.

(this can repeat multiple times)

For context, we plan to contribute code implementing algorithms that work in this fashion upstream at some point.

#8817 works towards having a composable architecture that lends itself to such flows, and will eliminate most major sources of code duplication concerning TreeNode. Similar to how DynTreeNode works, it introduces one generic object for ExecutionPlan-related trees, and one generic object for PhysicalExpr-related trees.

Once we have such a reusable/composable "infrastructure", we can discuss/ideate on various flavors of the transform APIs and find a design that works well.

[peter-toth]

The main concern I have is composability. Downstream, we have many use cases where the flow goes something like this:

1. Visit/transform a given tree, and save per-node data during the process.
2. Make a decision based depending on the final result of the visitation.
3. Visit/transform the tree again, re-use the saved per-node data and only modify when necessary.

(this can repeat multiple times)

This is actually interresting use case and I don't see that requirement in the current codebase. Are you adding this optimization in #8817 to somewhere?
Also, I'm not sure how often is this required so it might not be a good idea to transform all usecases to this direction.

#8817 works towards having a composable architecture that lends itself to such flows, and will eliminate most major sources of code duplication concerning TreeNode. Similar to how DynTreeNode works, it introduces one generic object for ExecutionPlan-related trees, and one generic object for PhysicalExpr-related trees.

I don't think that code duplication is the major issue here. I think the main issue with the current 11 derived trees (and also with the proposed generic PlanContext and ExprContext) is that the algorithms they use need to maintain both the base tree and the derived tree and keeping them in sync. Also the current closures (and #8817 too) are mutating the encoded state in tree nodes so it is very hard to reason about the correctness of the algorithms.
IMO in most of the cases (I checked only 9 out of 11 in the current codebase in #8664) this is simply not neded. In my transform_with_payload solution the immutable state of nodes, that is propagated down and up during transformation, and the clear transition functions encoded into closures give much cleaner and simle algorithms.

What I would propose is proabaly to take a look at #8664 first (it is a much simpler PR than #8817 is now) and check where the new API can help. If there remain any usecase where transform_with_payload doesn't prove to be useful there is sill the option to do the PlanContext / ExprContext way of refactor there only.

[peter-toth]

@alamb If you don't mind I would suggest a bit different list of tasks as some of the API inconsistencies can be fixed without transform_with_payload.

1. For example Consolidate TreeNode transform and rewrite APIs #8891 tries to fix 2 things that don't need a new base API.
   The 2 questions there is more like if / to what extent we want to:

   • changerewrite / TreeNodeRewriter signature / semantics
   • combine VisitRecursion + RewriteRecursion into TreeNodeRecursion
   • change transform signature / semantics

   to make the APIs more consistent.

2. For transform_with_payload I already have a PR: Transform with payload #8664, it is not exactly the above "ideal" form of the API, but already shows its capabilities. The question here is more like if it make sense / when can we move towards that base API. cc @ozankabak

3. For tranform / rewrite early termination I want to open a new PR.

4. For the Transformed enum issue I have a PR also: Remove Transformed enum #8835. If the decision is to remove the enum then it is independent of other PRs as the conflict it causes is minimal and straightforward to fix. But if we want to keep and fix it then then it will be huge change and only make sense to do after all the above.

[ozankabak]

This is actually interresting use case and I don't see that requirement in the current codebase. Are you adding this optimization in #8817 to somewhere?

Use cases where composability/reusability is critical are not yet in the current codebase (upstream) yet but at least we at Synnada have downstream use cases.

Also, I'm not sure how often is this required so it might not be a good idea to transform all usecases to this direction.

Having worked with most of the physical optimization rules as well as some logical rules for a long time now (upstream and downstream), I think I can safely that this is necessary. As I mentioned before, it is also on our roadmap to contribute more physical optimizations upstream in the near future and they need this kind of composability/reusability.

I don't think that code duplication is the major issue here. I think the main issue with the current 11 derived trees (and also with the proposed generic PlanContext and ExprContext) is that the algorithms they use need to maintain both the base tree and the derived tree and keeping them in sync.

This is not really true though -- once you have the generics, you do not need to maintain anything with regards to apply_children, map_children and others.

All in all, the composability/reusability we will get via #8817 is really critical to certain use cases and upcoming contributions. Let's get that in and then let's discuss refining the transform signatures and fixing inefficiencies pertaining to Recursion enums. I actually like some of your core ideas but I get the sense that you may not have the full picture on this part of the code, its current/possible downstream uses, possible future extensions and subtle logical dependencies.

In any case, I agree that parts of #8891 would be a good next step to iterate on and improve after #8817 merges. FYI, I hope to finish #8817 in a few days so that I can help with the rest of the work.

[peter-toth]

All in all, the composability/reusability we will get via #8817 is really critical to certain use cases and upcoming contributions. Let's get that in and then let's discuss refining the transform signatures and fixing inefficiencies pertaining to Recursion enums. I actually like some of your core ideas but I get the sense that you may not have the full picture on this part of the code, its current/possible downstream uses, possible future extensions and subtle logical dependencies.

I can't argue with that since I don't know your downstream code, but I'm looking forward to see a usecase in a future contribution where PlanContext or ExprContext is actually needed.

In any case, I agree that parts of #8891 would be a good next step to iterate on and improve after #8817 merges. FYI, I hope to finish #8817 in a few days so that I can help with the rest of the work.

Thanks for the update!

[alamb]

@alamb If you don't mind I would suggest a bit different list of tasks as some of the API inconsistencies can be fixed without transform_with_payload.

Thank you -- I have removed that suggestion from the description here and linked the PRs you mentioned to this ticket

[alamb]

Update here is that we are making good progress

See

• Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit  #9913
• Simplify TreeNode recursions #9965
• Consolidate LogicalPlan tree node walking / rewriting code in one place #9994

[alamb]

Here is some thoughts from @peter-toth on what a unified API might look like from #10035 (comment)

BTW, if we wanted to consolidate the TreeNode API terminonlogy I would suggest the following:

traversal order	inspecting	transforming
order agnostic	for_each() (new API, an alias to for_each_up(), but its docs shouldn't specify any order)	transform() (remains an alias to transform_up() but its docs shouldn't specify any order)
top-down	for_each_down() (new API), apply() (deprecated, from now just an alias to for_each_down())	transform_down() (changes to f: &mut F where F: FnMut(Self) -> Result<Transformed<Self>>, which is a breaking change but requires a trivial fix from the users), transform_down_mut() (deprecated, from now just an alias to transform_down())
bottom-up	for_each_up() (new API)	transform_up() (changes to f: &mut F where F: FnMut(Self) -> Result<Transformed<Self>>, which is a breaking change but requires a trivial fix from the users), transform_up_mut() (deprecated, from now just an alias to transform_up())
combined with separete f_down and f_up closures		transform_down_up()
combined with incorporated f_down() and f_up() methods into an object	visit() + TreeNodeVisitor	rewrite() + TreeNodeRewriter

+ Maybe rename apply_children() to for_each_children() and map_children() to transform_children(). Although renaming these would be a breaking change if something that is built on DataFusion defines a new kind of TreeNode, but fixing them would also be trivial.

+ LogicalPlan::..._with_subqueries() APIs should be alligned with the above TreeNode ones, but as they are fairly new (#9913) they haven't landed in any release yet.

[peter-toth]

I think the epic description is a bit misleading as the transform_with_payload API hasn't been implemented but we managed to unify the APIs and decrease the number of TreeNodes types in different PRs.

BTW, I'm still investigating transform_with_payload (#8664) idea, if it can help in other areas of the code...

[alamb]

Ok, given we largely implemented what this epic was designed to do, let's close it and track additional work in a new ticket. I filed #10121 to track

[alamb]

if can help in other areas of the code...

Well, since you asked @peter-toth 😄 migrating other passes listed on to use the TreeNode API would certainly be helpful. I know CSE is a big one. #9873 might also be an interesting exercise

[peter-toth]

Well, since you asked @peter-toth 😄 migrating other passes listed on to use the TreeNode API would certainly be helpful. I know CSE is a big one. #9873 might also be an interesting exercise

Sure. I want to take a look at #10097 next, but then I can check CSE: #9873