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

costin · costin · commit b29a2cb360a0 · 2019-09-26T22:32:42.000+03:00
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 f4c2964)
diff --git a/x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/FetchSizeTestCase.java b/x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/jdbc/FetchSizeTestCase.java
@@ -156,20 +156,7 @@ private void assertNestedDocuments(ResultSet rs, int i) throws SQLException {
      * page size affects the result not the intermediate query.
      */
     public void testPivotPaging() 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());
+        addPivotData();
         
         try (Connection c = esJdbc();
              Statement s = c.createStatement()) {
@@ -204,4 +191,50 @@ public void testPivotPaging() throws Exception {
         }
         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());
+    }
 }
diff --git a/x-pack/plugin/sql/qa/src/main/resources/pivot.csv-spec b/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
diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/PivotRowSet.java b/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;
