Skip to content

Commit bdbd15e

Browse files
alex-spiesalex-spies
authored
ESQL: Push down filters even in case of renames in Evals (#114411) (#114612)
Optimize queries like ... | EVAL b = a, c = b | WHERE c > 2 to ... | WHERE a > 2 | EVAL b = a, c = b
1 parent 915d7cc commit bdbd15e

File tree

3 files changed

+184
-14
lines changed

3 files changed

+184
-14
lines changed

‎docs/changelog/114411.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 114411
2+
summary: "ESQL: Push down filters even in case of renames in Evals"
3+
area: ES|QL
4+
type: enhancement
5+
issues: []

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

Lines changed: 47 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,13 @@
77

88
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
99

10+
import org.elasticsearch.xpack.esql.core.expression.Alias;
1011
import org.elasticsearch.xpack.esql.core.expression.Attribute;
12+
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
1113
import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
1214
import org.elasticsearch.xpack.esql.core.expression.Expression;
1315
import org.elasticsearch.xpack.esql.core.expression.Expressions;
16+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
1417
import org.elasticsearch.xpack.esql.core.expression.predicate.Predicates;
1518
import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
1619
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
@@ -25,6 +28,7 @@
2528

2629
import java.util.ArrayList;
2730
import java.util.List;
31+
import java.util.function.Function;
2832
import java.util.function.Predicate;
2933

3034
public final class PushDownAndCombineFilters extends OptimizerRules.OptimizerRule<Filter> {
@@ -43,20 +47,37 @@ protected LogicalPlan rule(Filter filter) {
4347
filter,
4448
agg,
4549
e -> e instanceof Attribute && agg.output().contains(e) && agg.groupings().contains(e) == false
46-
|| e instanceof AggregateFunction
50+
|| e instanceof AggregateFunction,
51+
NO_OP
4752
);
4853
} else if (child instanceof Eval eval) {
49-
// Don't push if Filter (still) contains references of Eval's fields.
50-
var attributes = new AttributeSet(Expressions.asAttributes(eval.fields()));
51-
plan = maybePushDownPastUnary(filter, eval, attributes::contains);
54+
// Don't push if Filter (still) contains references to Eval's fields.
55+
// Account for simple aliases in the Eval, though - these shouldn't stop us.
56+
AttributeMap.Builder<Expression> aliasesBuilder = AttributeMap.builder();
57+
for (Alias alias : eval.fields()) {
58+
aliasesBuilder.put(alias.toAttribute(), alias.child());
59+
}
60+
AttributeMap<Expression> evalAliases = aliasesBuilder.build();
61+
62+
Function<Expression, Expression> resolveRenames = expr -> expr.transformDown(ReferenceAttribute.class, r -> {
63+
Expression resolved = evalAliases.resolve(r, null);
64+
// Avoid resolving to an intermediate attribute that only lives inside the Eval - only replace if the attribute existed
65+
// before the Eval.
66+
if (resolved instanceof Attribute && eval.inputSet().contains(resolved)) {
67+
return resolved;
68+
}
69+
return r;
70+
});
71+
72+
plan = maybePushDownPastUnary(filter, eval, evalAliases::containsKey, resolveRenames);
5273
} else if (child instanceof RegexExtract re) {
5374
// Push down filters that do not rely on attributes created by RegexExtract
5475
var attributes = new AttributeSet(Expressions.asAttributes(re.extractedFields()));
55-
plan = maybePushDownPastUnary(filter, re, attributes::contains);
76+
plan = maybePushDownPastUnary(filter, re, attributes::contains, NO_OP);
5677
} else if (child instanceof Enrich enrich) {
5778
// Push down filters that do not rely on attributes created by Enrich
5879
var attributes = new AttributeSet(Expressions.asAttributes(enrich.enrichFields()));
59-
plan = maybePushDownPastUnary(filter, enrich, attributes::contains);
80+
plan = maybePushDownPastUnary(filter, enrich, attributes::contains, NO_OP);
6081
} else if (child instanceof Project) {
6182
return PushDownUtils.pushDownPastProject(filter);
6283
} else if (child instanceof OrderBy orderBy) {
@@ -67,21 +88,35 @@ protected LogicalPlan rule(Filter filter) {
6788
return plan;
6889
}
6990

70-
private static LogicalPlan maybePushDownPastUnary(Filter filter, UnaryPlan unary, Predicate<Expression> cannotPush) {
91+
private static Function<Expression, Expression> NO_OP = expression -> expression;
92+
93+
private static LogicalPlan maybePushDownPastUnary(
94+
Filter filter,
95+
UnaryPlan unary,
96+
Predicate<Expression> cannotPush,
97+
Function<Expression, Expression> resolveRenames
98+
) {
7199
LogicalPlan plan;
72100
List<Expression> pushable = new ArrayList<>();
73101
List<Expression> nonPushable = new ArrayList<>();
74102
for (Expression exp : Predicates.splitAnd(filter.condition())) {
75-
(exp.anyMatch(cannotPush) ? nonPushable : pushable).add(exp);
103+
Expression resolvedExp = resolveRenames.apply(exp);
104+
if (resolvedExp.anyMatch(cannotPush)) {
105+
// Add the original expression to the non-pushables.
106+
nonPushable.add(exp);
107+
} else {
108+
// When we can push down, we use the resolved expression.
109+
pushable.add(resolvedExp);
110+
}
76111
}
77112
// Push the filter down even if it might not be pushable all the way to ES eventually: eval'ing it closer to the source,
78113
// potentially still in the Exec Engine, distributes the computation.
79-
if (pushable.size() > 0) {
80-
if (nonPushable.size() > 0) {
81-
Filter pushed = new Filter(filter.source(), unary.child(), Predicates.combineAnd(pushable));
114+
if (pushable.isEmpty() == false) {
115+
Filter pushed = filter.with(unary.child(), Predicates.combineAnd(pushable));
116+
if (nonPushable.isEmpty() == false) {
82117
plan = filter.with(unary.replaceChild(pushed), Predicates.combineAnd(nonPushable));
83118
} else {
84-
plan = unary.replaceChild(filter.with(unary.child(), filter.condition()));
119+
plan = unary.replaceChild(pushed);
85120
}
86121
} else {
87122
plan = filter;

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

Lines changed: 132 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,14 @@
99

1010
import org.elasticsearch.index.IndexMode;
1111
import org.elasticsearch.test.ESTestCase;
12+
import org.elasticsearch.xpack.esql.core.expression.Alias;
13+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
14+
import org.elasticsearch.xpack.esql.core.expression.Expression;
1215
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
16+
import org.elasticsearch.xpack.esql.core.expression.predicate.Predicates;
1317
import org.elasticsearch.xpack.esql.core.expression.predicate.logical.And;
1418
import org.elasticsearch.xpack.esql.expression.function.aggregate.Count;
19+
import org.elasticsearch.xpack.esql.expression.function.scalar.math.Pow;
1520
import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
1621
import org.elasticsearch.xpack.esql.expression.function.scalar.string.WildcardLike;
1722
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThan;
@@ -20,17 +25,23 @@
2025
import org.elasticsearch.xpack.esql.index.EsIndex;
2126
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
2227
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
28+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
2329
import org.elasticsearch.xpack.esql.plan.logical.Filter;
30+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
31+
import org.elasticsearch.xpack.esql.plan.logical.Project;
2432
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
2533

34+
import java.util.ArrayList;
2635
import java.util.List;
2736

2837
import static java.util.Collections.emptyList;
2938
import static java.util.Collections.emptyMap;
3039
import static java.util.Collections.singletonList;
40+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.FOUR;
3141
import static org.elasticsearch.xpack.esql.EsqlTestUtils.ONE;
3242
import static org.elasticsearch.xpack.esql.EsqlTestUtils.THREE;
3343
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TWO;
44+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
3445
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getFieldAttribute;
3546
import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOf;
3647
import static org.elasticsearch.xpack.esql.EsqlTestUtils.greaterThanOrEqualOf;
@@ -77,6 +88,115 @@ public void testPushDownFilter() {
7788
assertEquals(new EsqlProject(EMPTY, combinedFilter, projections), new PushDownAndCombineFilters().apply(fb));
7889
}
7990

91+
public void testPushDownFilterPastRenamingProject() {
92+
FieldAttribute a = getFieldAttribute("a");
93+
FieldAttribute b = getFieldAttribute("b");
94+
EsRelation relation = relation(List.of(a, b));
95+
96+
Alias aRenamed = new Alias(EMPTY, "a_renamed", a);
97+
Alias aRenamedTwice = new Alias(EMPTY, "a_renamed_twice", aRenamed.toAttribute());
98+
Alias bRenamed = new Alias(EMPTY, "b_renamed", b);
99+
100+
Project project = new Project(EMPTY, relation, List.of(aRenamed, aRenamedTwice, bRenamed));
101+
102+
GreaterThan aRenamedTwiceGreaterThanOne = greaterThanOf(aRenamedTwice.toAttribute(), ONE);
103+
LessThan bRenamedLessThanTwo = lessThanOf(bRenamed.toAttribute(), TWO);
104+
Filter filter = new Filter(EMPTY, project, Predicates.combineAnd(List.of(aRenamedTwiceGreaterThanOne, bRenamedLessThanTwo)));
105+
106+
LogicalPlan optimized = new PushDownAndCombineFilters().apply(filter);
107+
108+
Project optimizedProject = as(optimized, Project.class);
109+
assertEquals(optimizedProject.projections(), project.projections());
110+
Filter optimizedFilter = as(optimizedProject.child(), Filter.class);
111+
assertEquals(optimizedFilter.condition(), Predicates.combineAnd(List.of(greaterThanOf(a, ONE), lessThanOf(b, TWO))));
112+
EsRelation optimizedRelation = as(optimizedFilter.child(), EsRelation.class);
113+
assertEquals(optimizedRelation, relation);
114+
}
115+
116+
// ... | eval a_renamed = a, a_renamed_twice = a_renamed, a_squared = pow(a, 2)
117+
// | where a_renamed > 1 and a_renamed_twice < 2 and a_squared < 4
118+
// ->
119+
// ... | where a > 1 and a < 2 | eval a_renamed = a, a_renamed_twice = a_renamed, non_pushable = pow(a, 2) | where a_squared < 4
120+
public void testPushDownFilterOnAliasInEval() {
121+
FieldAttribute a = getFieldAttribute("a");
122+
FieldAttribute b = getFieldAttribute("b");
123+
EsRelation relation = relation(List.of(a, b));
124+
125+
Alias aRenamed = new Alias(EMPTY, "a_renamed", a);
126+
Alias aRenamedTwice = new Alias(EMPTY, "a_renamed_twice", aRenamed.toAttribute());
127+
Alias bRenamed = new Alias(EMPTY, "b_renamed", b);
128+
Alias aSquared = new Alias(EMPTY, "a_squared", new Pow(EMPTY, a, TWO));
129+
Eval eval = new Eval(EMPTY, relation, List.of(aRenamed, aRenamedTwice, aSquared, bRenamed));
130+
131+
// We'll construct a Filter after the Eval that has conditions that can or cannot be pushed before the Eval.
132+
List<Expression> pushableConditionsBefore = List.of(
133+
greaterThanOf(a.toAttribute(), TWO),
134+
greaterThanOf(aRenamed.toAttribute(), ONE),
135+
lessThanOf(aRenamedTwice.toAttribute(), TWO),
136+
lessThanOf(aRenamedTwice.toAttribute(), bRenamed.toAttribute())
137+
);
138+
List<Expression> pushableConditionsAfter = List.of(
139+
greaterThanOf(a.toAttribute(), TWO),
140+
greaterThanOf(a.toAttribute(), ONE),
141+
lessThanOf(a.toAttribute(), TWO),
142+
lessThanOf(a.toAttribute(), b.toAttribute())
143+
);
144+
List<Expression> nonPushableConditions = List.of(
145+
lessThanOf(aSquared.toAttribute(), FOUR),
146+
greaterThanOf(aRenamedTwice.toAttribute(), aSquared.toAttribute())
147+
);
148+
149+
// Try different combinations of pushable and non-pushable conditions in the filter while also randomizing their order a bit.
150+
for (int numPushable = 0; numPushable <= pushableConditionsBefore.size(); numPushable++) {
151+
for (int numNonPushable = 0; numNonPushable <= nonPushableConditions.size(); numNonPushable++) {
152+
if (numPushable == 0 && numNonPushable == 0) {
153+
continue;
154+
}
155+
156+
List<Expression> conditions = new ArrayList<>();
157+
158+
int pushableIndex = 0, nonPushableIndex = 0;
159+
// Loop and add either a pushable or non-pushable condition to the filter.
160+
boolean addPushable;
161+
while (pushableIndex < numPushable || nonPushableIndex < numNonPushable) {
162+
if (pushableIndex == numPushable) {
163+
addPushable = false;
164+
} else if (nonPushableIndex == numNonPushable) {
165+
addPushable = true;
166+
} else {
167+
addPushable = randomBoolean();
168+
}
169+
170+
if (addPushable) {
171+
conditions.add(pushableConditionsBefore.get(pushableIndex++));
172+
} else {
173+
conditions.add(nonPushableConditions.get(nonPushableIndex++));
174+
}
175+
}
176+
177+
Filter filter = new Filter(EMPTY, eval, Predicates.combineAnd(conditions));
178+
179+
LogicalPlan plan = new PushDownAndCombineFilters().apply(filter);
180+
181+
if (numNonPushable > 0) {
182+
Filter optimizedFilter = as(plan, Filter.class);
183+
assertEquals(optimizedFilter.condition(), Predicates.combineAnd(nonPushableConditions.subList(0, numNonPushable)));
184+
plan = optimizedFilter.child();
185+
}
186+
Eval optimizedEval = as(plan, Eval.class);
187+
assertEquals(optimizedEval.fields(), eval.fields());
188+
plan = optimizedEval.child();
189+
if (numPushable > 0) {
190+
Filter pushedFilter = as(plan, Filter.class);
191+
assertEquals(pushedFilter.condition(), Predicates.combineAnd(pushableConditionsAfter.subList(0, numPushable)));
192+
plan = pushedFilter.child();
193+
}
194+
EsRelation optimizedRelation = as(plan, EsRelation.class);
195+
assertEquals(optimizedRelation, relation);
196+
}
197+
}
198+
}
199+
80200
public void testPushDownLikeRlikeFilter() {
81201
EsRelation relation = relation();
82202
org.elasticsearch.xpack.esql.core.expression.predicate.regex.RLike conditionA = rlike(getFieldAttribute("a"), "foo");
@@ -125,7 +245,17 @@ public void testSelectivelyPushDownFilterPastFunctionAgg() {
125245
assertEquals(expected, new PushDownAndCombineFilters().apply(fb));
126246
}
127247

128-
private EsRelation relation() {
129-
return new EsRelation(EMPTY, new EsIndex(randomAlphaOfLength(8), emptyMap()), randomFrom(IndexMode.values()), randomBoolean());
248+
private static EsRelation relation() {
249+
return relation(List.of());
250+
}
251+
252+
private static EsRelation relation(List<Attribute> fieldAttributes) {
253+
return new EsRelation(
254+
EMPTY,
255+
new EsIndex(randomAlphaOfLength(8), emptyMap()),
256+
fieldAttributes,
257+
randomFrom(IndexMode.values()),
258+
randomBoolean()
259+
);
130260
}
131261
}

0 commit comments

Comments
 (0)