SQL: Fix issue with optimization on queries with ORDER BY/LIMIT

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: elastic#40211

Author: matriv

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

@@ -1909,7 +1909,7 @@ static class SkipQueryOnLimitZero extends OptimizerRule<Limit> {
         @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,37 @@ protected LogicalPlan rule(Limit limit) {
     static class SkipQueryIfFoldingProjection extends OptimizerRule<LogicalPlan> {
         @Override
         protected LogicalPlan rule(LogicalPlan plan) {
-            if (plan instanceof Project) {
-                Project p = (Project) plan;
+            LogicalPlan optimized = plan.transformDown(p -> {
                 List<Object> 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()));
                 }
+                return p;
+            }, Project.class);
+
+            if (optimized.equals(plan) == false) {
+                LocalRelation lr = findLocalRelation(optimized);
+                if (lr != null) {
+                    return lr;
+                }
             }
-            if (plan instanceof Aggregate) {
-                Aggregate a = (Aggregate) plan;
+
+            optimized = optimized.transformDown(a -> {
                 List<Object> values = extractConstants(a.aggregates());
                 if (values.size() == a.aggregates().size() && isNotQueryWithFromClauseAndFilterFoldedToFalse(a)) {
                     return new LocalRelation(a.source(), new SingletonExecutable(a.output(), values.toArray()));
                 }
+                return a;
+            }, Aggregate.class);
+
+            if (optimized.equals(plan) == false) {
+                LocalRelation lr = findLocalRelation(optimized);
+                if (lr != null) {
+                    return lr;
+                }
             }
+
             return plan;
         }
 
@@ -1959,6 +1975,19 @@ private static boolean isNotQueryWithFromClauseAndFilterFoldedToFalse(UnaryPlan
             return (!(plan.child() instanceof LocalRelation) || (plan.child() instanceof LocalRelation &&
                 !(((LocalRelation) plan.child()).executable() instanceof EmptyExecutable)));
         }
+
+        private static LocalRelation findLocalRelation(LogicalPlan plan) {
+            Holder<LocalRelation> logicalPlanHolder = new Holder<>();
+            plan.transformDown(l -> {
+                logicalPlanHolder.set(l);
+                return l;
+            }, LocalRelation.class);
+
+            if (logicalPlanHolder.get() != null) {
+                return logicalPlanHolder.get();
+            }
+            return logicalPlanHolder.get();
+        }
     }
 
 

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

@@ -1226,7 +1226,7 @@ public void testTranslateMinToFirst() {
         List<Order> 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 void testTranslateMaxToLast() {
         assertTrue(result instanceof OrderBy);
         List<Order> 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 void testSortAggregateOnOrderByWithTwoFields() {
         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 void testSortAggregateOnOrderByOnlyAliases() {
         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));
     }
 }

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

@@ -70,6 +70,36 @@ public void testFoldingToLocalExecWithProject() {
         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 void testLocalExecWithPrunedFilterWithFunctionAndAggregation() {
         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());