introduce PushDownAndCombineOrderBy optimization

x-pack/plugin/sql/qa/server/src/main/resources/select.sql-spec

@@ -153,6 +153,8 @@ selectOrderByOrderByOrderByLimit
 SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY emp_no DESC) ORDER BY emp_no ASC) ORDER BY emp_no DESC LIMIT 5;
 selectOrderByOrderByOrderByLimitLimit
 SELECT * FROM (SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY emp_no DESC) ORDER BY emp_no ASC) ORDER BY emp_no DESC LIMIT 12) LIMIT 6;
+selectOrderByLimitSameOrderBy
+SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY emp_no LIMIT 10)) ORDER BY emp_no LIMIT 5;
 selectGroupByOrderByLimit
 SELECT * FROM (SELECT max(salary) AS max, languages FROM test_emp GROUP BY languages) ORDER BY max DESC LIMIT 3;
 selectGroupByOrderByLimitNulls
@@ -165,3 +167,21 @@ selectGroupByWithAliasedSubQuery
 SELECT max, languages FROM (SELECT max(salary) AS max, languages FROM test_emp GROUP BY languages ORDER BY max ASC NULLS LAST) AS subquery;
 selectConstantFromSubQuery
 SELECT * FROM (SELECT * FROM (SELECT 1));
+
+//
+// SELECT with multiple ORDER BY
+//
+selectTwoDistinctOrderBy
+SELECT * FROM (SELECT * FROM test_emp ORDER BY gender NULLS FIRST) ORDER BY languages NULLS FIRST;
+selectThreeDistinctOrderBy
+SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY languages NULLS FIRST) ORDER BY gender NULLS FIRST;
+selectTwoDistinctOrderByWithDuplicate1
+SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY languages NULLS FIRST) ORDER BY salary;
+selectTwoDistinctOrderByWithDuplicate2
+SELECT * FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY salary) ORDER BY languages NULLS FIRST;
+selectTwoOrderByWithDistinctOrder
+SELECT * FROM (SELECT * FROM test_emp ORDER BY gender ASC NULLS FIRST) ORDER BY gender DESC NULLS LAST;
+selectTwoOrderByWithShadowedOrder
+SELECT * FROM (SELECT * FROM test_emp ORDER BY salary, gender ASC, languages NULLS LAST) ORDER BY languages NULLS FIRST, gender DESC NULLS LAST;
+selectTwoOrderByWithWhere
+SELECT * FROM (SELECT * FROM test_emp ORDER BY gender NULLS FIRST) WHERE salary > 10000 ORDER BY languages NULLS FIRST;

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.ql.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.ql.expression.UnresolvedAttribute;
 import org.elasticsearch.xpack.ql.expression.function.Function;
+import org.elasticsearch.xpack.ql.expression.function.Functions;
 import org.elasticsearch.xpack.ql.expression.function.aggregate.AggregateFunction;
 import org.elasticsearch.xpack.ql.expression.function.aggregate.Count;
 import org.elasticsearch.xpack.ql.expression.function.aggregate.InnerAggregate;
@@ -93,6 +94,7 @@
 import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -165,6 +167,9 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 new PushDownAndCombineFilters()
         );
 
+        Batch pushDownOrderBy = new Batch("Push Down Order By",
+            new PushDownAndCombineOrderBy());
+
         Batch aggregate = new Batch("Aggregation Rewrite",
                 new ReplaceMinMaxWithTopHits(),
                 new ReplaceAggsWithMatrixStats(),
@@ -189,7 +194,7 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 CleanAliases.INSTANCE,
                 new SetAsOptimized());
 
