Skip to content

Commit a036535

Browse files
astefanastefan
committed
Address reviews
1 parent 16311c2 commit a036535

File tree

4 files changed

+196
-23
lines changed

4 files changed

+196
-23
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,13 +140,15 @@ protected static Batch<LogicalPlan> substitutions() {
140140
// translate metric aggregates after surrogate substitution and replace nested expressions with eval (again)
141141
new TranslateMetricsAggregate(),
142142
new ReplaceAggregateNestedExpressionWithEval(),
143+
// this one needs to be placed before ReplaceAliasingEvalWithProject, so that any potential aliasing eval (eval x = y)
144+
// is not replaced with a Project before the eval to be copied on the left hand side of an InlineJoin
145+
new PropagateInlineEvals(),
143146
new ReplaceRegexMatch(),
144147
new ReplaceTrivialTypeConversions(),
145-
new ReplaceOrderByExpressionWithEval(),
146-
new PropagateInlineEvals(),
147148
new ReplaceAliasingEvalWithProject(),
148149
new SkipQueryOnEmptyMappings(),
149-
new SubstituteSpatialSurrogates()
150+
new SubstituteSpatialSurrogates(),
151+
new ReplaceOrderByExpressionWithEval()
150152
// new NormalizeAggregate(), - waits on https://github.com/elastic/elasticsearch/issues/100634
151153
);
152154
}

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

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ protected LogicalPlan rule(InlineJoin plan) {
3939

4040
// grouping references
4141
List<Alias> groupingAlias = new ArrayList<>();
42+
// TODO: replace this with AttributeSet
4243
Map<String, ReferenceAttribute> groupingRefs = new LinkedHashMap<>();
4344

4445
// perform only one iteration that does two things
@@ -59,23 +60,21 @@ protected LogicalPlan rule(InlineJoin plan) {
5960
}
6061

6162
// find their declaration and remove it
62-
// TODO: this doesn't take into account aliasing
6363
if (p instanceof Eval eval) {
64-
if (groupingRefs.size() > 0) {
65-
List<Alias> fields = eval.fields();
66-
List<Alias> remainingEvals = new ArrayList<>(fields.size());
67-
for (Alias f : fields) {
68-
if (groupingRefs.remove(f.name()) != null) {
69-
groupingAlias.add(f);
70-
} else {
71-
remainingEvals.add(f);
72-
}
73-
}
74-
if (remainingEvals.size() != fields.size()) {
75-
// if all fields are moved, replace the eval
76-
p = remainingEvals.size() == 0 ? eval.child() : new Eval(eval.source(), eval.child(), remainingEvals);
64+
List<Alias> fields = eval.fields();
65+
List<Alias> remainingEvals = new ArrayList<>(fields.size());
66+
for (Alias f : fields) {
67+
// TODO: look into identifying refs by their NameIds instead
68+
if (groupingRefs.remove(f.name()) != null) {
69+
groupingAlias.add(f);
70+
} else {
71+
remainingEvals.add(f);
7772
}
7873
}
74+
if (remainingEvals.size() != fields.size()) {
75+
// if all fields are moved, replace the eval
76+
p = remainingEvals.size() == 0 ? eval.child() : new Eval(eval.source(), eval.child(), remainingEvals);
77+
}
7978
}
8079
return p;
8180
});

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -214,8 +214,8 @@ public class LogicalPlanOptimizerTests extends ESTestCase {
214214
private static EnrichResolution enrichResolution;
215215
private static final LiteralsOnTheRight LITERALS_ON_THE_RIGHT = new LiteralsOnTheRight();
216216

217-
private static class SubstitutionOnlyOptimizer extends LogicalPlanOptimizer {
218-
static SubstitutionOnlyOptimizer INSTANCE = new SubstitutionOnlyOptimizer(unboundLogicalOptimizerContext());
217+
public static class SubstitutionOnlyOptimizer extends LogicalPlanOptimizer {
218+
public static SubstitutionOnlyOptimizer INSTANCE = new SubstitutionOnlyOptimizer(unboundLogicalOptimizerContext());
219219

220220
SubstitutionOnlyOptimizer(LogicalOptimizerContext optimizerContext) {
221221
super(optimizerContext);
@@ -6078,10 +6078,6 @@ public void testSimplifyComparisonArithmeticWithFloatsAndDirectionChange() {
60786078
doTestSimplifyComparisonArithmetics("float * -2 < 4", "float", GT, -2d);
60796079
}
60806080

6081-
private void assertNullLiteral(Expression expression) {
6082-
assertNull(as(expression, Literal.class).value());
6083-
}
6084-
60856081
@AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/108519")
60866082
public void testSimplifyComparisonArithmeticSkippedOnIntegerArithmeticalOverflow() {
60876083
assertNotSimplified("integer - 1 " + randomBinaryComparison() + " " + Long.MAX_VALUE);

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

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
9+
10+
import org.elasticsearch.index.IndexMode;
11+
import org.elasticsearch.test.ESTestCase;
12+
import org.elasticsearch.xpack.esql.EsqlTestUtils;
13+
import org.elasticsearch.xpack.esql.analysis.Analyzer;
14+
import org.elasticsearch.xpack.esql.analysis.AnalyzerContext;
15+
import org.elasticsearch.xpack.esql.analysis.EnrichResolution;
16+
import org.elasticsearch.xpack.esql.core.expression.Expressions;
17+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
18+
import org.elasticsearch.xpack.esql.core.type.EsField;
19+
import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
20+
import org.elasticsearch.xpack.esql.index.EsIndex;
21+
import org.elasticsearch.xpack.esql.index.IndexResolution;
22+
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
23+
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizerTests;
24+
import org.elasticsearch.xpack.esql.parser.EsqlParser;
25+
import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
26+
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
27+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
28+
import org.elasticsearch.xpack.esql.plan.logical.Limit;
29+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
30+
import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
31+
import org.elasticsearch.xpack.esql.plan.logical.join.StubRelation;
32+
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
33+
import org.junit.BeforeClass;
34+
35+
import java.util.List;
36+
import java.util.Map;
37+
38+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_VERIFIER;
39+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
40+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.loadMapping;
41+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
42+
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.defaultLookupResolution;
43+
import static org.hamcrest.Matchers.contains;
44+
import static org.hamcrest.Matchers.hasSize;
45+
import static org.hamcrest.Matchers.is;
46+
47+
public class PropagateInlineEvalsTests extends ESTestCase {
48+
49+
private static EsqlParser parser;
50+
private static Map<String, EsField> mapping;
51+
private static Analyzer analyzer;
52+
53+
@BeforeClass
54+
public static void init() {
55+
parser = new EsqlParser();
56+
mapping = loadMapping("mapping-basic.json");
57+
EsIndex test = new EsIndex("test", mapping, Map.of("test", IndexMode.STANDARD));
58+
IndexResolution getIndexResult = IndexResolution.valid(test);
59+
analyzer = new Analyzer(
60+
new AnalyzerContext(
61+
EsqlTestUtils.TEST_CFG,
62+
new EsqlFunctionRegistry(),
63+
getIndexResult,
64+
defaultLookupResolution(),
65+
new EnrichResolution()
66+
),
67+
TEST_VERIFIER
68+
);
69+
}
70+
71+
/**
72+
* Expects after running the {@link LogicalPlanOptimizer#substitutions()}:
73+
*
74+
* Limit[1000[INTEGER],false]
75+
* \_InlineJoin[LEFT,[y{r}#10],[y{r}#10],[y{r}#10]]
76+
* |_Eval[[gender{f}#13 AS y]]
77+
* | \_EsqlProject[[emp_no{f}#11, languages{f}#14, gender{f}#13]]
78+
* | \_EsRelation[test][_meta_field{f}#17, emp_no{f}#11, first_name{f}#12, ..]
79+
* \_Aggregate[STANDARD,[y{r}#10],[MAX(languages{f}#14,true[BOOLEAN]) AS max_lang, y{r}#10]]
80+
* \_StubRelation[[emp_no{f}#11, languages{f}#14, gender{f}#13, y{r}#10]]
81+
*/
82+
public void testGroupingAliasingMoved_To_LeftSideOfJoin() {
83+
var plan = plan("""
84+
from test
85+
| keep emp_no, languages, gender
86+
| inlinestats max_lang = MAX(languages) BY y = gender
87+
""", LogicalPlanOptimizerTests.SubstitutionOnlyOptimizer.INSTANCE);
88+
89+
var limit = as(plan, Limit.class);
90+
var inline = as(limit.child(), InlineJoin.class);
91+
var leftEval = as(inline.left(), Eval.class);
92+
var project = as(leftEval.child(), EsqlProject.class);
93+
94+
assertThat(Expressions.names(project.projections()), contains("emp_no", "languages", "gender"));
95+
96+
as(project.child(), EsRelation.class);
97+
var rightAgg = as(inline.right(), Aggregate.class);
98+
var stubRelation = as(rightAgg.child(), StubRelation.class);
99+
assertThat(Expressions.names(stubRelation.expressions()), contains("emp_no", "languages", "gender", "y"));
100+
101+
var groupings = rightAgg.groupings();
102+
var aggs = rightAgg.aggregates();
103+
var ref = as(groupings.get(0), ReferenceAttribute.class);
104+
assertThat(aggs.get(1), is(ref));
105+
assertThat(leftEval.fields(), hasSize(1));
106+
assertThat(leftEval.fields().get(0).toAttribute(), is(ref)); // the only grouping is passed as eval on the join's left hand side
107+
assertThat(leftEval.fields().get(0).name(), is("y"));
108+
}
109+
110+
/**
111+
* Expects after running the {@link LogicalPlanOptimizer#substitutions()}:
112+
* Limit[1000[INTEGER],false]
113+
* \_InlineJoin[LEFT,[f{r}#18, g{r}#21, first_name_l{r}#9],[f{r}#18, g{r}#21, first_name_l{r}#9],[f{r}#18, g{r}#21, first_name_l{
114+
* r}#9]]
115+
* |_Eval[[LEFT(last_name{f}#27,1[INTEGER]) AS f, gender{f}#25 AS g]]
116+
* | \_Eval[[LEFT(first_name{f}#24,1[INTEGER]) AS first_name_l]]
117+
* | \_EsqlProject[[emp_no{f}#23, languages{f}#26, gender{f}#25, last_name{f}#27, first_name{f}#24]]
118+
* | \_EsRelation[test][_meta_field{f}#29, emp_no{f}#23, first_name{f}#24, ..]
119+
* \_Aggregate[STANDARD,[f{r}#18, g{r}#21, first_name_l{r}#9],[MAX(languages{f}#26,true[BOOLEAN]) AS max_lang, MIN(languages{f}
120+
* #26,true[BOOLEAN]) AS min_lang, f{r}#18, g{r}#21, first_name_l{r}#9]]
121+
* \_StubRelation[[emp_no{f}#23, languages{f}#26, gender{f}#25, last_name{f}#27, first_name{f}#24, first_name_l{r}#9, f{r}#18, g
122+
* {r}#21]]
123+
*/
124+
public void testGroupingAliasingMoved_To_LeftSideOfJoin_WithExpression() {
125+
var plan = plan("""
126+
from test
127+
| keep emp_no, languages, gender, last_name, first_name
128+
| eval first_name_l = left(first_name, 1)
129+
| inlinestats max_lang = MAX(languages), min_lang = MIN(languages) BY f = left(last_name, 1), g = gender, first_name_l
130+
""", LogicalPlanOptimizerTests.SubstitutionOnlyOptimizer.INSTANCE);
131+
132+
var limit = as(plan, Limit.class);
133+
var inline = as(limit.child(), InlineJoin.class);
134+
var leftEval1 = as(inline.left(), Eval.class);
135+
var leftEval2 = as(leftEval1.child(), Eval.class);
136+
var project = as(leftEval2.child(), EsqlProject.class);
137+
138+
assertThat(Expressions.names(project.projections()), contains("emp_no", "languages", "gender", "last_name", "first_name"));
139+
140+
as(project.child(), EsRelation.class);
141+
var rightAgg = as(inline.right(), Aggregate.class);
142+
var stubRelation = as(rightAgg.child(), StubRelation.class);
143+
assertThat(
144+
Expressions.names(stubRelation.expressions()),
145+
contains("emp_no", "languages", "gender", "last_name", "first_name", "first_name_l", "f", "g")
146+
);
147+
148+
var groupings = rightAgg.groupings();
149+
assertThat(groupings, hasSize(3));
150+
var aggs = rightAgg.aggregates();
151+
var ref1 = as(groupings.get(0), ReferenceAttribute.class); // f = left(last_name, 1)
152+
var ref2 = as(groupings.get(1), ReferenceAttribute.class); // g = gender
153+
var ref3 = as(groupings.get(2), ReferenceAttribute.class); // first_name_l
154+
assertThat(aggs.get(2), is(ref1));
155+
assertThat(aggs.get(3), is(ref2));
156+
assertThat(leftEval1.fields(), hasSize(2));
157+
assertThat(leftEval1.fields().get(0).toAttribute(), is(ref1)); // f = left(last_name, 1)
158+
assertThat(leftEval1.fields().get(0).name(), is("f"));
159+
assertThat(leftEval1.fields().get(1).toAttribute(), is(ref2)); // g = gender
160+
assertThat(leftEval1.fields().get(1).name(), is("g"));
161+
assertThat(leftEval2.fields(), hasSize(1));
162+
assertThat(leftEval2.fields().get(0).toAttribute(), is(ref3)); // first_name_l is in the second eval (the one the user added)
163+
assertThat(leftEval2.fields().get(0).name(), is("first_name_l"));
164+
}
165+
166+
private LogicalPlan plan(String query, LogicalPlanOptimizer optimizer) {
167+
var analyzed = analyzer.analyze(parser.createStatement(query));
168+
var optimized = optimizer.optimize(analyzed);
169+
return optimized;
170+
}
171+
172+
@Override
173+
protected List<String> filteredWarnings() {
174+
return withDefaultLimitWarning(super.filteredWarnings());
175+
}
176+
}

0 commit comments

Comments
 (0)