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,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/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/CanPaginateVisitor.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.executor;
+
+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.
+ * The purpose of this visitor is to activate legacy engine fallback mechanism.
+ * 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.
+ * So it accepts only queries like `select * from $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;
+    }
+
+    return Boolean.TRUE;
+  }
+
+  @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

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

@@ -0,0 +1,162 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.executor;
+
+import com.google.common.hash.HashCode;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.zip.GZIPInputStream;
+import java.util.zip.GZIPOutputStream;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.expression.NamedExpression;
+import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer;
+import org.opensearch.sql.opensearch.executor.Cursor;
+import org.opensearch.sql.planner.physical.PaginateOperator;
+import org.opensearch.sql.planner.physical.PhysicalPlan;
+import org.opensearch.sql.planner.physical.ProjectOperator;
+import org.opensearch.sql.storage.StorageEngine;
+import org.opensearch.sql.storage.TableScanOperator;
+
+/**
+ * This class is entry point to paged requests. It is responsible to cursor serialization
+ * and deserialization.
+ */
+@RequiredArgsConstructor
+public class PaginatedPlanCache {
+  public static final String CURSOR_PREFIX = "n:";
+  private final StorageEngine storageEngine;
+
+  public boolean canConvertToCursor(UnresolvedPlan plan) {
+    return plan.accept(new CanPaginateVisitor(), null);
+  }
+
+  /**
+   * Converts a physical plan tree to a cursor. May cache plan related data somewhere.
+   */
+  public Cursor convertToCursor(PhysicalPlan plan) throws IOException {
+    if (plan instanceof PaginateOperator) {
+      var cursor = plan.toCursor();
+      if (cursor == null) {
+        return Cursor.None;
+      }
+      var raw = CURSOR_PREFIX + compress(cursor);
+      return new Cursor(raw.getBytes());
+    }
+    return Cursor.None;
+  }
+
+  /**
+   * Compress serialized query plan.
+   * @param str string representing a query plan
+   * @return str compressed with gzip.
+   */
+  String compress(String str) throws IOException {
+    if (str == null || str.length() == 0) {
+      return "";
+    }
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    GZIPOutputStream gzip = new GZIPOutputStream(out);
+    gzip.write(str.getBytes());
+    gzip.close();
+    return HashCode.fromBytes(out.toByteArray()).toString();
+  }
+
+  /**
+   * Decompresses a query plan that was compress with {@link PaginatedPlanCache#compress}.
+   * @param input compressed query plan
+   * @return decompressed string
+   */
+  String decompress(String input) throws IOException {
+    if (input == null || input.length() == 0) {
+      return "";
+    }
+    GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(
+        HashCode.fromString(input).asBytes()));
+    return new String(gzip.readAllBytes());
+  }
+
+  /**
+   * Parse `NamedExpression`s from cursor.
+   * @param listToFill List to fill with data.
+   * @param cursor Cursor to parse.
+   * @return Remaining part of the cursor.
+   */
+  private String parseNamedExpressions(List<NamedExpression> listToFill, String cursor) {
+    var serializer = new DefaultExpressionSerializer();
+    if (cursor.startsWith(")")) { //empty list
+      return cursor.substring(cursor.indexOf(',') + 1);
+    }
+    while (!cursor.startsWith("(")) {
+      listToFill.add((NamedExpression)
+          serializer.deserialize(cursor.substring(0,
+              Math.min(cursor.indexOf(','), cursor.indexOf(')')))));
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+    }
+    return cursor;
+  }
+
+  /**
+    * Converts a cursor to a physical plan tree.
+    */
+  public PhysicalPlan convertToPlan(String cursor) {
+    if (!cursor.startsWith(CURSOR_PREFIX)) {
+      throw new UnsupportedOperationException("Unsupported cursor");
+    }
+    try {
+      cursor = cursor.substring(CURSOR_PREFIX.length());
+      cursor = decompress(cursor);
+
+      // TODO Parse with ANTLR or serialize as JSON/XML
+      if (!cursor.startsWith("(Paginate,")) {
+        throw new UnsupportedOperationException("Unsupported cursor");
+      }
+      // TODO add checks for > 0
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      final int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10);
+
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      final int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10);
+
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      if (!cursor.startsWith("(Project,")) {
+        throw new UnsupportedOperationException("Unsupported cursor");
+      }
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      if (!cursor.startsWith("(namedParseExpressions,")) {
+        throw new UnsupportedOperationException("Unsupported cursor");
+      }
+
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      List<NamedExpression> namedParseExpressions = new ArrayList<>();
+      cursor = parseNamedExpressions(namedParseExpressions, cursor);
+
+      List<NamedExpression> projectList = new ArrayList<>();
+      if (!cursor.startsWith("(projectList,")) {
+        throw new UnsupportedOperationException("Unsupported cursor");
+      }
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      cursor = parseNamedExpressions(projectList, cursor);
+
+      if (!cursor.startsWith("(OpenSearchPagedIndexScan,")) {
+        throw new UnsupportedOperationException("Unsupported cursor");
+      }
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      var indexName = cursor.substring(0, cursor.indexOf(','));
+      cursor = cursor.substring(cursor.indexOf(',') + 1);
+      var scrollId = cursor.substring(0, cursor.indexOf(')'));
+      TableScanOperator scan = storageEngine.getTableScan(indexName, scrollId);
+
+      return new PaginateOperator(new ProjectOperator(scan, projectList, namedParseExpressions),
+          pageSize, currentPageIndex);
+    } catch (Exception e) {
+      throw new UnsupportedOperationException("Unsupported cursor", e);
+    }
+  }
+}

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.
    */