Support pagination in V2 engine, phase 1 (#226)

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;
@@ -83,6 +84,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;
@@ -520,6 +522,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

@@ -47,6 +47,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;
@@ -289,4 +290,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/datasource/model/DataSourceMetadata.java

@@ -12,16 +12,13 @@
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableMap;
-import com.google.gson.Gson;
-import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.NoArgsConstructor;
-import lombok.RequiredArgsConstructor;
 import lombok.Setter;
 import org.opensearch.sql.datasource.DataSourceService;
 

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/QueryId.java

@@ -16,6 +16,7 @@
  * Query id of {@link AbstractPlan}.
  */
 public class QueryId {
+  public static final QueryId None = new QueryId("");
   /**
    * Query id.
    */

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

@@ -15,7 +15,9 @@
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.planner.PlanContext;
 import org.opensearch.sql.planner.Planner;
+import org.opensearch.sql.planner.logical.LogicalPaginate;
 import org.opensearch.sql.planner.logical.LogicalPlan;
+import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
 /**
@@ -28,7 +30,15 @@ public class QueryService {
 
   private final ExecutionEngine executionEngine;
 
+  /**
+   * There are two planners, one - to handle pagination requests (cursor/scroll) only and
+   * another one for everything else.
+   * @see OpenSearchPluginModule#queryPlanFactory (:plugin module)
+   * @see LogicalPlanOptimizer#paginationCreate
+   * @see QueryService
+   */
   private final Planner planner;
+  private final Planner paginationPlanner;
 
   /**
    * Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.
@@ -46,6 +56,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.
@@ -97,6 +115,6 @@ public LogicalPlan analyze(UnresolvedPlan plan) {
    * Translate {@link LogicalPlan} to {@link PhysicalPlan}.
    */
   public PhysicalPlan plan(LogicalPlan plan) {
-    return planner.plan(plan);
+    return plan instanceof LogicalPaginate ? paginationPlanner.plan(plan) : planner.plan(plan);
   }
 }

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.PaginatedPlanCache;
+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 PaginatedPlanCache paginatedPlanCache;
+
+  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,
+                               PaginatedPlanCache planCache,
+                               ResponseListener<ExecutionEngine.QueryResponse>
+                                   queryResponseListener) {
+    super(queryId);
+    this.cursor = cursor;
+    this.paginatedPlanCache = planCache;
+    this.queryService = queryService;
+    this.queryResponseListener = queryResponseListener;
+  }
+
+  @Override
+  public void execute() {
+    try {
+      PhysicalPlan plan = paginatedPlanCache.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."));
+  }
+}