SQL: Fix nested ORDER BY (#76277) (#76335)

* introduce PushDownAndCombineOrderBy optimization

* move merging of order by clauses to QueryFolder

* address review comments

* fix iteration

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

@@ -165,3 +165,27 @@ 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 emp_no FROM (SELECT * FROM test_emp ORDER BY gender NULLS FIRST) ORDER BY languages NULLS FIRST;
+selectThreeDistinctOrderBy
+SELECT emp_no FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY languages NULLS FIRST) ORDER BY gender NULLS FIRST;
+selectThreeDistinctOrderBy2
+SELECT emp_no FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY languages NULLS FIRST, gender NULLS FIRST;
+selectThreeDistinctOrderBy3
+SELECT emp_no FROM (SELECT * FROM test_emp ORDER BY salary, languages NULLS FIRST) ORDER BY gender NULLS FIRST;
+selectTwoDistinctOrderByWithDuplicate1
+SELECT emp_no FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY languages NULLS FIRST) ORDER BY salary;
+selectTwoDistinctOrderByWithDuplicate2
+SELECT emp_no FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY salary) ORDER BY salary) ORDER BY languages NULLS FIRST;
+selectTwoOrderByWithDistinctOrder
+SELECT emp_no FROM (SELECT * FROM test_emp ORDER BY gender ASC NULLS FIRST) ORDER BY gender DESC NULLS LAST;
+selectTwoOrderByWithShadowedOrder
+SELECT emp_no FROM (SELECT * FROM test_emp ORDER BY salary, gender ASC, languages NULLS LAST) ORDER BY languages NULLS FIRST, gender DESC NULLS LAST;
+selectTwoOrderByWithWhere
+SELECT emp_no FROM (SELECT * FROM test_emp ORDER BY gender NULLS FIRST) WHERE salary > 10000 ORDER BY languages NULLS FIRST;
+selectThreeDerivedOrderBy
+SELECT emp_no FROM (SELECT * FROM (SELECT * FROM test_emp ORDER BY DAY(hire_date)) ORDER BY emp_no + 100) ORDER BY salary / 10;

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

@@ -90,6 +90,7 @@
 import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.ListIterator;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.atomic.AtomicReference;
