Merge OpenSearchPagedIndexScan and OpenSearchIndexScan

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -42,6 +42,7 @@
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.tree.AD;
 import org.opensearch.sql.ast.tree.Aggregation;
+import org.opensearch.sql.ast.tree.Cursor;
 import org.opensearch.sql.ast.tree.Dedupe;
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.Filter;
@@ -64,7 +65,6 @@
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.data.model.ExprMissingValue;
 import org.opensearch.sql.data.type.ExprCoreType;
-import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
@@ -80,6 +80,7 @@
 import org.opensearch.sql.expression.function.FunctionName;
 import org.opensearch.sql.expression.function.TableFunctionImplementation;
 import org.opensearch.sql.expression.parse.ParseExpression;
+import org.opensearch.sql.planner.LogicalCursor;
 import org.opensearch.sql.planner.logical.LogicalAD;
 import org.opensearch.sql.planner.logical.LogicalAggregation;
 import org.opensearch.sql.planner.logical.LogicalDedupe;
@@ -211,6 +212,10 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex
         tableFunctionImplementation.applyArguments());
   }
 
+  @Override
+  public LogicalPlan visitCursor(Cursor cursor, AnalysisContext context) {
+    return new LogicalCursor(cursor.getCursor());
+  }
 
   @Override
   public LogicalPlan visitLimit(Limit node, AnalysisContext context) {

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -41,6 +41,7 @@
 import org.opensearch.sql.ast.statement.Statement;
 import org.opensearch.sql.ast.tree.AD;
 import org.opensearch.sql.ast.tree.Aggregation;
+import org.opensearch.sql.ast.tree.Cursor;
 import org.opensearch.sql.ast.tree.Dedupe;
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.Filter;
@@ -299,4 +300,8 @@ public T visitExplain(Explain node, C context) {
   public T visitPaginate(Paginate paginate, C context) {
     return visitChildren(paginate, context);
   }
+
+  public T visitCursor(Cursor cursor, C context) {
+    return visit(cursor, context);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java

@@ -0,0 +1,21 @@
+package org.opensearch.sql.ast.tree;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+@RequiredArgsConstructor
+public class Cursor extends UnresolvedPlan {
+  @Getter
+  final String cursor;
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitCursor(this, context);
+  }
+
+  @Override
+  public UnresolvedPlan attach(UnresolvedPlan child) {
+    throw new UnsupportedOperationException("Cursor unresolved plan does not support children");
+  }
+}

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

@@ -17,11 +17,14 @@
 import org.opensearch.sql.ast.statement.Explain;
 import org.opensearch.sql.ast.statement.Query;
 import org.opensearch.sql.ast.statement.Statement;
+import org.opensearch.sql.ast.tree.Cursor;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.exception.UnsupportedCursorRequestException;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
 import org.opensearch.sql.executor.QueryService;
+import org.opensearch.sql.executor.pagination.CanPaginateVisitor;
 import org.opensearch.sql.executor.pagination.PlanSerializer;
 
 /**
@@ -65,7 +68,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,15 +78,18 @@ 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);
+    var plan = new QueryPlan(queryId, new Cursor(cursor), queryService, queryResponseListener);
     return isExplain ? new ExplainPlan(queryId, plan, explainListener) : plan;
   }
 
+   boolean canConvertToCursor(UnresolvedPlan plan) {
+    return plan.accept(new CanPaginateVisitor(), null);
+  }
+
   @Override
   public AbstractPlan visitQuery(
       Query node,
@@ -94,7 +100,7 @@ public AbstractPlan visitQuery(
         context.getLeft().isPresent(), "[BUG] query listener must be not null");
 
     if (node.getFetchSize() > 0) {
-      if (planSerializer.canConvertToCursor(node.getPlan())) {
+      if (canConvertToCursor(node.getPlan())) {
         return new QueryPlan(QueryId.queryId(), node.getPlan(), node.getFetchSize(),
             queryService,
             context.getLeft().get());
@@ -119,7 +125,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/executor/pagination/PlanSerializer.java

@@ -18,9 +18,8 @@
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
 import lombok.RequiredArgsConstructor;
-import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.exception.NoCursorException;
-import org.opensearch.sql.planner.SerializablePlan;
+import org.opensearch.sql.planner.ExternalizablePlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.storage.StorageEngine;
 
@@ -34,17 +33,14 @@ public class PlanSerializer {
 
   private final StorageEngine engine;
 
-  public boolean canConvertToCursor(UnresolvedPlan plan) {
-    return plan.accept(new CanPaginateVisitor(), null);
-  }
 
   /**
    * Converts a physical plan tree to a cursor.
    */
   public Cursor convertToCursor(PhysicalPlan plan) {
     try {
       return new Cursor(CURSOR_PREFIX
-          + serialize(((SerializablePlan) plan).getPlanForSerialization()));
+          + serialize(((ExternalizablePlan) plan).getPlanForSerialization()));
       // ClassCastException thrown when a plan in the tree doesn't implement SerializablePlan
     } catch (NotSerializableException | ClassCastException | NoCursorException e) {
       return Cursor.None;

core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java → core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java

@@ -23,11 +23,11 @@
  *   </li>
  *   <li>
  *     Overwrite {@link #getPlanForSerialization} to return
- *     another instance of {@link SerializablePlan}.
+ *     another instance of {@link ExternalizablePlan}.
  *   </li>
  * </ul>
  */
-public interface SerializablePlan extends Externalizable {
+public interface ExternalizablePlan extends Externalizable {
 
   /**
    * Argument is an instance of {@link PlanSerializer.CursorDeserializationStream}.
@@ -57,7 +57,7 @@ public interface SerializablePlan extends Externalizable {
    * It is needed to skip a `ResourceMonitorPlan` instance only, actually.
    * @return Next plan for serialization.
    */
-  default SerializablePlan getPlanForSerialization() {
+  default ExternalizablePlan getPlanForSerialization() {
     return this;
   }
 }

core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java

@@ -0,0 +1,21 @@
+package org.opensearch.sql.planner;
+
+import java.util.List;
+import lombok.Getter;
+import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
+
+public class LogicalCursor extends LogicalPlan {
+  @Getter
+  final String cursor;
+
+  public LogicalCursor(String cursor) {
+    super(List.of());
+    this.cursor = cursor;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitCursor(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java

@@ -6,6 +6,7 @@
 
 package org.opensearch.sql.planner.logical;
 
+import org.opensearch.sql.planner.LogicalCursor;
 import org.opensearch.sql.storage.read.TableScanBuilder;
 import org.opensearch.sql.storage.write.TableWriteBuilder;
 
@@ -108,4 +109,8 @@ public R visitAD(LogicalAD plan, C context) {
   public R visitPaginate(LogicalPaginate plan, C context) {
     return visitNode(plan, context);
   }
+
+  public R visitCursor(LogicalCursor plan, C context) {
+    return visitNode(plan, context);
+  }
 }

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,55 @@
+/*
+ * 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;
+
+/**
+ * A {@link LogicalPlanOptimizer} rule that pushes down page size
+ * to table scan builder.
+ */
+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();
+  }
+}