Merge OpenSearchPagedIndexScan and OpenSearchIndexScan

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java

@@ -65,7 +65,7 @@ public void onFailure(Exception e) {
   /**
    * Create QueryExecution from Statement.
    */
-  public AbstractPlan createContinuePaginatedPlan(
+  public AbstractPlan create(
       Statement statement,
       Optional<ResponseListener<ExecutionEngine.QueryResponse>> queryListener,
       Optional<ResponseListener<ExecutionEngine.ExplainResponse>> explainListener) {
@@ -75,9 +75,9 @@ public AbstractPlan createContinuePaginatedPlan(
   /**
    * Creates a ContinuePaginatedPlan from a cursor.
    */
-  public AbstractPlan createContinuePaginatedPlan(String cursor, boolean isExplain,
-      ResponseListener<ExecutionEngine.QueryResponse> queryResponseListener,
-      ResponseListener<ExecutionEngine.ExplainResponse> explainListener) {
+  public AbstractPlan create(String cursor, boolean isExplain,
+                             ResponseListener<ExecutionEngine.QueryResponse> queryResponseListener,
+                             ResponseListener<ExecutionEngine.ExplainResponse> explainListener) {
     QueryId queryId = QueryId.queryId();
     var plan = new ContinuePaginatedPlan(queryId, cursor, queryService,
             planSerializer, queryResponseListener);
@@ -119,7 +119,7 @@ public AbstractPlan visitExplain(
 
     return new ExplainPlan(
         QueryId.queryId(),
-        createContinuePaginatedPlan(node.getStatement(),
+        create(node.getStatement(),
             Optional.of(NO_CONSUMER_RESPONSE_LISTENER), Optional.empty()),
         context.getRight().get());
   }

core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java

@@ -13,7 +13,6 @@
 import java.util.List;
 import java.util.stream.Collectors;
 import org.opensearch.sql.planner.logical.LogicalPlan;
-import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder;
 import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter;
 import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
 import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
@@ -52,7 +51,7 @@ public static LogicalPlanOptimizer create() {
          * Phase 2: Transformations that rely on data source push down capability
          */
         new CreateTableScanBuilder(),
-        new CreatePagingTableScanBuilder(),
+        new PushDownPageSize(),
         TableScanPushDown.PUSH_DOWN_FILTER,
         TableScanPushDown.PUSH_DOWN_AGGREGATION,
         TableScanPushDown.PUSH_DOWN_SORT,

core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.optimizer;
+
+import com.facebook.presto.matching.Captures;
+import com.facebook.presto.matching.Pattern;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.Optional;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
+import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.storage.read.TableScanBuilder;
+
+public class PushDownPageSize implements Rule<LogicalPaginate> {
+  @Override
+  public Pattern<LogicalPaginate> pattern() {
+    return Pattern.typeOf(LogicalPaginate.class)
+      .matching(lp -> findTableScanBuilder(lp).isPresent());
+  }
+
+  @Override
+  public LogicalPlan apply(LogicalPaginate plan, Captures captures) {
+
+    var builder = findTableScanBuilder(plan).orElseThrow();
+    if (!builder.pushDownPageSize(plan)) {
+      throw new IllegalStateException("Failed to push down LogicalPaginate");
+    }
+    return plan.getChild().get(0);
+  }
+
+  private Optional<TableScanBuilder> findTableScanBuilder(LogicalPaginate logicalPaginate) {
+    Deque<LogicalPlan> plans = new ArrayDeque<>();
+    plans.add(logicalPaginate);
+    do {
+      final var plan = plans.removeFirst();
+      final var children = plan.getChild();
+      if (children.stream().anyMatch(TableScanBuilder.class::isInstance)) {
+        if (children.size() > 1) {
+          throw new UnsupportedOperationException(
+            "Unsupported plan: relation operator cannot have siblings");
+        }
+        return Optional.of((TableScanBuilder) children.get(0));
+      }
+      plans.addAll(children);
+    } while (!plans.isEmpty());
+    return Optional.empty();
+  }
+}

core/src/main/java/org/opensearch/sql/storage/Table.java

@@ -100,8 +100,4 @@ default StreamingSource asStreamingSource() {
     throw new UnsupportedOperationException();
   }
 
-  default TableScanBuilder createPagedScanBuilder(int pageSize) {
-    var error = String.format("'%s' does not support pagination", getClass().toString());
-    throw new UnsupportedOperationException(error);
-  }
 }

core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java

@@ -11,6 +11,7 @@
 import org.opensearch.sql.planner.logical.LogicalHighlight;
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalNested;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
 import org.opensearch.sql.planner.logical.LogicalProject;
@@ -28,7 +29,7 @@ public abstract class TableScanBuilder extends LogicalPlan {
   /**
    * Construct and initialize children to empty list.
    */
-  public TableScanBuilder() {
+  protected TableScanBuilder() {
     super(Collections.emptyList());
   }
 
@@ -116,6 +117,10 @@ public boolean pushDownNested(LogicalNested nested) {
     return false;
   }
 
+  public boolean pushDownPageSize(LogicalPaginate paginate) {
+    return false;
+  }
+
   @Override
   public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
     return visitor.visitTableScanBuilder(this, context);

core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java

@@ -62,23 +62,23 @@ void init() {
   public void createFromQueryShouldSuccess() {
     Statement query = new Query(plan, 0);
     AbstractPlan queryExecution =
-        factory.createContinuePaginatedPlan(query, Optional.of(queryListener), Optional.empty());
+        factory.create(query, Optional.of(queryListener), Optional.empty());
     assertTrue(queryExecution instanceof QueryPlan);
   }
 
   @Test
   public void createFromExplainShouldSuccess() {
     Statement query = new Explain(new Query(plan, 0));
     AbstractPlan queryExecution =
-        factory.createContinuePaginatedPlan(query, Optional.empty(), Optional.of(explainListener));
+        factory.create(query, Optional.empty(), Optional.of(explainListener));
     assertTrue(queryExecution instanceof ExplainPlan);
   }
 
   @Test
   public void createFromCursorShouldSuccess() {
-    AbstractPlan queryExecution = factory.createContinuePaginatedPlan("", false,
+    AbstractPlan queryExecution = factory.create("", false,
         queryListener, explainListener);
-    AbstractPlan explainExecution = factory.createContinuePaginatedPlan("", true,
+    AbstractPlan explainExecution = factory.create("", true,
         queryListener, explainListener);
     assertAll(
         () -> assertTrue(queryExecution instanceof ContinuePaginatedPlan),
@@ -91,7 +91,7 @@ public void createFromQueryWithoutQueryListenerShouldThrowException() {
     Statement query = new Query(plan, 0);
 
     IllegalArgumentException exception =
-        assertThrows(IllegalArgumentException.class, () -> factory.createContinuePaginatedPlan(
+        assertThrows(IllegalArgumentException.class, () -> factory.create(
             query, Optional.empty(), Optional.empty()));
     assertEquals("[BUG] query listener must be not null", exception.getMessage());
   }
@@ -101,7 +101,7 @@ public void createFromExplainWithoutExplainListenerShouldThrowException() {
     Statement query = new Explain(new Query(plan, 0));
 
     IllegalArgumentException exception =
-        assertThrows(IllegalArgumentException.class, () -> factory.createContinuePaginatedPlan(
+        assertThrows(IllegalArgumentException.class, () -> factory.create(
             query, Optional.empty(), Optional.empty()));
     assertEquals("[BUG] explain listener must be not null", exception.getMessage());
   }
@@ -129,7 +129,7 @@ public void createQueryWithFetchSizeWhichCanBePaged() {
     factory = new QueryPlanFactory(queryService, planSerializer);
     Statement query = new Query(plan, 10);
     AbstractPlan queryExecution =
-        factory.createContinuePaginatedPlan(query, Optional.of(queryListener), Optional.empty());
+        factory.create(query, Optional.of(queryListener), Optional.empty());
     assertTrue(queryExecution instanceof QueryPlan);
   }
 
@@ -139,7 +139,7 @@ public void createQueryWithFetchSizeWhichCannotBePaged() {
     factory = new QueryPlanFactory(queryService, planSerializer);
     Statement query = new Query(plan, 10);
     assertThrows(UnsupportedCursorRequestException.class,
-        () -> factory.createContinuePaginatedPlan(query,
+        () -> factory.create(query,
             Optional.of(queryListener), Optional.empty()));
   }
 }