SQL: Fix ORDER BY on aggregates and GROUPed BY fields (#51894)

Previously, in the in-memory sorting module
`LocalAggregationSorterListener` only the aggregate functions where used
(grabbed by the `sortingColumns`). As a consequence, if the ORDER BY
was also using columns of the GROUP BY clause, (especially in the case
of higher priority - before the aggregate functions) wrong results were
produced. E.g.:
```
SELECT gender, MAX(salary) AS max FROM test_emp
GROUP BY gender
ORDER BY gender, max
```

Add all columns of the ORDER BY to the `sortingColumns` so that the
`LocalAggregationSorterListener` can use the correct comparators in
the underlying PriorityQueue used to implement the in-memory sorting.

Fixes: #50355
(cherry picked from commit be680af11c823292c2d115bff01658f7b75abd76)

x-pack/plugin/sql/qa/src/main/resources/agg-ordering.csv-spec

@@ -23,4 +23,50 @@ g:s  | gender:s  | s3:i  | SUM(salary):i | s5:i
 M    |M          |2671054|2671054        |2671054
 F    |F          |1666196|1666196        |1666196
 null |null       |487605 |487605         |487605
-;
+;
+
+histogramDateTimeWithCountAndOrder_1
+schema::h:ts|c:l
+SELECT HISTOGRAM(birth_date, INTERVAL 1 YEAR) AS h, COUNT(*) as c FROM test_emp GROUP BY h ORDER BY h DESC, c ASC;
+
+           h            |       c
+------------------------+---------------
+1965-01-01T00:00:00.000Z|1
+1964-01-01T00:00:00.000Z|4
+1963-01-01T00:00:00.000Z|7
+1962-01-01T00:00:00.000Z|6
+1961-01-01T00:00:00.000Z|8
+1960-01-01T00:00:00.000Z|8
+1959-01-01T00:00:00.000Z|9
+1958-01-01T00:00:00.000Z|7
+1957-01-01T00:00:00.000Z|4
+1956-01-01T00:00:00.000Z|5
+1955-01-01T00:00:00.000Z|4
+1954-01-01T00:00:00.000Z|8
+1953-01-01T00:00:00.000Z|11
+1952-01-01T00:00:00.000Z|8
+null                    |10
+;
+
+histogramDateTimeWithCountAndOrder_2
+schema::h:ts|c:l
+SELECT HISTOGRAM(birth_date, INTERVAL 1 YEAR) AS h, COUNT(*) as c FROM test_emp GROUP BY h ORDER BY c DESC, h ASC;
+
+           h            |       c
+------------------------+---------------
+1953-01-01T00:00:00.000Z|11
+null                    |10
+1959-01-01T00:00:00.000Z|9
+1952-01-01T00:00:00.000Z|8
+1954-01-01T00:00:00.000Z|8
+1960-01-01T00:00:00.000Z|8
+1961-01-01T00:00:00.000Z|8
+1958-01-01T00:00:00.000Z|7
+1963-01-01T00:00:00.000Z|7
+1962-01-01T00:00:00.000Z|6
+1956-01-01T00:00:00.000Z|5
+1955-01-01T00:00:00.000Z|4
+1957-01-01T00:00:00.000Z|4
+1964-01-01T00:00:00.000Z|4
+1965-01-01T00:00:00.000Z|1
+;

x-pack/plugin/sql/qa/src/main/resources/agg-ordering.sql-spec

@@ -80,7 +80,7 @@ aggNotSpecifiedInTheAggregateAndGroupWithHavingWithLimitAndDirection
 SELECT gender, MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY MAX(salary) ASC, c DESC LIMIT 5;
 
 groupAndAggNotSpecifiedInTheAggregateWithHaving
-SELECT gender, MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY gender, MAX(salary);
+SELECT gender, MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY gender NULLS FIRST, MAX(salary);
 
 multipleAggsThatGetRewrittenWithAliasOnAMediumGroupBy
 SELECT languages, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY languages ORDER BY max;
@@ -136,5 +136,26 @@ SELECT gender AS g, first_name AS f, last_name AS l FROM test_emp GROUP BY f, ge
 multipleGroupingsAndOrderingByGroups_8
 SELECT gender AS g, first_name, last_name FROM test_emp GROUP BY g, last_name, first_name ORDER BY gender ASC, first_name DESC, last_name ASC;
 
-multipleGroupingsAndOrderingByGroupsWithFunctions
+multipleGroupingsAndOrderingByGroupsAndAggs_1
-SELECT first_name f, last_name l, gender g, CONCAT(first_name, last_name) c FROM test_emp GROUP BY gender, l, f, c ORDER BY gender, c DESC, first_name, last_name ASC;
+SELECT gender, MIN(salary) AS min, COUNT(*) AS c, MAX(salary) AS max FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY gender ASC NULLS FIRST, MAX(salary) DESC;
+
+multipleGroupingsAndOrderingByGroupsAndAggs_2
+SELECT gender, MIN(salary) AS min, COUNT(*) AS c, MAX(salary) AS max FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY gender DESC NULLS LAST, MAX(salary) ASC;
+
+multipleGroupingsAndOrderingByGroupsWithFunctions_1
+SELECT first_name f, last_name l, gender g, CONCAT(first_name, last_name) c FROM test_emp GROUP BY gender, l, f, c ORDER BY gender NULLS FIRST, c DESC, first_name, last_name ASC;
+
+multipleGroupingsAndOrderingByGroupsWithFunctions_2
+SELECT first_name f, last_name l, gender g, CONCAT(first_name, last_name) c FROM test_emp GROUP BY gender, l, f, c ORDER BY c DESC, gender DESC NULLS LAST, first_name, last_name ASC;
+
+multipleGroupingsAndOrderingByGroupsAndAggregatesWithFunctions_1
+SELECT CONCAT('foo', gender) g, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY g ORDER BY 1 NULLS FIRST, 2, 3;
+
+multipleGroupingsAndOrderingByGroupsAndAggregatesWithFunctions_2
+SELECT CONCAT('foo', gender) g, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY g ORDER BY 1 DESC NULLS LAST, 2, 3;
+
+multipleGroupingsAndOrderingByGroupsAndAggregatesWithFunctions_3
+SELECT CONCAT('foo', gender) g, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY g ORDER BY 2, 1 NULLS FIRST, 3;
+
+multipleGroupingsAndOrderingByGroupsAndAggregatesWithFunctions_4
+SELECT CONCAT('foo', gender) g, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY g ORDER BY 3 DESC, 1 NULLS FIRST, 2;

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/Querier.java

@@ -593,36 +593,51 @@ public class Querier {
             this.sortingColumns = sortingColumns;
         }
 
-        // compare row based on the received attribute sort
+        /**
-        // if a sort item is not in the list, it is assumed the sorting happened in ES
+         * Compare row based on the received attribute sort
-        // and the results are left as is (by using the row ordering), otherwise it is sorted based on the given criteria.
+         * <ul>
-        //
+         *     <li>
-        // Take for example ORDER BY a, x, b, y
+         *         If a tuple in {@code sortingColumns} has a null comparator, it is assumed the sorting
-        // a, b - are sorted in ES
+         *         happened in ES and the results are left as is (by using the row ordering), otherwise it is
-        // x, y - need to be sorted client-side
+         *         sorted based on the given criteria.
-        // sorting on x kicks in, only if the values for a are equal.
+         *     </li>
-
+         *     <li>
+         *         If no tuple exists in {@code sortingColumns} for an output column, it means this column
+         *         is not included at all in the ORDER BY
+         *     </li>
+         *</ul>
+         *
+         * Take for example ORDER BY a, x, b, y
+         * a, b - are sorted in ES
+         * x, y - need to be sorted client-side
+         * sorting on x kicks in only if the values for a are equal.
+         * sorting on y kicks in only if the values for a, x and b are all equal
+         *
+         */
         // thanks to @jpountz for the row ordering idea as a way to preserve ordering
         @SuppressWarnings("unchecked")
         @Override
         protected boolean lessThan(Tuple<List<?>, Integer> l, Tuple<List<?>, Integer> r) {
             for (Tuple<Integer, Comparator> tuple : sortingColumns) {
-                int i = tuple.v1().intValue();
+                int columnIdx = tuple.v1().intValue();
                 Comparator comparator = tuple.v2();
 
-                Object vl = l.v1().get(i);
+                // Get the values for left and right rows at the current column index
-                Object vr = r.v1().get(i);
+                Object vl = l.v1().get(columnIdx);
+                Object vr = r.v1().get(columnIdx);
                 if (comparator != null) {
                     int result = comparator.compare(vl, vr);
-                    // if things are equals, move to the next comparator
+                    // if things are not equal: return the comparison result,
+                    // otherwise: move to the next comparator to solve the tie.
                     if (result != 0) {
                         return result > 0;
                     }
                 }
-                // no comparator means the existing order needs to be preserved
+                // no comparator means the rows are pre-ordered by ES for the column at
+                // the current index and the existing order needs to be preserved
                 else {
-                    // check the values - if they are equal move to the next comparator
+                    // check the values - if they are not equal return the row order
-                    // otherwise return the row order
+                    // otherwise: move to the next comparator to solve the tie.
                     if (Objects.equals(vl, vr) == false) {
                         return l.v2().compareTo(r.v2()) > 0;
                     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/SourceGenerator.java

@@ -110,7 +110,7 @@ public abstract class SourceGenerator {
             source.sort("_doc");
             return;
         }
-        for (Sort sortable : container.sort()) {
+        for (Sort sortable : container.sort().values()) {
             SortBuilder<?> sortBuilder = null;
 
             if (sortable instanceof AttributeSort) {

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

@@ -67,6 +67,7 @@ import org.elasticsearch.xpack.sql.querydsl.container.ComputedRef;
 import org.elasticsearch.xpack.sql.querydsl.container.GlobalCountRef;
 import org.elasticsearch.xpack.sql.querydsl.container.GroupByRef;
 import org.elasticsearch.xpack.sql.querydsl.container.GroupByRef.Property;
+import org.elasticsearch.xpack.sql.querydsl.container.GroupingFunctionSort;
 import org.elasticsearch.xpack.sql.querydsl.container.MetricAggRef;
 import org.elasticsearch.xpack.sql.querydsl.container.PivotColumnRef;
 import org.elasticsearch.xpack.sql.querydsl.container.QueryContainer;
@@ -682,37 +683,36 @@ class QueryFolder extends RuleExecutor<PhysicalPlan> {
 
                     // TODO: might need to validate whether the target field or group actually exist
                     if (group != null && group != Aggs.IMPLICIT_GROUP_KEY) {
-                        // check whether the lookup matches a group
+                        qContainer = qContainer.updateGroup(group.with(direction));
-                        if (group.id().equals(lookup)) {
-                            qContainer = qContainer.updateGroup(group.with(direction));
-                        }
-                        // else it's a leafAgg
-                        else {
-                            qContainer = qContainer.updateGroup(group.with(direction));
-                        }
                     }
+
+                    // field
+                    if (orderExpression instanceof FieldAttribute) {
+                        qContainer = qContainer.addSort(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));
+                    }
+                    // histogram
+                    else if (orderExpression instanceof Histogram) {
+                        qContainer = qContainer.addSort(lookup, new GroupingFunctionSort(direction, missing));
+                    }
+                    // score
+                    else if (orderExpression instanceof Score) {
+                        qContainer = qContainer.addSort(lookup, new ScoreSort(direction, missing));
+                    }
+                    // agg function
+                    else if (orderExpression instanceof AggregateFunction) {
+                        qContainer = qContainer.addSort(lookup,
+                                new AggregateSort((AggregateFunction) orderExpression, direction, missing));
+                    }
+                    // unknown
                     else {
-                        // scalar functions typically require script ordering
+                        throw new SqlIllegalArgumentException("unsupported sorting expression {}", orderExpression);
-                        if (orderExpression instanceof ScalarFunction) {
-                            ScalarFunction sf = (ScalarFunction) orderExpression;
-                            // nope, use scripted sorting
-                            qContainer = qContainer.addSort(new ScriptSort(sf.asScript(), direction, missing));
-                        }
-                        // score
-                        else if (orderExpression instanceof Score) {
-                            qContainer = qContainer.addSort(new ScoreSort(direction, missing));
-                        }
-                        // field
-                        else if (orderExpression instanceof FieldAttribute) {
-                            qContainer = qContainer.addSort(new AttributeSort((FieldAttribute) orderExpression, direction, missing));
-                        }
-                        // agg function
-                        else if (orderExpression instanceof AggregateFunction) {
-                            qContainer = qContainer.addSort(new AggregateSort((AggregateFunction) orderExpression, direction, missing));
-                        } else {
-                            // unknown
-                            throw new SqlIllegalArgumentException("unsupported sorting expression {}", orderExpression);
-                        }
                     }
                 }
 

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.sql.querydsl.container;
+
+import java.util.Objects;
+
+public class GroupingFunctionSort extends Sort {
+
+    public GroupingFunctionSort(Direction direction, Missing missing) {
+        super(direction, missing);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(direction(), missing());
+    }
+    
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+        
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+        
+        GroupingFunctionSort other = (GroupingFunctionSort) obj;
+        return Objects.equals(direction(), other.direction())
+                && Objects.equals(missing(), other.missing());
+    }
+}

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

@@ -17,7 +17,6 @@ import org.elasticsearch.xpack.ql.expression.AttributeMap;
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.FieldAttribute;
-import org.elasticsearch.xpack.ql.expression.function.aggregate.AggregateFunction;
 import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.ql.expression.gen.pipeline.ConstantInput;
 import org.elasticsearch.xpack.ql.expression.gen.pipeline.Pipe;
@@ -43,15 +42,12 @@ import java.util.BitSet;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptyMap;
-import static java.util.Collections.emptySet;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.xpack.ql.util.CollectionUtils.combine;
 
@@ -82,7 +78,7 @@ public class QueryContainer {
     // at scrolling, their inputs (leaves) get updated
     private final AttributeMap<Pipe> scalarFunctions;
 
-    private final Set<Sort> sort;
+    private final Map<String, Sort> sort;
     private final int limit;
     private final boolean trackHits;
     private final boolean includeFrozen;
@@ -106,7 +102,7 @@ public class QueryContainer {
             AttributeMap<Expression> aliases,
             Map<String, GroupByKey> pseudoFunctions,
             AttributeMap<Pipe> scalarFunctions,
-            Set<Sort> sort,
+            Map<String, Sort> sort,
             int limit,
             boolean trackHits,
             boolean includeFrozen,
@@ -117,7 +113,7 @@ public class QueryContainer {
         this.aliases = aliases == null || aliases.isEmpty() ? AttributeMap.emptyAttributeMap() : aliases;
         this.pseudoFunctions = pseudoFunctions == null || pseudoFunctions.isEmpty() ? emptyMap() : pseudoFunctions;
         this.scalarFunctions = scalarFunctions == null || scalarFunctions.isEmpty() ? AttributeMap.emptyAttributeMap() : scalarFunctions;
-        this.sort = sort == null || sort.isEmpty() ? emptySet() : sort;
+        this.sort = sort == null || sort.isEmpty() ? emptyMap() : sort;
         this.limit = limit;
         this.trackHits = trackHits;
         this.includeFrozen = includeFrozen;
@@ -134,45 +130,48 @@ public class QueryContainer {
             return emptyList();
         }
 
-        List<Tuple<Integer, Comparator>> sortingColumns = new ArrayList<>(sort.size());
+        for (Sort s : sort.values()) {
-
-        boolean aggSort = false;
-        for (Sort s : sort) {
-            Tuple<Integer, Comparator> tuple = new Tuple<>(Integer.valueOf(-1), null);
-            
             if (s instanceof AggregateSort) {
-                AggregateSort as = (AggregateSort) s;
+                customSort = Boolean.TRUE;
-                // find the relevant column of each aggregate function
+                break;
-                AggregateFunction af = as.agg();
-
-                aggSort = true;
-                int atIndex = -1;
-                String id = Expressions.id(af);
-
-                for (int i = 0; i < fields.size(); i++) {
-                    Tuple<FieldExtraction, String> field = fields.get(i);
-                    if (field.v2().equals(id)) {
-                        atIndex = i;
-                        break;
-                    }
-                }
-                if (atIndex == -1) {
-                    throw new SqlIllegalArgumentException("Cannot find backing column for ordering aggregation [{}]", s);
-                }
-                // assemble a comparator for it
-                Comparator comp = s.direction() == Sort.Direction.ASC ? Comparator.naturalOrder() : Comparator.reverseOrder();
-                comp = s.missing() == Sort.Missing.FIRST ? Comparator.nullsFirst(comp) : Comparator.nullsLast(comp);
-
-                tuple = new Tuple<>(Integer.valueOf(atIndex), comp);
             }
-            sortingColumns.add(tuple);
+        }
+
+        // If no custom sort is used break early
+        if (customSort == null) {
+            customSort = Boolean.FALSE;
+            return emptyList();
+        }
+
+        List<Tuple<Integer, Comparator>> sortingColumns = new ArrayList<>(sort.size());
+        for (Map.Entry<String, Sort> entry : sort.entrySet()) {
+            String expressionId = entry.getKey();
+            Sort s = entry.getValue();
+
+            int atIndex = -1;
+            for (int i = 0; i < fields.size(); i++) {
+                Tuple<FieldExtraction, String> field = fields.get(i);
+                if (field.v2().equals(expressionId)) {
+                    atIndex = i;
+                    break;
+                }
+            }
+            if (atIndex == -1) {
+                throw new SqlIllegalArgumentException("Cannot find backing column for ordering aggregation [{}]", s);
+            }
+
+            // assemble a comparator for it, if it's not an AggregateSort
+            // then it's pre-sorted by ES so use null
+            Comparator comp = null;
+            if (s instanceof AggregateSort) {
+                comp = s.direction() == Sort.Direction.ASC ? Comparator.naturalOrder() : Comparator.reverseOrder();
+                comp = s.missing() == Sort.Missing.FIRST ? Comparator.nullsFirst(comp) : Comparator.nullsLast(comp);
+            }
+
+            sortingColumns.add(new Tuple<>(Integer.valueOf(atIndex), comp));
         }
         
-        if (customSort == null) {
+        return sortingColumns;
-            customSort = Boolean.valueOf(aggSort);
-        }
-
-        return aggSort ? sortingColumns : emptyList();
     }
 
     /**
@@ -230,7 +229,7 @@ public class QueryContainer {
         return pseudoFunctions;
     }
 
-    public Set<Sort> sort() {
+    public Map<String, Sort> sort() {
         return sort;
     }
 
@@ -304,10 +303,10 @@ public class QueryContainer {
         return new QueryContainer(query, aggs, fields, aliases, pseudoFunctions, procs, sort, limit, trackHits, includeFrozen, minPageSize);
     }
 
-    public QueryContainer addSort(Sort sortable) {
+    public QueryContainer addSort(String expressionId, Sort sortable) {
-        Set<Sort> sort = new LinkedHashSet<>(this.sort);
+        Map<String, Sort> newSort = new LinkedHashMap<>(this.sort);
-        sort.add(sortable);
+        newSort.put(expressionId, sortable);
-        return new QueryContainer(query, aggs, fields, aliases, pseudoFunctions, scalarFunctions, sort, limit, trackHits, includeFrozen,
+        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/QuerierTests.java

@@ -66,6 +66,70 @@ public class QuerierTests extends ESTestCase {
         }
     }
 
+    @SuppressWarnings("rawtypes")
+    public void testAggSorting_TwoFields_One_Presorted() {
+        List<Tuple<Integer, Comparator>> tuples = new ArrayList<>(2);
+        tuples.add(new Tuple<>(0, null));
+        tuples.add(new Tuple<>(1, Comparator.reverseOrder()));
+        Querier.AggSortingQueue queue = new AggSortingQueue(20, tuples);
+
+        for (int i = 1; i <= 100; i++) {
+            queue.insertWithOverflow(new Tuple<>(Arrays.asList(i <= 5 ? null : 100 - i + 1, i), i));
+        }
+        List<List<?>> results = queue.asList();
+
+        assertEquals(20, results.size());
+        for (int i = 0; i < 20; i++) {
+            assertEquals(i < 5 ? null : 100 - i, results.get(i).get(0));
+            assertEquals(i < 5 ? 5 - i : i + 1, results.get(i).get(1));
+        }
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    public void testAggSorting_FourFields() {
+        List<Comparator> comparators = Arrays.asList(
+                Comparator.naturalOrder(),
+                Comparator.naturalOrder(),
+                Comparator.reverseOrder(),
+                Comparator.naturalOrder()
+        );
+        List<Tuple<Integer, Comparator>> tuples = new ArrayList<>(4);
+        tuples.add(new Tuple<>(0, null));
+        tuples.add(new Tuple<>(1, comparators.get(1)));
+        tuples.add(new Tuple<>(2, null));
+        tuples.add(new Tuple<>(3, comparators.get(3)));
+        Querier.AggSortingQueue queue = new AggSortingQueue(35, tuples);
+
+        List<List<Integer>> expected = new ArrayList<>(128);
+        for (int i = 0; i < 128; i++) {
+            int col1 = i / 16;
+            int col2 = 15 - (i / 8);
+            int col3 = 32 - (i / 4);
+            int col4 = 127 - i;
+
+            expected.add(Arrays.asList(col1, col2, col3, col4));
+            queue.insertWithOverflow(new Tuple<>(Arrays.asList(col1, col2, col3, col4), i));
+        }
+
+        expected.sort((o1, o2) -> {
+            for (int i = 0; i < 4; i++) {
+                int result = comparators.get(i).compare(o1.get(i), o2.get(i));
+                if (result != 0) {
+                    return result;
+                }
+            }
+            return 0;
+        });
+        List<List<?>> results = queue.asList();
+
+        assertEquals(35, results.size());
+        for (int i = 0; i < 35; i++) {
+            for (int j = 0; j < 4; j++) {
+                assertEquals(expected.get(i).get(j), results.get(i).get(j));
+            }
+        }
+    }
+
     @SuppressWarnings("rawtypes")
     public void testAggSorting_Randomized() {
         // Initialize comparators for fields (columns)
@@ -76,7 +140,7 @@ public class QuerierTests extends ESTestCase {
             boolean order = randomBoolean();
             ordering[j] = order;
             Comparator comp = order ? Comparator.naturalOrder() : Comparator.reverseOrder();
-            tuples.add(new Tuple<Integer, Comparator>(j, comp));
+            tuples.add(new Tuple<>(j, comp));
         }
 
         // Insert random no of documents (rows) with random 0/1 values for each field

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

@@ -97,7 +97,7 @@ public class SourceGeneratorTests extends ESTestCase {
 
     public void testSortScoreSpecified() {
         QueryContainer container = new QueryContainer()
-                .addSort(new ScoreSort(Direction.DESC, null));
+                .addSort("id", new ScoreSort(Direction.DESC, null));
         SearchSourceBuilder sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertEquals(singletonList(scoreSort()), sourceBuilder.sorts());
     }
@@ -106,14 +106,14 @@ public class SourceGeneratorTests extends ESTestCase {
         FieldSortBuilder sortField = fieldSort("test").unmappedType("keyword");
         
         QueryContainer container = new QueryContainer()
-                .addSort(new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")), Direction.ASC,
+                .addSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
-                        Missing.LAST));
+                        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(new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")), Direction.DESC,
+                .addSort("id", new AttributeSort(new FieldAttribute(Source.EMPTY, "test", new KeywordEsField("test")),
-                        Missing.FIRST));
+                        Direction.DESC, Missing.FIRST));
         sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertEquals(singletonList(sortField.order(SortOrder.DESC).missing("_first")), sourceBuilder.sorts());
     }
@@ -137,4 +137,4 @@ public class SourceGeneratorTests extends ESTestCase {
         SearchSourceBuilder sourceBuilder = SourceGenerator.sourceBuilder(container, null, randomIntBetween(1, 10));
         assertNull(sourceBuilder.sorts());
     }
-}
+}