SQL: Check case where the pivot limit is reached

In some cases, the fetch size affects the way the groups are returned
causing the last page to beyond the limit. Add dedicated check to
prevent extra data from being returned.

Fix #47002

Author: costin

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

@@ -114,18 +114,16 @@ null           |48396.28571428572|62140.666666666664
 1              |49767.22222222222|47073.25  
 ;
 
-// AwaitsFix https://github.com/elastic/elasticsearch/issues/47002
-// averageWithOneValueAndOrder
-// schema::languages:bt|'F':d
-// SELECT * FROM (SELECT languages, gender, salary FROM test_emp) PIVOT (AVG(salary) FOR gender IN ('F')) ORDER BY languages DESC LIMIT 4;
-//
-//    languages   |       'F'
-// ---------------+------------------
-// 5              |46705.555555555555
-// 4              |49291.5
-// 3              |53660.0
-// 2              |50684.4
-// ;
+averageWithOneValueAndOrder
+schema::languages:bt|'F':d
+SELECT * FROM (SELECT languages, gender, salary FROM test_emp) PIVOT (AVG(salary) FOR gender IN ('F')) ORDER BY languages DESC LIMIT 4;
+   languages   |       'F'
+---------------+------------------
+5              |46705.555555555555
+4              |49291.5
+3              |53660.0
+2              |50684.4
+;
 
 averageWithTwoValuesAndOrderDesc
 schema::languages:bt|'M':d|'F':d
@@ -165,20 +163,18 @@ null           |48396.28571428572|62140.666666666664
 5              |39052.875        |46705.555555555555
 ;
 
-// AwaitsFix https://github.com/elastic/elasticsearch/issues/47002
-// sumWithoutSubquery
-// schema::birth_date:ts|emp_no:i|first_name:s|gender:s|hire_date:ts|last_name:s|1:i|2:i
-// SELECT * FROM test_emp PIVOT (SUM(salary) FOR languages IN (1, 2)) LIMIT 5;
-//
-//      birth_date      |    emp_no     |  first_name   |    gender     |      hire_date      |   last_name   |       1       |       2
-// ---------------------+---------------+---------------+---------------+---------------------+---------------+---------------+---------------
-// null                 |10041          |Uri            |F              |1989-11-12 00:00:00.0|Lenart         |56415          |null
-// null                 |10043          |Yishay         |M              |1990-10-20 00:00:00.0|Tzvieli        |34341          |null
-// null                 |10044          |Mingsen        |F              |1994-05-21 00:00:00.0|Casley         |39728          |null
-// 1952-04-19 00:00:00.0|10009          |Sumant         |F              |1985-02-18 00:00:00.0|Peac           |66174          |null
-// 1953-01-07 00:00:00.0|10067          |Claudi         |M              |1987-03-04 00:00:00.0|Stavenow       |null           |52044
-// 1953-01-23 00:00:00.0|10019          |Lillian        |null           |1999-04-30 00:00:00.0|Haddadi        |73717          |null
-// ;
+sumWithoutSubquery
+schema::birth_date:ts|emp_no:i|first_name:s|gender:s|hire_date:ts|last_name:s|1:i|2:i
+SELECT * FROM test_emp PIVOT (SUM(salary) FOR languages IN (1, 2)) LIMIT 5;
+
+     birth_date      |    emp_no     |  first_name   |    gender     |      hire_date      |   last_name   |       1       |       2
+---------------------+---------------+---------------+---------------+---------------------+---------------+---------------+---------------
+null                 |10041          |Uri            |F              |1989-11-12 00:00:00.0|Lenart         |56415          |null
+null                 |10043          |Yishay         |M              |1990-10-20 00:00:00.0|Tzvieli        |34341          |null
+null                 |10044          |Mingsen        |F              |1994-05-21 00:00:00.0|Casley         |39728          |null
+1952-04-19 00:00:00.0|10009          |Sumant         |F              |1985-02-18 00:00:00.0|Peac           |66174          |null
+1953-01-07 00:00:00.0|10067          |Claudi         |M              |1987-03-04 00:00:00.0|Stavenow       |null           |52044
+;
 
 averageWithOneValueAndMath
 schema::languages:bt|'F':d

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

@@ -60,6 +60,10 @@ class PivotRowSet extends SchemaCompositeAggRowSet {
                 currentRowGroupKey = key;
                 // save the data
                 data.add(currentRow);
+
+                if (limit > 0 && data.size() == limit) {
+                    break;
+                }
                 // create a new row
                 currentRow = new Object[columnCount()];
             }
@@ -76,19 +80,23 @@ class PivotRowSet extends SchemaCompositeAggRowSet {
             }
         }
 
-        // add the last group if any of the following matches:
-        // a. the last key has been sent before (it's the last page)
-        if ((previousLastKey != null && sameCompositeKey(previousLastKey, currentRowGroupKey))) {
+        // check the last group using the following:
+        // a. limit has been reached, the rest of the data is ignored.
+        if (limit > 0 && data.size() == limit) {
+            afterKey = null;
+        }
+        // b. the last key has been sent before (it's the last page)
+        else if ((previousLastKey != null && sameCompositeKey(previousLastKey, currentRowGroupKey))) {
             data.add(currentRow);
             afterKey = null;
         }
-        // b. all the values are initialized (there might be another page but no need to ask for the group again)
-        // c. or no data was added (typically because there's a null value such as the group)
+        // c. all the values are initialized (there might be another page but no need to ask for the group again)
+        // d. or no data was added (typically because there's a null value such as the group)
         else if (hasNull(currentRow) == false || data.isEmpty()) {
             data.add(currentRow);
             afterKey = currentRowGroupKey;
         }
-        //otherwise we can't tell whether it's complete or not
+        // otherwise we can't tell whether it's complete or not
         // so discard the last group and ask for it on the next page
         else {
             afterKey = lastCompletedGroupKey;