-        return Arrays.asList(substitutions, refs, operators, aggregate, local, label);
+        return Arrays.asList(substitutions, refs, operators, pushDownOrderBy, aggregate, local, label);
     }
 
     static class RewritePivot extends OptimizerRule<Pivot> {
@@ -1255,6 +1260,58 @@ private boolean foldable(Expression e) {
 
     }
 
+    static class PushDownAndCombineOrderBy extends OptimizerRule<OrderBy> {
+
+        @Override
+        protected LogicalPlan rule(OrderBy plan) {
+            if (plan.child() instanceof OrderBy) {
+                OrderBy child = (OrderBy) plan.child();
+                List<Order> combinedOrder = combineOrders(plan.order(), child.order());
+
+                return new OrderBy(child.source(), child.child(), combinedOrder);
+
+            } else if (plan.child() instanceof Filter || plan.child() instanceof Project) {
+                UnaryPlan child = (UnaryPlan) plan.child();
+                Holder<Boolean> hasAggregate = new Holder<>(false);
+                if (child instanceof Filter) {
+                    child.forEachExpressionDown(e -> hasAggregate.set(hasAggregate.get() || Functions.isAggregate(e)));
+                }
+
+                if (hasAggregate.get() == false) {
+                    return child.replaceChildrenSameSize(
+                        Collections.singletonList(plan.replaceChildrenSameSize(Collections.singletonList(child.child())))
+                    );
+                }
+
+            } else if (plan.child() instanceof Limit) {
+                // OrderBy can be pushed through a Limit if the Limit's child is also an OrderBy and the two sort orders are equal
+                Limit child = (Limit) plan.child();
+                if (child.child() instanceof OrderBy) {
+                    OrderBy innerOrderBy = (OrderBy) child.child();
+                    List<Order> combinedOrder = combineOrders(plan.order(), innerOrderBy.order());
+
+                    if (combinedOrder.equals(innerOrderBy.order())) {
+                        return child.replaceChildrenSameSize(Collections.singletonList(innerOrderBy));
+                    }
+                }
+            }
+
+            return plan;
+        }
+
+        private List<Order> combineOrders(List<Order> parentOrders, List<Order> childOrders) {
+            List<Order> combined = parentOrders;
+
+            for (Order childOrder : childOrders) {
+                if(combined.stream().noneMatch(parentOrder -> childOrder.child().equals(parentOrder.child()))){
+                    combined.add(childOrder);
+                }
+            }
+
+            return combined;
+        }
+    }
+
     abstract static class OptimizerBasicRule extends Rule<LogicalPlan, LogicalPlan> {
 
         @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java

@@ -701,6 +701,10 @@ protected PhysicalPlan rule(OrderExec plan) {
                 EsQueryExec exec = (EsQueryExec) plan.child();
                 QueryContainer qContainer = exec.queryContainer();
 
+                if (qContainer.sort().isEmpty() == false) {
+                    throw new SqlIllegalArgumentException("QueryContainer already defines sort order");
+                }
+
                 for (Order order : plan.order()) {
                     Direction direction = Direction.from(order.direction());
                     Missing missing = Missing.from(order.nullsPosition());

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

@@ -12,6 +12,8 @@
 import org.elasticsearch.xpack.ql.expression.function.aggregate.AggregateFunction;
 import org.elasticsearch.xpack.ql.index.EsIndex;
 import org.elasticsearch.xpack.ql.index.IndexResolution;
+import org.elasticsearch.xpack.ql.querydsl.container.AttributeSort;
+import org.elasticsearch.xpack.ql.querydsl.container.Sort;
 import org.elasticsearch.xpack.ql.type.EsField;
 import org.elasticsearch.xpack.sql.SqlTestUtils;
 import org.elasticsearch.xpack.sql.analysis.analyzer.Analyzer;
@@ -534,6 +536,33 @@ public void testPivotHasSameQueryAsGroupBy() {
         }
     }
 
+    public void testFoldRedundantOrderBy() {
+        PhysicalPlan p = plan("SELECT * FROM (SELECT * FROM test ORDER BY int LIMIT 10) ORDER BY int");
+        assertEquals(EsQueryExec.class, p.getClass());
+        EsQueryExec ee = (EsQueryExec) p;
+        assertEquals(10, ee.queryContainer().limit());
+        assertEquals(1, ee.queryContainer().sort().size());
+        AttributeSort as = (AttributeSort) ee.queryContainer().sort().values().toArray()[0];
+        assertEquals("test.int", as.attribute().qualifiedName());
+    }
+
+    public void testFoldShadowedOrderBy() {
+        PhysicalPlan p = plan(
+            "SELECT * FROM (SELECT * FROM test ORDER BY int ASC, keyword NULLS LAST) ORDER BY keyword NULLS FIRST, int DESC"
+        );
+        assertEquals(EsQueryExec.class, p.getClass());
+        EsQueryExec ee = (EsQueryExec) p;
+        assertEquals(2, ee.queryContainer().sort().size());
+
+        AttributeSort as1 = (AttributeSort) ee.queryContainer().sort().values().toArray()[0];
+        assertEquals("test.keyword", as1.attribute().qualifiedName());
+        assertEquals(Sort.Missing.FIRST, as1.missing());
+
+        AttributeSort as2 = (AttributeSort) ee.queryContainer().sort().values().toArray()[1];
+        assertEquals("test.int", as2.attribute().qualifiedName());
+        assertEquals(Sort.Direction.DESC, as2.direction());
+    }
+
     private static String randomOrderByAndLimit(int noOfSelectArgs) {
         return SqlTestUtils.randomOrderByAndLimit(noOfSelectArgs, random());
     }