DEMO: Introduce TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 10% faster planning time
#9780
Conversation
github-actions
bot
added
logical-expr
github-actions
bot
added
logical-expr
alamb
force-pushed
the
alamb/optimizer_tree_node
branch
from
da634a2
to
233d8e9
Compare
|
Update here is that I think this PR now has the basic code to do in-place updates of TreeNodes in general and TODOs:
Follow ons: |
alamb
changed the title
TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place
| ------Inner Join: Filter: t1.t1_id = t2.t2_id AND t1.t1_name = outer_ref(t0.t0_name) | ||
| --------TableScan: t1 | ||
| --------TableScan: t2 | ||
| LeftSemi Join: t0.t0_name = __correlated_sq_2.t1_name |
There was a problem hiding this comment.
this actually seems like an improvement to me -- as the subquery expressions are now properly visited by the rewriter
|
Updated description with performance results (about a 10% improvement for this PR alone). I think we can squeeze a bunch more out of this approach with the additional changes to rewrite in place (will file follow on tickets). Proceeding to get this PR ready for review |
| @@ -88,7 +88,7 @@ mod tests { | |||
|
|
|||
| use crate::test::*; | |||
|
|
|||
| fn assert_optimized_plan_equal(plan: &LogicalPlan, expected: &str) -> Result<()> { | |||
| fn assert_optimized_plan_equal(plan: LogicalPlan, expected: &str) -> Result<()> { | |||
There was a problem hiding this comment.
A large number of test changes below are due to taking an owned LogicalPlan (to avoid having to copy it)
|
|
||
| let formatted_plan = format!("{optimized_plan:?}"); | ||
| assert_eq!(formatted_plan, expected); | ||
| assert_eq!(plan.schema(), optimized_plan.schema()); |
There was a problem hiding this comment.
The test calls Optimizer::optimze now, which already checks the schemas are the same after each rule, so there is no need to check that again.
| @@ -651,8 +651,8 @@ explain SELECT t1_id, t1_name FROM t1 WHERE t1_id in (SELECT t2_id FROM t2 where | |||
| logical_plan | |||
| Filter: t1.t1_id IN (<subquery>) | |||
| --Subquery: | |||
| ----Limit: skip=0, fetch=10 | |||
| ------Projection: t2.t2_id | |||
| ----Projection: t2.t2_id | |||
There was a problem hiding this comment.
The limit is now pushed under the subquery
alamb
changed the title
TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in placeTreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 15% faster planning time
alamb
changed the title
TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 15% faster planning timeTreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 10% faster planning time
| @@ -59,7 +59,7 @@ pub fn main() -> Result<()> { | |||
|
|
|||
| // then run the optimizer with our custom rule | |||
| let optimizer = Optimizer::with_rules(vec![Arc::new(MyOptimizerRule {})]); | |||
| let optimized_plan = optimizer.optimize(&analyzed_plan, &config, observe)?; | |||
| let optimized_plan = optimizer.optimize(analyzed_plan, &config, observe)?; | |||
There was a problem hiding this comment.
This is the only API change exposed to users -- Optimizer::optimize takes an owned LogicalPlan rather than a &LogicalPlan. This is to avoid the need to copy over and over again
| @@ -174,6 +182,70 @@ pub trait TreeNode: Sized { | |||
| }) | |||
| } | |||
|
|
|||
| /// Implements the [visitor pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for | |||
There was a problem hiding this comment.
Here is the new TreeNode API to rewrite the nodes in place
I also tried Transformed<&mut Node>, but this got caught up in the borrow checker. Returning Transformed<()> seemed to work and made using this API feasible
@peter-toth and @berkaysynnada I would love any feedback on this API that you may have
| fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>( | ||
| &self, | ||
| f: &mut F, | ||
| ) -> Result<TreeNodeRecursion>; | ||
|
|
||
| /// Apply transform `F` to the node's children. Note that the transform `F` | ||
| /// might have a direction (pre-order or post-order). | ||
| /// Rewrite the node's children in place using `F`. |
There was a problem hiding this comment.
I chose to leave mutate_children unimplemented for Expr and the other TreeNodes initially to reduce the size of this PR (and thus the review burden). I will file a follow on ticket to track doing so
| /// | ||
| /// * If `f` returns Ok, sets `self.transformed` to `true` if either self or | ||
| /// the result of `f` were transformed. | ||
| pub fn and_then<F>(self, f: F) -> Result<Transformed<()>> |
There was a problem hiding this comment.
Introducing this function allowed the implementation of mutate and mutate children to be more concise. I couldn't figure out how to implement it in general for two Transformed as there are two different payloads
| use std::sync::{Arc, OnceLock}; | ||
|
|
||
| impl LogicalPlan { | ||
| /// applies `f` to each expression of this node, potentially rewriting it in |
There was a problem hiding this comment.
This file contains the implementation to allow rewriting the LogicalPlan in place. rewrite_exprs is needed initially to rewrite subquery references. I also expect we will be able to use this to implement more performant optimizer passes in the future (e.g. ExprSimplifier)
| .collect(); | ||
|
|
||
| let exprs = plan.expressions(); | ||
| plan.with_new_exprs(exprs, new_inputs).map(Some) |
There was a problem hiding this comment.
this used to copy all expressions and the entire LogicalPlan at least once
| let optimize_self_opt = self.optimize_node(rule, plan, config)?; | ||
| let optimize_inputs_opt = match &optimize_self_opt { | ||
| Some(optimized_plan) => { | ||
| self.optimize_inputs(rule, optimized_plan, config)? |
There was a problem hiding this comment.
each of these calls also copies the plan
| assert_eq!( | ||
| "Optimizer rule 'get table_scan rule' failed\ncaused by\nget table_scan rule\ncaused by\n\ | ||
| Internal error: Failed due to a difference in schemas, \ | ||
| original schema: DFSchema { fields: [\ |
There was a problem hiding this comment.
this test actually shows the output was reversed -- the original schema was [] not the table schema (which is what is changed by the optimizer)
| &OptimizerContext::new(), | ||
| )? | ||
| .unwrap_or_else(|| plan.clone()); | ||
| let optimized_plan = |
There was a problem hiding this comment.
This now uses the Optimizer directly rather than calling optimize_recursively (which the Optimizer does internally)
| fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {} | ||
|
|
||
| // in tests we are applying only one rule once | ||
| let opt_context = OptimizerContext::new().with_max_passes(1); |
There was a problem hiding this comment.
it turns out if you apply this rule multiple times, it changes the output. I explicitly made the test only run it a single time to match the previous behavior
Use TreeNode API in Optimizer
alamb
force-pushed
the
alamb/optimizer_tree_node
branch
from
1011559
to
1bedfec
Compare
| .clone(); | ||
|
|
||
| // take the old value out of the Arc | ||
| std::mem::swap(node, &mut new_node); |
There was a problem hiding this comment.
I think we can use std::mem::take here.
There was a problem hiding this comment.
There is no default impl for Arc<LogicalPlan> so I am not sure how to use the take API here 🤔
There was a problem hiding this comment.
I think we can have a default implementation with empty relation and empty schema, like the placeholder here
There was a problem hiding this comment.
I agree it would be nice -- maybe like adding a Default impl for LogicalPlan (that is empty relation?)
| /// Applies `f` to rewrite the existing node, while avoiding `clone`'ing as much | ||
| /// as possiblw. | ||
| /// | ||
| /// TODO eventually remove `Arc<LogicalPlan>` from `LogicalPlan` and have it own |
There was a problem hiding this comment.
Do you mean we can have LogicalPlan eventually? Not even with Box<LogicalPlan>? I think we may need Box for recursive plan
There was a problem hiding this comment.
Yeah, sorry I meant use Box eventually
| // we make a temporary Expr to rewrite and then put it back | ||
|
|
||
| let mut swap_column = Column::new_unqualified("TEMP_unnest_column"); | ||
| std::mem::swap(column, &mut swap_column); |
There was a problem hiding this comment.
I think we can have mem::take here too with a default implementation for Column
There was a problem hiding this comment.
I tried using take and it seems like there is no default impl for Column 🤔
error[E0277]: the trait bound `datafusion_common::Column: std::default::Default` is not satisfied
--> datafusion/expr/src/logical_plan/mutate.rs:288:41
|
288 | let mut swap_column= std::mem::take(column);
| -------------- ^^^^^^ the trait `std::default::Default` is not implemented for `datafusion_common::Column`
| |
| required by a bound introduced by this call
There was a problem hiding this comment.
yes, I would prefer "" as default column if it makes sense
There was a problem hiding this comment.
I see -- the reason I left a TEMP... name here was so that on error it was obvious that the Column was partially rewritten.
Given the comments below, however, it seems like we may not even need this API
| // rewrite children recursively with mutator | ||
| self.mutate_children(|c| c.mutate(mutator))? | ||
| .try_transform_node_with( | ||
| |_: ()| mutator.f_up(self), |
There was a problem hiding this comment.
It seems |_| is also fine. what is () for?
There was a problem hiding this comment.
The () is the type of the Node in Transform<()> which is somewhat confusing
|
I wonder if it would make sense to split this PR into 2 parts as it seem to mix 2 very different ideas.
So I don't think we need a new API for #9768. Don't get me wrong, I favour the new API but I would rather move that to a completely separate PR and slowly transform the code to using that in the future. As for the new |
|
Thank you for the feedback @peter-toth
I can definitely split it into two PRs. My rationale for one PR was that a PR for
I agree
The only difference I see in the APIs is that using
Changing
This is also an excellent call -- I actually tried initially to use So how about this for a plan:
We can then compare the performance of the TreeNodeRewriter to this PR with TreeNodeMutator and decide and see if the additional TreeNodeMutator API is worth the extra API surface / migration. |
Sounds very good to me, I'm really curious how much improvement the changes can bring. |
| /// applies `f` to LogicalPlans in any subquery expressions | ||
| /// | ||
| /// If Err is returned, the plan may be left in a partially modified state | ||
| fn rewrite_subqueries<F>(&mut self, mut f: F) -> Result<Transformed<()>> |
There was a problem hiding this comment.
I'm a bit concerned about this part. I mean the TreeNode APIs are currently inconsistent about handling "inner" trees. I'm refferring to LogicalPlan for example. It overrides TreeNode::apply() and TreeNode::visit() to apply f on subquery's logical plans too.
- But first of all, none of the other trees do that,
- secodly,
LogicalPlandoesn't overriderewrite()ortransform...()so transformations are applied only the "current" tree and not on "whole" logical plan.
I'm not sure that the current state is right or wrong as I can imagine both usescases in which
- a transformations needs to be applied on the current tree only
- a transformations need to change the whole tree (subplans trees included)
Now, this code seems to apply f on subquery logical plans too. If we keep it that way mutate() will be inconsistent with other APIs. But I'm also not sure that why it applies f on the root node only? Shouldn't it call subquery.subquery.mutate(f) to start a traversal on the subquery plan tree?
Is it needed because of the Optimizer? Does it optimize the whole logical plan at once? If so, I think we should introduce new APIs on LogicalPlan that is different to the existing TreeNode APIs. In Spark for example we have TreeNode.transform() and QueryPlan.transformWithSubqueries() for the different purposes.
There was a problem hiding this comment.
Shall I open a PR that adds LogicalPlan::rewriteWithSubqueries() and LogicalPlan::transform...WithSubqueries()?
There was a problem hiding this comment.
Now, this code seems to apply f on subquery logical plans too. If we keep it that way mutate() will be inconsistent with other APIs
I actually found I needed to handle recursing into subqueries to get some optimizer tests tests passing. Interestingly, recursing into subqueries also seems to have improved planning in some cases like #9780 (comment)
Is it needed because of the Optimizer? Does it optimize the whole logical plan at once? If so, I think we should introduce new APIs on LogicalPlan that is different to the existing TreeNode APIs. In Spark for example we have TreeNode.transform() and QueryPlan.transformWithSubqueries() for the different purposes.
OptimizerRules can currently specify if they want to handle recursion themselves or not: https://docs.rs/datafusion/latest/datafusion/optimizer/trait.OptimizerRule.html#method.apply_order
I'm not sure that the current state is right or wrong as I can imagine both usescases in which
- a transformations needs to be applied on the current tree only
- a transformations need to change the whole tree (subplans trees included)
I agree I can see usecases for both. I think the optimizer rules that handle subqueries specially currently do their own walk of the tree. For example https://github.com/apache/arrow-datafusion/blob/cc1db8a2043c73bda7adec309b42c08d88defab8/datafusion/optimizer/src/decorrelate_predicate_subquery.rs doesn't specify apply_order and walks down the tree
Shall I open a PR that adds
LogicalPlan::rewriteWithSubqueries()andLogicalPlan::transform...WithSubqueries()?
It seems like a good idea to me. If you you are going to do this, however, I would recommend finding some manual recursion / optimizer pass (maybe decorrelate_predicate_subquery.rs?) that you can use for a usecase of the new API.
Thank you again for all your help
There was a problem hiding this comment.
I've opened a PR here: #9913.
I think the new LogicalPlan methods like LogicalPlan::transform_expressions() can be useful for 2. in #9780 (comment)
There was a problem hiding this comment.
Update here I have created two PRs as discussed
- Refactor
Optimizerto use owned plans andTreeNodeAPI (10% faster planning) #9948 - Avoid copying (so much) for LogicalPlan::map_children #9946
Unfortunately, currently individually they either don't improve performance or actually slow things down.
I am going to test what happens when I run them both together
There was a problem hiding this comment.
When I combined both #9948 and #9946 tpch planning time improves by 12% (basically the same as this PR)
Thus I conclude we do not need the tree mutator API quite yet. Let's continue the discussions on those other PRs
Details
group both_tests main optimizer_tree_node2
----- ---------- ---- --------------------
logical_aggregate_with_join 1.01 1203.3±30.57µs ? ?/sec 1.00 1197.3±17.89µs ? ?/sec 1.00 1194.4±22.14µs ? ?/sec
logical_plan_tpcds_all 1.00 157.6±2.32ms ? ?/sec 1.00 157.3±3.68ms ? ?/sec 1.00 157.6±3.67ms ? ?/sec
logical_plan_tpch_all 1.01 16.9±0.34ms ? ?/sec 1.00 16.7±0.22ms ? ?/sec 1.04 17.5±0.37ms ? ?/sec
logical_select_all_from_1000 1.01 19.6±0.17ms ? ?/sec 1.00 19.4±0.17ms ? ?/sec 1.00 19.4±0.17ms ? ?/sec
logical_select_one_from_700 1.01 790.8±12.36µs ? ?/sec 1.00 784.0±12.23µs ? ?/sec 1.00 786.1±12.87µs ? ?/sec
logical_trivial_join_high_numbered_columns 1.00 740.8±12.72µs ? ?/sec 1.00 737.2±28.09µs ? ?/sec 1.00 739.5±12.25µs ? ?/sec
logical_trivial_join_low_numbered_columns 1.00 719.9±14.19µs ? ?/sec 1.02 731.2±11.31µs ? ?/sec 1.00 721.4±19.28µs ? ?/sec
physical_plan_tpcds_all 1.00 1666.1±16.04ms ? ?/sec 1.12 1872.6±16.17ms ? ?/sec 1.20 2.0±0.02s ? ?/sec
physical_plan_tpch_all 1.00 110.2±1.86ms ? ?/sec 1.10 121.3±1.90ms ? ?/sec 1.21 133.3±2.21ms ? ?/sec
physical_plan_tpch_q1 1.00 6.2±0.11ms ? ?/sec 1.19 7.4±0.12ms ? ?/sec 1.23 7.6±0.07ms ? ?/sec
physical_plan_tpch_q10 1.00 5.1±0.07ms ? ?/sec 1.11 5.6±0.11ms ? ?/sec 1.20 6.1±0.10ms ? ?/sec
physical_plan_tpch_q11 1.00 4.5±0.09ms ? ?/sec 1.10 5.0±0.07ms ? ?/sec 1.20 5.4±0.10ms ? ?/sec
physical_plan_tpch_q12 1.00 3.6±0.06ms ? ?/sec 1.11 4.0±0.07ms ? ?/sec 1.21 4.4±0.11ms ? ?/sec
physical_plan_tpch_q13 1.00 2.4±0.04ms ? ?/sec 1.10 2.7±0.06ms ? ?/sec 1.18 2.9±0.08ms ? ?/sec
physical_plan_tpch_q14 1.00 3.1±0.05ms ? ?/sec 1.09 3.4±0.06ms ? ?/sec 1.19 3.7±0.09ms ? ?/sec
physical_plan_tpch_q16 1.00 4.5±0.06ms ? ?/sec 1.10 4.9±0.08ms ? ?/sec 1.24 5.6±0.10ms ? ?/sec
physical_plan_tpch_q17 1.00 4.3±0.07ms ? ?/sec 1.09 4.7±0.05ms ? ?/sec 1.21 5.2±0.10ms ? ?/sec
physical_plan_tpch_q18 1.00 4.6±0.09ms ? ?/sec 1.10 5.1±0.09ms ? ?/sec 1.22 5.6±0.10ms ? ?/sec
physical_plan_tpch_q19 1.00 9.0±0.11ms ? ?/sec 1.06 9.5±0.11ms ? ?/sec 1.23 11.1±0.14ms ? ?/sec
physical_plan_tpch_q2 1.00 9.6±0.12ms ? ?/sec 1.11 10.6±0.18ms ? ?/sec 1.23 11.7±0.14ms ? ?/sec
physical_plan_tpch_q20 1.00 5.5±0.08ms ? ?/sec 1.11 6.1±0.08ms ? ?/sec 1.24 6.9±0.14ms ? ?/sec
physical_plan_tpch_q21 1.00 7.6±0.11ms ? ?/sec 1.12 8.5±0.14ms ? ?/sec 1.22 9.3±0.13ms ? ?/sec
physical_plan_tpch_q22 1.00 4.1±0.10ms ? ?/sec 1.10 4.5±0.10ms ? ?/sec 1.23 5.0±0.12ms ? ?/sec
physical_plan_tpch_q3 1.00 3.6±0.07ms ? ?/sec 1.08 3.9±0.06ms ? ?/sec 1.19 4.3±0.11ms ? ?/sec
physical_plan_tpch_q4 1.00 2.7±0.06ms ? ?/sec 1.11 3.0±0.05ms ? ?/sec 1.21 3.2±0.11ms ? ?/sec
physical_plan_tpch_q5 1.00 5.3±0.08ms ? ?/sec 1.08 5.7±0.12ms ? ?/sec 1.17 6.2±0.08ms ? ?/sec
physical_plan_tpch_q6 1.00 1905.1±30.83µs ? ?/sec 1.06 2.0±0.03ms ? ?/sec 1.15 2.2±0.05ms ? ?/sec
physical_plan_tpch_q7 1.00 7.0±0.11ms ? ?/sec 1.08 7.6±0.12ms ? ?/sec 1.19 8.4±0.14ms ? ?/sec
physical_plan_tpch_q8 1.00 8.9±0.14ms ? ?/sec 1.10 9.7±0.18ms ? ?/sec 1.19 10.6±0.16ms ? ?/sec
physical_plan_tpch_q9 1.00 6.7±0.11ms ? ?/sec 1.10 7.3±0.10ms ? ?/sec 1.21 8.0±0.10ms ? ?/sec
physical_select_all_from_1000 1.00 114.1±0.77ms ? ?/sec 1.13 129.2±1.14ms ? ?/sec 1.12 128.3±0.79ms ? ?/sec
physical_select_one_from_700 1.00 3.9±0.08ms ? ?/sec 1.03 4.1±0.05ms ? ?/sec 1.04 4.1±0.04ms ? ?/sec
There was a problem hiding this comment.
Very nice improvement! Thanks @alamb!
|
Converting to draft as we work on this API so someone doesn't merge it accidentally |
|
I will review this PR and the conversation here tomorrow. Thanks @alamb, the results are so promising 🚀 |
|
Per #9780 (comment) I think we can accomplish the same thing without a mutator at this time, so closing this PR |
alamb
changed the title
TreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 10% faster planning timeTreeNodeMutator for rewriting TreeNodes in place, change optimizer to rewrite LogicalPlan in place - 10% faster planning time
Note this PR looks large, but it is mostly comments and changes of
&plantoplanin testsWhich issue does this PR close?
Part of #9637 (based on some ideas from #9708). 🙏 @jayzhan211
closes #9768
Rationale for this change
See #9637
TLDR is that the optimizer currently copies
LogicalPlans many times unnecessarily (I think it is quadratic if not worse in the number of tree nodes). This is both slow and requires many memory allocations.Note we can do even better, but I want to keep the PR sizes as small as possible.
What changes are included in this PR?
mutate/MutatorAPIs toTreeNodemutateforLogicalPlan(aka&mut LogicalPlan)Optimizerto useTreeNodeAPII am trying to keep this PR relatively small, and I have other improvements in mind for follow on PRs (see below). This PR is now finally feasible thanks to the really nice TreeNode cleanup from @peter-toth and @berkaysynnada in #8891
Are these changes tested?
Functionally covered by existing tests, and I measured performance using the benchmarks -- I see an mprovement of about 10% across TPCH queries.
Details
Are there any user-facing changes?
LogicalPlanin place usingTreeNodeAPIPlanned Follow Ons:
Here are follow on tickets that will improve things even more
OptimizerRuleto rewrite in place (and avoid yet another copy)Exprs in place in addition toLogicalPlan(implementTreeNodeMutator-- specifically ExprRewriter should be a big win)UserDefinedLogicalPlanTreeNodeRewritertoTreeNodeMutator(e.g.ExprRewriter, etc)