From f37f2b5d39676d45706efd29b2fab9650158d5a7 Mon Sep 17 00:00:00 2001 From: Marios Trivyzas Date: Wed, 20 Mar 2019 23:51:15 +0100 Subject: [PATCH] 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 --- .../xpack/sql/optimizer/Optimizer.java | 23 ++++--- .../xpack/sql/optimizer/OptimizerTests.java | 12 ++-- .../xpack/sql/planner/QueryFolderTests.java | 60 +++++++++++++++++++ 3 files changed, 82 insertions(+), 13 deletions(-) diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java index 4b59db68063..c0cee88bb06 100644 --- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java +++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java @@ -1909,7 +1909,7 @@ public class Optimizer extends RuleExecutor { @Override protected LogicalPlan rule(Limit limit) { if (limit.limit() instanceof Literal) { - if (Integer.valueOf(0).equals((((Literal) limit.limit()).fold()))) { + if (Integer.valueOf(0).equals((limit.limit().fold()))) { return new LocalRelation(limit.source(), new EmptyExecutable(limit.output())); } } @@ -1920,21 +1920,30 @@ public class Optimizer extends RuleExecutor { static class SkipQueryIfFoldingProjection extends OptimizerRule { @Override protected LogicalPlan rule(LogicalPlan plan) { - if (plan instanceof Project) { - Project p = (Project) plan; + Holder optimizedPlan = new Holder<>(); + plan.forEachDown(p -> { List values = extractConstants(p.projections()); if (values.size() == p.projections().size() && !(p.child() instanceof EsRelation) && isNotQueryWithFromClauseAndFilterFoldedToFalse(p)) { - return new LocalRelation(p.source(), new SingletonExecutable(p.output(), values.toArray())); + optimizedPlan.set(new LocalRelation(p.source(), new SingletonExecutable(p.output(), values.toArray()))); } + }, Project.class); + + if (optimizedPlan.get() != null) { + return optimizedPlan.get(); } - if (plan instanceof Aggregate) { - Aggregate a = (Aggregate) plan; + + plan.forEachDown(a -> { List values = extractConstants(a.aggregates()); if (values.size() == a.aggregates().size() && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) { - return new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray())); + optimizedPlan.set(new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray()))); } + }, Aggregate.class); + + if (optimizedPlan.get() != null) { + return optimizedPlan.get(); } + return plan; } diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java index 0d9df66f378..7c01fd8ff15 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java +++ b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java @@ -1226,7 +1226,7 @@ public class OptimizerTests extends ESTestCase { List order = ((OrderBy) result).order(); assertEquals(2, order.size()); assertEquals(First.class, order.get(0).child().getClass()); - assertEquals(min2, order.get(1).child());; + assertEquals(min2, order.get(1).child()); First first = (First) order.get(0).child(); assertTrue(((OrderBy) result).child() instanceof Aggregate); @@ -1249,7 +1249,7 @@ public class OptimizerTests extends ESTestCase { assertTrue(result instanceof OrderBy); List order = ((OrderBy) result).order(); assertEquals(Last.class, order.get(0).child().getClass()); - assertEquals(max2, order.get(1).child());; + assertEquals(max2, order.get(1).child()); Last last = (Last) order.get(0).child(); assertTrue(((OrderBy) result).child() instanceof Aggregate); @@ -1288,8 +1288,8 @@ public class OptimizerTests extends ESTestCase { assertEquals(2, groupings.size()); assertTrue(groupings.get(0) instanceof FieldAttribute); assertTrue(groupings.get(1) instanceof FieldAttribute); - assertEquals(firstField, ((FieldAttribute) groupings.get(0))); - assertEquals(secondField, ((FieldAttribute) groupings.get(1))); + assertEquals(firstField, groupings.get(0)); + assertEquals(secondField, groupings.get(1)); } public void testSortAggregateOnOrderByOnlyAliases() { @@ -1320,7 +1320,7 @@ public class OptimizerTests extends ESTestCase { assertEquals(2, groupings.size()); assertTrue(groupings.get(0) instanceof Alias); assertTrue(groupings.get(1) instanceof Alias); - assertEquals(firstAlias, ((Alias) groupings.get(0))); - assertEquals(secondAlias, ((Alias) groupings.get(1))); + assertEquals(firstAlias, groupings.get(0)); + assertEquals(secondAlias, groupings.get(1)); } } diff --git a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java index 6e538bb0e7a..c94da662151 100644 --- a/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java +++ b/x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryFolderTests.java @@ -70,6 +70,36 @@ public class QueryFolderTests extends ESTestCase { assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#")); } + public void testFoldingToLocalExecWithProjectAndLimit() { + PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 LIMIT 10"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(EmptyExecutable.class, le.executable().getClass()); + EmptyExecutable ee = (EmptyExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#")); + } + + public void testFoldingToLocalExecWithProjectAndOrderBy() { + PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 ORDER BY 1"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(EmptyExecutable.class, le.executable().getClass()); + EmptyExecutable ee = (EmptyExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#")); + } + + public void testFoldingToLocalExecWithProjectAndOrderByAndLimit() { + PhysicalPlan p = plan("SELECT keyword FROM test WHERE 1 = 2 ORDER BY 1 LIMIT 10"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(EmptyExecutable.class, le.executable().getClass()); + EmptyExecutable ee = (EmptyExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("keyword{f}#")); + } + public void testLocalExecWithPrunedFilterWithFunction() { PhysicalPlan p = plan("SELECT E() FROM test WHERE PI() = 5"); assertEquals(LocalExec.class, p.getClass()); @@ -90,6 +120,36 @@ public class QueryFolderTests extends ESTestCase { assertThat(ee.output().get(0).toString(), startsWith("E(){c}#")); } + public void testFoldingToLocalExecWithAggregationAndLimit() { + PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 LIMIT 10"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(SingletonExecutable.class, le.executable().getClass()); + SingletonExecutable ee = (SingletonExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#")); + } + + public void testFoldingToLocalExecWithAggregationAndOrderBy() { + PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 ORDER BY 1"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(SingletonExecutable.class, le.executable().getClass()); + SingletonExecutable ee = (SingletonExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#")); + } + + public void testFoldingToLocalExecWithAggregationAndOrderByAndLimit() { + PhysicalPlan p = plan("SELECT 'foo' FROM test GROUP BY 1 ORDER BY 1 LIMIT 10"); + assertEquals(LocalExec.class, p.getClass()); + LocalExec le = (LocalExec) p; + assertEquals(SingletonExecutable.class, le.executable().getClass()); + SingletonExecutable ee = (SingletonExecutable) le.executable(); + assertEquals(1, ee.output().size()); + assertThat(ee.output().get(0).toString(), startsWith("'foo'{c}#")); + } + public void testLocalExecWithoutFromClause() { PhysicalPlan p = plan("SELECT E(), 'foo', abs(10)"); assertEquals(LocalExec.class, p.getClass());