SQL: Fix LIMIT bug in agg sorting (#41258)

When specifying a limit over an agg sorting, the limit will be pushed
down to the grouping which affects the custom sorting. This commit fixes
that and restricts the limit only to sorting.

Fix #40984

(cherry picked from commit da37265)

Author: costin

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

@@ -29,9 +29,24 @@ SELECT MAX(salary) AS max, MIN(salary) AS min FROM test_emp HAVING MIN(salary) >
 aggWithoutAlias
 SELECT MAX(salary) AS max FROM test_emp GROUP BY gender ORDER BY MAX(salary);
 
+aggWithoutAliasWithLimit
+SELECT MAX(salary) AS max FROM test_emp GROUP BY gender ORDER BY MAX(salary) LIMIT 3;
+
+aggWithoutAliasWithLimitDesc
+SELECT MAX(salary) AS max FROM test_emp GROUP BY gender ORDER BY MAX(salary) DESC LIMIT 3;
+
 aggWithAlias
 SELECT MAX(salary) AS m FROM test_emp GROUP BY gender ORDER BY m;
 
+aggOrderByCountWithLimit
+SELECT MAX(salary) AS max, COUNT(*) AS c FROM test_emp GROUP BY gender ORDER BY c LIMIT 3;
+
+aggOrderByCountWithLimitDescAndGrouping
+SELECT gender, COUNT(*) AS c FROM test_emp GROUP BY gender ORDER BY c DESC LIMIT 5;
+
+aggOrderByCountWithLimitDesc
+SELECT MAX(salary) AS max, COUNT(*) AS c FROM test_emp GROUP BY gender ORDER BY c DESC LIMIT 3;
+
 multipleAggsThatGetRewrittenWithoutAlias
 SELECT MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY gender ORDER BY MAX(salary);
 
@@ -56,12 +71,21 @@ SELECT MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c
 aggNotSpecifiedInTheAggregateAndGroupWithHaving
 SELECT gender, MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY MAX(salary), gender;
 
+aggNotSpecifiedInTheAggregateAndGroupWithHavingWithLimit
+SELECT gender, MIN(salary) AS min, COUNT(*) AS c FROM test_emp GROUP BY gender HAVING c > 1 ORDER BY MAX(salary), c LIMIT 5;
+
+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);
 
 multipleAggsThatGetRewrittenWithAliasOnAMediumGroupBy
 SELECT languages, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY languages ORDER BY max;
 
+multipleAggsThatGetRewrittenWithAliasOnAMediumGroupByWithLimit
+SELECT languages, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY languages ORDER BY max DESC LIMIT 5;
+
 multipleAggsThatGetRewrittenWithAliasOnALargeGroupBy
 SELECT emp_no, MAX(salary) AS max, MIN(salary) AS min FROM test_emp GROUP BY emp_no ORDER BY max;
 

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

@@ -80,7 +80,8 @@ public static SearchSourceBuilder sourceBuilder(QueryContainer container, QueryB
             if (source.size() == -1) {
                 source.size(sz);
             }
-            if (aggBuilder instanceof CompositeAggregationBuilder) {
+            // limit the composite aggs only for non-local sorting
+            if (aggBuilder instanceof CompositeAggregationBuilder && container.sortingColumns().isEmpty()) {
                 ((CompositeAggregationBuilder) aggBuilder).size(sz);
             }
         }