Improve constant folding for associative operations#23215
Open
hadrian-reppas wants to merge 6 commits into
Open
Improve constant folding for associative operations#23215hadrian-reppas wants to merge 6 commits into
hadrian-reppas wants to merge 6 commits into
Conversation
nathanb9
reviewed
Jun 27, 2026
| impl<'n> TreeNodeVisitor<'n> for AdjacentLiteralVisitor { | ||
| type Node = Expr; | ||
|
|
||
| fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> { |
Contributor
There was a problem hiding this comment.
last_expr_was_literal is not reset when visiting a nonliteral. For (lit(1) + col("c1")) + lit(2), the rewrite appears to rebuild the same expression but still results in Transformed::yes
Instead you could make has_adjacent_literals inspect the flattened operand sequence for the same operator, and reset adjacency whenever a nonliteral operand appears.
Author
There was a problem hiding this comment.
I kept the AdjacentLiteralVisitor and changed it so that it updates last_expr_was_literal correctly. is_associative_with_adjacent_literals needs to take the expression by reference because we call it from the big match statement in the Simplifier. I'm not sure I can construct the flattened operand sequence with just a reference to the expression. I suppose I could construct a Vec<&Expr>, but I'm not sure how much better that would be. Unless I could somehow write one function that does &Expr -> Vec<&Expr> and Expr -> Vec<Expr>.
Author
|
Also I realized that you have to check that the types match because otherwise reassociating the expression can change when casts happen which can change the answer if there is overflow. For example: |
Author
|
@nathanb9, I made some changes and the checks are now green. What do you think? |
Author
|
Bumping this @nathanb9 |
Contributor
|
To me this makes sense @hadrian-reppas |
nathanb9
approved these changes
Jul 16, 2026
Contributor
|
run benchmarks sql_planner sql_planner_extended |
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
This comment was marked as duplicate.
Jefffrey
reviewed
Jul 16, 2026
Comment on lines
+2371
to
+2383
| match &expr { | ||
| Expr::BinaryExpr(binary) | ||
| if binary.op == op | ||
| && matches!(info.get_data_type(&expr), Ok(dt) if &dt == datatype) => | ||
| { | ||
| let Expr::BinaryExpr(binary) = expr else { | ||
| unreachable!() | ||
| }; | ||
| flatten(op, datatype, info, *binary.left, out); | ||
| flatten(op, datatype, info, *binary.right, out); | ||
| } | ||
| _ => out.push(expr), | ||
| } |
Contributor
There was a problem hiding this comment.
Suggested change
| match &expr { | |
| Expr::BinaryExpr(binary) | |
| if binary.op == op | |
| && matches!(info.get_data_type(&expr), Ok(dt) if &dt == datatype) => | |
| { | |
| let Expr::BinaryExpr(binary) = expr else { | |
| unreachable!() | |
| }; | |
| flatten(op, datatype, info, *binary.left, out); | |
| flatten(op, datatype, info, *binary.right, out); | |
| } | |
| _ => out.push(expr), | |
| } | |
| match expr { | |
| Expr::BinaryExpr(binary) | |
| if binary.op == op | |
| && matches!(info.get_data_type(&expr), Ok(dt) if &dt == datatype) => | |
| { | |
| flatten(op, datatype, info, *binary.left, out); | |
| flatten(op, datatype, info, *binary.right, out); | |
| } | |
| _ => out.push(expr), | |
| } |
minor refactor to avoid the duplicate unstructuring
| } | ||
|
|
||
| // (A + 1) + 2 -> A + (1 + 2) | ||
| expr if is_associative_with_adjacent_literals(&expr, self.info) => { |
Contributor
There was a problem hiding this comment.
i do worry about this guard; it seems in the worst case, with Simplifier already doing a traversal (up) of the expr tree, this guard then does another traversal (down) of the tree, repeatedly as the simplifier goes up 🤔
ill see if benchmarks show this as an actual impact
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
This comment was marked as outdated.
This comment was marked as outdated.
Contributor
|
run benchmark sql_planner_extended |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
Benchmark for this request hit the 7200s job deadline before finishing. Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
tohuya6
pushed a commit
to tohuya6/datafusion
that referenced
this pull request
Jul 18, 2026
… size to 5 (apache#23659) on my M4 Macbook pro, it estimates it would take 50 minutes (!) to run the benchmark ```sh datafusion (main)$ cargo bench -p datafusion --bench sql_planner_extended logical_plan Compiling datafusion-common v54.0.0 (/Users/jeffrey/Code/datafusion/datafusion/common) ... Compiling datafusion v54.0.0 (/Users/jeffrey/Code/datafusion/datafusion/core) Finished `bench` profile [optimized] target(s) in 8m 11s Running benches/sql_planner_extended.rs (/Users/jeffrey/.cargo_target_cache/release/deps/sql_planner_extended-c026520b80bba716) Gnuplot not found, using plotters backend Benchmarking logical_plan_optimize: Warming up for 3.0000 s Warning: Unable to complete 100 samples in 5.0s. You may wish to increase target time to 3077.4s, or reduce sample count to 10. Benchmarking logical_plan_optimize: Collecting 100 samples in estimated 3077.4 s (100 iterations) ``` and running it via the bot just times out: apache#23215 (comment) so might actually be useful to reduce the sample size so it can actually run and be of use to us
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
4 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?
Closes #17158
Rationale for this change
The query planner currently does not optimize expressions like
c1 || 'a' || 'b' || c2(which it could simplify it toc1 || 'ab' || c2). This is because theConstEvaluatorlooks forBinaryExprs where both arguments areLiterals. The expression above looks likeso it does not get optimized. Since string concatenation is associative, we can rewrite the expression to look like
which the
ConstEvaluatorcan optimize toWe start by flattening nested occurrences of a the same associative operator into a list. Then we rebuild the expression tree from this list so that adjacent
Literals appear together in the same subtree. This rewrite happens in theexpr_simplifier::Simplifierso that theConstEvaluatorcan do its work on a subsequent pass.What changes are included in this PR?
This PR implements the
is_associative_with_adjacent_literalsandreassociate_literalsfunctions.There are some remaining questions:
Are these changes tested?
Yes
Are there any user-facing changes?
No