SQL: Check case where the pivot limit is reached (#47121)

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

Fix #47002

(cherry picked from commit f4c29646f097bbd29855300342823ef4cef61c05)

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

@@ -156,20 +156,7 @@ public class FetchSizeTestCase extends JdbcIntegrationTestCase {
      * page size affects the result not the intermediate query.
      */
     public void testPivotPaging() throws Exception {
-        Request request = new Request("PUT", "/test_pivot/_bulk");
+        addPivotData();
-        request.addParameter("refresh", "true");
-        StringBuilder bulk = new StringBuilder();
-        String[] continent = new String[] { "AF", "AS", "EU", "NA", "SA", "AQ", "AU" };
-        for (int i = 0; i <= 100; i++) {
-            bulk.append("{\"index\":{}}\n");
-            bulk.append("{\"item\":").append(i % 10)
-                .append(", \"entry\":").append(i)
-                .append(", \"amount\" : ").append(randomInt(999))
-                .append(", \"location\" : \"").append(continent[i % (continent.length)]).append("\"")
-                .append("}\n");
-        }
-        request.setJsonEntity(bulk.toString());
-        assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
         
         try (Connection c = esJdbc();
              Statement s = c.createStatement()) {
@@ -204,4 +191,50 @@ public class FetchSizeTestCase extends JdbcIntegrationTestCase {
         }
         assertNoSearchContexts();
     }
+    
+    
+    public void testPivotPagingWithLimit() throws Exception {
+        addPivotData();
+
+        try (Connection c = esJdbc();
+             Statement s = c.createStatement()) {
+            
+            // run a query with a limit that is not a multiple of the fetch size
+            String query = "SELECT * FROM "
+                    + "(SELECT item, amount, location FROM test_pivot)"
+                    + " PIVOT (AVG(amount) FOR location IN ( 'EU', 'NA' ) ) LIMIT 5";
+            // set size smaller than an agg page
+            s.setFetchSize(20);
+            try (ResultSet rs = s.executeQuery(query)) {
+                assertEquals(3, rs.getMetaData().getColumnCount());
+                for (int i = 0; i < 4; i++) {
+                    assertTrue(rs.next());
+                    assertEquals(2, rs.getFetchSize());
+                    assertEquals(Long.valueOf(i), rs.getObject("item"));
+                }
+                // last entry
+                assertTrue(rs.next());
+                assertEquals(1, rs.getFetchSize());
+                assertFalse("LIMIT should be reached", rs.next());
+            }
+        }
+        assertNoSearchContexts();
+    }
+
+    private void addPivotData() throws Exception {
+        Request request = new Request("PUT", "/test_pivot/_bulk");
+        request.addParameter("refresh", "true");
+        StringBuilder bulk = new StringBuilder();
+        String[] continent = new String[] { "AF", "AS", "EU", "NA", "SA", "AQ", "AU" };
+        for (int i = 0; i <= 100; i++) {
+            bulk.append("{\"index\":{}}\n");
+            bulk.append("{\"item\":").append(i % 10)
+                .append(", \"entry\":").append(i)
+                .append(", \"amount\" : ").append(randomInt(999))
+                .append(", \"location\" : \"").append(continent[i % (continent.length)]).append("\"")
+                .append("}\n");
+        }
+        request.setJsonEntity(bulk.toString());
+        assertEquals(200, client().performRequest(request).getStatusLine().getStatusCode());
+    }
 }

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
-// averageWithOneValueAndOrder
+schema::languages:bt|'F':d
-// 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;
-// SELECT * FROM (SELECT languages, gender, salary FROM test_emp) PIVOT (AVG(salary) FOR gender IN ('F')) ORDER BY languages DESC LIMIT 4;
+   languages   |       'F'
-//
+---------------+------------------
-//    languages   |       'F'
+5              |46705.555555555555
-// ---------------+------------------
+4              |49291.5
-// 5              |46705.555555555555
+3              |53660.0
-// 4              |49291.5
+2              |50684.4
-// 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
-// sumWithoutSubquery
+schema::birth_date:ts|emp_no:i|first_name:s|gender:s|hire_date:ts|last_name:s|1:i|2:i
-// 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;
-// 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
-//      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                 |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                 |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
-// 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
-// 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-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
-// ;
 
 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:
+        // check the last group using the following:
-        // a. the last key has been sent before (it's the last page)
+        // a. limit has been reached, the rest of the data is ignored.
-        if ((previousLastKey != null && sameCompositeKey(previousLastKey, currentRowGroupKey))) {
+        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. 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)
+        // 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;