Implement reverse performance optimization

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

@@ -38,6 +38,7 @@
 import java.util.stream.Stream;
 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;
@@ -567,19 +568,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();
   }
 

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

@@ -8,6 +8,7 @@
 import com.google.common.collect.ImmutableList;
 import java.util.List;
 import org.apache.calcite.plan.RelOptRule;
+import org.opensearch.sql.calcite.rule.SortReverseOptimizationRule;
 
 public class OpenSearchRules {
   private static final PPLAggregateConvertRule AGGREGATE_CONVERT_RULE =
@@ -16,6 +17,11 @@ public class OpenSearchRules {
   public static final List<RelOptRule> OPEN_SEARCH_OPT_RULES =
       ImmutableList.of(AGGREGATE_CONVERT_RULE);
 
+  public static final List<RelOptRule> OPTIMIZATION_RULES = ImmutableList.of();
+
+  public static final List<RelOptRule> OPEN_SEARCH_POST_AGG_RULES =
+      ImmutableList.of(SortReverseOptimizationRule.INSTANCE);
+
   // prevent instantiation
   private OpenSearchRules() {}
 }

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

@@ -33,6 +33,16 @@ public void register(RelOptPlanner planner) {
       planner.addRule(rule);
     }
 
+    // Register optimization rules
+    for (RelOptRule rule : OpenSearchRules.OPTIMIZATION_RULES) {
+      planner.addRule(rule);
+    }
+
+    // Register post-aggregation rules (run after aggregate optimizations)
+    for (RelOptRule rule : OpenSearchRules.OPEN_SEARCH_POST_AGG_RULES) {
+      planner.addRule(rule);
+    }
+
     // remove this rule otherwise opensearch can't correctly interpret approx_count_distinct()
     // it is converted to cardinality aggregation in OpenSearch
     planner.removeRule(CoreRules.AGGREGATE_EXPAND_DISTINCT_AGGREGATES);

core/src/main/java/org/opensearch/sql/calcite/rule/SortReverseOptimizationRule.java

@@ -0,0 +1,117 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.rule;
+
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
+
+/** Combines sort then reverse into 1 sort. */
+public class SortReverseOptimizationRule extends RelOptRule {
+  private static final Logger LOG = LogManager.getLogger(SortReverseOptimizationRule.class);
+
+  public static final SortReverseOptimizationRule INSTANCE = new SortReverseOptimizationRule();
+
+  private SortReverseOptimizationRule() {
+    super(
+        operand(LogicalSort.class, operand(LogicalSort.class, any())),
+        "SortReverseOptimizationRule");
+  }
+
+  @Override
+  public boolean matches(RelOptRuleCall call) {
+    LogicalSort outerSort = call.rel(0);
+    LogicalSort innerSort = call.rel(1);
+
+    LOG.debug("SortReverseOptimizationRule.matches() called");
+    LOG.debug("Outer sort: {}", outerSort);
+    LOG.debug("Inner sort: {}", innerSort);
+    LOG.debug("Inner sort input: {}", innerSort.getInput());
+
+    // Don't optimize if outer sort is a LogicalSystemLimit
+    if (call.rel(0) instanceof LogicalSystemLimit) {
+      LOG.debug("Skipping: outer sort is LogicalSystemLimit");
+      return false;
+    }
+
+    // Don't optimize if inner sort has a fetch limit (head/limit before sort)
+    // This preserves limit-then-sort semantics
+    if (innerSort.fetch != null) {
+      LOG.debug("Skipping: inner sort has fetch limit: {}", innerSort.fetch);
+      return false;
+    }
+
+    // Must be same field with opposite directions (sort | reverse pattern)
+    boolean matches = hasSameFieldWithOppositeDirection(outerSort, innerSort);
+    LOG.debug("Same field with opposite direction: {}", matches);
+    return matches;
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    LogicalSort outerSort = call.rel(0);
+    LogicalSort innerSort = call.rel(1);
+
+    LOG.debug("SortReverseOptimizationRule.onMatch() applying transformation");
+    LOG.debug("Transforming from: {} -> {}", outerSort, innerSort);
+
+    LogicalSort optimizedSort =
+        LogicalSort.create(
+            innerSort.getInput(), outerSort.getCollation(), outerSort.offset, outerSort.fetch);
+
+    LOG.debug("Transformed to: {}", optimizedSort);
+    call.transformTo(optimizedSort);
+  }
+
+  private boolean hasSameFieldWithOppositeDirection(LogicalSort outerSort, LogicalSort innerSort) {
+    var outerFields = outerSort.getCollation().getFieldCollations();
+    var innerFields = innerSort.getCollation().getFieldCollations();
+
+    if (outerFields.isEmpty() || innerFields.isEmpty()) {
+      LOG.debug("No field collations found");
+      return false;
+    }
+
+    // Must have same number of fields
+    if (outerFields.size() != innerFields.size()) {
+      LOG.debug(
+          "Different number of sort fields: outer={}, inner={}",
+          outerFields.size(),
+          innerFields.size());
+      return false;
+    }
+
+    // Check all fields have same index but opposite directions
+    for (int i = 0; i < outerFields.size(); i++) {
+      var outerField = outerFields.get(i);
+      var innerField = innerFields.get(i);
+
+      LOG.debug(
+          "Field {}: outer(index={}, direction={}), inner(index={}, direction={})",
+          i,
+          outerField.getFieldIndex(),
+          outerField.getDirection(),
+          innerField.getFieldIndex(),
+          innerField.getDirection());
+
+      if (outerField.getFieldIndex() != innerField.getFieldIndex()
+          || outerField.getDirection() == innerField.getDirection()) {
+        LOG.debug(
+            "Field {} mismatch: same index={}, opposite direction={}",
+            i,
+            outerField.getFieldIndex() == innerField.getFieldIndex(),
+            outerField.getDirection() != innerField.getDirection());
+        return false;
+      }
+    }
+
+    LOG.debug("All fields match with opposite directions");
+    return true;
+  }
+}

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

@@ -16,6 +16,9 @@
 import java.util.List;
 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;
@@ -391,4 +394,35 @@ public Void visitInputRef(RexInputRef inputRef) {
     visitor.visitEach(rexNodes);
     return selectedColumns;
   }
+
+  /**
+   * 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);
+  }
 }

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

@@ -206,8 +206,7 @@ public void testFilterFunctionScriptPushDownExplain() throws Exception {
   @Test
   public void testExplainWithReverse() throws IOException {
     String result =
-        executeWithReplace(
-            "explain source=opensearch-sql_test_index_account | sort age | reverse | head 5");
+        executeWithReplace("explain source=opensearch-sql_test_index_account | reverse | head 5");
 
     // Verify that the plan contains a LogicalSort with fetch (from head 5)
     assertTrue(result.contains("LogicalSort") && result.contains("fetch=[5]"));
@@ -217,6 +216,30 @@ public void testExplainWithReverse() throws IOException {
     assertTrue(result.contains("dir0=[DESC]"));
   }
 
+  @Test
+  public void testExplainWithReversePushdown() throws IOException {
+    String query = "source=opensearch-sql_test_index_account | sort - age | reverse";
+    var result = explainQueryToString(query);
+    String expected =
+        isPushdownEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_reverse_pushdown_single.json")
+            : loadFromFile(
+                "expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_single.json");
+    assertJsonEqualsIgnoreId(expected, result);
+  }
+
+  @Test
+  public void testExplainWithReversePushdownMultipleFields() throws IOException {
+    String query = "source=opensearch-sql_test_index_account | sort - age, + firstname | reverse";
+    var result = explainQueryToString(query);
+    String expected =
+        isPushdownEnabled()
+            ? loadFromFile("expectedOutput/calcite/explain_reverse_pushdown_multiple.json")
+            : loadFromFile(
+                "expectedOutput/calcite_no_pushdown/explain_reverse_pushdown_multiple.json");
+    assertJsonEqualsIgnoreId(expected, result);
+  }
+
   @Test
   public void testExplainWithTimechartAvg() throws IOException {
     var result = explainQueryToString("source=events | timechart span=1m avg(cpu_usage) by host");

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

@@ -97,14 +97,70 @@ public void testReverseWithComplexPipeline() throws IOException {
   }
 
   @Test
-  public void testReverseWithMultipleSorts() throws IOException {
-    // Use the existing BANK data but with a simpler, more predictable query
+  public void testReverseWithDescendingSort() throws IOException {
+    // Test reverse with descending sort (- age)
     JSONObject result =
         executeQuery(
             String.format(
-                "source=%s | sort account_number | fields account_number | reverse | head 3",
+                "source=%s | sort - account_number | fields account_number | reverse",
                 TEST_INDEX_BANK));
     verifySchema(result, schema("account_number", "bigint"));
-    verifyDataRowsInOrder(result, rows(32), rows(25), rows(20));
+    verifyDataRowsInOrder(
+        result, rows(1), rows(6), rows(13), rows(18), rows(20), rows(25), rows(32));
+  }
+
+  @Test
+  public void testReverseWithMixedSortDirections() throws IOException {
+    // Test reverse with mixed sort directions (- age, + firstname)
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort - account_number, + firstname | fields account_number, firstname"
+                    + " | reverse",
+                TEST_INDEX_BANK));
+    verifySchema(result, schema("account_number", "bigint"), schema("firstname", "string"));
+    verifyDataRowsInOrder(
+        result,
+        rows(1, "Amber JOHnny"),
+        rows(6, "Hattie"),
+        rows(13, "Nanette"),
+        rows(18, "Dale"),
+        rows(20, "Elinor"),
+        rows(25, "Virginia"),
+        rows(32, "Dillard"));
+  }
+
+  @Test
+  public void testDoubleReverseWithDescendingSort() throws IOException {
+    // Test double reverse with descending sort (- age)
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort - account_number | fields account_number | reverse | reverse",
+                TEST_INDEX_BANK));
+    verifySchema(result, schema("account_number", "bigint"));
+    verifyDataRowsInOrder(
+        result, rows(32), rows(25), rows(20), rows(18), rows(13), rows(6), rows(1));
+  }
+
+  @Test
+  public void testDoubleReverseWithMixedSortDirections() throws IOException {
+    // Test double reverse with mixed sort directions (- age, + firstname)
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | sort - account_number, + firstname | fields account_number, firstname"
+                    + " | reverse | reverse",
+                TEST_INDEX_BANK));
+    verifySchema(result, schema("account_number", "bigint"), schema("firstname", "string"));
+    verifyDataRowsInOrder(
+        result,
+        rows(32, "Dillard"),
+        rows(25, "Virginia"),
+        rows(20, "Elinor"),
+        rows(18, "Dale"),
+        rows(13, "Nanette"),
+        rows(6, "Hattie"),
+        rows(1, "Amber JOHnny"));
   }
 }