Implement reverse performance optimization

ahkcs · ahkcs · commit f1ed013f879b · 2025-11-12T12:07:11.000-08:00
This commit optimizes the `reverse` command in the Calcite planner by intelligently reversing existing sort collations instead of always using the ROW_NUMBER() approach. Key changes: - Added PlanUtils.reverseCollation() method to flip sort directions and null directions - Updated CalciteRelNodeVisitor.visitReverse() to: - Check for existing sort collations - Reverse them if present (more efficient) - Fall back to ROW_NUMBER() when no sort exists - Added comprehensive integration test expected outputs for: - Single field reverse pushdown - Multiple field reverse pushdown - Reverse fallback cases - Double reverse no-op optimizations This optimization significantly improves performance when reversing already-sorted data by leveraging database-native sort reversal. Based on PR opensearch-project#4056 by @selsong Signed-off-by: Kai Huang <ahkcs@amazon.com>
diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java
@@ -45,6 +45,7 @@
 import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.ViewExpanders;
+import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.JoinRelType;
@@ -674,19 +675,30 @@ public RelNode visitHead(Head node, CalcitePlanContext context) {
   public RelNode visitReverse(
       org.opensearch.sql.ast.tree.Reverse node, CalcitePlanContext context) {
     visitChildren(node, context);
-    // Add ROW_NUMBER() column
-    RexNode rowNumber =
-        context
-            .relBuilder
-            .aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
-            .over()
-            .rowsTo(RexWindowBounds.CURRENT_ROW)
-            .as(REVERSE_ROW_NUM);
-    context.relBuilder.projectPlus(rowNumber);
-    // Sort by row number descending
-    context.relBuilder.sort(context.relBuilder.desc(context.relBuilder.field(REVERSE_ROW_NUM)));
-    // Remove row number column
-    context.relBuilder.projectExcept(context.relBuilder.field(REVERSE_ROW_NUM));
+
+    // Check if there's an existing sort to reverse
+    List<RelCollation> collations =
+        context.relBuilder.getCluster().getMetadataQuery().collations(context.relBuilder.peek());
+    RelCollation collation = collations != null && !collations.isEmpty() ? collations.get(0) : null;
+
+    if (collation != null && !collation.getFieldCollations().isEmpty()) {
+      // If there's an existing sort, reverse its direction
+      RelCollation reversedCollation = PlanUtils.reverseCollation(collation);
+      context.relBuilder.sort(reversedCollation);
+    } else {
+      // Fallback: use ROW_NUMBER approach when no existing sort
+      RexNode rowNumber =
+          context
+              .relBuilder
+              .aggregateCall(SqlStdOperatorTable.ROW_NUMBER)
+              .over()
+              .rowsTo(RexWindowBounds.CURRENT_ROW)
+              .as(REVERSE_ROW_NUM);
+      context.relBuilder.projectPlus(rowNumber);
+      context.relBuilder.sort(context.relBuilder.desc(context.relBuilder.field(REVERSE_ROW_NUM)));
+      context.relBuilder.projectExcept(context.relBuilder.field(REVERSE_ROW_NUM));
+    }
+
     return context.relBuilder.peek();
   }
 
diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java b/core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java
@@ -22,6 +22,9 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelHomogeneousShuttle;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelShuttle;
@@ -568,6 +571,37 @@ public Void visitCorrelVariable(RexCorrelVariable correlVar) {
     }
   }
 
+  /**
+   * Reverses the direction of a RelCollation.
+   *
+   * @param original The original collation to reverse
+   * @return A new RelCollation with reversed directions
+   */
+  public static RelCollation reverseCollation(RelCollation original) {
+    if (original == null || original.getFieldCollations().isEmpty()) {
+      return original;
+    }
+
+    List<RelFieldCollation> reversedFields = new ArrayList<>();
+    for (RelFieldCollation field : original.getFieldCollations()) {
+      RelFieldCollation.Direction reversedDirection = field.direction.reverse();
+
+      // Handle null direction properly - reverse it as well
+      RelFieldCollation.NullDirection reversedNullDirection =
+          field.nullDirection == RelFieldCollation.NullDirection.FIRST
+              ? RelFieldCollation.NullDirection.LAST
+              : field.nullDirection == RelFieldCollation.NullDirection.LAST
+                  ? RelFieldCollation.NullDirection.FIRST
+                  : field.nullDirection;
+
+      RelFieldCollation reversedField =
+          new RelFieldCollation(field.getFieldIndex(), reversedDirection, reversedNullDirection);
+      reversedFields.add(reversedField);
+    }
+
+    return RelCollations.of(reversedFields);
+  }
+
   /** Adds a rel node to the top of the stack while preserving the field names and aliases. */
   static void replaceTop(RelBuilder relBuilder, RelNode relNode) {
     try {
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_double_reverse_no_op.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_double_reverse_no_op.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "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",
+    "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"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_double_reverse_sort_no_op.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_double_reverse_sort_no_op.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalSystemLimit(sort0=[$8], sort1=[$1], dir0=[DESC-nulls-last], dir1=[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], sort1=[$1], dir0=[DESC-nulls-last], dir1=[ASC-nulls-first])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "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\" : \"desc\",\n    \"missing\" : \"_last\"\n  }\n}, {\n  \"firstname.keyword\" : {\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\":\"desc\",\"missing\":\"_last\"}},{\"firstname.keyword\":{\"order\":\"asc\",\"missing\":\"_first\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_fallback.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_fallback.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalSystemLimit(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=[$17], dir0=[DESC], fetch=[5])\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], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], __reverse_row_num__=[ROW_NUMBER() OVER ()])\n        CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableCalc(expr#0..11=[{inputs}], proj#0..10=[{exprs}])\n    EnumerableLimit(fetch=[5])\n      EnumerableSort(sort0=[$11], dir0=[DESC])\n        EnumerableWindow(window#0=[window(rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])\n          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_multiple.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_multiple.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "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",
+    "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"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_single.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_reverse_pushdown_single.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "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",
+    "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"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_fallback.json b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_fallback.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "logical": "LogicalSystemLimit(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=[$17], dir0=[DESC], fetch=[5])\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], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], __reverse_row_num__=[ROW_NUMBER() OVER ()])\n        CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableCalc(expr#0..17=[{inputs}], proj#0..10=[{exprs}])\n    EnumerableLimit(fetch=[5])\n      EnumerableSort(sort0=[$17], dir0=[DESC])\n        EnumerableWindow(window#0=[window(rows between UNBOUNDED PRECEDING and CURRENT ROW aggs [ROW_NUMBER()])])\n          CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_multiple.json b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_multiple.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "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",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableSort(sort0=[$8], sort1=[$1], dir0=[ASC-nulls-first], dir1=[DESC-nulls-last])\n    EnumerableCalc(expr#0..16=[{inputs}], proj#0..10=[{exprs}])\n      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n"
+  }
+}
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_single.json b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_single.json
@@ -0,0 +1,6 @@
+{
+  "calcite": {
+    "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",
+    "physical": "EnumerableLimit(fetch=[10000])\n  EnumerableSort(sort0=[$8], dir0=[ASC-nulls-first])\n    EnumerableCalc(expr#0..16=[{inputs}], proj#0..10=[{exprs}])\n      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n"
+  }
+}
