Transform with payload

[peter-toth]

Which issue does this PR close?

Part of #8913, #8663.

Rationale for this change

After #8817 the derived TreeNode trees got unified into 2 generic trees: PlanContext and ExprContext. These generic trees still have the drawback that they contain derived nodes (contains some data and a link to a node of the original tree) for each of the original tree nodes to store some additional information / state of the node, which means transformations needs to keep the 2 trees in sync. While this is done automaticaly, it comes with additional costs and make tranformation algorightms less readable as the generic derived nodes has only a data field that encodes different information in different trees.

This PR proposes an alternative approach to many of the transformations that uses derived trees currently. The new API that PR adds to TreeNode looks like the following:

    /// 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<PD>` type payload to propagate down to the node's children
    ///    (one `PD` element is propagated down to each child).
    ///
    /// The `f_up` closure takes
    /// - a `PC` type payload from `f_down` and
    /// - a `Vec<PU>` type payload collected from the node's children
    /// and returns a tuple made of:
    /// - a possibly modified node,
    /// - a `PU` type payload to propagate up to the node's parent.
    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<(Transformed<Self>, Vec<PD>, PC)>,
        FU: FnMut(Self, PC, Vec<PU>) -> Result<(Transformed<Self>, PU)>,

as well as the TreeNode.transform_down_with_payload() and TreeNode.transform_up_with_payload() variants that do one direction transformations.

With the help of above new API this PR refactors 9 of the 11 derived trees and makes ExprContext obsolete.

Besides the cleaner state transition functions, the new API is capable to describe some transformations that currently require 2 pass with 1 pass. With the new API, the example given in #8817 with RequiringExec that requires a top-down transformation to store minimum sort requirement in derived nodes and then a bottom-up transformation to insert SortExecss can be done in 1 pass without creating any additional trees.

There are 2 examples for this advanced usecase in this PR:

• PlanWithCorrespondingCoalescePartitions:
  Currently parallelize_sorts is a bottom-up transformation while at some nodes it kicks off remove_corresponding_coalesce_in_sub_plan to do a top-down transformation of the subtree.
  After this PR propagate_unnecessary_coalesce_connections_down top-down and parallelize_sorts_up bottom-up transformations are incorporated into one pass using transform_down_with_payload.
• OrderPreservationContext:
  Currently replace_with_order_preserving_variants is a bottom-up transformation while at some nodes it starts a top-down plan_with_order_preserving_variants transformation.
  After this PR propagate_order_maintaining_connections_down top-down and replace_with_order_preserving_variants_up bottom-up transformations are incorporated into one pass using transform_down_with_payload.

Please note that the new API can't replace or non-trivially can replace derived trees in some cases like:

• Node level states needs to be stored between 2 consecutive "same direction" transformations, but in this case the 2 transformations could be unified into 1 one.
• Node level states needs to be stored between a bottom-up and a consecutive top-down transformation (the other order, a top-down followed by a bottom-up transformation can be trivially rewritten), but in this case there might be better algorightms that need only 1 pass traversal. Please find 2 advanced examples above.

This PR doesn't necessary want to remove PlanContext and ExprContext as there might be downstream usescases of those. The purpose of this PR is to offer better API (more effective and readable transformations) for usescases where it make sense to use.

This PR implements idea 4. from #7942.

What changes are included in this PR?

This PR:

• Adds TreeNode.transform_with_payload(), TreeNode.transform_down_with_payload() and TreeNode.transform_up_with_payload(),
• Refactors SortPushDown and PlanWithRequitements using TreeNode.transform_down_with_payload(),
• Refactors ExprOrdering, ExprTreeNode and PipelineStatePropagator using TreeNode.transform_up_with_payload(),
• Refactors OrderPreservationContext and PlanWithCorrespondingCoalescePartitions using TreeNode.transform_with_payload().

The remaining 2 derived trees can be refactored in follow-up PRs if possible and needed.

Are these changes tested?

Using exinsting tests.

Are there any user-facing changes?

No.

[peter-toth]

cc @berkaysynnada, @alamb, @ozankabak

[alamb]

Thank you @peter-toth -- this seems like an improvement to me. Perhaps @viirya and @berkaysynnada can take a look too prior to merging this

[alamb]

We could potentially reduce redundancy by implementing this function terms of transform_with_payload and empty f_up and PU=() -- it probably doesn't matter but I figured I would bring it up

[peter-toth]

You are right, both the down and up versions can be implemented using transform_with_payload. But I would keep them separated in this PR and check later what happens if we we use empty f_up and PU=(). Does the "up logic" get optimized out by the compiler at "down invocations"? I.e. payload_up vec is not allocated?

[alamb]

It might make this code simpler to make this struct RequiredKeyOrder and give its methods names rather than dealing with Vec<Vec<...>> -- but that can be a refactor for a different PR perhaps

[peter-toth]

Thank you @peter-toth -- this seems like an improvement to me. Perhaps @viirya and @berkaysynnada can take a look too prior to merging this

