Skip to content

Commit 7a6a3a5

Browse files
penghuopenghuo
authored
Fix #5114: preserve head/TopK semantics for sort-expression pushdown (#5135)
* Fix bug 5114 Signed-off-by: Peng Huo <penghuo@gmail.com> * Update Signed-off-by: Peng Huo <penghuo@gmail.com> * Address comments Signed-off-by: Peng Huo <penghuo@gmail.com> * Fix spotless import ordering for EnumerableLimitSort Signed-off-by: Peng Huo <penghuo@gmail.com> --------- Signed-off-by: Peng Huo <penghuo@gmail.com>
1 parent d18c514 commit 7a6a3a5

4 files changed

Lines changed: 123 additions & 0 deletions

File tree

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2033,6 +2033,17 @@ public void testTopKThenSortExplain() throws IOException {
20332033
+ "| fields age"));
20342034
}
20352035

2036+
@Test
2037+
public void testIssue5114SortExprHeadExplain() throws IOException {
2038+
enabledOnlyWhenPushdownIsEnabled();
2039+
String query =
2040+
"source=opensearch-sql_test_index_account | eval a = rand() | sort a | fields"
2041+
+ " account_number | head 5";
2042+
var result = explainQueryYaml(query);
2043+
String expected = loadExpectedPlan("explain_issue_5114_sort_expr_head_push.yaml");
2044+
assertYamlEqualsIgnoreId(expected, result);
2045+
}
2046+
20362047
@Test
20372048
public void testGeoIpPushedInAgg() throws IOException {
20382049
// This explain IT verifies that externally registered UDF can be properly pushed down

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(account_number=[$0])
5+
LogicalSort(sort0=[$17], dir0=[ASC-nulls-first], fetch=[5])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], a=[RAND()])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
8+
physical: |
9+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number], SORT_EXPR->[RAND() ASCENDING NULLS_FIRST], LIMIT->5, LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":5,"timeout":"1m","_source":{"includes":["account_number"],"excludes":[]},"sort":[{"_script":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXQAbnsKICAib3AiOiB7CiAgICAibmFtZSI6ICJSQU5EIiwKICAgICJraW5kIjogIk9USEVSX0ZVTkNUSU9OIiwKICAgICJzeW50YXgiOiAiRlVOQ1RJT04iCiAgfSwKICAib3BlcmFuZHMiOiBbXQp9\"}","lang":"opensearch_compounded_script","params":{"MISSING_MAX":false,"utcTimestamp": 0,"SOURCES":[],"DIGESTS":[]}},"type":"number","order":"asc"}}]}, requestedTotalSize=5, pageSize=null, startFrom=0)])

‎integ-test/src/yamlRestTest/resources/rest-api-spec/test/issues/5114.yml‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
setup:
2+
- do:
3+
query.settings:
4+
body:
5+
transient:
6+
plugins.calcite.enabled: true
7+
8+
- do:
9+
indices.create:
10+
index: issue5114
11+
body:
12+
settings:
13+
number_of_shards: 1
14+
number_of_replicas: 0
15+
mappings:
16+
properties:
17+
id:
18+
type: integer
19+
name:
20+
type: keyword
21+
reportsTo:
22+
type: keyword
23+
24+
- do:
25+
bulk:
26+
refresh: true
27+
body:
28+
- '{"index": {"_index": "issue5114", "_id": "1"}}'
29+
- '{"id": 1, "name": "Dev", "reportsTo": "Eliot"}'
30+
- '{"index": {"_index": "issue5114", "_id": "2"}}'
31+
- '{"id": 2, "name": "Eliot", "reportsTo": "Ron"}'
32+
- '{"index": {"_index": "issue5114", "_id": "3"}}'
33+
- '{"id": 3, "name": "Ron", "reportsTo": "Andrew"}'
34+
- '{"index": {"_index": "issue5114", "_id": "4"}}'
35+
- '{"id": 4, "name": "Andrew", "reportsTo": null}'
36+
- '{"index": {"_index": "issue5114", "_id": "5"}}'
37+
- '{"id": 5, "name": "Asya", "reportsTo": "Ron"}'
38+
- '{"index": {"_index": "issue5114", "_id": "6"}}'
39+
- '{"id": 6, "name": "Dan", "reportsTo": "Andrew"}'
40+
41+
---
42+
teardown:
43+
- do:
44+
indices.delete:
45+
index: issue5114
46+
ignore_unavailable: true
47+
- do:
48+
query.settings:
49+
body:
50+
transient:
51+
plugins.calcite.enabled: false
52+
53+
---
54+
"Issue 5114: head should be preserved for non-order-equivalent deterministic sort expression":
55+
- skip:
56+
features:
57+
- headers
58+
- do:
59+
headers:
60+
Content-Type: 'application/json'
61+
ppl:
62+
body:
63+
query: source=issue5114 | eval a = abs(id) + 1 | sort a | fields id | head 5
64+
65+
- match: { total: 5 }
66+
- match: { schema: [ { name: id, type: int } ] }
67+
68+
---
69+
"Issue 5114: head should be preserved for non-deterministic sort expression":
70+
- skip:
71+
features:
72+
- headers
73+
- do:
74+
headers:
75+
Content-Type: 'application/json'
76+
ppl:
77+
body:
78+
query: source=issue5114 | eval a = rand() | sort a | fields id | head 5
79+
80+
- match: { total: 5 }
81+
- match: { schema: [ { name: id, type: int } ] }
82+
83+
---
84+
"Issue 5114 control: order-equivalent expression remains correct":
85+
- skip:
86+
features:
87+
- headers
88+
- do:
89+
headers:
90+
Content-Type: 'application/json'
91+
ppl:
92+
body:
93+
query: source=issue5114 | eval a = id + 1 | sort a | fields id | head 5
94+
95+
- match: { total: 5 }
96+
- match: { schema: [ { name: id, type: int } ] }
97+
- match: { datarows: [ [ 1 ], [ 2 ], [ 3 ], [ 4 ], [ 5 ] ] }

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/rules/SortExprIndexScanRule.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import java.util.Map;
1212
import java.util.Optional;
1313
import java.util.function.Predicate;
14+
import org.apache.calcite.adapter.enumerable.EnumerableLimitSort;
1415
import org.apache.calcite.plan.RelOptRuleCall;
1516
import org.apache.calcite.rel.RelFieldCollation;
1617
import org.apache.calcite.rel.RelFieldCollation.Direction;
@@ -47,6 +48,11 @@ protected SortExprIndexScanRule(SortExprIndexScanRule.Config config) {
4748
@Override
4849
protected void onMatchImpl(RelOptRuleCall call) {
4950
final Sort sort = call.rel(0);
51+
// EnumerableLimitSort carries fetch semantics; this rule doesn't preserve it on physical
52+
// scans because limit pushdown path is logical-only.
53+
if (sort instanceof EnumerableLimitSort) {
54+
return;
55+
}
5056
final Project project = call.rel(1);
5157
final AbstractCalciteIndexScan scan = call.rel(2);
5258

0 commit comments

Comments
 (0)