Support pagination in V2 engine, phase 1

core/build.gradle

@@ -56,6 +56,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

@@ -50,6 +50,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;
@@ -85,6 +86,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;
@@ -529,6 +531,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,45 @@
+/*
+ * 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;
+
+@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) {
+    assert this.child == null;
+    this.child = child;
+    return this;
+  }
+}

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

@@ -0,0 +1,9 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.exception;
+
+public class UnsupportedCursorRequestException extends RuntimeException {
+}

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

@@ -0,0 +1,104 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.executor;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
+import org.opensearch.sql.ast.expression.AllFields;
+import org.opensearch.sql.ast.tree.Project;
+import org.opensearch.sql.ast.tree.Relation;
+
+/**
+ * Use this unresolved plan visitor to check if a plan can be serialized by PaginatedPlanCache.
+ * If plan.accept(new CanPaginateVisitor(...)) returns true,
+ * then PaginatedPlanCache.convertToCursor will succeed.
+ * Otherwise, it will fail.
+ * Currently, the conditions are:
+ * - only projection of a relation is supported.
+ * - projection only has * (a.k.a. allFields).
+ * - Relation only scans one table
+ * - The table is an open search index.
+ * See PaginatedPlanCache.canConvertToCursor for usage.
+ */
+public class CanPaginateVisitor extends AbstractNodeVisitor<Boolean, Object> {
+
+  @Override
+  public Boolean visitRelation(Relation node, Object context) {
+    if (!node.getChild().isEmpty()) {
+      // Relation instance should never have a child, but check just in case.
+      return Boolean.FALSE;
+    }
+
+    // TODO use storageEngine from the calling PaginatedPlanCache to determine if
+    // node.getTableQualifiedName is provided by the storage engine. Return false if it's
+    // not the case.
+    return Boolean.TRUE;
+  }
+
+  /*
+  private Boolean canPaginate(Node node, Object context) {
+    AtomicBoolean result = new AtomicBoolean(true);
+    node.getChild().forEach(n -> result.set(result.get() && n.accept(this, context)));
+    return result.get();
+  }
+
+  For queries without `FROM` clause.
+  Required to overload `toCursor` function in `ValuesOperator` and modify cursor parsing.
+  @Override
+  public Boolean visitValues(Values node, Object context) {
+    return canPaginate(node, context);
+  }
+
+  For queries with LIMIT clause:
+  Required to overload `toCursor` function in `LimitOperator` and modify cursor parsing.
+  @Override
+  public Boolean visitLimit(Limit node, Object context) {
+    return canPaginate(node, context);
+  }
+
+  For queries with ORDER BY clause:
+  Required to overload `toCursor` function in `SortOperator` and modify cursor parsing.
+  @Override
+  public Boolean visitSort(Sort node, Object context) {
+    return canPaginate(node, context);
+  }
+
+  For queries with WHERE clause:
+  Required to overload `toCursor` function in `FilterOperator` and modify cursor parsing.
+  @Override
+  public Boolean visitFilter(Filter node, Object context) {
+    return canPaginate(node, context);
+  }
+  */
+
+  @Override
+  public Boolean visitChildren(Node node, Object context) {
+    return Boolean.FALSE;
+  }
+
+  @Override
+  public Boolean visitProject(Project node, Object context) {
+    // Allow queries with 'SELECT *' only. Those restriction could be removed, but consider
+    // in-memory aggregation performed by window function (see WindowOperator).
+    // SELECT max(age) OVER (PARTITION BY city) ...
+    var projections = node.getProjectList();
+    if (projections.size() != 1) {
+      return Boolean.FALSE;
+    }
+
+    if (!(projections.get(0) instanceof AllFields)) {
+      return Boolean.FALSE;
+    }
+
+    var children = node.getChild();
+    if (children.size() != 1) {
+      return Boolean.FALSE;
+    }
+
+    return children.get(0).accept(this, context);
+  }
+}

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.opensearch.executor.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