Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit #9913

peter-toth · 2024-04-02T17:45:25Z

Which issue does this PR close?

Part of #8913.

Rationale for this change

TreeNode APIs are inconsistent on LogicaPlan because LogicalPlan:apply() and LogicalPlan::visit() are recursing into subquery plans too (these methods override the original TreeNode::apply and TreeNode::visit() methods).
But there is no LogicalPlan:rewite() or LogicalPlan::transform() overrides defined so calling rewrite() and transform() on a LogicalPlan does rewrite/transform only the current level of the plan and doesn't touch subquery plans.

What changes are included in this PR?

This PR:

Moves the code of LogicalPlan::apply() to LogicalPlan::apply_with_subqueries(), and LogicalPlan::visit() to LogicalPlan::visit_with_subqueries() to separate functionality from original the TreeNode::apply() and TreeNode::visit(). So from now the new explicit LogicalPlan::..._with_subqueries() APIs do traverse into subquery plans, while the original TreeNode APIs don't.
Adds LogicalPlan::rewrite_with_subqueries(), LogicalPlan::transform_with_subqueries(), LogicalPlan::transform_down_with_subqueries(), LogicalPlan::transform_down_mut_with_subqueries(), LogicalPlan::transform_up_with_subqueries(), LogicalPlan::transform_up_mut_with_subqueries() and LogicalPlan::transform_down_up_with_subqueries() to be able to rewrite LogicalPlans with their subquery plans.
Fixes the issue that subquery plans were not treated as child of the node. In the LogicalPlan::..._with_subqueries()s APIs the subquery plans behave like the first children of a node (siblings of the real children of the node).

Are these changes tested?

Yes, with Existing UTs.

Are there any user-facing changes?

Yes - LogicalPlan::apply() no longer recurses into subqueries. The API user needs to change the code to LogicalPlan::apply_with_subqueries() to get the old behaviour.

alamb

Thank you (again) @peter-toth -- this looks really nice as well

I think there are is some more documentation we could add, but overall the only concern I have about this PR is the lack of test changes (like it seems to change the semantics of the analyzer / correlated expressions but no tests are changed). So this perhaps suggests we have a lack of test coverage in this area 🤔

alamb · 2024-04-02T19:53:19Z

datafusion/expr/src/logical_plan/plan.rs

@@ -370,6 +373,221 @@ impl LogicalPlan {
        }
    }

