Skip to content

Commit d4a2d19

Browse files
qianheng-awsqianheng-aws
authored
Support push down sort after limit (#4657)
* Support push down sort after limit Signed-off-by: Heng Qian <[email protected]> * Fix IT Signed-off-by: Heng Qian <[email protected]> * Fix IT after merging main Signed-off-by: Heng Qian <[email protected]> * spotless apply Signed-off-by: Heng Qian <[email protected]> --------- Signed-off-by: Heng Qian <[email protected]>
1 parent 018fc59 commit d4a2d19

File tree

13 files changed

+89
-32
lines changed

13 files changed

+89
-32
lines changed

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1402,4 +1402,18 @@ public void testNestedAggregationsExplain() throws IOException {
14021402
+ " timestamp, value_range, category",
14031403
TEST_INDEX_TIME_DATA)));
14041404
}
1405+
1406+
@Test
1407+
public void testTopKThenSortExplain() throws IOException {
1408+
enabledOnlyWhenPushdownIsEnabled();
1409+
String expected = loadExpectedPlan("explain_top_k_then_sort_push.yaml");
1410+
assertYamlEqualsIgnoreId(
1411+
expected,
1412+
explainQueryYaml(
1413+
"source=opensearch-sql_test_index_account"
1414+
+ "| sort balance"
1415+
+ "| head 5 "
1416+
+ "| sort age "
1417+
+ "| fields age"));
1418+
}
14051419
}

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteSortCommandIT.java‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,6 @@ public void init() throws Exception {
2424
enableCalcite();
2525
}
2626

27-
// TODO: Move this test to SortCommandIT once head-then-sort is fixed in v2.
28-
@Test
29-
public void testHeadThenSort() throws IOException {
30-
JSONObject result =
31-
executeQuery(String.format("source=%s | head 2 | sort age | fields age", TEST_INDEX_BANK));
32-
verifyOrder(result, rows(32), rows(36));
33-
}
34-
3527
@Test
3628
public void testPushdownSortPlusExpression() throws IOException {
3729
String ppl =

‎integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -265,14 +265,8 @@ public void testSortThenLimitExplain() throws IOException {
265265
+ "| fields age"));
266266
}
267267

268-
/**
269-
* Push down LIMIT only Sort should NOT be pushed down since DSL process limit before sort when
270-
* they coexist
271-
*/
272268
@Test
273269
public void testLimitThenSortExplain() throws IOException {
274-
// TODO: Fix the expected output in expectedOutput/ppl/explain_limit_then_sort_push.json (v2)
275-
// limit-then-sort should not be pushed down.
276270
String expected = loadExpectedPlan("explain_limit_then_sort_push.yaml");
277271
assertYamlEqualsIgnoreId(
278272
expected,

‎integ-test/src/test/java/org/opensearch/sql/ppl/SortCommandIT.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,4 +253,16 @@ public void testSortWithAscMultipleFields() throws IOException {
253253
rows(36, 20),
254254
rows(39, 25));
255255
}
256+
257+
@Test
258+
public void testHeadThenSort() throws IOException {
259+
JSONObject result =
260+
executeQuery(String.format("source=%s | head 2 | sort age | fields age", TEST_INDEX_BANK));
261+
if (isPushdownDisabled()) {
262+
// Pushdown is disabled, it will retrieve the first 2 docs since there's only 1 shard.
263+
verifyOrder(result, rows(32), rows(36));
264+
} else {
265+
verifyOrder(result, rows(28), rows(32));
266+
}
267+
}
256268
}

