SQL: Fix wrong results when sorting on aggregate (#43154)

- Previously, when shorting on an aggregate function the bucket
processing ended early when the explicit (LIMIT XXX) or the impliciti
limit of 512 was reached. As a consequence, only a set of grouping
buckets was processed and the results returned didn't reflect the global
ordering.

- Previously, the priority queue shorting method had an inverse
comparison check and the final response from the priority queue was also
returned in the inversed order because of the calls to the `pop()`
method.

Fixes: #42851

(cherry picked from commit 19909ed)

Author: matriv

x-pack/plugin/sql/qa/single-node/src/test/java/org/elasticsearch/xpack/sql/qa/single_node/JdbcCsvSpecIT.java

@@ -32,7 +32,7 @@ public JdbcCsvSpecIT(String fileName, String groupName, String testName, Integer
     @Override
     protected int fetchSize() {
         // using a smaller fetchSize for nested documents' tests to uncover bugs
-        // similar with https://github.com/elastic/elasticsearch/issues/35176 quicker
+        // similar to https://github.com/elastic/elasticsearch/issues/35176 quicker
         return fileName.startsWith("nested") && randomBoolean() ? randomIntBetween(1,5) : super.fetchSize();
     }
 }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/ErrorsTestCase.java

@@ -21,4 +21,5 @@ public interface ErrorsTestCase {
     void testSelectGroupByScore() throws Exception;
     void testSelectScoreSubField() throws Exception;
     void testSelectScoreInScalar() throws Exception;
+    void testHardLimitForSortOnAggregate() throws Exception;
 }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/cli/CliIntegrationTestCase.java

@@ -56,8 +56,8 @@ protected SecurityConfig securityConfig() {
         return null;
     }
 
-    protected void index(String index, CheckedConsumer<XContentBuilder, IOException> body) throws IOException {
-        Request request = new Request("PUT", "/" + index + "/doc/1");
+    protected void index(String index, int docId, CheckedConsumer<XContentBuilder, IOException> body) throws IOException {
+        Request request = new Request("PUT", "/" + index + "/doc/" + docId);
         request.addParameter("refresh", "true");
         XContentBuilder builder = JsonXContent.contentBuilder().startObject();
         body.accept(builder);
@@ -66,6 +66,10 @@ protected void index(String index, CheckedConsumer<XContentBuilder, IOException>
         client().performRequest(request);
     }
 
+    protected void index(String index, CheckedConsumer<XContentBuilder, IOException> body) throws IOException {
+        index(index, 1, body);
+    }
+
     public String command(String command) throws IOException {
         return cli.command(command);
     }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/cli/ErrorsTestCase.java

@@ -97,8 +97,15 @@ public void testSelectScoreInScalar() throws Exception {
         assertEquals("line 1:12: [SCORE()] cannot be an argument to a function" + END, readLine());
     }
 
+    @Override
+    public void testHardLimitForSortOnAggregate() throws Exception {
+        index("test", body -> body.field("a", 1).field("b", 2));
+        String commandResult = command("SELECT max(a) max FROM test GROUP BY b ORDER BY max LIMIT 10000");
+        assertEquals(START + "Bad request [[3;33;22mThe maximum LIMIT for aggregate sorting is [512], received [10000]" + END,
+            commandResult);
+    }
+
     public static void assertFoundOneProblem(String commandResult) {
         assertEquals(START + "Bad request [[3;33;22mFound 1 problem(s)", commandResult);
     }
-
 }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/cli/FetchSizeTestCase.java

@@ -49,4 +49,32 @@ public void testInvalidFetchSize() throws IOException {
         assertEquals(ErrorsTestCase.START + "Invalid fetch size [[3;33;22m" + Long.MAX_VALUE + ErrorsTestCase.END,
             command("fetch size = " + Long.MAX_VALUE));
     }
+
+    // Test for issue: https://github.com/elastic/elasticsearch/issues/42851
+    // Even though fetch size and limit are smaller than the noRows, all buckets
+    // should be processed to achieve the global ordering of the aggregate function.
+    public void testOrderingOnAggregate() throws IOException {
+        Request request = new Request("PUT", "/test/doc/_bulk");
+        request.addParameter("refresh", "true");
+        StringBuilder bulk = new StringBuilder();
+        for (int i = 1; i <= 100; i++) {
+            bulk.append("{\"index\":{}}\n");
+            bulk.append("{\"a\":").append(i).append(", \"b\" : ").append(i).append("}\n");
+        }
+        request.setJsonEntity(bulk.toString());
+        client().performRequest(request);
+
+        assertEquals("[?1l>[?1000l[?2004lfetch size set to [90m4[0m", command("fetch size = 4"));
+        assertEquals("[?1l>[?1000l[?2004lfetch separator set to \"[90m -- fetch sep -- [0m\"",
+            command("fetch separator = \" -- fetch sep -- \""));
+        assertThat(command("SELECT max(b) FROM test GROUP BY a ORDER BY max(b) DESC LIMIT 20"), containsString("max(b)"));
+        assertThat(readLine(), containsString("----------"));
+        for (int i = 100; i > 80; i--) {
+            if (i < 100 && i % 4 == 0) {
+                assertThat(readLine(), containsString(" -- fetch sep -- "));
+            }
+            assertThat(readLine(), containsString(Integer.toString(i)));
+        }
+        assertEquals("", readLine());
+    }
 }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/ErrorsTestCase.java

@@ -116,4 +116,14 @@ public void testSelectScoreInScalar() throws Exception {
             assertThat(e.getMessage(), startsWith("Found 1 problem(s)\nline 1:12: [SCORE()] cannot be an argument to a function"));
         }
     }
