Merge multiple sort when occurs

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchSortIndexScanRule.java

@@ -43,7 +43,7 @@ public interface Config extends RelRule.Config {
                         .oneInput(
                             b1 ->
                                 b1.operand(CalciteLogicalIndexScan.class)
-                                    .predicate(OpenSearchIndexScanRule::test)
+                                    .predicate(OpenSearchIndexScanRule::noAggregatePushed)
                                     .noInputs()));
 
     @Override

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java

@@ -85,7 +85,7 @@ public double estimateRowCount(RelMetadataQuery mq) {
             (rowCount, action) ->
                 switch (action.type) {
                       case AGGREGATION -> mq.getRowCount((RelNode) action.digest);
-                      case PROJECT -> rowCount;
+                      case PROJECT, SORT -> rowCount;
                       case FILTER -> NumberUtil.multiply(
                           rowCount, RelMdUtil.guessSelectivity((RexNode) action.digest));
                       case LIMIT -> (Integer) action.digest;

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java

@@ -7,16 +7,20 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import lombok.Getter;
 import org.apache.calcite.plan.Convention;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.core.Aggregate;
@@ -199,8 +203,15 @@ public CalciteLogicalIndexScan pushDownLimit(Integer limit, Integer offset) {
 
   public CalciteLogicalIndexScan pushDownSort(List<RelFieldCollation> collations) {
     try {
+      // Merge with existing sort if any
+      RelCollation existingCollation = getTraitSet().getCollation();
+      List<RelFieldCollation> existingFieldCollations =
+          existingCollation == null ? List.of() : existingCollation.getFieldCollations();
+      List<RelFieldCollation> mergedCollations =
+          mergeCollations(existingFieldCollations, collations);
+
       // Propagate the sort to the new scan
-      RelTraitSet newTraitSet = getTraitSet().plus(RelCollations.of(collations));
+      RelTraitSet newTraitSet = getTraitSet().plus(RelCollations.of(mergedCollations));
       CalciteLogicalIndexScan newScan =
           new CalciteLogicalIndexScan(
               getCluster(),
@@ -212,7 +223,7 @@ public CalciteLogicalIndexScan pushDownSort(List<RelFieldCollation> collations)
               pushDownContext.clone());
 
       List<SortBuilder<?>> builders = new ArrayList<>();
-      for (RelFieldCollation collation : collations) {
+      for (RelFieldCollation collation : mergedCollations) {
         int index = collation.getFieldIndex();
         String fieldName = this.getRowType().getFieldNames().get(index);
         RelFieldCollation.Direction direction = collation.getDirection();
@@ -245,4 +256,32 @@ public CalciteLogicalIndexScan pushDownSort(List<RelFieldCollation> collations)
     }
     return null;
   }
+
+  /**
+   * Merges existing and new collations, ensuring that the last occurrence of each field index takes
+   * precedence.
+   *
+   * @param existingCollations Existing collation list.
+   * @param newCollations New collation list to be merged.
+   * @return Merged list of collations.
+   */
+  private static List<RelFieldCollation> mergeCollations(
+      List<RelFieldCollation> existingCollations, List<RelFieldCollation> newCollations) {
+    List<RelFieldCollation> concatenatedCollations = new ArrayList<>(existingCollations);
+    concatenatedCollations.addAll(newCollations);
+    LinkedList<RelFieldCollation> mergedCollations = new LinkedList<>();
+    for (RelFieldCollation collation : concatenatedCollations) {
+      // If the collation is already in the merged list, remove it from the list before adding
+      // This is because the sort that comes later in the list should take precedence
+      OptionalInt index =
+          IntStream.range(0, mergedCollations.size())
+              .filter(i -> mergedCollations.get(i).getFieldIndex() == collation.getFieldIndex())
+              .findFirst();
+      if (index.isPresent()) {
+        mergedCollations.remove(index.getAsInt());
+      }
+      mergedCollations.add(collation);
+    }
+    return mergedCollations;
+  }
 }