@@ -701,7 +702,13 @@ protected PhysicalPlan rule(OrderExec plan) {
                 EsQueryExec exec = (EsQueryExec) plan.child();
                 QueryContainer qContainer = exec.queryContainer();
 
-                for (Order order : plan.order()) {
+                // Reverse traversal together with the upwards fold direction ensures that sort clauses are added in reverse order of
+                // precedence. E.g. for the plan `OrderBy[a desc,b](OrderBy[a asc,c](EsExec[...]))`, `prependSort` is called with the
+                // following sequence of arguments: `c`, `a asc`, `b`, `a desc`. The resulting sort order is `a desc,b,c`.
+                ListIterator<Order> it = plan.order().listIterator(plan.order().size());
+                while (it.hasPrevious()) {
+                    Order order = it.previous();
+
                     Direction direction = Direction.from(order.direction());
                     Missing missing = Missing.from(order.nullsPosition());
 
@@ -722,26 +729,26 @@ protected PhysicalPlan rule(OrderExec plan) {
 
                     // field
                     if (orderExpression instanceof FieldAttribute) {
-                        qContainer = qContainer.addSort(lookup,
+                        qContainer = qContainer.prependSort(lookup,
                                 new AttributeSort((FieldAttribute) orderExpression, direction, missing));
                     }
                     // scalar functions typically require script ordering
                     else if (orderExpression instanceof ScalarFunction) {
                         ScalarFunction sf = (ScalarFunction) orderExpression;
                         // nope, use scripted sorting
-                        qContainer = qContainer.addSort(lookup, new ScriptSort(sf.asScript(), direction, missing));
+                        qContainer = qContainer.prependSort(lookup, new ScriptSort(sf.asScript(), direction, missing));
                     }
                     // histogram
                     else if (orderExpression instanceof Histogram) {
-                        qContainer = qContainer.addSort(lookup, new GroupingFunctionSort(direction, missing));
+                        qContainer = qContainer.prependSort(lookup, new GroupingFunctionSort(direction, missing));
                     }
                     // score
                     else if (orderExpression instanceof Score) {
-                        qContainer = qContainer.addSort(lookup, new ScoreSort(direction, missing));
+                        qContainer = qContainer.prependSort(lookup, new ScoreSort(direction, missing));
                     }
                     // agg function
                     else if (orderExpression instanceof AggregateFunction) {
-                        qContainer = qContainer.addSort(lookup,
+                        qContainer = qContainer.prependSort(lookup,
                                 new AggregateSort((AggregateFunction) orderExpression, direction, missing));
                     }
                     // unknown

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/container/QueryContainer.java

@@ -306,9 +306,16 @@ public QueryContainer withScalarProcessors(AttributeMap<Pipe> procs) {
         return new QueryContainer(query, aggs, fields, aliases, pseudoFunctions, procs, sort, limit, trackHits, includeFrozen, minPageSize);
     }
 
-    public QueryContainer addSort(String expressionId, Sort sortable) {
-        Map<String, Sort> newSort = new LinkedHashMap<>(this.sort);
+    /**
+     * Adds a sort expression that takes precedence over all existing sort expressions. Expressions are prepended because the logical plan
+     * is folded from bottom up. So the most significant sort order will be added last.
+     */
+    public QueryContainer prependSort(String expressionId, Sort sortable) {
+        Map<String, Sort> newSort = new LinkedHashMap<>(this.sort.size() + 1);
         newSort.put(expressionId, sortable);
+        for (Map.Entry<String, Sort> entry : this.sort.entrySet()) {
+            newSort.putIfAbsent(entry.getKey(), entry.getValue());
+        }
         return new QueryContainer(query, aggs, fields, aliases, pseudoFunctions, scalarFunctions, newSort, limit, trackHits, includeFrozen,
                 minPageSize);
     }

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/execution/search/SourceGeneratorTests.java

@@ -92,7 +92,7 @@ public void testSelectScoreForcesTrackingScore() {
 
     public void testSortScoreSpecified() {
         QueryContainer container = new QueryContainer()
-                .addSort("id", new ScoreSort(Direction.DESC, null));
+                .prependSort("id", new ScoreSort(Direction.DESC, null));
         SearchSourceBuilder sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertEquals(singletonList(scoreSort()), sourceBuilder.sorts());
     }
@@ -101,13 +101,13 @@ public void testSortFieldSpecified() {
         FieldSortBuilder sortField = fieldSort("test").unmappedType("keyword");
 
         QueryContainer container = new QueryContainer()
-                .addSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
+                .prependSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
                         Direction.ASC, Missing.LAST));
         SearchSourceBuilder sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertEquals(singletonList(sortField.order(SortOrder.ASC).missing("_last")), sourceBuilder.sorts());
 
         container = new QueryContainer()
-                .addSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
+                .prependSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
                         Direction.DESC, Missing.FIRST));
         sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertEquals(singletonList(sortField.order(SortOrder.DESC).missing("_first")), sourceBuilder.sorts());

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,25 @@ public void testPivotHasSameQueryAsGroupBy() {
         }
     }
 
+    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;
+
+        Sort[] sort = ee.queryContainer().sort().values().toArray(new Sort[0]);
+        assertEquals(2, sort.length);
+
+        AttributeSort as1 = (AttributeSort) sort[0];
+        assertEquals("test.keyword", as1.attribute().qualifiedName());
+        assertEquals(Sort.Missing.FIRST, as1.missing());
+
+        AttributeSort as2 = (AttributeSort) sort[1];
+        assertEquals("test.int", as2.attribute().qualifiedName());
+        assertEquals(Sort.Direction.DESC, as2.direction());
+    }
+
     private static String randomOrderByAndLimit(int noOfSelectArgs) {
         return SqlTestUtils.randomOrderByAndLimit(noOfSelectArgs, random());
     }