+
+    @Override
+    public void testHardLimitForSortOnAggregate() throws Exception {
+        index("test", body -> body.field("a", 1).field("b", 2));
+        try (Connection c = esJdbc()) {
+            SQLException e = expectThrows(SQLException.class, () ->
+                c.prepareStatement("SELECT max(a) max FROM test GROUP BY b ORDER BY max LIMIT 10000").executeQuery());
+            assertEquals("The maximum LIMIT for aggregate sorting is [512], received [10000]", e.getMessage());
+        }
+    }
 }

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/SqlSpecTestCase.java

@@ -75,6 +75,13 @@ public SqlSpecTestCase(String fileName, String groupName, String testName, Integ
         this.query = query;
     }
 
+    @Override
+    protected int fetchSize() {
+        // using a smaller fetchSize for nested documents' tests to uncover bugs
+        // similar to https://github.com/elastic/elasticsearch/issues/42581
+        return randomIntBetween(1, 20);
+    }
+
     @Override
     protected final void doTest() throws Throwable {
         // we skip the tests in case of these locales because ES-SQL is Locale-insensitive for now

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/rest/RestSqlTestCase.java

@@ -281,6 +281,13 @@ public void testSelectScoreInScalar() throws Exception {
             containsString("line 1:12: [SCORE()] cannot be an argument to a function"));
     }
 
