Skip to content

Commit 4343b47

Browse files
author
Selina Song
committed
add opt rule for sort flip
Signed-off-by: Selina Song <[email protected]>
1 parent 3a44688 commit 4343b47

File tree

11 files changed

+102
-27
lines changed

11 files changed

+102
-27
lines changed

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -569,11 +569,18 @@ public RelNode visitHead(Head node, CalcitePlanContext context) {
569569
public RelNode visitReverse(
570570
org.opensearch.sql.ast.tree.Reverse node, CalcitePlanContext context) {
571571
visitChildren(node, context);
572-
573-
RelCollation collation = context.relBuilder.peek().getTraitSet().getCollation();
574-
if (collation == null || collation == RelCollations.EMPTY) {
575-
// If no collation exists, use the traditional row_number approach
576-
// Add ROW_NUMBER() column
572+
573+
// Check if there's an existing sort to reverse
574+
List<RelCollation> collations =
575+
context.relBuilder.getCluster().getMetadataQuery().collations(context.relBuilder.peek());
576+
RelCollation collation = collations != null && !collations.isEmpty() ? collations.get(0) : null;
577+
578+
if (collation != null && !collation.getFieldCollations().isEmpty()) {
579+
// If there's an existing sort, reverse its direction
580+
RelCollation reversedCollation = PlanUtils.reverseCollation(collation);
581+
context.relBuilder.sort(reversedCollation);
582+
} else {
583+
// Fallback: use ROW_NUMBER approach when no existing sort
577584
RexNode rowNumber =
578585
context
579586
.relBuilder
@@ -582,16 +589,10 @@ public RelNode visitReverse(
582589
.rowsTo(RexWindowBounds.CURRENT_ROW)
583590
.as(REVERSE_ROW_NUM);
584591
context.relBuilder.projectPlus(rowNumber);
585-
// Sort by row number descending
586592
context.relBuilder.sort(context.relBuilder.desc(context.relBuilder.field(REVERSE_ROW_NUM)));
587-
// Remove row number column
588593
context.relBuilder.projectExcept(context.relBuilder.field(REVERSE_ROW_NUM));
589-
} else {
590-
// If collation exists, reverse the sort direction
591-
RelCollation reversedCollation = PlanUtils.reverseCollation(collation);
592-
context.relBuilder.sort(reversedCollation);
593594
}
594-
595+
595596
return context.relBuilder.peek();
596597
}
597598

‎core/src/main/java/org/opensearch/sql/calcite/plan/OpenSearchRules.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,19 @@
88
import com.google.common.collect.ImmutableList;
99
import java.util.List;
1010
import org.apache.calcite.plan.RelOptRule;
11+
import org.apache.calcite.rel.convert.ConverterRule;
12+
import org.opensearch.sql.calcite.rule.SortReverseOptimizationRule;
1113

1214
public class OpenSearchRules {
1315
private static final PPLAggregateConvertRule AGGREGATE_CONVERT_RULE =
1416
PPLAggregateConvertRule.Config.SUM_CONVERTER.toRule();
1517

16-
public static final List<RelOptRule> OPEN_SEARCH_OPT_RULES =
18+
public static final List<ConverterRule> OPEN_SEARCH_OPT_RULES = ImmutableList.of();
19+
20+
public static final List<RelOptRule> OPTIMIZATION_RULES =
21+
ImmutableList.of(SortReverseOptimizationRule.INSTANCE);
22+
23+
public static final List<RelOptRule> OPEN_SEARCH_OPT_RULES_OLD =
1724
ImmutableList.of(AGGREGATE_CONVERT_RULE);
1825

1926
// prevent instantiation

‎core/src/main/java/org/opensearch/sql/calcite/plan/OpenSearchTableScan.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ public void register(RelOptPlanner planner) {
3232
for (RelOptRule rule : OpenSearchRules.OPEN_SEARCH_OPT_RULES) {
3333
planner.addRule(rule);
3434
}
35+
36+
// Register optimization rules
37+
for (RelOptRule rule : OpenSearchRules.OPTIMIZATION_RULES) {
38+
planner.addRule(rule);
39+
}
3540

3641
// remove this rule otherwise opensearch can't correctly interpret approx_count_distinct()
3742
// it is converted to cardinality aggregation in OpenSearch

‎core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -372,9 +372,18 @@ public static RelCollation reverseCollation(RelCollation original) {
372372
List<RelFieldCollation> reversedFields = new ArrayList<>();
373373
for (RelFieldCollation field : original.getFieldCollations()) {
374374
RelFieldCollation.Direction reversedDirection = field.direction.reverse();
375+
376+
// Handle null direction properly - reverse it as well
377+
RelFieldCollation.NullDirection reversedNullDirection = field.nullDirection;
378+
if (reversedNullDirection == RelFieldCollation.NullDirection.FIRST) {
379+
reversedNullDirection = RelFieldCollation.NullDirection.LAST;
380+
} else if (reversedNullDirection == RelFieldCollation.NullDirection.LAST) {
381+
reversedNullDirection = RelFieldCollation.NullDirection.FIRST;
382+
}
383+
// UNSPECIFIED remains UNSPECIFIED
375384

376385
RelFieldCollation reversedField =
377-
new RelFieldCollation(field.getFieldIndex(), reversedDirection, field.nullDirection);
386+
new RelFieldCollation(field.getFieldIndex(), reversedDirection, reversedNullDirection);
378387
reversedFields.add(reversedField);
379388
}
380389

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,7 @@ public void testExplainWithReversePushdownMultipleFields() throws IOException {
232232
assertJsonEqualsIgnoreId(expected, result);
233233
}
234234

235+
<<<<<<< HEAD
235236
@Test
236237
public void testExplainWithTimechartAvg() throws IOException {
237238
var result = explainQueryToString("source=events | timechart span=1m avg(cpu_usage) by host");
@@ -260,6 +261,35 @@ public void testDoubleReverseWithSortNoOp() throws IOException {
260261
String expected = loadFromFile("expectedOutput/calcite/explain_double_reverse_sort_no_op.json");
261262
assertJsonEqualsIgnoreId(expected, result);
262263
}
264+
=======
265+
// @Test
266+
// public void testDoubleReverseNoOp() throws IOException {
267+
// String query =
268+
// "source=opensearch-sql_test_index_account | fields account_number | reverse | reverse";
269+
// var result = explainQueryToString(query);
270+
// String expected = loadFromFile("expectedOutput/calcite/explain_double_reverse_no_op.json");
271+
// assertJsonEqualsIgnoreId(expected, result);
272+
// }
273+
274+
// @Test
275+
// public void testTripleReverseOneOp() throws IOException {
276+
// String query =
277+
// "source=opensearch-sql_test_index_account | fields account_number | reverse | reverse |"
278+
// + " reverse";
279+
// var result = explainQueryToString(query);
280+
// String expected = loadFromFile("expectedOutput/calcite/explain_triple_reverse_one_op.json");
281+
// assertJsonEqualsIgnoreId(expected, result);
282+
// }
283+
//
284+
// @Test
285+
// public void testDoubleReverseWithSortNoOp() throws IOException {
286+
// String query =
287+
// "source=opensearch-sql_test_index_account | sort - age, + firstname | reverse | reverse";
288+
// var result = explainQueryToString(query);
289+
// String expected = loadFromFile("expectedOutput/calcite/explain_double_reverse_sort_no_op.json");
290+
// assertJsonEqualsIgnoreId(expected, result);
291+
// }
292+
>>>>>>> f6a6803f3 (add opt rule for sort flip)
263293

264294
@Test
265295
public void noPushDownForAggOnWindow() throws IOException {

‎integ-test/src/test/resources/expectedOutput/calcite/explain_double_reverse_no_op.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"calcite": {
3+
"logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(cpu_usage=[$0], @timestamp=[$1])\n LogicalSort(sort0=[$8], dir0=[DESC])\n LogicalProject(cpu_usage=[$0], @timestamp=[$1], _id=[$2], _index=[$3], _score=[$4], _maxscore=[$5], _sort=[$6], _routing=[$7], __reverse_row_num__=[ROW_NUMBER() OVER ()])\n LogicalSort(sort0=[$8], dir0=[DESC])\n LogicalProject(cpu_usage=[$0], @timestamp=[$1], _id=[$2], _index=[$3], _score=[$4], _maxscore=[$5], _sort=[$6], _routing=[$7], __reverse_row_num__=[ROW_NUMBER() OVER ()])\n CalciteLogicalIndexScan(table=[[OpenSearch, events]])\n",
4+
"physical": "EnumerableCalc(expr#0..3=[{inputs}], proj#0..1=[{exprs}])\n EnumerableLimit(fetch=[10000])\n EnumerableSort(sort0=[$3], dir0=[DESC])\n EnumerableWindow(window#0=[window(rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])\n EnumerableSort(sort0=[$2], dir0=[DESC])\n EnumerableWindow(window#0=[window(rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])\n CalciteEnumerableIndexScan(table=[[OpenSearch, events]], PushDownContext=[[PROJECT->[cpu_usage, @timestamp]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"cpu_usage\",\"@timestamp\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
5+
}
6+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"calcite": {
33
"logical": "LogicalSystemLimit(sort0=[$0], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(age=[$8])\n LogicalSort(sort0=[$8], dir0=[ASC-nulls-first])\n LogicalSort(fetch=[5])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4-
"physical": "EnumerableLimit(fetch=[10000])\n EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[LIMIT->5, PROJECT->[age]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
4+
"physical": "EnumerableSort(sort0=[$0], dir0=[ASC-nulls-first])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[LIMIT->5, PROJECT->[age], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"age\"],\"excludes\":[]}}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
55
}
66
}

‎integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_multiple.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"calcite": {
3-
"logical": "LogicalSystemLimit(sort0=[$8], sort1=[$1], dir0=[ASC], dir1=[DESC], fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10])\n LogicalSort(sort0=[$8], sort1=[$1], dir0=[ASC], dir1=[DESC])\n LogicalSort(sort0=[$8], sort1=[$1], dir0=[DESC-nulls-last], dir1=[ASC-nulls-first])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4-
"physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], SORT->[{\n \"age\" : {\n \"order\" : \"asc\",\n \"missing\" : \"_last\"\n }\n}, {\n \"firstname.keyword\" : {\n \"order\" : \"desc\",\n \"missing\" : \"_first\"\n }\n}], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_last\"}},{\"firstname.keyword\":{\"order\":\"desc\",\"missing\":\"_first\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
3+
"logical": "LogicalSystemLimit(sort0=[$8], sort1=[$1], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last], fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10])\n LogicalSort(sort0=[$8], sort1=[$1], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last])\n LogicalSort(sort0=[$8], sort1=[$1], dir0=[DESC-nulls-last], dir1=[ASC-nulls-first])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4+
"physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], SORT->[{\n \"age\" : {\n \"order\" : \"asc\",\n \"missing\" : \"_first\"\n }\n}, {\n \"firstname.keyword\" : {\n \"order\" : \"desc\",\n \"missing\" : \"_last\"\n }\n}], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}},{\"firstname.keyword\":{\"order\":\"desc\",\"missing\":\"_last\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
55
}
66
}

‎integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_single.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"calcite": {
3-
"logical": "LogicalSystemLimit(sort0=[$8], dir0=[ASC], fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10])\n LogicalSort(sort0=[$8], dir0=[ASC])\n LogicalSort(sort0=[$8], dir0=[DESC-nulls-last])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4-
"physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], SORT->[{\n \"age\" : {\n \"order\" : \"asc\",\n \"missing\" : \"_last\"\n }\n}], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_last\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
3+
"logical": "LogicalSystemLimit(sort0=[$8], dir0=[ASC-nulls-first], fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10])\n LogicalSort(sort0=[$8], dir0=[ASC-nulls-first])\n LogicalSort(sort0=[$8], dir0=[DESC-nulls-last])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
4+
"physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], SORT->[{\n \"age\" : {\n \"order\" : \"asc\",\n \"missing\" : \"_first\"\n }\n}], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]},\"sort\":[{\"age\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
55
}
66
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,10 @@ public void register(RelOptPlanner planner) {
7373
for (RelOptRule rule : OpenSearchRules.OPEN_SEARCH_OPT_RULES) {
7474
planner.addRule(rule);
7575
}
76+
77+
for (RelOptRule rule : OpenSearchRules.OPTIMIZATION_RULES) {
78+
planner.addRule(rule);
79+
}
7680

7781
// remove this rule otherwise opensearch can't correctly interpret approx_count_distinct()
7882
// it is converted to cardinality aggregation in OpenSearch

0 commit comments

Comments
 (0)