Thanks @alamb for the review. Just a note that in this discussion: #8672 (comment) some concerns came up regarding the new transform_with_payload() methods might not be able to cover all usecases and so can't replace all derived trees. cc @ozankabak
I might need a few days to check that with a more complex derived tree like OrderPreservationContext and extend this PR or open a follow-up PR if refactor is possible.
But, even if it turns out that not all derived trees can be refactored, IMO the new methods are still useful as the usecase of SortPushDown, PlanWithRequitements and ExprOrdering proved.

[ozankabak]

Yes, I have some concerns about this PR (and getting rid of derived trees in general), so let's not merge this for now.

I gave a brief explanation before, but @berkaysynnada and I will give more details on our reasoning after the holiday (next week).

Let's not merge until we have consensus and more discussion

[peter-toth]

@ozankabak, @berkaysynnada, I've refactored OrderPreservationContext using TreeNode.transform_with_payload() in 269d8ba commit of this PR.

IMO the commit not just removes OrderPreservationContext, but also makes the algorighm simpler and more effective.
I was able to eliminate the inner top-down transformation of get_updated_plan() that was kicked off at SortExec nodes during the bottom-up transformation of transform_up(). The new algorighm, that uses TreeNode.transform_with_payload(), does only one traversal. Please find the details here: 269d8ba#diff-7894863f285ed3ad976c94603d639b99955d4a581e5eaee5f5946960fcbc21e1R52-R67

Please let me know if you still have concerns.

[ozankabak]

@peter-toth thank you for working on this. We were a little busy but we will work on this and reply in a couple days

[ozankabak]

Two pieces of news from our side:

1. I think I figured out a possible way to simplify the overall TreeNode-related code and reduce code duplication significantly. I plan to post a PR on Monday.
2. I think the integrated up/down idea may indeed be useful. I will try to figure out the trade-off between using the existing Visitor/pre_visit/post_visit mechanism and having a new kind of transform function.

In any case, 2 (which is the main concern of this PR) will benefit from 1 so I will go step by step finish off 1 first, and then continue working on this PR (and 2 in general).

[peter-toth]

Thanks for the update @ozankabak.

As regards to 1., #7942 and our discussion with @berkaysynnada at #8663 (comment) might be relevant. Although that PR didn't try to eliminate code duplication, it might be useful to design better tree visit/transform functions in terms of effectiveness (self mutation of nodes where possible) and API usability (common TreeNodeRecursion enum results to control recursions).

I think 2. is somewhat orthogonal to 1., the new transformWithPayload functions of this PR depend on map_children only. I beleive whatever will be the alternative of map_children after you simplification the code of this PR can be adjusted easily. (Actually, the original version in #7942 did use mutable references to nodes but this separate PR alligns with the current consume and produce way of other TreeNode APIs.)

You might have seen that yesterday I pushed another commit to this PR: 88a43c1 to refactor and remove PlanWithCorrespondingCoalescePartitions using transform_with_payload(). The change is very similar to the previous OrderPreservationContext related one. The new algorightm not just removes the derived tree, but also offers a more effective one pass transformation.

If you have some time after your simplification a review of this PR would be greatly appreciated as I would like to continue working on refactoring the remaining derived trees (either in this PR or follow-up ones).

[ozankabak]

If you have some time after your simplification a review of this PR would be greatly appreciated as I would like to continue working on refactoring the remaining derived trees (either in this PR or follow-up ones).

Certainly, I will do a deep dive on this after I finish the simplification PR. Thank you for the awesome collaboration

[ozankabak]

FYI, the draft PR is #8817. I will focus on this once we are done with it 🚀

[peter-toth]

I will update this PR once #8817 gets merged.

[alamb]

Marking as draft as we work on #8817 -- I am trying to clean up the list of PRs that need review

[peter-toth]

Marking as draft as we work on #8817 -- I am trying to clean up the list of PRs that need review

Thanks @alamb, I will rebase this PR and #8891 once that one gets merged.

[peter-toth]

@alamb, @ozankabak, @berkaysynnada I rebased this PR and updated the PR description after #8817 got merged. Please share your thoughts.

[berkaysynnada]

@peter-toth what do you think about first focusing on #8891 and then on this PR? I believe it holds higher importance and could contribute more to the simplification process. It would be easier if we concentrate these 2 PR's one by one, rather than proceeding with both at the same time.

[peter-toth]

@peter-toth what do you think about first focusing on #8891 and then on this PR? I believe it holds higher importance and could contribute more to the simplification process. It would be easier if we concentrate these 2 PR's one by one, rather than proceeding with both at the same time.

The 2 PRs are not connected but I'm fine with any order. #8891 is a small PR so I think we can probably conclude on something easily. And actually there will be a few more API fixing PRs in #8913.

[alamb]

I wonder what the current state / plan for this PR is? (I am going through the PR review backlog)

[peter-toth]

I wonder what the current state / plan for this PR is? (I am going through the PR review backlog)

Thanks @alamb fot the question! I plan to rebase this PR in a few days, maybe a week. Let me move this to draft till then.

[github-actions]

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.