‎integ-test/src/test/resources/expectedOutput/calcite/explain_join_with_fields.yaml‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,15 @@ calcite:
1212
EnumerableCalc(expr#0..13=[{inputs}], account_number=[$t1], firstname=[$t2], address=[$t3], birthdate=[$t4], gender=[$t5], city=[$t6], lastname=[$t7], balance=[$t8], employer=[$t9], state=[$t10], age=[$t11], email=[$t12], male=[$t13])
1313
EnumerableLimit(fetch=[10000])
1414
EnumerableMergeJoin(condition=[=($0, $1)], joinType=[left])
15-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number], SORT->[{
15+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number], LIMIT->10000, SORT->[{
1616
"account_number" : {
1717
"order" : "asc",
1818
"missing" : "_last"
1919
}
20-
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","_source":{"includes":["account_number"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
21-
EnumerableSort(sort0=[$0], dir0=[ASC])
22-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], LIMIT->50000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":50000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]}}, requestedTotalSize=50000, pageSize=null, startFrom=0)])
20+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["account_number"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
21+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], LIMIT->50000, SORT->[{
22+
"account_number" : {
23+
"order" : "asc",
24+
"missing" : "_last"
25+
}
26+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":50000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=50000, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_limit_then_sort_push.yaml‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@ calcite:
66
LogicalSort(fetch=[5])
77
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
88
physical: |
9-
EnumerableLimit(fetch=[10000])
10-
EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])
11-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[age], LIMIT->5], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["age"],"excludes":[]}}, requestedTotalSize=5, pageSize=null, startFrom=0)])
9+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[age], LIMIT->5, SORT->[{
10+
"age" : {
11+
"order" : "asc",
12+
"missing" : "_first"
13+
}
14+
}], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["age"],"excludes":[]},"sort":[{"age":{"order":"asc","missing":"_first"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_merge_join_sort_push.yaml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,5 +17,9 @@ calcite:
1717
"missing" : "_last"
1818
}
1919
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
20-
EnumerableSort(sort0=[$0], dir0=[ASC])
21-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], LIMIT->50000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":50000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]}}, requestedTotalSize=50000, pageSize=null, startFrom=0)])
20+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[PROJECT->[account_number, firstname, address, birthdate, gender, city, lastname, balance, employer, state, age, email, male], LIMIT->50000, SORT->[{
21+
"account_number" : {
22+
"order" : "asc",
23+
"missing" : "_last"
24+
}
25+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":50000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","birthdate","gender","city","lastname","balance","employer","state","age","email","male"],"excludes":[]},"sort":[{"account_number":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=50000, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_scalar_correlated_subquery_in_select.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@ calcite:
1313
EnumerableCalc(expr#0..3=[{inputs}], expr#4=[IS NULL($t3)], expr#5=[0:BIGINT], expr#6=[CASE($t4, $t5, $t3)], id=[$t1], name=[$t0], count_dept=[$t6])
1414
EnumerableLimit(fetch=[10000])
1515
EnumerableMergeJoin(condition=[=($1, $2)], joinType=[left])
16-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_worker]], PushDownContext=[[PROJECT->[name, id], SORT->[{
16+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_worker]], PushDownContext=[[PROJECT->[name, id], LIMIT->10000, SORT->[{
1717
"id" : {
1818
"order" : "asc",
1919
"missing" : "_last"
2020
}
21-
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","_source":{"includes":["name","id"],"excludes":[]},"sort":[{"id":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
22-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_work_information]], PushDownContext=[[PROJECT->[name, uid], FILTER->AND(IS NOT NULL($1), IS NOT NULL($0)), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count(name)=COUNT($1)), SORT->[0]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"must":[{"exists":{"field":"uid","boost":1.0}},{"exists":{"field":"name","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["name","uid"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"uid":{"terms":{"field":"uid","missing_bucket":true,"missing_order":"last","order":"asc"}}}]},"aggregations":{"count(name)":{"value_count":{"field":"name"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
21+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["name","id"],"excludes":[]},"sort":[{"id":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
22+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_work_information]], PushDownContext=[[PROJECT->[name, uid], FILTER->AND(IS NOT NULL($1), IS NOT NULL($0)), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count(name)=COUNT($1)), SORT->[0]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":0,"timeout":"1m","query":{"bool":{"must":[{"exists":{"field":"uid","boost":1.0}},{"exists":{"field":"name","boost":1.0}}],"adjust_pure_negative":true,"boost":1.0}},"_source":{"includes":["name","uid"],"excludes":[]},"aggregations":{"composite_buckets":{"composite":{"size":1000,"sources":[{"uid":{"terms":{"field":"uid","missing_bucket":true,"missing_order":"last","order":"asc"}}}]},"aggregations":{"count(name)":{"value_count":{"field":"name"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_top_k_then_sort_push.yaml‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(sort0=[$0], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(age=[$8])
5+
LogicalSort(sort0=[$8], dir0=[ASC-nulls-first])
6+
LogicalSort(sort0=[$3], dir0=[ASC-nulls-first], fetch=[5])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
8+
physical: |
9+
EnumerableLimit(fetch=[10000])
10+
EnumerableCalc(expr#0..1=[{inputs}], age=[$t1])
11+
EnumerableSort(sort0=[$1], dir0=[ASC-nulls-first])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[balance, age], SORT->[{
13+
"balance" : {
14+
"order" : "asc",
15+
"missing" : "_first"
16+
}
17+
}], LIMIT->5], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["balance","age"],"excludes":[]},"sort":[{"balance":{"order":"asc","missing":"_first"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])

‎integ-test/src/test/resources/expectedOutput/calcite/explain_trendline_sort_push.yaml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,5 +11,9 @@ calcite:
1111
EnumerableCalc(expr#0..3=[{inputs}], expr#4=[1], expr#5=[>($t1, $t4)], expr#6=[CAST($t3):DOUBLE NOT NULL], expr#7=[/($t2, $t6)], expr#8=[null:NULL], expr#9=[CASE($t5, $t7, $t8)], ageTrend=[$t9])
1212
EnumerableWindow(window#0=[window(rows between $1 PRECEDING and CURRENT ROW aggs [COUNT(), $SUM0($0), COUNT($0)])])
1313
EnumerableCalc(expr#0=[{inputs}], expr#1=[IS NOT NULL($t0)], age=[$t0], $condition=[$t1])
14-
EnumerableSort(sort0=[$0], dir0=[ASC])
15-
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[age], LIMIT->5], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["age"],"excludes":[]}}, requestedTotalSize=5, pageSize=null, startFrom=0)])
14+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[age], LIMIT->5, SORT->[{
15+
"age" : {
16+
"order" : "asc",
17+
"missing" : "_last"
18+
}
19+
}]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["age"],"excludes":[]},"sort":[{"age":{"order":"asc","missing":"_last"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])

0 commit comments

Comments
 (0)