Support pagination in V2 engine, phase 1

core/build.gradle

@@ -57,6 +57,7 @@ dependencies {
     testImplementation('org.junit.jupiter:junit-jupiter:5.6.2')
     testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: '2.1'
     testImplementation group: 'org.mockito', name: 'mockito-core', version: '3.12.4'
+    testImplementation group: 'org.mockito', name: 'mockito-inline', version: '3.12.4'
     testImplementation group: 'org.mockito', name: 'mockito-junit-jupiter', version: '3.12.4'
 }
 

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

@@ -49,6 +49,7 @@
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.ML;
+import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -84,6 +85,7 @@
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalML;
 import org.opensearch.sql.planner.logical.LogicalMLCommons;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalProject;
 import org.opensearch.sql.planner.logical.LogicalRareTopN;
@@ -535,6 +537,12 @@ public LogicalPlan visitML(ML node, AnalysisContext context) {
     return new LogicalML(child, node.getArguments());
   }
 
+  @Override
+  public LogicalPlan visitPaginate(Paginate paginate, AnalysisContext context) {
+    LogicalPlan child = paginate.getChild().get(0).accept(this, context);
+    return new LogicalPaginate(paginate.getPageSize(), List.of(child));
+  }
+
   /**
    * The first argument is always "asc", others are optional.
    * Given nullFirst argument, use its value. Otherwise just use DEFAULT_ASC/DESC.

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

@@ -48,6 +48,7 @@
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.ML;
+import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Project;
 import org.opensearch.sql.ast.tree.RareTopN;
@@ -294,4 +295,8 @@ public T visitQuery(Query node, C context) {
   public T visitExplain(Explain node, C context) {
     return visitStatement(node, context);
   }
+
+  public T visitPaginate(Paginate paginate, C context) {
+    return visitChildren(paginate, context);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/statement/Query.java

@@ -27,6 +27,7 @@
 public class Query extends Statement {
 
   protected final UnresolvedPlan plan;
+  protected final int fetchSize;
 
   @Override
   public <R, C> R accept(AbstractNodeVisitor<R, C> visitor, C context) {

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
+
+/**
+ * AST node to represent pagination operation.
+ * Actually a wrapper to the AST.
+ */
+@RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@ToString
+public class Paginate extends UnresolvedPlan {
+  @Getter
+  private final int pageSize;
+  private UnresolvedPlan child;
+
+  public Paginate(int pageSize, UnresolvedPlan child) {
+    this.pageSize = pageSize;
+    this.child = child;
+  }
+
+  @Override
+  public List<? extends Node> getChild() {
+    return List.of(child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitPaginate(this, context);
+  }
+
+  @Override
+  public UnresolvedPlan attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+}

core/src/main/java/org/opensearch/sql/exception/NoCursorException.java

@@ -0,0 +1,13 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.exception;
+
+/**
+ * This should be thrown on serialization of a PhysicalPlan tree if paging is finished.
+ * Processing of such exception should outcome of responding no cursor to the user.
+ */
+public class NoCursorException extends RuntimeException {
+}

core/src/main/java/org/opensearch/sql/exception/UnsupportedCursorRequestException.java

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.exception;
+
+/**
+ * This should be thrown by V2 engine to support fallback scenario.
+ */
+public class UnsupportedCursorRequestException extends RuntimeException {
+}

core/src/main/java/org/opensearch/sql/executor/ExecutionEngine.java

@@ -14,6 +14,7 @@
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.executor.pagination.Cursor;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 /**
@@ -53,6 +54,8 @@ void execute(PhysicalPlan plan, ExecutionContext context,
   class QueryResponse {
     private final Schema schema;
     private final List<ExprValue> results;
+    private final long total;
+    private final Cursor cursor;
   }
 
   @Data

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -46,6 +46,14 @@ public void execute(UnresolvedPlan plan,
     }
   }
 
+  /**
+   * Execute a physical plan without analyzing or planning anything.
+   */
+  public void executePlan(PhysicalPlan plan,
+                          ResponseListener<ExecutionEngine.QueryResponse> listener) {
+    executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener);
+  }
+
   /**
    * Execute the {@link UnresolvedPlan}, with {@link PlanContext} and using {@link ResponseListener}
    * to get response.

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

@@ -0,0 +1,59 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.executor.execution;
+
+import org.opensearch.sql.common.response.ResponseListener;
+import org.opensearch.sql.executor.ExecutionEngine;
+import org.opensearch.sql.executor.QueryId;
+import org.opensearch.sql.executor.QueryService;
+import org.opensearch.sql.executor.pagination.PlanSerializer;
+import org.opensearch.sql.planner.physical.PhysicalPlan;
+
+/**
+ * ContinuePaginatedPlan represents cursor a request.
+ * It returns subsequent pages to the user (2nd page and all next).
+ * {@link PaginatedPlan}
+ */
+public class ContinuePaginatedPlan extends AbstractPlan {
+
+  private final String cursor;
+  private final QueryService queryService;
+  private final PlanSerializer planSerializer;
+
+  private final ResponseListener<ExecutionEngine.QueryResponse> queryResponseListener;
+
+
+  /**
+   * Create an abstract plan that can continue paginating a given cursor.
+   */
+  public ContinuePaginatedPlan(QueryId queryId, String cursor, QueryService queryService,
+                               PlanSerializer planCache,
+                               ResponseListener<ExecutionEngine.QueryResponse>
+                                   queryResponseListener) {
+    super(queryId);
+    this.cursor = cursor;
+    this.planSerializer = planCache;
+    this.queryService = queryService;
+    this.queryResponseListener = queryResponseListener;
+  }
+
+  @Override
+  public void execute() {
+    try {
+      PhysicalPlan plan = planSerializer.convertToPlan(cursor);
+      queryService.executePlan(plan, queryResponseListener);
+    } catch (Exception e) {
+      queryResponseListener.onFailure(e);
+    }
+  }
+
+  @Override
+  public void explain(ResponseListener<ExecutionEngine.ExplainResponse> listener) {
+    listener.onFailure(new UnsupportedOperationException(
+        "Explain of a paged query continuation is not supported. "
+        + "Use `explain` for the initial query request."));
+  }
+}

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

