Skip to content

Commit 8832d25

Browse files
matrivmatriv
authored
SQL: Fix issue with optimization on queries with ORDER BY/LIMIT (#40256)
Previously, when a trival plain `SELECT` or a trivial `SELECT` with aggregations has also an `ORDER BY` or a `LIMIT` or both, then the optimization to convert it to a `LocalRelation` was skipped resulting in exception thrown. E.g.:: ``` SELECT 'foo' FROM test LIMIT 10 ``` or ``` SELECT 'foo' FROM test GROUP BY 1 ORDER BY 1 ``` Fixes: #40211
1 parent 2945b8f commit 8832d25

File tree

3 files changed

+82
-13
lines changed

3 files changed

+82
-13
lines changed

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1909,7 +1909,7 @@ static class SkipQueryOnLimitZero extends OptimizerRule<Limit> {
19091909
@Override
19101910
protected LogicalPlan rule(Limit limit) {
19111911
if (limit.limit() instanceof Literal) {
1912-
if (Integer.valueOf(0).equals((((Literal) limit.limit()).fold()))) {
1912+
if (Integer.valueOf(0).equals((limit.limit().fold()))) {
19131913
return new LocalRelation(limit.source(), new EmptyExecutable(limit.output()));
19141914
}
19151915
}
@@ -1920,21 +1920,30 @@ protected LogicalPlan rule(Limit limit) {
19201920
static class SkipQueryIfFoldingProjection extends OptimizerRule<LogicalPlan> {
19211921
@Override
19221922
protected LogicalPlan rule(LogicalPlan plan) {
1923-
if (plan instanceof Project) {
1924-
Project p = (Project) plan;
1923+
Holder<LocalRelation> optimizedPlan = new Holder<>();
1924+
plan.forEachDown(p -> {
19251925
List<Object> values = extractConstants(p.projections());
19261926
if (values.size() == p.projections().size() && !(p.child() instanceof EsRelation) &&
19271927
isNotQueryWithFromClauseAndFilterFoldedToFalse(p)) {
1928-
return new LocalRelation(p.source(), new SingletonExecutable(p.output(), values.toArray()));
1928+
optimizedPlan.set(new LocalRelation(p.source(), new SingletonExecutable(p.output(), values.toArray())));
19291929
}
1930+
}, Project.class);
1931+
1932+
if (optimizedPlan.get() != null) {
1933+
return optimizedPlan.get();
19301934
}
1931-
if (plan instanceof Aggregate) {
1932-
Aggregate a = (Aggregate) plan;
1935+
1936+
plan.forEachDown(a -> {
19331937
List<Object> values = extractConstants(a.aggregates());
19341938
if (values.size() == a.aggregates().size() && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
1935-
return new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray()));
1939+
optimizedPlan.set(new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray())));
19361940
}
1941+
}, Aggregate.class);
1942+
1943+
if (optimizedPlan.get() != null) {
1944+
return optimizedPlan.get();
19371945
}
1946+
19381947
return plan;
19391948
}
19401949

‎x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1226,7 +1226,7 @@ public void testTranslateMinToFirst() {
12261226
List<Order> order = ((OrderBy) result).order();
12271227
assertEquals(2, order.size());
12281228
assertEquals(First.class, order.get(0).child().getClass());
1229-
assertEquals(min2, order.get(1).child());;
1229+
assertEquals(min2, order.get(1).child());
12301230
First first = (First) order.get(0).child();
12311231

12321232
assertTrue(((OrderBy) result).child() instanceof Aggregate);
@@ -1249,7 +1249,7 @@ public void testTranslateMaxToLast() {
12491249
assertTrue(result instanceof OrderBy);
12501250
List<Order> order = ((OrderBy) result).order();
12511251
assertEquals(Last.class, order.get(0).child().getClass());
1252-
assertEquals(max2, order.get(1).child());;
1252+
assertEquals(max2, order.get(1).child());
12531253
Last last = (Last) order.get(0).child();
12541254

12551255
assertTrue(((OrderBy) result).child() instanceof Aggregate);
@@ -1288,8 +1288,8 @@ public void testSortAggregateOnOrderByWithTwoFields() {
12881288
assertEquals(2, groupings.size());
12891289
assertTrue(groupings.get(0) instanceof FieldAttribute);
12901290
assertTrue(groupings.get(1) instanceof FieldAttribute);
1291-
assertEquals(firstField, ((FieldAttribute) groupings.get(0)));
1292-
assertEquals(secondField, ((FieldAttribute) groupings.get(1)));
1291+
assertEquals(firstField, groupings.get(0));
1292+
assertEquals(secondField, groupings.get(1));
12931293
}
12941294

12951295
public void testSortAggregateOnOrderByOnlyAliases() {
@@ -1320,7 +1320,7 @@ public void testSortAggregateOnOrderByOnlyAliases() {
13201320
assertEquals(2, groupings.size());
13211321
assertTrue(groupings.get(0) instanceof Alias);
13221322
assertTrue(groupings.get(1) instanceof Alias);
1323-
assertEquals(firstAlias, ((Alias) groupings.get(0)));
1324-
assertEquals(secondAlias, ((Alias) groupings.get(1)));
1323+
assertEquals(firstAlias, groupings.get(0));
1324+
assertEquals(secondAlias, groupings.get(1));
13251325
}
13261326
}

‎x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,36 @@ public void testFoldingToLocalExecWithProject() {
7070
assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
7171
}
7272

73+
public void testFoldingToLocalExecWithProjectAndLimit() {
74+
PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 LIMIT 10");
75+
assertEquals(LocalExec.class, p.getClass());
76+
LocalExec le = (LocalExec) p;
77+
assertEquals(EmptyExecutable.class, le.executable().getClass());
78+
EmptyExecutable ee = (EmptyExecutable) le.executable();
79+
assertEquals(1, ee.output().size());
80+
assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
81+
}
82+
83+
public void testFoldingToLocalExecWithProjectAndOrderBy() {
84+
PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 ORDER BY 1");
85+
assertEquals(LocalExec.class, p.getClass());
86+
LocalExec le = (LocalExec) p;
87+
assertEquals(EmptyExecutable.class, le.executable().getClass());
88+
EmptyExecutable ee = (EmptyExecutable) le.executable();
89+
assertEquals(1, ee.output().size());
90+
assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
91+
}
92+
93+
public void testFoldingToLocalExecWithProjectAndOrderByAndLimit() {
94+
PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 ORDER BY 1 LIMIT 10");
95+
assertEquals(LocalExec.class, p.getClass());
96+
LocalExec le = (LocalExec) p;
97+
assertEquals(EmptyExecutable.class, le.executable().getClass());
98+
EmptyExecutable ee = (EmptyExecutable) le.executable();
99+
assertEquals(1, ee.output().size());
100+
assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#"));
101+
}
102+
73103
public void testLocalExecWithPrunedFilterWithFunction() {
74104
PhysicalPlan p = plan("SELECT E() FROM test WHERE PI() = 5");
75105
assertEquals(LocalExec.class, p.getClass());
@@ -90,6 +120,36 @@ public void testLocalExecWithPrunedFilterWithFunctionAndAggregation() {
90120
assertThat(ee.output().get(0).toString(), startsWith("E(){c}#"));
91121
}
92122

123+
public void testFoldingToLocalExecWithAggregationAndLimit() {
124+
PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 LIMIT 10");
125+
assertEquals(LocalExec.class, p.getClass());
126+
LocalExec le = (LocalExec) p;
127+
assertEquals(SingletonExecutable.class, le.executable().getClass());
128+
SingletonExecutable ee = (SingletonExecutable) le.executable();
129+
assertEquals(1, ee.output().size());
130+
assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#"));
131+
}
132+
133+
public void testFoldingToLocalExecWithAggregationAndOrderBy() {
134+
PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 ORDER BY 1");
135+
assertEquals(LocalExec.class, p.getClass());
136+
LocalExec le = (LocalExec) p;
137+
assertEquals(SingletonExecutable.class, le.executable().getClass());
138+
SingletonExecutable ee = (SingletonExecutable) le.executable();
139+
assertEquals(1, ee.output().size());
140+
assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#"));
141+
}
142+
143+
public void testFoldingToLocalExecWithAggregationAndOrderByAndLimit() {
144+
PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 ORDER BY 1 LIMIT 10");
145+
assertEquals(LocalExec.class, p.getClass());
146+
LocalExec le = (LocalExec) p;
147+
assertEquals(SingletonExecutable.class, le.executable().getClass());
148+
SingletonExecutable ee = (SingletonExecutable) le.executable();
149+
assertEquals(1, ee.output().size());
150+
assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#"));
151+
}
152+
93153
public void testLocalExecWithoutFromClause() {
94154
PhysicalPlan p = plan("SELECT E(), 'foo', abs(10)");
95155
assertEquals(LocalExec.class, p.getClass());

0 commit comments

Comments
 (0)