Allow pushing down multiple times

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

Author: yuancu

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

@@ -43,7 +43,8 @@ public RelNode convert(RelNode rel) {
     final CalciteLogicalIndexScan scan = (CalciteLogicalIndexScan) rel;
     return new CalciteEnumerableIndexScan(
         scan.getCluster(),
-        rel.getTraitSet().getCollation(),
+        // Retains RelDistribution and RelCollation but replaces Convention
+        scan.getTraitSet().plus(EnumerableConvention.INSTANCE),
         scan.getHints(),
         scan.getTable(),
         scan.getOsIndex(),

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

@@ -18,8 +18,4 @@ static boolean test(CalciteLogicalIndexScan scan) {
   static boolean sortByFieldsOnly(LogicalSort sort) {
     return !sort.getCollation().getFieldCollations().isEmpty() && sort.fetch == null;
   }
-
-  static boolean isSortPushed(CalciteLogicalIndexScan scan) {
-    return scan.getPushDownContext().isSortPushed();
-  }
 }

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

@@ -5,7 +5,6 @@
 
 package org.opensearch.sql.opensearch.planner.physical;
 
-import java.util.function.Predicate;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelRule;
 import org.apache.calcite.rel.logical.LogicalSort;
@@ -44,9 +43,7 @@ public interface Config extends RelRule.Config {
                         .oneInput(
                             b1 ->
                                 b1.operand(CalciteLogicalIndexScan.class)
-                                    .predicate(
-                                        Predicate.not(OpenSearchIndexScanRule::isSortPushed)
-                                            .and(OpenSearchIndexScanRule::test))
+                                    .predicate(OpenSearchIndexScanRule::test)
                                     .noInputs()));
 
     @Override

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

@@ -6,7 +6,6 @@
 package org.opensearch.sql.opensearch.storage.scan;
 
 import java.util.List;
-import org.apache.calcite.adapter.enumerable.EnumerableConvention;
 import org.apache.calcite.adapter.enumerable.EnumerableRel;
 import org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;
 import org.apache.calcite.adapter.enumerable.PhysType;
@@ -21,7 +20,7 @@
 import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptTable;
-import org.apache.calcite.plan.RelTrait;
+import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.rel.type.RelDataType;
@@ -45,21 +44,13 @@ public class CalciteEnumerableIndexScan extends CalciteIndexScan implements Enum
    */
   public CalciteEnumerableIndexScan(
       RelOptCluster cluster,
-      RelTrait relTrait,
+      RelTraitSet traitSet,
       List<RelHint> hints,
       RelOptTable table,
       OpenSearchIndex osIndex,
       RelDataType schema,
       PushDownContext pushDownContext) {
-
-    super(
-        cluster,
-        cluster.traitSetOf(EnumerableConvention.INSTANCE, relTrait),
-        hints,
-        table,
-        osIndex,
-        schema,
-        pushDownContext);
+    super(cluster, traitSet, hints, table, osIndex, schema, pushDownContext);
   }
 
   @Override

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

@@ -65,7 +65,6 @@ 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() {
@@ -76,13 +75,9 @@ 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);
     }
 
@@ -91,12 +86,6 @@ 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 {