Support Sort pushdown

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

Author: yuancu

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

@@ -16,9 +16,12 @@ public class OpenSearchIndexRules {
       OpenSearchFilterIndexScanRule.Config.DEFAULT.toRule();
   private static final OpenSearchAggregateIndexScanRule AGGREGATE_INDEX_SCAN =
       OpenSearchAggregateIndexScanRule.Config.DEFAULT.toRule();
+  private static final OpenSearchSortIndexScanRule SORT_INDEX_SCAN =
+      OpenSearchSortIndexScanRule.Config.DEFAULT.toRule();
 
   public static final List<RelOptRule> OPEN_SEARCH_INDEX_SCAN_RULES =
-      ImmutableList.of(PROJECT_INDEX_SCAN, FILTER_INDEX_SCAN, AGGREGATE_INDEX_SCAN);
+      ImmutableList.of(
+          PROJECT_INDEX_SCAN, FILTER_INDEX_SCAN, AGGREGATE_INDEX_SCAN, SORT_INDEX_SCAN);
 
   // prevent instantiation
   private OpenSearchIndexRules() {}

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.planner.physical;
+
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.immutables.value.Value;
+import org.opensearch.sql.opensearch.storage.scan.CalciteLogicalIndexScan;
+
+@Value.Enclosing
+public class OpenSearchSortIndexScanRule extends RelRule<OpenSearchSortIndexScanRule.Config> {
+
+  protected OpenSearchSortIndexScanRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final LogicalSort sort = call.rel(0);
+    final CalciteLogicalIndexScan scan = call.rel(1);
+
+    if (!sort.getCollation().getFieldCollations().isEmpty() && sort.fetch == null) {
+      var collations = sort.collation.getFieldCollations();
+      CalciteLogicalIndexScan newScan = scan.pushDownSort(collations);
+      if (newScan != null) {
+        call.transformTo(newScan);
+      }
+    }
+  }
+
+  /** Rule configuration. */
+  @Value.Immutable
+  public interface Config extends RelRule.Config {
+    OpenSearchSortIndexScanRule.Config DEFAULT =
+        ImmutableOpenSearchSortIndexScanRule.Config.builder()
+            .build()
+            .withOperandSupplier(
+                b0 ->
+                    b0.operand(LogicalSort.class)
+                        .oneInput(
+                            b1 ->
+                                b1.operand(CalciteLogicalIndexScan.class)
+                                    .predicate(OpenSearchIndexScanRule::test)
+                                    .noInputs()));
+
+    @Override
+    default OpenSearchSortIndexScanRule toRule() {
+      return new OpenSearchSortIndexScanRule(this);
+    }
+  }
+}

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

@@ -65,6 +65,7 @@ public RelWriter explainTerms(RelWriter pw) {
   // TODO: should we consider equivalent among PushDownContexts with different push down sequence?
   public static class PushDownContext extends ArrayDeque<PushDownAction> {
     private boolean isAggregatePushed = false;
+    private boolean isSortPushed = false;
 
     @Override
     public PushDownContext clone() {
@@ -75,9 +76,13 @@ public PushDownContext clone() {
     public boolean add(PushDownAction pushDownAction) {
       // Defense check. It should never do push down to this context after aggregate push-down.
       assert !isAggregatePushed : "Aggregate has already been pushed!";
+      assert !isSortPushed : "Sort has already been pushed!";
       if (pushDownAction.type == PushDownType.AGGREGATION) {
         isAggregatePushed = true;
       }
+      if (pushDownAction.type == PushDownType.SORT) {
+        isSortPushed = true;
+      }
       return super.add(pushDownAction);
     }
 
@@ -86,13 +91,19 @@ public boolean isAggregatePushed() {
       isAggregatePushed = !isEmpty() && super.peekLast().type == PushDownType.AGGREGATION;
       return isAggregatePushed;
     }
+
+    public boolean isSortPushed() {
+      if (isSortPushed) return true;
+      isSortPushed = !isEmpty() && super.peekLast().type == PushDownType.SORT;
+      return isSortPushed;
+    }
   }
 
   protected enum PushDownType {
     FILTER,
     PROJECT,
     AGGREGATION,
-    // SORT,
+    SORT,
     // LIMIT,
     // HIGHLIGHT,
     // NESTED

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.opensearch.storage.scan;
 
 import com.google.common.collect.ImmutableList;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -16,6 +17,7 @@
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.hint.RelHint;
@@ -27,6 +29,10 @@
 import org.apache.logging.log4j.Logger;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.search.aggregations.AggregationBuilder;
+import org.opensearch.search.sort.ScoreSortBuilder;
+import org.opensearch.search.sort.SortBuilder;
+import org.opensearch.search.sort.SortBuilders;
+import org.opensearch.search.sort.SortOrder;
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
@@ -163,4 +169,45 @@ public CalciteLogicalIndexScan pushDownAggregate(Aggregate aggregate) {
     }
     return null;
   }
+
+  public CalciteLogicalIndexScan pushDownSort(List<RelFieldCollation> collations) {
+    if (getPushDownContext().isSortPushed()) {
+      return null;
+    }
+    try {
+      CalciteLogicalIndexScan newScan = this.copyWithNewSchema(this.getRowType());
+      List<SortBuilder<?>> builders = new ArrayList<>();
+      for (RelFieldCollation collation : collations) {
+        int index = collation.getFieldIndex();
+        String fieldName = this.getRowType().getFieldNames().get(index);
+        RelFieldCollation.Direction direction = collation.getDirection();
+        // Default sort order is ASCENDING
+        SortOrder order =
+            RelFieldCollation.Direction.DESCENDING.equals(direction)
+                ? SortOrder.DESC
+                : SortOrder.ASC;
+        // TODO: support script sort and distance sort
+        SortBuilder<?> sortBuilder;
+        if (ScoreSortBuilder.NAME.equals(fieldName)) {
+          sortBuilder = SortBuilders.scoreSort();
+        } else {
+          sortBuilder = SortBuilders.fieldSort(fieldName);
+        }
+        builders.add(sortBuilder.order(order));
+      }
+      newScan.pushDownContext.add(
+          PushDownAction.of(
+              PushDownType.SORT,
+              builders.toString(),
+              requestBuilder -> requestBuilder.pushDownSort(builders)));
+      return newScan;
+    } catch (Exception e) {
+      if (LOG.isDebugEnabled()) {
+        LOG.debug("Cannot pushdown the sort {}", collations, e);
+      } else {
+        LOG.warn("Cannot pushdown the sort {}, ", collations);
+      }
+    }
+    return null;
+  }
 }