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 index f12f069a3b33f..61cd6e93c1831 100644 --- 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()); + } } \ No newline at end of file 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 index ae761b6432edf..0baa18765ff66 100644 --- 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 index 6839e7275ae63..3d7e12b3d9b09 100644 --- 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;