+    pub fn transform_expressions<F: FnMut(Expr) -> Result<Transformed<Expr>>>(


Can we please add some comments explaining what this is doing (namely that it applying a function to effectively rewrite the expressions of the LogicalPlan node in place)? It is a very cool API to see actually, as it will hopefully save a bunch of copying

alamb · 2024-04-02T19:54:02Z

datafusion/expr/src/logical_plan/plan.rs

+            }) => expr
+                .into_iter()
+                .map_until_stop_and_collect(f)?
+                .update_data(|expr| {


😍 -- that is quite a neat way to express this pattern and get the ownership sorted out 👍

datafusion/expr/src/logical_plan/plan.rs

alamb · 2024-04-02T19:56:23Z

datafusion/expr/src/logical_plan/plan.rs

@@ -1080,8 +1298,72 @@ impl LogicalPlan {
 }

 impl LogicalPlan {
+    pub fn apply_with_subqueries<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(


Likewise here, can we please add docstrings?

alamb · 2024-04-02T19:58:35Z

datafusion/expr/src/logical_plan/plan.rs

+                visitor.f_up(self)
+            }
+            TreeNodeRecursion::Jump => {
+                self.visit_subqueries(visitor)?;


Does this treat subqueries as though they were siblings of the node (like if TNR::Jump is returned subquery children will still be visited?)

If so I think I would find that unexpected (I would expect that the subqueries are treated like additional children)

I just moved the code here from LogicalPlan::visit() and in this PR I didn't want to change how subqueries are handled. The new code I added in LogicalPlan::rewrite_with_subqueries() are in sync with this behaviour.

But this is a good question and I think you are right and I too would treat them as additional children. But that would be an API behavior change so I would rather change that in a separate PR.

alamb · 2024-04-02T20:00:18Z

datafusion/optimizer/src/analyzer/mod.rs

@@ -155,7 +155,7 @@ impl Analyzer {

 /// Do necessary check and fail the invalid plan
 fn check_plan(plan: &LogicalPlan) -> Result<()> {
-    plan.apply(&mut |plan: &LogicalPlan| {
+    plan.apply_with_subqueries(&mut |plan: &LogicalPlan| {


It is interesting that this didn't change any plans / tests.

Basically I just moved the code of LogicalPlan:: apply() (that overrided TreeNode:: apply()) to LogicalPlan:: apply_with_subqueries() and changed the LogicalPlan:: apply() calls to LogicalPlan:: apply_with_subqueries(). I updated the PR description to be clear.

peter-toth · 2024-04-03T18:30:26Z

Thank you (again) @peter-toth -- this looks really nice as well

I think there are is some more documentation we could add, but overall the only concern I have about this PR is the lack of test changes (like it seems to change the semantics of the analyzer / correlated expressions but no tests are changed). So this perhaps suggests we have a lack of test coverage in this area 🤔

I rephrased the PR description to be clear. This PR doesn't change the analyzer or anything that uses LogicalPlan APIs. This PR just moves the overriden TreeNode APIs (apply()/visit()) in LogicalPlan to under new names so as to be clear they recurse into subqueries. Also, this PR adds the missing LogicalPlan::rewrite_with_subqueries() that is capable to rewrite LogicalPlans with their subqueries.

peter-toth · 2024-04-04T07:14:38Z

Let me update this PR with the latest main (including #9876) and address review feedbacks today. Until then I'm marking this PR as draft.

peter-toth · 2024-04-05T14:58:41Z

@alamb , while I was working on this I realized that currently TreeNode recursion definitions are a bit hard to follow and could be simplified, which will be useful this PR as well. I've opened #9965 to do that. Can we please review/merge #9965 first, then I can wrap up this PR.

peter-toth · 2024-04-05T18:58:38Z

@alamb, I've updated the PR. I will try to add tests for the LogicalPlan::..._with_subqueries() APIs in this PR or a in follow-up PR on Monday.

…n-subquery-handling

alamb · 2024-04-05T21:05:40Z

@alamb, I've updated the PR. I will try to add tests for the LogicalPlan::..._with_subqueries() APIs in this PR or a in follow-up PR on Monday.

I think it would be great to do this as a follow up PR on Monday, however I think the utility of unit tests on this code is somewhat limited. It will be tested end to end once I switch the optimizers to using the TreeNode API in #9954.

I took the liberty of pushing a commit to this branch to fix clippy and plan to merge it in when CI passes

alamb

Thank you @peter-toth . Love it

alamb · 2024-04-05T21:16:34Z

datafusion/expr/src/logical_plan/plan.rs

-    where
-        F: FnMut(&Expr) -> Result<(), E>,
-    {
+    pub fn apply_expressions<F: FnMut(&Expr) -> Result<TreeNodeRecursion>>(


This is an api change as it renames a public function. I couldn't find a way to get the error type to line up, so I think it is unavoidable.

I'll mark this PR as API change

Actually, I found another pattern here: https://github.com/apache/arrow-datafusion/blob/d201ec742b08fa2a0549fc27f0b39622d67391ab/datafusion/expr/src/utils.rs#L684

Maybe we can pull this part (add apply_expressions into a separate PR)? I think it would be valuable in its own -- for example it could be used to void the expressions() copy here:

https://github.com/apache/arrow-datafusion/blob/27bdf3c8c976874aae2417cb2aa435448832ecc0/datafusion/optimizer/src/analyzer/subquery.rs#L236-L241

It could also make this PR smaller

Yes, we can do that. Let me file a new PR with this change tomorrow and then look into the infinite recursion issue of this PR (peter-toth#3). Until then I'm marking this PR as draft.

I've opened #9984 as I think we don't need that inspect_expr_pre(). This PR together with #9984 will remove all usecases of it.

As for https://github.com/apache/arrow-datafusion/blob/27bdf3c8c976874aae2417cb2aa435448832ecc0/datafusion/optimizer/src/analyzer/subquery.rs#L236-L241 I plan to check all .expressions() calls and see if it can be transformed to the better .apply_expressions() once this PR is merged.

Sounds good

I've opened #9984 as I think we don't need that inspect_expr_pre(). This PR together with #9984 will remove all usecases of it.

I merged #9984

In order to make this PR easier to merge (so it isn't an API change) I

pushed 7c90e70 to restore inspect_expressions (by calling apply_expressions) and deprecated it.

Removed the API change

I plan to check all .expressions() calls and see if it can be transformed to the better .apply_expressions() once this PR is merged.

This would be good -- I think there are some more (though #9974 did get merged)

alamb · 2024-04-06T09:15:06Z

🤔 it appears this PR has some sort of infinite loop -- I will investigate

peter-toth · 2024-04-06T09:36:39Z

🤔 it appears this PR has some sort of infinite loop -- I will investigate

Thanks @alamb . Maybe because now we treat subqueries as additional children? I can also look into it later today or tomorrow...

…es()`

peter-toth · 2024-04-07T14:45:47Z

This is now ready for review again. The infinite recursion was due to an extra accidental .apply() call. Fixed here: c723492

…n-subquery-handling

alamb

Thank you @peter-toth -- this PR looks great to me.

I took the liberty of restring inspect_expressions and marking it deprecated and merging up from main.

As I have several PRs backed up on this one, I would like to merge it asap so plan to do so first thing tomorrow morning unless anyone else would like time to review

alamb · 2024-04-07T16:03:56Z

datafusion/expr/src/logical_plan/plan.rs

-    where
-        F: FnMut(&Expr) -> Result<(), E>,
-    {
+    pub fn apply_expressions<F: FnMut(&Expr) -> Result<TreeNodeRecursion>>(


Sounds good

alamb · 2024-04-07T16:13:35Z

datafusion/expr/src/logical_plan/plan.rs

-    where
-        F: FnMut(&Expr) -> Result<(), E>,
-    {
+    pub fn apply_expressions<F: FnMut(&Expr) -> Result<TreeNodeRecursion>>(


I've opened #9984 as I think we don't need that inspect_expr_pre(). This PR together with #9984 will remove all usecases of it.

I merged #9984

In order to make this PR easier to merge (so it isn't an API change) I

pushed 7c90e70 to restore inspect_expressions (by calling apply_expressions) and deprecated it.

Removed the API change

alamb · 2024-04-07T16:23:25Z

I am not sure what you are planning on working on next @peter-toth but one thing I noticed is that keeping track of all the LogicalPlan tree node walking / rewriting is somewhat hard for me (and datafusion/expr/src/logical_plan/plan.rs is getting pretty large).

I wonder if we should move all the code into a single module (maybe alongside the impl TreeNode for LogicalPlan) ? Maybe somethign like datafusion/expr/src/logical_plan/tree_node.rs or datafusion/expr/src/logical_plan/apply_visit.rs 🤔

alamb · 2024-04-07T18:05:53Z

datafusion/expr/src/logical_plan/plan.rs

-                        synthetic_plan.apply(op)?;
-                    }
-                    _ => {}
+    pub fn visit_with_subqueries<V: TreeNodeVisitor<Node = Self>>(


as a follow on PR I also plan to add documentation to these functions. Perhaps I can do so as part of moving the code into its own module

Docs: #9996

peter-toth · 2024-04-07T18:31:41Z

I took the liberty of restring inspect_expressions and marking it deprecated and merging up from main.

Sure, thanks for doing that!

Removed the API change

I was thinking about this a bit more and I unfortunately think this PR is still an API change. If some external code calls .apply() on a LogicalPlan then this PR might break it. This is because the implementation of .apply()will be de default implementation of TreeNode::apply() after this PR and no longer recurse into subqueries. The API user needs to change the code to .apply_with_subqueries() to get the old behaviour.
But I have no other idea to fix the current inconsistency between .apply()/.visit(), that do recurse into subqueries and .rewrite()/.transform(), that don't. (Well we could change .rewrite()/.transform() to do recurse into subqueries, but that too would be a breaking change.)

I am not sure what you are planning on working on next @peter-toth but one thing I noticed is that keeping track of all the LogicalPlan tree node walking / rewriting is somewhat hard for me (and datafusion/expr/src/logical_plan/plan.rs is getting pretty large).

I agree. I'm happy to take that task, but I will be travelling from Tuesday till the end of the week, so most likely I can only do that once I'm back.

alamb · 2024-04-08T09:28:42Z

I was thinking about this a bit more and I unfortunately think this PR is still an API change.

That makes sense -- I have added the label back.

But I have no other idea to fix the current inconsistency between

I think it is inevitable that there are different calls for LogicalPlan for traversing with subqueries and without subqueries

I agree. I'm happy to take that task,

Awesome -- thank you. I filed #9994 to track the idea. I may make a few small PRs in this direction to start the process too

alamb · 2024-04-08T09:28:55Z

Thanks again @peter-toth

alamb · 2024-04-08T10:20:01Z

datafusion/expr/src/logical_plan/plan.rs

@@ -1079,57 +1317,178 @@ impl LogicalPlan {
    }
 }

+/// This macro is used to determine continuation during combined transforming


I think this may be dead code -- can somone check #9997 ?

Actually, I deliberately added ..._with_subqueries() versions of all base TreeNode APIs to LogicalPlan. These 3 macros are used in LogicalPlan::rewrite_with_subqueries(), LogicalPlan::transform_down_with_subqueries() and LogicalPlan::transform_up_with_subqueries().

Cool -- thanks -- I think that will be clearer as we consolidate the tree rewriting code

peter-toth · 2024-04-08T11:06:58Z

Thanks @alamb for the review!

github-actions bot added logical-expr Logical plan and expressions optimizer Optimizer rules core Core DataFusion crate labels Apr 2, 2024

peter-toth changed the title ~~Consistent LogicalPlan Subquery handling~~ Consistent LogicalPlan subquery handling Apr 2, 2024

peter-toth mentioned this pull request Apr 2, 2024

DEMO: Introduce TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 10% faster planning time #9780

Closed

5 tasks

alamb mentioned this pull request Apr 2, 2024

DataFusion weekly project plan (Andrew Lamb) - April 1, 2024 #9899

Closed

7 tasks

peter-toth force-pushed the consistent-logicalplan-subquery-handling branch 2 times, most recently from f223ee8 to ea53b77 Compare April 2, 2024 18:58

alamb reviewed Apr 2, 2024

View reviewed changes

alamb changed the title ~~Consistent LogicalPlan subquery handling~~ Consistent LogicalPlan subquery handling in TreeNode::apply Apr 2, 2024

alamb changed the title ~~Consistent LogicalPlan subquery handling in TreeNode::apply~~ Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit Apr 2, 2024

peter-toth marked this pull request as draft April 4, 2024 07:14

This was referenced Apr 4, 2024

Refactor Optimizer to use owned plans and TreeNode API (10% faster planning) #9948

Merged

Avoid copying (so much) for LogicalPlan::map_children #9946

Closed

peter-toth mentioned this pull request Apr 5, 2024

Simplify TreeNode recursions #9965

Merged

fix

20f3d65

peter-toth force-pushed the consistent-logicalplan-subquery-handling branch from ea53b77 to 20f3d65 Compare April 5, 2024 18:50

peter-toth marked this pull request as ready for review April 5, 2024 18:53

alamb added 2 commits April 5, 2024 17:00

Merge remote-tracking branch 'apache/main' into consistent-logicalpla…

9206cdf

…n-subquery-handling

clippy

caa0d67

alamb added the api change Changes the API exposed to users of the crate label Apr 5, 2024

alamb changed the title ~~Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit~~ Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit. Rename inspect_expressions to apply_expressions Apr 5, 2024

alamb previously approved these changes Apr 5, 2024

View reviewed changes

peter-toth added 2 commits April 7, 2024 10:03

fix LogicalPlan::apply_expressions() and `LogicalPlan::map_subqueri…

6d4712f

…es()`

fix LogicalPlan::visit_with_subqueries()

90e7941

peter-toth mentioned this pull request Apr 7, 2024

Minor: Use Expr::apply() instead of inspect_expr_pre() #9984

Merged

peter-toth marked this pull request as ready for review April 7, 2024 14:44

alamb added 2 commits April 7, 2024 12:11

Add deprecated LogicalPlan::inspect_expressions

7c90e70

Merge remote-tracking branch 'apache/main' into consistent-logicalpla…

bd355a3

…n-subquery-handling

alamb approved these changes Apr 7, 2024

View reviewed changes

alamb removed the api change Changes the API exposed to users of the crate label Apr 7, 2024

alamb changed the title ~~Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit. Rename inspect_expressions to apply_expressions~~ Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit Apr 7, 2024

alamb reviewed Apr 7, 2024

View reviewed changes

alamb added the api change Changes the API exposed to users of the crate label Apr 8, 2024

alamb mentioned this pull request Apr 8, 2024

Consolidate LogicalPlan tree node walking / rewriting code in one place #9994

Closed

alamb merged commit 0a4d9a6 into apache:main Apr 8, 2024
26 checks passed

This was referenced Apr 8, 2024

Epic: Unified TreeNode rewrite API #8913

Closed

Minor: Improve documentation on LogicalPlan::apply* and LogicalPlan::map* #9996

Merged

Minor: adjust visibility of LogicalPlan:: apply_subqueries and LogicalPlan::map_subueries #9997

Closed

alamb reviewed Apr 8, 2024

View reviewed changes

This was referenced Apr 8, 2024

Minor: Make LogicalPlan::apply_subqueries and LogicalPlan::map_subqueries pub #9998

Merged

Avoid LogicalPlan::clone() in LogicalPlan::map_children when possible #9999

Merged

DataFusion weekly project plan (Andrew Lamb) - April 8, 2024 #10002

Closed

peter-toth mentioned this pull request Apr 9, 2024

Introduce TreeNode::exists() API, avoid copying expressions #10008

Merged

alamb mentioned this pull request Apr 10, 2024

Avoid copies in InlineTableScan via TreeNode API #10038

Merged

peter-toth mentioned this pull request Apr 12, 2024

Improve documentation on TreeNode #10035

Merged

alamb mentioned this pull request Apr 16, 2024

Deprecate TreeNode transform_xx_mut methods #10097

Closed

This pull request was closed.

Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit #9913

Consistent LogicalPlan subquery handling in TreeNode::apply and TreeNode::visit #9913

Conversation

peter-toth commented Apr 2, 2024 • edited by alamb 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?

alamb left a comment • 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

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 commented Apr 3, 2024

peter-toth commented Apr 4, 2024

peter-toth commented Apr 5, 2024 • edited Loading

peter-toth commented Apr 5, 2024 • edited Loading

alamb commented Apr 5, 2024

alamb left a comment

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 Apr 7, 2024 • 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

Choose a reason for hiding this comment

alamb commented Apr 6, 2024

peter-toth commented Apr 6, 2024

peter-toth commented Apr 7, 2024 • edited Loading

alamb 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 commented Apr 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peter-toth commented Apr 7, 2024 • edited Loading

alamb commented Apr 8, 2024

alamb commented Apr 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peter-toth commented Apr 8, 2024

peter-toth commented Apr 2, 2024 •

edited by alamb

Loading

alamb left a comment •

edited

Loading

peter-toth commented Apr 5, 2024 •

edited

Loading

peter-toth commented Apr 5, 2024 •

edited

Loading

peter-toth Apr 7, 2024 •

edited

Loading

peter-toth commented Apr 7, 2024 •

edited

Loading

peter-toth commented Apr 7, 2024 •

edited

Loading