+    @Override
+    public void testHardLimitForSortOnAggregate() throws Exception {
+        index("{\"a\": 1, \"b\": 2}");
+        expectBadRequest(() -> runSql(randomMode(), "SELECT max(a) max FROM test GROUP BY b ORDER BY max LIMIT 10000"),
+            containsString("The maximum LIMIT for aggregate sorting is [512], received [10000]"));
+    }
+
     protected void expectBadRequest(CheckedSupplier<Map<String, Object>, Exception> code, Matcher<String> errorMessageMatcher) {
         try {
             Map<String, Object> result = code.get();

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

@@ -99,13 +99,13 @@ aggNotSpecifiedWithHavingOnLargeGroupBy
 SELECT MAX(salary) AS max FROM test_emp GROUP BY emp_no HAVING AVG(salary) > 1000 ORDER BY MIN(salary);
 
 aggWithTieBreakerDescAsc
-SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MIN(languages) DESC NULLS FIRST, emp_no ASC;
+SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MIN(languages) DESC NULLS FIRST, emp_no ASC LIMIT 50;
 
 aggWithTieBreakerDescDesc
-SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MIN(languages) DESC NULLS FIRST, emp_no DESC;
+SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MIN(languages) DESC NULLS FIRST, emp_no DESC LIMIT 50;
 
 aggWithTieBreakerAscDesc
-SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MAX(languages) ASC NULLS FIRST, emp_no DESC;
+SELECT emp_no, MIN(languages) AS min FROM test_emp GROUP BY emp_no ORDER BY MAX(languages) ASC NULLS FIRST, emp_no DESC LIMIT 50;
 
 aggWithMixOfOrdinals
 SELECT gender AS g, MAX(salary) AS m FROM test_emp GROUP BY gender ORDER BY 2 DESC LIMIT 3;

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

@@ -133,8 +133,9 @@ public void onResponse(SearchResponse r) {
                         return;
                     }
 
-                    updateCompositeAfterKey(r, query);
-                    CompositeAggsRowSet rowSet = new CompositeAggsRowSet(extractors, mask, r, limit, serializeQuery(query), indices);
+                    boolean hasAfterKey = updateCompositeAfterKey(r, query);
+                    CompositeAggsRowSet rowSet = new CompositeAggsRowSet(extractors, mask, r, limit,
+                            hasAfterKey ? serializeQuery(query) : null, indices);
                     listener.onResponse(rowSet);
                 } catch (Exception ex) {
                     listener.onFailure(ex);
@@ -167,7 +168,7 @@ static CompositeAggregation getComposite(SearchResponse response) {
         throw new SqlIllegalArgumentException("Unrecognized root group found; {}", agg.getClass());
     }
 
-    static void updateCompositeAfterKey(SearchResponse r, SearchSourceBuilder next) {
+    static boolean updateCompositeAfterKey(SearchResponse r, SearchSourceBuilder next) {
         CompositeAggregation composite = getComposite(r);
 
         if (composite == null) {
@@ -176,22 +177,25 @@ static void updateCompositeAfterKey(SearchResponse r, SearchSourceBuilder next)
 
         Map<String, Object> afterKey = composite.afterKey();
         // a null after-key means done
-        if (afterKey != null) {
-            AggregationBuilder aggBuilder = next.aggregations().getAggregatorFactories().iterator().next();
-            // update after-key with the new value
-            if (aggBuilder instanceof CompositeAggregationBuilder) {
-                CompositeAggregationBuilder comp = (CompositeAggregationBuilder) aggBuilder;
-                comp.aggregateAfter(afterKey);
-            } else {
-                throw new SqlIllegalArgumentException("Invalid client request; expected a group-by but instead got {}", aggBuilder);
-            }
+        if (afterKey == null) {
+            return false;
+        }
+
+        AggregationBuilder aggBuilder = next.aggregations().getAggregatorFactories().iterator().next();
+        // update after-key with the new value
+        if (aggBuilder instanceof CompositeAggregationBuilder) {
+            CompositeAggregationBuilder comp = (CompositeAggregationBuilder) aggBuilder;
+            comp.aggregateAfter(afterKey);
+            return true;
+        } else {
+            throw new SqlIllegalArgumentException("Invalid client request; expected a group-by but instead got {}", aggBuilder);
         }
     }
 
     /**
      * Deserializes the search source from a byte array.
      */
-    static SearchSourceBuilder deserializeQuery(NamedWriteableRegistry registry, byte[] source) throws IOException {
+    private static SearchSourceBuilder deserializeQuery(NamedWriteableRegistry registry, byte[] source) throws IOException {
         try (NamedWriteableAwareStreamInput in = new NamedWriteableAwareStreamInput(StreamInput.wrap(source), registry)) {
             return new SearchSourceBuilder(in);
         }