Skip to content

Commit 3bc9d62

Browse files
songkant-awssongkant-aws
committed
Simplify some code and add more explain tests
Signed-off-by: Songkan Tang <[email protected]>
1 parent cc61524 commit 3bc9d62

File tree

14 files changed

+154
-104
lines changed

14 files changed

+154
-104
lines changed

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -867,6 +867,35 @@ public void testSimpleSortExpressionPushDownWithOnlyExprProjected() throws Excep
867867
assertJsonEqualsIgnoreId(expected, result);
868868
}
869869

870+
@Test
871+
public void testComplexSortExpressionPushDownExplain() throws Exception {
872+
String query =
873+
"source=opensearch-sql_test_index_bank| eval age2 = age + balance | sort age2 | fields age,"
874+
+ " age2";
875+
var result = explainQueryYaml(query);
876+
String expected = loadExpectedPlan("explain_complex_sort_expr_push.yaml");
877+
assertYamlEqualsIgnoreId(expected, result);
878+
}
879+
880+
@Test
881+
public void testComplexSortExpressionPushDownWithOnlyExprProjected() throws Exception {
882+
String query =
883+
"source=opensearch-sql_test_index_bank| eval age2 = age + balance | sort age2 | fields"
884+
+ " age2";
885+
var result = explainQueryYaml(query);
886+
String expected = loadExpectedPlan("explain_complex_sort_expr_single_expr_output_push.yaml");
887+
assertYamlEqualsIgnoreId(expected, result);
888+
}
889+
890+
@Test
891+
public void testComplexSortExpressionPushDownWithoutExprProjected() throws Exception {
892+
String query =
893+
"source=opensearch-sql_test_index_bank| eval age2 = age + balance | sort age2 | fields age";
894+
var result = explainQueryYaml(query);
895+
String expected = loadExpectedPlan("explain_complex_sort_expr_no_expr_output_push.yaml");
896+
assertYamlEqualsIgnoreId(expected, result);
897+
}
898+
870899
@Test
871900
public void testRexExplain() throws IOException {
872901
String query =

‎integ-test/src/test/resources/expectedOutput/calcite/explain_complex_sort_expr_no_expr_output_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(age=[$10])
5+
LogicalSort(sort0=[$19], dir0=[ASC-nulls-first])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _id=[$13], _index=[$14], _score=[$15], _maxscore=[$16], _sort=[$17], _routing=[$18], age2=[+($10, $7)])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
8+
physical: |
9+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]], PushDownContext=[[SORT_EXPR->[+($10, $7) ASCENDING FIRST], LIMIT->10000, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["age"],"excludes":[]},"sort":[{"_script":{"script":{"source":"{\"langType\":\"calcite\",\"script\":\"rO0ABXNyABFqYXZhLnV0aWwuQ29sbFNlcleOq7Y6G6gRAwABSQADdGFneHAAAAADdwQAAAAGdAAHcm93VHlwZXQA0HsKICAiZmllbGRzIjogWwogICAgewogICAgICAidHlwZSI6ICJJTlRFR0VSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgIm5hbWUiOiAiYWdlIgogICAgfSwKICAgIHsKICAgICAgInR5cGUiOiAiQklHSU5UIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgIm5hbWUiOiAiYmFsYW5jZSIKICAgIH0KICBdLAogICJudWxsYWJsZSI6IGZhbHNlCn10AARleHBydADFewogICJvcCI6IHsKICAgICJuYW1lIjogIisiLAogICAgImtpbmQiOiAiUExVUyIsCiAgICAic3ludGF4IjogIkJJTkFSWSIKICB9LAogICJvcGVyYW5kcyI6IFsKICAgIHsKICAgICAgImlucHV0IjogMCwKICAgICAgIm5hbWUiOiAiJDAiCiAgICB9LAogICAgewogICAgICAiaW5wdXQiOiAxLAogICAgICAibmFtZSI6ICIkMSIKICAgIH0KICBdCn10AApmaWVsZFR5cGVzc3IAEWphdmEudXRpbC5IYXNoTWFwBQfawcMWYNEDAAJGAApsb2FkRmFjdG9ySQAJdGhyZXNob2xkeHA/QAAAAAAADHcIAAAAEAAAAAJ0AAdiYWxhbmNlfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAETE9OR3QAA2FnZX5xAH4ACnQAB0lOVEVHRVJ4eA==\"}","lang":"opensearch_compounded_script","params":{"utcTimestamp": 0}},"type":"number","order":"asc"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])

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

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

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

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

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_complex_sort_expr_no_expr_output_push.yaml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(age=[$10])
5+
LogicalSort(sort0=[$19], dir0=[ASC-nulls-first])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _id=[$13], _index=[$14], _score=[$15], _maxscore=[$16], _sort=[$17], _routing=[$18], age2=[+($10, $7)])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
8+
physical: |
9+
EnumerableCalc(expr#0..1=[{inputs}], age=[$t0])
10+
EnumerableLimit(fetch=[10000])
11+
EnumerableSort(sort0=[$1], dir0=[ASC-nulls-first])
12+
EnumerableCalc(expr#0..18=[{inputs}], expr#19=[+($t10, $t7)], age=[$t10], age2=[$t19])
13+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_complex_sort_expr_push.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(sort0=[$1], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(age=[$10], age2=[$19])
5+
LogicalSort(sort0=[$19], dir0=[ASC-nulls-first])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _id=[$13], _index=[$14], _score=[$15], _maxscore=[$16], _sort=[$17], _routing=[$18], age2=[+($10, $7)])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
8+
physical: |
9+
EnumerableLimit(fetch=[10000])
10+
EnumerableSort(sort0=[$1], dir0=[ASC-nulls-first])
11+
EnumerableCalc(expr#0..18=[{inputs}], expr#19=[+($t10, $t7)], age=[$t10], age2=[$t19])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])

‎integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_complex_sort_expr_single_expr_output_push.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
calcite:
2+
logical: |
3+
LogicalSystemLimit(sort0=[$0], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])
4+
LogicalProject(age2=[$19])
5+
LogicalSort(sort0=[$19], dir0=[ASC-nulls-first])
6+
LogicalProject(account_number=[$0], firstname=[$1], address=[$2], birthdate=[$3], gender=[$4], city=[$5], lastname=[$6], balance=[$7], employer=[$8], state=[$9], age=[$10], email=[$11], male=[$12], _id=[$13], _index=[$14], _score=[$15], _maxscore=[$16], _sort=[$17], _routing=[$18], age2=[+($10, $7)])
7+
CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])
8+
physical: |
9+
EnumerableLimit(fetch=[10000])
10+
EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])
11+
EnumerableCalc(expr#0..18=[{inputs}], expr#19=[+($t10, $t7)], age2=[$t19])
12+
CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_bank]])

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

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232
* takes effect, the input collation is changed to a sort over field instead of original sort over
3333
* expression. It changes the collation requirement of the whole query.
3434
*
35+
* <p>Another problem is if sort expression is pushed down to scan, the Enumerable project doesn't
36+
* know the collation is already satisfied.
37+
*
3538
* <p>AbstractConverter physical node is supposed to resolve the problem of inconsistent collation
3639
* requirement between physical node input and output. This optimization rule finds equivalent
3740
* output expression collations and input field collations. If their collation traits are satisfied,
@@ -53,12 +56,12 @@ public void onMatch(RelOptRuleCall call) {
5356
final RelCollation toCollation = toTraits.getTrait(RelCollationTraitDef.INSTANCE);
5457

5558
// Branch 1: Check if complex expressions are already sorted by scan
56-
if (handleComplexExpressionsSortedByScan(call, converter, project, toTraits, toCollation)) {
59+
if (handleComplexExpressionsSortedByScan(call, project, toTraits, toCollation)) {
5760
return;
5861
}
5962

6063
// Branch 2: Handle simple expressions that can be transformed to field sorts
61-
handleSimpleExpressionFieldSorts(call, converter, project, toTraits, toCollation);
64+
handleSimpleExpressionFieldSorts(call, project, toTraits, toCollation);
6265
}
6366

6467
/**
@@ -68,19 +71,15 @@ public void onMatch(RelOptRuleCall call) {
6871
* @return true if handled, false if not applicable
6972
*/
7073
private boolean handleComplexExpressionsSortedByScan(
71-
RelOptRuleCall call,
72-
AbstractConverter converter,
73-
Project project,
74-
RelTraitSet toTraits,
75-
RelCollation toCollation) {
74+
RelOptRuleCall call, Project project, RelTraitSet toTraits, RelCollation toCollation) {
7675

7776
// Check if toCollation is null or not a simple RelCollation with field collations
7877
if (toCollation == null || toCollation.getFieldCollations().isEmpty()) {
7978
return false;
8079
}
8180

8281
// Extract the actual enumerable scan from the input, handling RelSubset case
83-
CalciteEnumerableIndexScan scan = extractScanFromInput(project.getInput());
82+
CalciteEnumerableIndexScan scan = extractEnumerableScanFromInput(project.getInput());
8483
if (scan == null) {
8584
return false;
8685
}
@@ -102,11 +101,7 @@ private boolean handleComplexExpressionsSortedByScan(
102101
* getOrderEquivalentInputInfo.
103102
*/
104103
private void handleSimpleExpressionFieldSorts(
105-
RelOptRuleCall call,
106-
AbstractConverter converter,
107-
Project project,
108-
RelTraitSet toTraits,
109-
RelCollation toCollation) {
104+
RelOptRuleCall call, Project project, RelTraitSet toTraits, RelCollation toCollation) {
110105

111106
RelTrait fromTrait = project.getInput().getTraitSet().getTrait(RelCollationTraitDef.INSTANCE);
112107

@@ -161,7 +156,7 @@ private void handleSimpleExpressionFieldSorts(
161156
* @param input The input RelNode to extract scan from
162157
* @return CalciteEnumerableIndexScan if found, null otherwise
163158
*/
164-
private static CalciteEnumerableIndexScan extractScanFromInput(RelNode input) {
159+
private static CalciteEnumerableIndexScan extractEnumerableScanFromInput(RelNode input) {
165160

166161
// Case 1: Direct CalciteEnumerableIndexScan (physical scan)
167162
if (input instanceof CalciteEnumerableIndexScan) {
@@ -174,14 +169,7 @@ private static CalciteEnumerableIndexScan extractScanFromInput(RelNode input) {
174169
RelNode bestPlan = subset.getBest();
175170
if (bestPlan != null) {
176171
// Recursively check the best plan
177-
return extractScanFromInput(bestPlan);
178-
}
179-
180-
// During physical optimization, we should have the best plan. But if not available yet,
181-
// we can check the original node (though it's less likely to be CalciteEnumerableIndexScan)
182-
RelNode original = subset.getOriginal();
183-
if (original != null) {
184-
return extractScanFromInput(original);
172+
return extractEnumerableScanFromInput(bestPlan);
185173
}
186174
}
187175

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

Lines changed: 22 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -126,38 +126,26 @@ private List<SortExpressionInfo> extractSortExpressionInfos(
126126
*/
127127
private SortExpressionInfo mapThroughProject(
128128
RexNode sortKey, Project project, CalciteLogicalIndexScan scan, RelFieldCollation collation) {
129-
130-
if (sortKey instanceof RexInputRef) {
131-
RexInputRef inputRef = (RexInputRef) sortKey;
132-
int index = inputRef.getIndex();
133-
if (index < project.getProjects().size()) {
134-
RexNode projectExpression = project.getProjects().get(index);
135-
136-
// If the project expression is a simple RexInputRef pointing to a scan field,
137-
// store the field name for stability
138-
if (projectExpression instanceof RexInputRef) {
139-
RexInputRef scanInputRef = (RexInputRef) projectExpression;
140-
int scanFieldIndex = scanInputRef.getIndex();
141-
142-
// Get the field name from the scan's row type
143-
List<String> scanFieldNames = scan.getRowType().getFieldNames();
144-
if (scanFieldIndex < scanFieldNames.size()) {
145-
String fieldName = scanFieldNames.get(scanFieldIndex);
146-
147-
// Create SortExpressionInfo with field name (stable reference)
148-
return new SortExpressionInfo(
149-
fieldName, collation.getDirection(), collation.nullDirection);
150-
}
151-
}
152-
153-
// For complex expressions, store the RexNode
154-
return new SortExpressionInfo(
155-
projectExpression, collation.getDirection(), collation.nullDirection);
156-
}
129+
assert sortKey instanceof RexInputRef : "sort key should be always RexInputRef";
130+
131+
RexInputRef inputRef = (RexInputRef) sortKey;
132+
RexNode projectExpression = project.getProjects().get(inputRef.getIndex());
133+
134+
// If the project expression is a simple RexInputRef pointing to a scan field,
135+
// store the field name for stability
136+
if (projectExpression instanceof RexInputRef) {
137+
RexInputRef scanInputRef = (RexInputRef) projectExpression;
138+
int scanFieldIndex = scanInputRef.getIndex();
139+
// Get the field name from the scan's row type
140+
List<String> scanFieldNames = scan.getRowType().getFieldNames();
141+
// Create SortExpressionInfo with field name (stable reference)
142+
return new SortExpressionInfo(
143+
scanFieldNames.get(scanFieldIndex), collation.getDirection(), collation.nullDirection);
157144
}
158145

159-
// Fallback: store the sort key as-is
160-
return new SortExpressionInfo(sortKey, collation.getDirection(), collation.nullDirection);
146+
// For complex expressions, store the RexNode
147+
return new SortExpressionInfo(
148+
projectExpression, collation.getDirection(), collation.nullDirection);
161149
}
162150

163151
/**
@@ -172,7 +160,7 @@ private boolean canPushDownSortExpressionInfos(List<SortExpressionInfo> sortExpr
172160
RexNode expr = info.getExpression();
173161
if (expr == null && StringUtils.isEmpty(info.getFieldName())) {
174162
return false;
175-
} else if (expr == null && !StringUtils.isEmpty(info.getFieldName())) {
163+
} else if (info.isSimpleFieldReference()) {
176164
continue;
177165
}
178166
// Reject literals or constant expression - they don't provide meaningful sorting
@@ -194,11 +182,13 @@ private boolean canPushDownSortExpressionInfos(List<SortExpressionInfo> sortExpr
194182
*/
195183
private boolean isSupportedSortScriptType(SqlTypeName sqlTypeName) {
196184
switch (sqlTypeName) {
185+
case TINYINT:
186+
case SMALLINT:
197187
case INTEGER:
198188
case BIGINT:
199189
case FLOAT:
190+
case REAL:
200191
case DOUBLE:
201-
case DECIMAL:
202192
case VARCHAR:
203193
case CHAR:
204194
return true;

‎opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ public void pushDownSortExpression(
232232
Map<String, ExprType> fieldTypes,
233233
RelDataType rowType,
234234
RelOptCluster cluster) {
235-
if (!org.apache.commons.lang3.StringUtils.isEmpty(sortExprInfo.getFieldName())) {
235+
if (sortExprInfo.isSimpleFieldReference()) {
236236
sourceBuilder.sort(SortBuilders.fieldSort(sortExprInfo.getFieldName()));
237237
return;
238238
}
@@ -266,7 +266,7 @@ public void pushDownLimit(Integer limit, Integer offset) {
266266
int newStartFrom = startFrom + offset;
267267

268268
if (newStartFrom >= maxResultWindow) {
269-
throw new OpenSearchRequestBuilder.PushDownUnSupportedException(
269+
throw new PushDownUnSupportedException(
270270
String.format(
271271
"Requested offset %d should be less than the max result window %d",
272272
newStartFrom, maxResultWindow));
@@ -468,7 +468,7 @@ private BoolQueryBuilder query() {
468468
* @param relDataType the return type of the expression
469469
* @return the appropriate ScriptSortType
470470
*/
471-
private ScriptSortType getScriptSortType(org.apache.calcite.rel.type.RelDataType relDataType) {
471+
private ScriptSortType getScriptSortType(RelDataType relDataType) {
472472
switch (relDataType.getSqlTypeName()) {
473473
case TINYINT:
474474
case SMALLINT:
@@ -477,14 +477,13 @@ private ScriptSortType getScriptSortType(org.apache.calcite.rel.type.RelDataType
477477
case FLOAT:
478478
case REAL:
479479
case DOUBLE:
480-
case DECIMAL:
481480
return ScriptSortType.NUMBER;
482481
case CHAR:
483482
case VARCHAR:
484483
return ScriptSortType.STRING;
485484
default:
486-
// Default to STRING for unknown types
487-
return ScriptSortType.STRING;
485+
throw new PushDownUnSupportedException(
486+
"Unsupported type for sort expression pushdown: " + relDataType);
488487
}
489488
}
490489
}

0 commit comments

Comments
 (0)