@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.executor.execution;
+
+import org.apache.commons.lang3.NotImplementedException;
+import org.opensearch.sql.ast.tree.Paginate;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.common.response.ResponseListener;
+import org.opensearch.sql.executor.ExecutionEngine;
+import org.opensearch.sql.executor.QueryId;
+import org.opensearch.sql.executor.QueryService;
+
+/**
+ * PaginatedPlan represents a page request. Dislike a regular QueryPlan,
+ * it returns paged response to the user and cursor, which allows to query
+ * next page.
+ * {@link ContinuePaginatedPlan}
+ */
+public class PaginatedPlan extends AbstractPlan {
+  final UnresolvedPlan plan;
+  final int fetchSize;
+  final QueryService queryService;
+  final ResponseListener<ExecutionEngine.QueryResponse>
+      queryResponseResponseListener;
+
+  /**
+   * Create an abstract plan that can start paging a query.
+   */
+  public PaginatedPlan(QueryId queryId, UnresolvedPlan plan, int fetchSize,
+                       QueryService queryService,
+                       ResponseListener<ExecutionEngine.QueryResponse>
+                           queryResponseResponseListener) {
+    super(queryId);
+    this.plan = plan;
+    this.fetchSize = fetchSize;
+    this.queryService = queryService;
+    this.queryResponseResponseListener = queryResponseResponseListener;
+  }
+
+  @Override
+  public void execute() {
+    queryService.execute(new Paginate(fetchSize, plan), queryResponseResponseListener);
+  }
+
+  @Override
+  public void explain(ResponseListener<ExecutionEngine.ExplainResponse> listener) {
+    listener.onFailure(new NotImplementedException(
+        "`explain` feature for paginated requests is not implemented yet."));
+  }
+}

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

@@ -18,9 +18,11 @@
 import org.opensearch.sql.ast.statement.Query;
 import org.opensearch.sql.ast.statement.Statement;
 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.PlanSerializer;
 
 /**
  * QueryExecution Factory.
@@ -37,9 +39,10 @@ public class QueryPlanFactory
    * Query Service.
    */
   private final QueryService queryService;
+  private final PlanSerializer planSerializer;
 
   /**
-   * NO_CONSUMER_RESPONSE_LISTENER should never been called. It is only used as constructor
+   * NO_CONSUMER_RESPONSE_LISTENER should never be called. It is only used as constructor
    * parameter of {@link QueryPlan}.
    */
   @VisibleForTesting
@@ -62,39 +65,62 @@ public void onFailure(Exception e) {
   /**
    * Create QueryExecution from Statement.
    */
-  public AbstractPlan create(
+  public AbstractPlan createContinuePaginatedPlan(
       Statement statement,
       Optional<ResponseListener<ExecutionEngine.QueryResponse>> queryListener,
       Optional<ResponseListener<ExecutionEngine.ExplainResponse>> explainListener) {
     return statement.accept(this, Pair.of(queryListener, explainListener));
   }
 
+  /**
+   * Creates a ContinuePaginatedPlan from a cursor.
+   */
+  public AbstractPlan createContinuePaginatedPlan(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);
+    return isExplain ? new ExplainPlan(queryId, plan, explainListener) : plan;
+  }
+
   @Override
   public AbstractPlan visitQuery(
       Query node,
-      Pair<
-              Optional<ResponseListener<ExecutionEngine.QueryResponse>>,
-              Optional<ResponseListener<ExecutionEngine.ExplainResponse>>>
+      Pair<Optional<ResponseListener<ExecutionEngine.QueryResponse>>,
+           Optional<ResponseListener<ExecutionEngine.ExplainResponse>>>
           context) {
     Preconditions.checkArgument(
         context.getLeft().isPresent(), "[BUG] query listener must be not null");
 
-    return new QueryPlan(QueryId.queryId(), node.getPlan(), queryService, context.getLeft().get());
+    if (node.getFetchSize() > 0) {
+      if (planSerializer.canConvertToCursor(node.getPlan())) {
+        return new PaginatedPlan(QueryId.queryId(), node.getPlan(), node.getFetchSize(),
+            queryService,
+            context.getLeft().get());
+      } else {
+        // This should be picked up by the legacy engine.
+        throw new UnsupportedCursorRequestException();
+      }
+    } else {
+      return new QueryPlan(QueryId.queryId(), node.getPlan(), queryService,
+          context.getLeft().get());
+    }
   }
 
   @Override
   public AbstractPlan visitExplain(
       Explain node,
-      Pair<
-              Optional<ResponseListener<ExecutionEngine.QueryResponse>>,
-              Optional<ResponseListener<ExecutionEngine.ExplainResponse>>>
+      Pair<Optional<ResponseListener<ExecutionEngine.QueryResponse>>,
+           Optional<ResponseListener<ExecutionEngine.ExplainResponse>>>
           context) {
     Preconditions.checkArgument(
         context.getRight().isPresent(), "[BUG] explain listener must be not null");
 
     return new ExplainPlan(
         QueryId.queryId(),
-        create(node.getStatement(), Optional.of(NO_CONSUMER_RESPONSE_LISTENER), Optional.empty()),
+        createContinuePaginatedPlan(node.getStatement(),
+            Optional.of(NO_CONSUMER_RESPONSE_LISTENER), Optional.empty()),
         context.getRight().get());
   }
 }