WIP: sort flip collation

Signed-off-by: Selina Song <[email protected]>

Author: Selina Song

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

@@ -38,6 +38,8 @@
 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.RelCollations;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.JoinRelType;
@@ -567,19 +569,29 @@ 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));
+    
+    RelCollation collation = context.relBuilder.peek().getTraitSet().getCollation();
+    if (collation == null || collation == RelCollations.EMPTY) {
+      // If no collation exists, use the traditional row_number approach
+      // 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));
+    } else {
+      // If collation exists, reverse the sort direction
+      RelCollation reversedCollation = PlanUtils.reverseCollation(collation);
+      context.relBuilder.sort(reversedCollation);
+    }
+    
     return context.relBuilder.peek();
   }
 

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

@@ -0,0 +1,65 @@
+/*
+ * 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.opensearch.sql.calcite.plan.LogicalSystemLimit;
+
+/**
+ * Optimization rule that eliminates redundant consecutive sorts on the same field.
+ * Detects: LogicalSort(field, direction1) -> LogicalSort(field, direction2)
+ * Converts to: LogicalSort(field, direction1) (keeps outer sort)
+ */
+public class SortReverseOptimizationRule extends RelOptRule {
+
+  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);
+    
+    // Don't optimize if outer sort is a LogicalSystemLimit - we want to preserve system limits
+    if (call.rel(0) instanceof LogicalSystemLimit) {
+      return false;
+    }
+    
+    return hasSameField(outerSort, innerSort);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    LogicalSort outerSort = call.rel(0);
+    LogicalSort innerSort = call.rel(1);
+    
+    LogicalSort optimizedSort = LogicalSort.create(
+        innerSort.getInput(),
+        outerSort.getCollation(),
+        outerSort.offset,
+        outerSort.fetch);
+    
+    call.transformTo(optimizedSort);
+  }
+  
+  private boolean hasSameField(LogicalSort outerSort, LogicalSort innerSort) {
+    if (outerSort.getCollation().getFieldCollations().isEmpty() 
+        || innerSort.getCollation().getFieldCollations().isEmpty()) {
+      return false;
+    }
+    
+    int outerField = outerSort.getCollation().getFieldCollations().get(0).getFieldIndex();
+    int innerField = innerSort.getCollation().getFieldCollations().get(0).getFieldIndex();
+    return outerField == innerField;
+  }
+}

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;
@@ -355,40 +358,29 @@ static RexNode derefMapCall(RexNode rexNode) {
     return rexNode;
   }
 
-  /** Check if contains RexOver */
-  static boolean containsRowNumberDedup(LogicalProject project) {
-    return project.getProjects().stream()
-        .anyMatch(p -> p instanceof RexOver && p.getKind() == SqlKind.ROW_NUMBER);
-  }
+  /**
+   * 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;
+    }
 
-  /** Get all RexWindow list from LogicalProject */
-  static List<RexWindow> getRexWindowFromProject(LogicalProject project) {
-    final List<RexWindow> res = new ArrayList<>();
-    final RexVisitorImpl<Void> visitor =
-        new RexVisitorImpl<>(true) {
-          @Override
-          public Void visitOver(RexOver over) {
-            res.add(over.getWindow());
-            return null;
-          }
-        };
-    visitor.visitEach(project.getProjects());
-    return res;
-  }
+    List<RelFieldCollation> reversedFields = new ArrayList<>();
+    for (RelFieldCollation field : original.getFieldCollations()) {
+      RelFieldCollation.Direction reversedDirection = field.direction.reverse();
+      
+      RelFieldCollation reversedField = new RelFieldCollation(
+          field.getFieldIndex(),
+          reversedDirection,
+          field.nullDirection
+      );
+      reversedFields.add(reversedField);
+    }
 
-  static List<Integer> getSelectColumns(List<RexNode> rexNodes) {
-    final List<Integer> selectedColumns = new ArrayList<>();
-    final RexVisitorImpl<Void> visitor =
-        new RexVisitorImpl<Void>(true) {
-          @Override
-          public Void visitInputRef(RexInputRef inputRef) {
-            if (!selectedColumns.contains(inputRef.getIndex())) {
-              selectedColumns.add(inputRef.getIndex());
-            }
-            return null;
-          }
-        };
-    visitor.visitEach(rexNodes);
-    return selectedColumns;
+    return RelCollations.of(reversedFields);
   }
 }

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

@@ -207,7 +207,7 @@ public void testFilterFunctionScriptPushDownExplain() throws Exception {
   public void testExplainWithReverse() throws IOException {
     String result =
         executeWithReplace(
-            "explain source=opensearch-sql_test_index_account | sort age | reverse | head 5");
+            "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]"));
@@ -216,6 +216,37 @@ public void testExplainWithReverse() throws IOException {
     assertTrue(result.contains("ROW_NUMBER()"));
     assertTrue(result.contains("dir0=[DESC]"));
   }
+  
+  @Test
+  public void testExplainWithReversePushdown() throws IOException {
+    // Test with a sort operation that should use the reverse pushdown optimization
+    String result =
+        executeWithReplace(
+            "explain source=opensearch-sql_test_index_account | sort - age | reverse");
+
+    // Verify that the plan contains a LogicalSort with ascending direction (reversed from DESC)
+    assertTrue(result.contains("LogicalSort"));
+    assertTrue(result.contains("dir0=[ASC]"));
+    
+    // Verify that ROW_NUMBER is NOT used (since we're using the collation-based optimization)
+    assertFalse(result.contains("ROW_NUMBER()"));
+  }
+  
+  @Test
+  public void testExplainWithReversePushdownMultipleFields() throws IOException {
+    // Test with multiple sort fields that should use the reverse pushdown optimization
+    String result =
+        executeWithReplace(
+            "explain source=opensearch-sql_test_index_account | sort - age, + firstname | reverse");
+
+    // Verify that the plan contains a LogicalSort with reversed directions
+    assertTrue(result.contains("LogicalSort"));
+    assertTrue(result.contains("dir0=[ASC]")); // age was DESC, now ASC
+    assertTrue(result.contains("dir1=[DESC]")); // firstname was ASC, now DESC
+    
+    // Verify that ROW_NUMBER is NOT used (since we're using the collation-based optimization)
+    assertFalse(result.contains("ROW_NUMBER()"));
+  }
 
   @Test
   public void testExplainWithTimechartAvg() throws IOException {

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

@@ -97,14 +97,68 @@ 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"));
   }
 }