Support Limit pushdown

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

Author: yuancu

integ-test/src/test/java/org/opensearch/sql/ppl/ExplainIT.java

@@ -96,7 +96,6 @@ public void testSortPushDownExplain() throws Exception {
 
   @Test
   public void testLimitPushDownExplain() throws Exception {
-    // TODO fix after https://github.com/opensearch-project/sql/issues/3381
     String expected =
         isCalciteEnabled()
             ? loadFromFile("expectedOutput/calcite/explain_limit_push.json")

integ-test/src/test/resources/expectedOutput/calcite/explain_limit_push.json

@@ -1,6 +1,6 @@
 {
   "calcite":{
     "logical":"LogicalProject(ageMinus=[$17])\n  LogicalSort(fetch=[5])\n    LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _id=[$11], _index=[$12], _score=[$13], _maxscore=[$14], _sort=[$15], _routing=[$16], ageMinus=[-($8, 30)])\n      CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n",
-    "physical":"EnumerableCalc(expr#0..16=[{inputs}], expr#17=[30], expr#18=[-($t8, $t17)], ageMinus=[$t18])\n  EnumerableLimit(fetch=[5])\n    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n"
+    "physical":"EnumerableCalc(expr#0..16=[{inputs}], expr#17=[30], expr#18=[-($t8, $t17)], ageMinus=[$t18])\n  CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[LIMIT->5], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":5,\"timeout\":\"1m\"}, requestedTotalSize=5, pageSize=null, startFrom=0)])\n"
   }
 }

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 OpenSearchLimitIndexScanRule LIMIT_INDEX_SCAN =
+      OpenSearchLimitIndexScanRule.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, LIMIT_INDEX_SCAN);
 
   // prevent instantiation
   private OpenSearchIndexRules() {}

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

@@ -0,0 +1,89 @@
+package org.opensearch.sql.opensearch.planner.physical;
+
+import java.util.Objects;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
+import org.immutables.value.Value;
+import org.opensearch.sql.opensearch.storage.scan.CalciteLogicalIndexScan;
+
+@Value.Enclosing
+public class OpenSearchLimitIndexScanRule extends RelRule<OpenSearchLimitIndexScanRule.Config> {
+
+  protected OpenSearchLimitIndexScanRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public void onMatch(RelOptRuleCall call) {
+    final LogicalSort sort = call.rel(0);
+    final CalciteLogicalIndexScan scan = call.rel(1);
+
+    // The LogicalSort is a LIMIT that should be pushed down when its fetch field is not null and
+    // its collation is empty.
+    // For example: `sort name | head 5` should not be pushed down because it has a field collation.
+    if (sort.fetch != null && sort.getCollation().getFieldCollations().isEmpty()) {
+      Integer limitValue = extractLimitValue(sort.fetch);
+      Integer offsetValue = extractOffsetValue(sort.offset);
+      if (limitValue != null) {
+        CalciteLogicalIndexScan newScan = scan.pushDownLimit(limitValue, offsetValue);
+        if (newScan != null) {
+          call.transformTo(newScan);
+        }
+      }
+    }
+  }
+
+  private static Integer extractLimitValue(RexNode fetch) {
+    if (fetch instanceof RexLiteral) {
+      return ((RexLiteral) fetch).getValueAs(Integer.class);
+    }
+    return null;
+  }
+
+  /**
+   * Extracts the offset value from the given `RexNode`. If the offset is `null`, it defaults to 0.
+   * For example:
+   *
+   * <ul>
+   *   <li><code>source=people | head 1</code> will have a <code>null</code> offset, which is
+   *       converted to 0.
+   *   <li><code>source=people | head 1 from 2</code> will have an offset of 2.
+   * </ul>
+   *
+   * @param offset The `RexNode` representing the offset.
+   * @return The extracted offset value, or `null` if it cannot be determined.
+   */
+  private static Integer extractOffsetValue(RexNode offset) {
+    if (Objects.isNull(offset)) {
+      return 0;
+    }
+    if (offset instanceof RexLiteral) {
+      return ((RexLiteral) offset).getValueAs(Integer.class);
+    }
+    return null;
+  }
+
+  /** Rule configuration. */
+  @Value.Immutable
+  public interface Config extends RelRule.Config {
+    OpenSearchLimitIndexScanRule.Config DEFAULT =
+        ImmutableOpenSearchLimitIndexScanRule.Config.builder()
+            .build()
+            .withOperandSupplier(
+                b0 ->
+                    b0.operand(LogicalSort.class)
+                        .oneInput(
+                            b1 ->
+                                b1.operand(CalciteLogicalIndexScan.class)
+                                    .predicate(OpenSearchIndexScanRule::test)
+                                    .noInputs()));
+
+    @Override
+    default OpenSearchLimitIndexScanRule toRule() {
+      return new OpenSearchLimitIndexScanRule(this);
+    }
+  }
+}

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

@@ -93,7 +93,7 @@ protected enum PushDownType {
     PROJECT,
     AGGREGATION,
     // SORT,
-    // LIMIT,
+    LIMIT,
     // HIGHLIGHT,
     // NESTED
   }

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

@@ -163,4 +163,23 @@ public CalciteLogicalIndexScan pushDownAggregate(Aggregate aggregate) {
     }
     return null;
   }
+
+  public CalciteLogicalIndexScan pushDownLimit(Integer limit, Integer offset) {
+    try {
+      CalciteLogicalIndexScan newScan = this.copyWithNewSchema(getRowType());
+      newScan.pushDownContext.add(
+          PushDownAction.of(
+              PushDownType.LIMIT,
+              limit,
+              requestBuilder -> requestBuilder.pushDownLimit(limit, offset)));
+      return newScan;
+    } catch (Exception e) {
+      if (LOG.isDebugEnabled()) {
+        LOG.debug("Cannot pushdown the limit {}", limit, e);
+      } else {
+        LOG.warn("Cannot pushdown the limit {}, ", limit);
+      }
+    }
+    return null;
+  }
 }