Skip to content

Commit 984063f

Browse files
drempapisdrempapis
authored
Merge branch 'main' into multiproject/lightweight_stats_action_rrc
2 parents 2a8a7eb + 4e4b89e commit 984063f

File tree

11 files changed

+331
-56
lines changed

11 files changed

+331
-56
lines changed

‎benchmarks/src/main/java/org/elasticsearch/benchmark/esql/QueryPlanningBenchmark.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ private LogicalPlan plan(String query) {
123123
}
124124

125125
@Benchmark
126-
public void run(Blackhole blackhole) {
126+
public void manyFields(Blackhole blackhole) {
127127
blackhole.consume(plan("FROM test | LIMIT 10"));
128128
}
129129
}

‎docs/changelog/127524.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 127524
2+
summary: Resolve groupings in aggregate before resolving references to groupings in
3+
the aggregations
4+
area: ES|QL
5+
type: bug
6+
issues: []

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -184,16 +184,19 @@ protected void doCollectFirst(Predicate<? super T> predicate, List<T> matches) {
184184
public T transformDown(Function<? super T, ? extends T> rule) {
185185
T root = rule.apply((T) this);
186186
Node<T> node = this.equals(root) ? this : root;
187-
188187
return node.transformChildren(child -> child.transformDown(rule));
189188
}
190189

191190
@SuppressWarnings("unchecked")
192191
public <E extends T> T transformDown(Class<E> typeToken, Function<E, ? extends T> rule) {
193-
// type filtering function
194192
return transformDown((t) -> (typeToken.isInstance(t) ? rule.apply((E) t) : t));
195193
}
196194

195+
@SuppressWarnings("unchecked")
196+
public <E extends T> T transformDown(Predicate<Node<?>> nodePredicate, Function<E, ? extends T> rule) {
197+
return transformDown((t) -> (nodePredicate.test(t) ? rule.apply((E) t) : t));
198+
}
199+
197200
@SuppressWarnings("unchecked")
198201
public T transformUp(Function<? super T, ? extends T> rule) {
199202
T transformed = transformChildren(child -> child.transformUp(rule));
@@ -203,10 +206,14 @@ public T transformUp(Function<? super T, ? extends T> rule) {
203206

204207
@SuppressWarnings("unchecked")
205208
public <E extends T> T transformUp(Class<E> typeToken, Function<E, ? extends T> rule) {
206-
// type filtering function
207209
return transformUp((t) -> (typeToken.isInstance(t) ? rule.apply((E) t) : t));
208210
}
209211

212+
@SuppressWarnings("unchecked")
213+
public <E extends T> T transformUp(Predicate<Node<?>> nodePredicate, Function<E, ? extends T> rule) {
214+
return transformUp((t) -> (nodePredicate.test(t) ? rule.apply((E) t) : t));
215+
}
216+
210217
@SuppressWarnings("unchecked")
211218
protected <R extends Function<? super T, ? extends T>> T transformChildren(Function<T, ? extends T> traversalOperation) {
212219
boolean childrenChanged = false;

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/bucket.csv-spec‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -850,3 +850,37 @@ c:long | b:date
850850
11 | 1984-05-01T00:00:00.000Z
851851
11 | 1991-01-01T00:00:00.000Z
852852
;
853+
854+
resolveGroupingsBeforeResolvingImplicitReferencesToGroupings
855+
required_capability: resolve_groupings_before_resolving_references_to_groupings_in_aggregations
856+
857+
FROM employees
858+
| STATS c = count(emp_no), b = BUCKET(hire_date, "1 year") + 1 year BY yr = BUCKET(hire_date, "1 year")
859+
| SORT yr
860+
| LIMIT 5
861+
;
862+
863+
c:long | b:datetime | yr:datetime
864+
11 | 1986-01-01T00:00:00.000Z | 1985-01-01T00:00:00.000Z
865+
11 | 1987-01-01T00:00:00.000Z | 1986-01-01T00:00:00.000Z
866+
15 | 1988-01-01T00:00:00.000Z | 1987-01-01T00:00:00.000Z
867+
9 | 1989-01-01T00:00:00.000Z | 1988-01-01T00:00:00.000Z
868+
13 | 1990-01-01T00:00:00.000Z | 1989-01-01T00:00:00.000Z
869+
;
870+
871+
resolveGroupingsBeforeResolvingExplicitReferencesToGroupings
872+
required_capability: resolve_groupings_before_resolving_references_to_groupings_in_aggregations
873+
874+
FROM employees
875+
| STATS c = count(emp_no), b = yr + 1 year BY yr = BUCKET(hire_date, "1 year")
876+
| SORT yr
877+
| LIMIT 5
878+
;
879+
880+
c:long | b:datetime | yr:datetime
881+
11 | 1986-01-01T00:00:00.000Z | 1985-01-01T00:00:00.000Z
882+
11 | 1987-01-01T00:00:00.000Z | 1986-01-01T00:00:00.000Z
883+
15 | 1988-01-01T00:00:00.000Z | 1987-01-01T00:00:00.000Z
884+
9 | 1989-01-01T00:00:00.000Z | 1988-01-01T00:00:00.000Z
885+
13 | 1990-01-01T00:00:00.000Z | 1989-01-01T00:00:00.000Z
886+
;

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3097,3 +3097,27 @@ ROW a = [1,2,3], b = 5
30973097
STD_DEV(a):double | STD_DEV(b):double
30983098
0.816496580927726 | 0.0
30993099
;
3100+
3101+
resolveGroupingsBeforeResolvingImplicitReferencesToGroupings
3102+
required_capability: resolve_groupings_before_resolving_references_to_groupings_in_aggregations
3103+
3104+
FROM employees
3105+
| EVAL date = "2025-01-01"::datetime
3106+
| stats m = MAX(hire_date) BY d = (date == "2025-01-01")
3107+
;
3108+
3109+
m:datetime | d:boolean
3110+
1999-04-30T00:00:00.000Z | true
3111+
;
3112+
3113+
resolveGroupingsBeforeResolvingExplicitReferencesToGroupings
3114+
required_capability: resolve_groupings_before_resolving_references_to_groupings_in_aggregations
3115+
3116+
FROM employees
3117+
| EVAL date = "2025-01-01"::datetime
3118+
| stats m = MAX(hire_date), x = d::int + 1 BY d = (date == "2025-01-01")
3119+
;
3120+
3121+
m:datetime | x:integer | d:boolean
3122+
1999-04-30T00:00:00.000Z | 2 | true
3123+
;

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1066,7 +1066,10 @@ public enum Cap {
10661066
*/
10671067
FIRST_OVER_TIME(Build.current().isSnapshot()),
10681068

1069-
;
1069+
/**
1070+
* Resolve groupings before resolving references to groupings in the aggregations.
1071+
*/
1072+
RESOLVE_GROUPINGS_BEFORE_RESOLVING_REFERENCES_TO_GROUPINGS_IN_AGGREGATIONS;
10701073

10711074
private final boolean enabled;
10721075

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java‎

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -173,13 +173,8 @@ public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerCon
173173
),
174174
new Batch<>(
175175
"Resolution",
176-
/*
177-
* ImplicitCasting must be before ResolveRefs. Because a reference is created for a Bucket in Aggregate's aggregates,
178-
* resolving this reference before implicit casting may cause this reference to have customMessage=true, it prevents further
179-
* attempts to resolve this reference.
180-
*/
181-
new ImplicitCasting(),
182176
new ResolveRefs(),
177+
new ImplicitCasting(),
183178
new ResolveUnionTypes() // Must be after ResolveRefs, so union types can be found
184179
),
185180
new Batch<>("Finish Analysis", Limiter.ONCE, new AddImplicitLimit(), new AddImplicitForkLimit(), new UnionTypesCleanup())
@@ -574,7 +569,7 @@ private Aggregate resolveAggregate(Aggregate aggregate, List<Attribute> children
574569
}
575570
}
576571

577-
if (Resolvables.resolved(groupings) == false || (Resolvables.resolved(aggregates) == false)) {
572+
if (Resolvables.resolved(groupings) == false || Resolvables.resolved(aggregates) == false) {
578573
ArrayList<Attribute> resolved = new ArrayList<>();
579574
for (Expression e : groupings) {
580575
Attribute attr = Expressions.attribute(e);
@@ -585,17 +580,29 @@ private Aggregate resolveAggregate(Aggregate aggregate, List<Attribute> children
585580
List<Attribute> resolvedList = NamedExpressions.mergeOutputAttributes(resolved, childrenOutput);
586581

587582
List<NamedExpression> newAggregates = new ArrayList<>();
588-
for (NamedExpression ag : aggregate.aggregates()) {
589-
var agg = (NamedExpression) ag.transformUp(UnresolvedAttribute.class, ua -> {
590-
Expression ne = ua;
591-
Attribute maybeResolved = maybeResolveAttribute(ua, resolvedList);
592-
if (maybeResolved != null) {
593-
changed.set(true);
594-
ne = maybeResolved;
595-
}
596-
return ne;
597-
});
598-
newAggregates.add(agg);
583+
// If the groupings are not resolved, skip the resolution of the references to groupings in the aggregates, resolve the
584+
// aggregations that do not reference to groupings, so that the fields/attributes referenced by the aggregations can be
585+
// resolved, and verifier doesn't report field/reference/column not found errors for them.
586+
boolean groupingResolved = Resolvables.resolved(groupings);
587+
int size = groupingResolved ? aggregates.size() : aggregates.size() - groupings.size();
588+
for (int i = 0; i < aggregates.size(); i++) {
589+
NamedExpression maybeResolvedAgg = aggregates.get(i);
590+
if (i < size) { // Skip resolving references to groupings in the aggregations if the groupings are not resolved yet.
591+
maybeResolvedAgg = (NamedExpression) maybeResolvedAgg.transformUp(UnresolvedAttribute.class, ua -> {
592+
Expression ne = ua;
593+
Attribute maybeResolved = maybeResolveAttribute(ua, resolvedList);
594+
// An item in aggregations can reference to groupings explicitly, if groupings are not resolved yet and
595+
// maybeResolved is not resolved, return the original UnresolvedAttribute, so that it has another chance
596+
// to get resolved in the next iteration.
597+
// For example STATS c = count(emp_no), x = d::int + 1 BY d = (date == "2025-01-01")
598+
if (groupingResolved || maybeResolved.resolved()) {
599+
changed.set(true);
600+
ne = maybeResolved;
601+
}
602+
return ne;
603+
});
604+
}
605+
newAggregates.add(maybeResolvedAgg);
599606
}
600607

601608
// TODO: remove this when Stats interface is removed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/OptimizerRules.java‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,13 @@
77
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
88

99
import org.elasticsearch.xpack.esql.core.expression.Expression;
10+
import org.elasticsearch.xpack.esql.core.tree.Node;
1011
import org.elasticsearch.xpack.esql.core.util.ReflectionUtils;
1112
import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
13+
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
14+
import org.elasticsearch.xpack.esql.plan.logical.Limit;
1215
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
16+
import org.elasticsearch.xpack.esql.plan.logical.Project;
1317
import org.elasticsearch.xpack.esql.rule.ParameterizedRule;
1418
import org.elasticsearch.xpack.esql.rule.Rule;
1519

@@ -55,12 +59,26 @@ public OptimizerExpressionRule(TransformDirection direction) {
5559
@Override
5660
public final LogicalPlan apply(LogicalPlan plan, LogicalOptimizerContext ctx) {
5761
return direction == TransformDirection.DOWN
58-
? plan.transformExpressionsDown(expressionTypeToken, e -> rule(e, ctx))
59-
: plan.transformExpressionsUp(expressionTypeToken, e -> rule(e, ctx));
62+
? plan.transformExpressionsDown(this::shouldVisit, expressionTypeToken, e -> rule(e, ctx))
63+
: plan.transformExpressionsUp(this::shouldVisit, expressionTypeToken, e -> rule(e, ctx));
6064
}
6165

6266
protected abstract Expression rule(E e, LogicalOptimizerContext ctx);
6367

68+
/**
69+
* Defines if a node should be visited or not.
70+
* Allows to skip nodes that are not applicable for the rule even if they contain expressions.
71+
* By default that skips FROM, LIMIT, PROJECT, KEEP and DROP but this list could be extended or replaced in subclasses.
72+
*/
73+
protected boolean shouldVisit(Node<?> node) {
74+
return switch (node) {
75+
case EsRelation relation -> false;
76+
case Project project -> false;// this covers project, keep and drop
77+
case Limit limit -> false;
78+
default -> true;
79+
};
80+
}
81+
6482
public Class<E> expressionToken() {
6583
return expressionTypeToken;
6684
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QueryPlan.java‎

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import java.util.List;
1919
import java.util.function.Consumer;
2020
import java.util.function.Function;
21+
import java.util.function.Predicate;
2122

2223
/**
2324
* There are two main types of plans, {@code LogicalPlan} and {@code PhysicalPlan}
@@ -109,22 +110,36 @@ public <E extends Expression> PlanType transformExpressionsOnlyUp(Class<E> typeT
109110
return transformPropertiesOnly(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)));
110111
}
111112

112-
public PlanType transformExpressionsDown(Function<Expression, ? extends Expression> rule) {
113-
return transformExpressionsDown(Expression.class, rule);
114-
}
115-
116113
public <E extends Expression> PlanType transformExpressionsDown(Class<E> typeToken, Function<E, ? extends Expression> rule) {
117114
return transformPropertiesDown(Object.class, e -> doTransformExpression(e, exp -> exp.transformDown(typeToken, rule)));
118115
}
119116

120-
public PlanType transformExpressionsUp(Function<Expression, ? extends Expression> rule) {
121-
return transformExpressionsUp(Expression.class, rule);
117+
public <E extends Expression> PlanType transformExpressionsDown(
118+
Predicate<Node<?>> shouldVisit,
119+
Class<E> typeToken,
120+
Function<E, ? extends Expression> rule
121+
) {
122+
return transformDown(
123+
shouldVisit,
124+
t -> t.transformNodeProps(Object.class, e -> doTransformExpression(e, exp -> exp.transformDown(typeToken, rule)))
125+
);
122126
}
123127

124128
public <E extends Expression> PlanType transformExpressionsUp(Class<E> typeToken, Function<E, ? extends Expression> rule) {
125129
return transformPropertiesUp(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)));
126130
}
127131

132+
public <E extends Expression> PlanType transformExpressionsUp(
133+
Predicate<Node<?>> shouldVisit,
134+
Class<E> typeToken,
135+
Function<E, ? extends Expression> rule
136+
) {
137+
return transformUp(
138+
shouldVisit,
139+
t -> t.transformNodeProps(Object.class, e -> doTransformExpression(e, exp -> exp.transformUp(typeToken, rule)))
140+
);
141+
}
142+
128143
@SuppressWarnings("unchecked")
129144
private static Object doTransformExpression(Object arg, Function<Expression, ? extends Expression> traversal) {
130145
if (arg instanceof Expression exp) {
@@ -184,18 +199,10 @@ public <E extends Expression> void forEachExpression(Class<E> typeToken, Consume
184199
forEachPropertyOnly(Object.class, e -> doForEachExpression(e, exp -> exp.forEachDown(typeToken, rule)));
185200
}
186201

187-
public void forEachExpressionDown(Consumer<? super Expression> rule) {
188-
forEachExpressionDown(Expression.class, rule);
189-
}
190-
191202
public <E extends Expression> void forEachExpressionDown(Class<? extends E> typeToken, Consumer<? super E> rule) {
192203
forEachPropertyDown(Object.class, e -> doForEachExpression(e, exp -> exp.forEachDown(typeToken, rule)));
193204
}
194205

195-
public void forEachExpressionUp(Consumer<? super Expression> rule) {
196-
forEachExpressionUp(Expression.class, rule);
197-
}
198-
199206
public <E extends Expression> void forEachExpressionUp(Class<E> typeToken, Consumer<? super E> rule) {
200207
forEachPropertyUp(Object.class, e -> doForEachExpression(e, exp -> exp.forEachUp(typeToken, rule)));
201208
}

0 commit comments

Comments
 (0)