Repository navigation
Rewrite SUM(X+c) to SUM(X) + c*COUNT(X) - #145510
parkertimmins merged 66 commits into
Conversation
…field) Introduces MvSingleValueOrNull, an internal optimizer function that returns a field value unchanged if single-valued, or null if multi-valued. This enables the rewrite SUM(X + c) → SUM(sv_X) + c * COUNT(sv_X) via the existing surrogate mechanism, allowing multiple SUM(field + c_i) expressions in the same STATS to share a single SUM and COUNT computation. Correctness: multivalued fields are excluded from SUM(X + c) because X+c returns null when X is MV. SINGLE_VALUE_OR_NULL replicates that filter, so the rewrite is exact. Relates: elastic#140470
The previous commit used Sum.surrogate() to rewrite SUM(field + c). This replaces it with a dedicated RewriteSumFieldPlusConstant optimizer rule that processes the full Aggregate node at once. Key improvement: the rule only fires when 2+ SUM(field + c_i) expressions share the same base field in the same STATS. A single SUM(X+c) is left untouched since rewriting it would produce more aggregates, not fewer. The rule must run first in the substitutions() phase, before ReplaceAggregateNestedExpressionWithEval, which would extract field+c into a pre-agg EVAL and hide the pattern. Relates: elastic#140470
…SUM(X+c)" This reverts commit 8fe3b48.
- MvSingleValueOrNullTests: unit tests for all supported types verifying single-valued pass-through and null for multi-valued positions - stats.csv-spec: three CSV spec tests covering SUM(X+c) correctness with single-valued integer fields, double fields filtered to single-valued rows, and null field values
6a91c6f to
2f63625
Compare
|
Buildkite benchmark this with clickbench-columnar-mode please |
The surrogate approach in Sum.surrogate() cannot work: ReplaceAggregateNestedExpressionWithEval runs first in the substitutions phase, extracting SUM(field + c) into a pre-agg EVAL. By the time SubstituteSurrogateAggregations fires, SUM sees an Attribute, not an Add, so the pattern never matches. Replace it with RewriteSumFieldPlusConstant, a dedicated OptimizerRule<Aggregate> placed first in substitutions(), before ReplaceAggregateNestedExpressionWithEval. It sees the raw SUM(field + c) pattern and rewrites multiple expressions sharing the same base field into a single SUM(SINGLE_VALUE_OR_NULL(field)) + COUNT(SINGLE_VALUE_OR_NULL(field)) pair, deriving each original sum via a cheap post-agg eval. The rule only fires when 2+ SUM(field + c_i) expressions share the same field, since a single SUM(X+c) would produce more aggregates after rewrite, not fewer. Relates: elastic#140470
Remove the 2+ expression threshold — the rewrite now fires for any single SUM(field + c). This eliminates the two-pass grouping approach: one pass suffices since fieldToSvPair deduplicates the shared aggregates across multiple SUM(field + c_i) expressions over the same field.
|
Buildkite benchmark this with clickbench-columnar-mode please |
- Fix filter check: use hasFilter() == false instead of filter() == null (default filter is Literal.TRUE, not null) - Fix null children: use Literal.TRUE / AggregateFunction.NO_WINDOW when constructing Sum/Count inside the rule - Switch to ParameterizedOptimizerRule to get context.configuration() for Add construction (CONFIGURATION_MARKER is invalid in optimizer rules) - Add testSumOfFieldPlusConstant() verifying that SUM(x+1) and SUM(x+2) share a single SUM(MvSingleValueOrNull(x)) / COUNT(MvSingleValueOrNull(x)) pair in the optimized plan
04582e9 to
84ab47c
Compare
|
Buildkite benchmark this with clickbench-columnar-mode please |
- Handle SUM(field - c) via Sub matching in addition to Add - Use Block.keepMask for MvSingleValueOrNull evaluation - Fix import ordering in LogicalPlanOptimizer - Add csv-spec tests for multi-valued fields and subtraction - Add optimizer plan tests for single and mixed ±c rewrites
Drop Sub handling, redundant csv-spec tests, and extra optimizer plan tests to keep the optimization minimal.
|
Pinging @elastic/es-analytical-engine (Team:Analytics) |
alex-spies
left a comment
There was a problem hiding this comment.
Ok, did another pass. Need to do another one tomorrow, this is quite a bit of code.
| * For single-valued blocks the default pass-through in {@link AbstractEvaluator} is used | ||
| * ({@link AbstractEvaluator#evalSingleValuedNotNullable} returns the block as-is). | ||
| * For blocks that may contain multi-valued positions, a boolean mask is built where only | ||
| * single-valued positions are kept, then {@link Block#keepMask} applies it using | ||
| * type-specialized logic. |
There was a problem hiding this comment.
Nice, this sounds very reasonable and correct.
| try ( | ||
| BooleanVector.FixedBuilder maskBuilder = driverContext.blockFactory().newBooleanVectorFixedBuilder(block.getPositionCount()) | ||
| ) { | ||
| for (int p = 0; p < block.getPositionCount(); p++) { |
There was a problem hiding this comment.
We have inherently non-null single-valued blocks that we call vector blocks. Making a new mask by adding a keep mask may actually make things worse because we will then emit a non-vector block (and lose some knowledge about the block, esp. that it consists of non-null single values, only).
We don't have to actually check for the block type, there's a simpler method called Block#mayHaveMultivaluedFields(), which is always false for vector blocks.
We should probably check for this condition and perform a no-op then?
| w.registerException(IllegalArgumentException.class, "single-value function encountered multi-value"); | ||
| } | ||
| } | ||
| maskBuilder.appendBoolean(valueCount == 1); |
There was a problem hiding this comment.
We could track if any position was multivalued, and default to a no-op if not.
| } | ||
|
|
||
| public TestSubstitutionOnlyOptimizer(TransportVersion minimumVersion) { | ||
| super(unboundLogicalOptimizerContext(minimumVersion)); |
There was a problem hiding this comment.
Version-aware optimizations were only added in, I think, August 2025, and are generally a rare occurrence so far. The TestSubstitutionOnlyOptimizer was used only in very few tests in the past. I think you're adapting the test infrastructure exactly as required for this PR. Thank you!
| \_Eval[[$$salary_+_1$SUM$0{r$}#2 + $$salary_+_1$SUM$1{r$}#3 * 1[INTEGER] AS s1#0, $$salary_+_1$SUM$0{r$}#2 + $$salary_+_1$SUM$1{r$}#3 * 2[INTEGER] AS s2#1]] | ||
| \_Limit[1000[INTEGER],false,false] | ||
| \_Aggregate[[],[SUM($$salary_+_1$SUM$0{r$}#4,true[BOOLEAN],PT0S[TIME_DURATION],compensated[KEYWORD],long_overflow_warn[KEYWORD]) AS $$salary_+_1$SUM$0#2, COUNT($$salary_+_1$SUM$0{r$}#4,true[BOOLEAN],PT0S[TIME_DURATION]) AS $$salary_+_1$SUM$1#3]] | ||
| \_Eval[[MVSINGLEVALUEORNULL(salary{f}#5) AS $$salary_+_1$SUM$0#4]] |
There was a problem hiding this comment.
For debugging purposes, it's a bit weird that we call MVSINGLEVALUEORNULL(salary{f}#5) by the name $$salary_+_1$SUM$0. This seems to be derived from the first SUM expression that we encounter, but it's misleading - this is not salary + 1, this is just salary (sans multivalues). I'd have expected $$salary$SUM$0, instead.
| Project[[s1{r}#0, s2{r}#1]] | ||
| \_Eval[[$$salary_+_1$SUM$0{r$}#2 + $$salary_+_1$SUM$1{r$}#3 * 1[INTEGER] AS s1#0, $$salary_+_1$SUM$0{r$}#2 + $$salary_+_1$SUM$1{r$}#3 * 2[INTEGER] AS s2#1]] | ||
| \_Limit[1000[INTEGER],false,false] | ||
| \_Aggregate[[],[SUM($$salary_+_1$SUM$0{r$}#4,true[BOOLEAN],PT0S[TIME_DURATION],compensated[KEYWORD],long_overflow_warn[KEYWORD]) AS $$salary_+_1$SUM$0#2, COUNT($$salary_+_1$SUM$0{r$}#4,true[BOOLEAN],PT0S[TIME_DURATION]) AS $$salary_+_1$SUM$1#3]] |
There was a problem hiding this comment.
Similarly, the temporary names for the SUM and COUNT expressions are rather confusing. They're re-using the same name that we gave to MVSINGLEVALUEORNULL(salary{f}#5) (which we called $$salary_+_1$SUM$0) and differ just by incrementing the id at the end. This is not going to be easy to debug.
I suggest this alternative naming (not checking how our temp name tooling helps here, so please push back if I'm suggesting nonsense):
MVSINGLEVALUEORNULL(salary{f}#5)->$$salary$MVSINGLEVALUEORNULL$0SUMof 1. ->$$salary$MVSINGLEVALUEORNULL_SUM$1COUNTof 1. ->$$salary$MVSINGLEVALUEORNULL_COUNT$2(note the ever-incrementing counter at the end. Re-using a number might lead to name conflicts.)
There was a problem hiding this comment.
What do you think of $$salary$MVSINGLEVALUEORNULL$SUM$1 and $$salary$MVSINGLEVALUEORNULL$COUNT$2 ? So sticking with $instead ofbetweenMVSINGLEVALUEORNULLand the aggregate function name. This let's us use a bit more of the existing temp name tooling. Though this PR already includes some custom naming logic, so change back to` is easy enough if you prefer.
alex-spies
left a comment
There was a problem hiding this comment.
Done reviewing now. Thanks @parkertimmins !
I have a couple more comments that I think would be good to address before merging. The highest prio ones are:
- Serialization tests for the new function (already mentioned in my previous review)
- Additional tests, esp. with
INLINE STATS. I'm actually unsure if it's going to be affected by this rule or not, but we could easily mess it up on accident. - Ensuring we don't mess up
Sum's properties - see my comments on the optimizer rule.
| FROM employees | ||
| | STATS s1 = SUM(salary + 1), s2 = SUM(salary + 2) |
There was a problem hiding this comment.
Another interesting edge case would be when there is a WHERE filter before the STATS which removes all rows. (SUM of no rows is null, and we could potentially introduce a regression on accident.)
There was a problem hiding this comment.
Good point, I'll add a csv test sumExpressionWithConstantWhereFiltersAllRows
| super(OptimizerRules.TransformDirection.UP); | ||
| } | ||
|
|
||
| private record SvPair(Attribute sum, Attribute count) {} |
There was a problem hiding this comment.
Can we please add some javadoc to explain the purpose of the local records?
|
|
||
| private record Match(Alias alias, Expression dataExpr, Expression constant, Sum sum, boolean isSubtraction, boolean constantIsRight) { | ||
| Key key() { | ||
| return new Key(dataExpr.canonical(), sum.summationMode().canonical()); |
There was a problem hiding this comment.
I don't understand why we single out the summation mode here. Moreover, looking at the individual components of SUMs will make this a little brittle and hard to reason about IMHO. What if we enhance the Sum class and introduce another way how two sums with same field inside them are different?
This is partially duplicating the logic for equality/semantic equality of Sum instances. Equality/semantic equality should better be delegated to the Sum class.
(Additionally and sadly, the summation mode is actually an enum, but for some reason we pass it to Sum as a literal expression (maybe there were serialization concerns at the time)).
Let's see how we can achieve that. What we want to do is to look for SUM(+- foo +- c) and ignore just the constant inside the sum and the negation of foo. That is, we want to group this with any other SUM(+-foo +-c') expression. This is the case when we strip negations from foo, ignore c and end up with semantically equal SUMs.
This can actually be formalized like this:
- Inspect the
field(). - See that it is of the form
+-foo +-cand extract thefooexpression. (Check foldability ofc, unwrap negations aroundfoo). - Replace the field by just
foo. (I think we're missing a helper methodSum#withField()- we can add it!) - Use the canonicalized version of this as a key OR compare with
.semanticEquals()(which is equivalent).
There was a problem hiding this comment.
More test ideas:
- What about multiple levels of negations?
SUM(-(-x) + 1) - What about expressions instead of fields?
SUM((2*salary) + 1, SUM(-3 - (2*salary))?
| * | PROJECT s1, s2, g | ||
| * </pre> | ||
| * | ||
| * <p>{@code x} can be any non-foldable expression (a field reference, a function call, etc.). |
There was a problem hiding this comment.
Let's add: this does not apply when the SUM has a filter, as in SUM(x + 1) WHERE y > z.
| if (s.field() instanceof Add add) { | ||
| if (add.right().foldable() && add.left().foldable() == false) { | ||
| return new Match(alias, add.left(), add.right(), s, false, true); | ||
| } else if (add.left().foldable() && add.right().foldable() == false) { | ||
| return new Match(alias, add.right(), add.left(), s, false, false); | ||
| } | ||
| } else if (s.field() instanceof Sub sub) { | ||
| if (sub.right().foldable() && sub.left().foldable() == false) { | ||
| return new Match(alias, sub.left(), sub.right(), s, true, true); | ||
| } else if (sub.left().foldable() && sub.right().foldable() == false) { | ||
| return new Match(alias, sub.right(), sub.left(), s, true, false); | ||
| } |
There was a problem hiding this comment.
I don't think we handle negations here, so we actually don't group SUM(-x + 1) and SUM(1 - x), is that right?
There was a problem hiding this comment.
Correct, we don't handle that case. I don't think it'd be terribly difficult to extend this to handle ... but this PR is already unwieldy. I'm inclined to add a test documenting that this doesn't apply and leave it as is.
| Map<Match.Key, SvPair> exprToSvPair = new HashMap<>(); | ||
| List<NamedExpression> newAggs = new ArrayList<>(); | ||
| List<Alias> newEvals = new ArrayList<>(); | ||
| int[] counter = { 0 }; |
There was a problem hiding this comment.
nit: What are we counting? I think this could also be called temporaryNameCounter, no?
| List<Source> warningSources = exprFieldSources.get(k); | ||
| var sv = new MvSingleValueOrNull(warningSources.get(0), de, warningSources); | ||
| var svSumName = TemporaryNameGenerator.temporaryName(sv, fs, counter[0]++); | ||
| var svSumExpr = new Sum( |
There was a problem hiding this comment.
We really shouldn't explicitly construct a Sum here. That's bound to break in the future if Sum gains more properties.
What we want to do is just replace the field by sv. Let's add a helper method Sum#withField and use it here. That method should preserve all of the Sum's other properties, in particular the original window.
| source, | ||
| sv, | ||
| Literal.TRUE, | ||
| AggregateFunction.NO_WINDOW, |
There was a problem hiding this comment.
Speaking of the window! Does this work with window aggregations? We should either add tests for that, or exclude this case entirely.
We only generate matches when Sums have no filter. We can expand that. We can add a helper method Sum#isSimpleSum that ensures that we have AggregateFunction.NO_WINDOW and no filter and use it in tryMatch.
The isSimpleSum method should also have an assertion that calling sum.info().properties() has length 5, so that this method gets updated when someone adds a new property later.
There was a problem hiding this comment.
I think it makes sense to just exclude window aggregations entirely. Nice, I think the isSimpleSum along with Sum.withField is a much cleaner way to encapsulate this logic and keep it future-proof!
| import static org.hamcrest.Matchers.instanceOf; | ||
| import static org.hamcrest.Matchers.not; | ||
|
|
||
| public class RewriteSumOfExpressionPlusConstantTests extends AbstractLogicalPlanOptimizerTests { |
There was a problem hiding this comment.
nit: these are so hard to read :/
It'd be nicer if we could use golden tests here and just plug in the TestSubstitutionOnlyOptimizer.
There was a problem hiding this comment.
Haha, I think you may be a bit more experienced at reading the printed query planner output! But this file definitely is hard to read. I went ahead and moved these to golden tests using TestSubstitutionOnlyOptimizer in 8b400bc
|
Also, I saw that there were discussions on preserving warning behavior. The current solution seems to address that and preserve the behavior of "1 warning per 1 agg function evaluation that led to a IMHO, it'd be fine to just emit 1 warning. In the future, we may want to dedupe expressions, and we'll have the same problem when |
My only concern is that this optimization is hidden from the user and what the user actually typed is not consistently found in the warnings. Every time we optimize something and we change the original query (for internal consumption), imo we should always preserve the original user intention regarding |
|
@alex-spies Thanks for the in-depth and thoughtful review! While working on your feedback, I merged main and discovered a more comprehensive rule to do this same optimization was recently merged in #148070. So I am closing this PR. |
|
@parkertimmins , I would really like to re-open this after discussing with @costin . You correctly spotted that #148070 doesn't respect multi-value and warning behavior, while your PR addresses that (and we're reverting #148070 in #148291). This is a lovely PR that's missing only a little last push. The coverage in #148070 is more comprehensive, but I think we can include more aggregation functions in the rule over time, while this PR really sets the foundation and deals with the tricky MV-bits already. |
…onstant Production changes from review feedback: - Refactor the rule to pre-alias MvSingleValueOrNull(x) in a dedicated EVAL before the STATS, so ReplaceAggregateNestedExpressionWithEval leaves it unchanged. The grouping key now delegates to Sum.withField(dataExpr).canonical(), capturing filter/window/mode without duplicating Sum's equality logic. Windowed sums are now excluded in addition to filtered ones, via a new Sum.isSimpleSum(). - MvSingleValueOrNull: add warningSources() accessor; override hashCode/equals to include warningSources; add fast-paths in nullifyMultiValued for blocks where keepMask can be skipped entirely. Tests: - MvSingleValueOrNullSerializationTests: serialization round-trip. - stats.csv-spec / inlinestats.csv-spec: end-to-end CSV tests. - RewriteSumOfExpressionPlusConstantGoldenTests: 12 new full-optimizer golden tests covering INLINE STATS, WHERE before/filtering all rows, EVAL override, AVG alongside, filtered/mixed aggs, chained STATS, mixed int/double, double negation, expressions, and negation vs subtraction not grouped.
Replace the AST-traversal assertions in RewriteSumOfExpressionPlusConstantTests with RewriteSumOfExpressionPlusConstantSubstitutionOnlyGoldenTests, which uses TestSubstitutionOnlyOptimizer to produce readable golden output. Two tests with non-deterministic output remain as simple assertions in the original test class. GoldenTestCase.TestBuilder: add .optimizer() hook so subclasses can supply a custom optimizer for the logical-optimization stage.
# Conflicts: # server/src/main/resources/transport/upper_bounds/9.5.csv
…essionPlusConstantGoldenTests Drop testWhereFiltersAllRowsBeforeStats (covered by csv tests) and remove the ANALYSIS stage from the golden test class since the rule fires during logical optimization and the analysis snapshots test nothing rule-specific.
🔍 Preview links for changed docs⏳ Building and deploying preview... View progress This comment will be updated with preview links when the build is complete. |
ℹ️ Important: Docs version tagging👋 Thanks for updating the docs! Just a friendly reminder that our docs are now cumulative. This means all 9.x versions are documented on the same page and published off of the main branch, instead of creating separate pages for each minor version. We use applies_to tags to mark version-specific features and changes. Expand for a quick overviewWhen to use applies_to tags:✅ At the page level to indicate which products/deployments the content applies to (mandatory) What NOT to do:❌ Don't remove or replace information that applies to an older version 🤔 Need help?
|
Replaces underscore-concatenated suffixes (MVSINGLEVALUEORNULL_SUM) with proper $-separated segments (MVSINGLEVALUEORNULL$SUM) via a new private syntheticName helper. Also fixes COUNT source attribution to point at the data expression rather than the whole STATS clause.
|
@alex-spies Okay, this PR is now ready for re-review. I think I've incorporated all the feedback. I've responded to some of your comments, but I'll recap the comments here, along with some questions, so it's all in one place:
|
|
@alex-spies @parkertimmins Is there anything holding up this PR? Let me know if I can help move it along. |
|
@costin I spoke with Alex earlier today. He has another review pass to do, but has it in his queue. I believe all feedback has been addressed, so once he gives the go-ahead, I'll merge it. |
alex-spies
left a comment
There was a problem hiding this comment.
Awesome PR. Thanks a lot for the iterations, @parkertimmins !
|
Retroactively closes elastic/esql-planning#694 (the SUM(field+const) re-land). |
Adds an optimizer rule that rewrites
SUM(field + c)intoSUM(field) + c * COUNT(field). When a query contains multiple such expressions over the same field (e.g.SUM(x+1),SUM(x+2), ...), they share a single SUM/COUNT pair instead of computing independent aggregations. The field is wrapped in an internal MvSingleValueOrNull function to preserve correct null semantics for multi-valued inputs.Closes #140470