Optimize logical optimizer: skip map_subqueries + in-place rewriting#22298
Open
adriangb wants to merge 3 commits into
Open
Optimize logical optimizer: skip map_subqueries + in-place rewriting#22298adriangb wants to merge 3 commits into
adriangb wants to merge 3 commits into
Conversation
3 tasks
github-actions
Bot
added
logical-expr
Contributor
Author
|
run benchmark sql_planner |
Three optimizations that together yield ~23-25% faster optimization on TPC-H/TPC-DS and up to 33% on expression-heavy queries: 1. map_subqueries short-circuit: skip expression tree walks when no subquery expressions exist. Previously rewrite_with_subqueries called map_subqueries at every plan node, walking all expression trees via ownership-based transform_down even with no subqueries. 2. plan_has_subqueries per-pass check: when no subqueries exist in the plan, bypass rewrite_with_subqueries entirely and use the cheaper rewrite_plan_in_place path. 3. rewrite_plan_in_place with Arc::make_mut: new map_children_mut method that mutates children in-place, avoiding the Arc::unwrap_or_clone + Arc::new allocation cycle at every node. The owned-plan rule API is bridged with std::mem::take, which is allocation-free: LogicalPlan::default() is an EmptyRelation that shares the process-wide static empty schema. Also adds optimizer-only benchmarks that isolate optimizer performance from SQL parsing and analysis overhead. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
adriangb
force-pushed
the
optimize-logical-optimizer
branch
from
d405cdf to
7d5359f
Compare
|
🤖 Criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing optimize-logical-optimizer (d405cdf) to 102da39 (merge-base) diff File an issue against this benchmark runner |
|
🤖 Criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
Contributor
Author
|
@alamb I know you worry about planning performance, wdyt of this change? |
Contributor
😍 Ok you nerd sniped me -- going to review this one now |
alamb
approved these changes
May 18, 2026
Contributor
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
TLDR is I think this is great and a really nice idea -- It is also nice it is so localized
As I mentioned in my comments, I was thinking we could expand the scope of this optimization (as a follow on PR) -- maybe we could implement the same/similar thing for Expr, PhysicalExpr and ExecutionPlan (though maybe they already have it)
@neilconway I think was working on something similar in #21749
FYI @joroKr21 as I think this may be related to some work you have been doing
| /// | ||
| /// This avoids the `Arc::unwrap_or_clone` + `Arc::new` cycle that the | ||
| /// ownership-based `TreeNode::rewrite` performs at every child node. | ||
| /// When the `Arc` refcount is 1 (always true in the optimizer), |
Contributor
There was a problem hiding this comment.
Why is this always true in the optimizer?
Contributor
Author
There was a problem hiding this comment.
I think it's usually true except for stuff like DynamicFilterPhsyicalExpr. But I'll remove the comment, I don't think it's helpful.
| /// When the `Arc` refcount is 1 (always true in the optimizer), | ||
| /// `Arc::make_mut` is essentially free. | ||
| /// | ||
| /// The `rule.rewrite()` API takes ownership, so we bridge the `&mut` to an |
Contributor
There was a problem hiding this comment.
this is an implementation detail -- I think it belongs next to the code below
| let mut changed = false; | ||
| if apply_order == ApplyOrder::TopDown { | ||
| let owned = std::mem::take(plan); | ||
| let result = rule.rewrite(owned, config)?; |
Contributor
There was a problem hiding this comment.
If this returns early, the plan is left with the default LogicalPlan -- I think we should make it clear in the comments what the expected semantics are when the rule fails (aka the tree is left in an inconsistent / errored state)
| } | ||
|
|
||
| // Recurse into children using Arc::make_mut (zero-cost when refcount == 1) | ||
| changed |= plan.map_children_mut(|child| { |
Contributor
There was a problem hiding this comment.
I wonder if we could eventually consider making a TreeNode API that uses the same trick to mutate the plans in place, rather than forcing a copy
| /// traversal is needed. When the plan has no subqueries, the cheaper | ||
| /// `rewrite` traversal is sufficient since all plan nodes are reachable | ||
| /// via direct children. | ||
| fn plan_has_subqueries(plan: &LogicalPlan) -> bool { |
Contributor
There was a problem hiding this comment.
it is unfortunate to have to walk the tree twice, but that does seem better than copying it in non subquery cases
It isn't clear to me why we couldn't do a similar "rewrite in place" trick with Expr / PhysicalExprs as well, to avoid the owned rewrite path entirely 🤔
| /// When the refcount is >1, it clones the inner value first. | ||
| /// | ||
| /// Returns `Ok(true)` if any child was modified by `f`. | ||
| pub fn map_children_mut<F: FnMut(&mut Self) -> Result<bool>>( |
Contributor
There was a problem hiding this comment.
It would be great if we could move this into TreeNode eventually
Address review feedback on the in-place optimizer rewrite: - Document the error contract of `rewrite_plan_in_place`: on `Err` the plan is left in an unspecified state because `rule.rewrite()` consumes it by value, and explain why it cannot be recovered without the clone the function exists to avoid. Note how `Optimizer::optimize` handles it. - Move the `mem::take` bridge explanation from the doc comment into an inline comment next to the code it describes. - Drop the inaccurate "Arc refcount is 1 is always true" claim. - Document that `LogicalPlan::map_children_mut` does not roll back partial mutations when `f` fails. Comment-only changes; no behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@alamb I've addressed the docs side of things. Regarding moving into TreeNode: I'm not sure yet this will work for other cases. In particular the |
`map_children_mut` was added as a public method on `LogicalPlan` in `datafusion-expr` only because the optimizer, in a different crate, needed to call it. But the optimizer is its sole consumer, and the `Arc::make_mut` in-place trick does not generalize to the other tree types (`Expr` children are `Box`ed; `PhysicalExpr`/`ExecutionPlan` children are `Arc<dyn _>`, which `Arc::make_mut` cannot handle), so committing to it as public API is not warranted. Move it into `optimizer.rs` as a private free function next to `rewrite_plan_in_place`, its only caller. `datafusion-expr` is now untouched by this PR except for the `map_subqueries` short-circuit, and the optimizer adds no public API. If `TreeNode` ever grows an in-place traversal this logic can move there with no breaking change. No behavior change; the 713 datafusion-optimizer tests pass unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@alamb I think it won't be a problem but just FYI I did move |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Supersedes #20837. That PR's branch was rebased onto current
main; GitHub doesnot allow reopening a closed PR once its branch has been force-pushed, so this is
a fresh PR carrying the same work, rebased and re-benchmarked.
Rationale for this change
Optimizer::optimizeruns every logical optimizer rule over the whole plan,repeatedly, until a fixed point. Benchmarking the optimizer in isolation (SQL
parsing and analysis excluded) surfaced two avoidable costs:
rewrite_with_subqueriescallsmap_subqueriesat every plan node, walkingevery expression tree via ownership-based
transform_down— even for planswith no subquery expressions at all, which is the common case.
TreeNode::rewritetraversal performs anArc::unwrap_or_clone+Arc::newcycle at every child node, re-allocatingthe
Arcspine on every pass even when nothing changes.What changes are included in this PR?
Three optimizations:
map_subqueriesshort-circuit — skip the expression-tree walk when a nodehas no subquery expressions.
plan_has_subqueriesper-pass check — when the whole plan has nosubqueries, bypass
rewrite_with_subqueriesentirely and use the cheaperin-place traversal.
rewrite_plan_in_placewithArc::make_mut— a privatemap_children_muthelper in the optimizer crate mutatesArc<LogicalPlan>children in place (copy-on-write; free when the refcount is 1, the common
case in the optimizer), avoiding the
Arc::unwrap_or_clone+Arc::newcycle. The owned-plan rule API is bridged with
std::mem::take, which isallocation-free because
LogicalPlan::default()is anEmptyRelationthatshares the process-wide static empty schema.
Also adds optimizer-only benchmarks to
sql_planner.rsthat isolate optimizercost from SQL parsing/analysis.
Benchmark results (optimizer-only, criterion, this PR vs
main)Measured A/B on a single machine: for the baseline run the two optimizer rule
files were reverted to
main(keeping the new benchmark code), then restored, soonly the optimization itself is being measured. The
optimizer_tpcds_allnumberwas confirmed across multiple back-to-back runs after an initial reading was
distorted by machine interference.
Possible future work (not in this PR)
The
mem::takebridge inrewrite_plan_in_placeis allocation-free, but it stillextracts an owned plan for every
(node, rule)pair up front — before the ruledecides whether it will transform — and the overwhelming majority of those pairs
are no-ops. A dynamic "lazy handle" rule API (the rule receives a handle that
derefs to
&LogicalPlanfor free and only pays the copy-on-write / move when itcalls
make_mut/replace) would let the framework skip the extraction forno-op rule invocations entirely. That is a breaking change to the public
OptimizerRuletrait and is out of scope here; it is being prototyped separately.Are these changes tested?
Yes. The existing optimizer test suite (713 tests across
datafusion-optimizer)passes unchanged. The optimizations are behavior-preserving: the in-place
traversal produces the same plans as the ownership-based traversal, and the
subquery short-circuit only skips work that is provably a no-op.
Are there any user-facing changes?
No. No new public API — the in-place traversal helper is private to the
optimizer crate. Optimization is faster; output plans are unchanged.