poc: scroll query request

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

@@ -46,6 +46,7 @@ public interface ExecutionEngine {
   class QueryResponse {
     private final Schema schema;
     private final List<ExprValue> results;
+    private final String cursor;
   }
 
   @Data

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.planner;
+
+import lombok.Getter;
+import lombok.Setter;
+
+/**
+ * The context used for planner.
+ */
+public class PlanContext {
+  @Getter
+  @Setter
+  private int fetchSize;
+
+  public PlanContext() {
+    this.fetchSize = 0;
+  }
+}

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

@@ -10,6 +10,7 @@
 
 import java.util.List;
 import lombok.RequiredArgsConstructor;
+import lombok.extern.java.Log;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
 import org.opensearch.sql.planner.logical.LogicalRelation;
@@ -39,15 +40,20 @@ public class Planner {
    * @param plan logical plan
    * @return optimal physical plan
    */
-  public PhysicalPlan plan(LogicalPlan plan) {
+  public PhysicalPlan plan(LogicalPlan plan, PlanContext planContext) {
     String tableName = findTableName(plan);
     if (isNullOrEmpty(tableName)) {
       return plan.accept(new DefaultImplementor<>(), null);
     }
 
     Table table = storageEngine.getTable(tableName);
     return table.implement(
-        table.optimize(optimize(plan)));
+        table.optimize(optimize(plan)),
+        planContext);
+  }
+
+  public PhysicalPlan plan(LogicalPlan plan) {
+    return plan(plan, new PlanContext());
   }
 
   private String findTableName(LogicalPlan plan) {

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java

@@ -32,6 +32,14 @@ public void open() {
     getChild().forEach(PhysicalPlan::open);
   }
 
+  public String getCursor() {
+    try {
+      return getChild().get(0).getCursor();
+    } catch (IndexOutOfBoundsException e) {
+      return getCursor();
+    }
+  }
+
   public void close() {
     getChild().forEach(PhysicalPlan::close);
   }

core/src/main/java/org/opensearch/sql/storage/Table.java

@@ -8,6 +8,7 @@
 
 import java.util.Map;
 import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.planner.PlanContext;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 
@@ -29,6 +30,8 @@ public interface Table {
    */
   PhysicalPlan implement(LogicalPlan plan);
 
+  PhysicalPlan implement(LogicalPlan plan, PlanContext planContext);
+
   /**
    * Optimize the {@link LogicalPlan} by storage engine rule.
    * The default optimize solution is no optimization.

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java

@@ -29,7 +29,9 @@
 import org.opensearch.sql.executor.ExecutionEngine.ExplainResponse;
 import org.opensearch.sql.legacy.metrics.MetricName;
 import org.opensearch.sql.legacy.metrics.Metrics;
+import org.opensearch.sql.planner.PlanContext;
 import org.opensearch.sql.opensearch.security.SecurityAccess;
+import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
 import org.opensearch.sql.protocol.response.QueryResult;
 import org.opensearch.sql.protocol.response.format.CsvResponseFormatter;
@@ -101,9 +103,10 @@ public RestChannelConsumer prepareRequest(SQLQueryRequest request, NodeClient no
     try {
       // For now analyzing and planning stage may throw syntax exception as well
       // which hints the fallback to legacy code is necessary here.
-      plan = sqlService.plan(
-                sqlService.analyze(
-                    sqlService.parse(request.getQuery())));
+      LogicalPlan logicalPlan = sqlService.analyze(sqlService.parse(request.getQuery()));
+      PlanContext planContext = new PlanContext();
+      planContext.setFetchSize(request.getFetchSize());
+      plan = sqlService.plan(logicalPlan, planContext);
     } catch (SyntaxCheckException e) {
       // When explain, print info log for what unsupported syntax is causing fallback to old engine
       if (request.isExplainRequest()) {
@@ -167,7 +170,7 @@ private ResponseListener<QueryResponse> createQueryResponseListener(RestChannel
       @Override
       public void onResponse(QueryResponse response) {
         sendResponse(channel, OK,
-            formatter.format(new QueryResult(response.getSchema(), response.getResults())));
+            formatter.format(new QueryResult(response.getSchema(), response.getResults(), response.getCursor())));
       }
 
       @Override

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlAction.java

@@ -149,7 +149,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
 
             // Route request to new query engine if it's supported already
             SQLQueryRequest newSqlRequest = new SQLQueryRequest(sqlRequest.getJsonContent(),
-                sqlRequest.getSql(), request.path(), request.params());
+                sqlRequest.getSql(), request.path(), request.params(), sqlRequest.fetchSize());
             RestChannelConsumer result = newSqlQueryHandler.prepareRequest(newSqlRequest, client);
             if (result != RestSQLQueryAction.NOT_SUPPORTED_YET) {
                 LOG.info("[{}] Request is handled by new SQL query engine", LogUtils.getRequestId());

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java

@@ -40,7 +40,7 @@ public void execute(PhysicalPlan physicalPlan, ResponseListener<QueryResponse> l
               result.add(plan.next());
             }
 
-            QueryResponse response = new QueryResponse(physicalPlan.schema(), result);
+            QueryResponse response = new QueryResponse(physicalPlan.schema(), result, plan.getCursor());
             listener.onResponse(response);
           } catch (Exception e) {
             listener.onFailure(e);

opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java

@@ -69,6 +69,11 @@ public List<PhysicalPlan> getChild() {
     return delegate.getChild();
   }
 
+  @Override
+  public String getCursor() {
+    return delegate.getCursor();
+  }
+
   @Override
   public boolean hasNext() {
     return delegate.hasNext();

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollCursorRequest.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.opensearch.request;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.action.search.SearchScrollRequest;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
+import org.opensearch.sql.opensearch.response.OpenSearchResponse;
+
+import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+/**
+ * OpenSearch scroll cursor request.
+ */
+@EqualsAndHashCode
+@Getter
+@ToString
+public class OpenSearchScrollCursorRequest implements OpenSearchRequest {
+
+  /** Default scroll context timeout in minutes. */
+  public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1L);
+
+  /** Index name. */
+  @EqualsAndHashCode.Exclude
+  @ToString.Exclude
+  private final OpenSearchExprValueFactory exprValueFactory;
+
+  /**
+   * Scroll id.
+   */
+  private String scrollId;
+
+  /** Search request source builder. */
+  private final SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+
+  public OpenSearchScrollCursorRequest(String scrollId, OpenSearchExprValueFactory exprValueFactory) {
+    this.scrollId = scrollId;
+    this.exprValueFactory = exprValueFactory;
+  }
+
+  @Override
+  public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
+                                   Function<SearchScrollRequest, SearchResponse> scrollAction) {
+    SearchResponse openSearchResponse;
+    openSearchResponse = scrollAction.apply(scrollRequest());
+
+    return new OpenSearchResponse(openSearchResponse, exprValueFactory);
+  }
+
+  @Override
+  public void clean(Consumer<String> cleanAction) {
+    try {
+      cleanAction.accept(getScrollId());
+    }
+    finally {
+
+    }
+  }
+
+  /**
+   * Generate OpenSearch scroll request by scroll id maintained.
+   *
+   * @return scroll request
+   */
+  public SearchScrollRequest scrollRequest() {
+    Objects.requireNonNull(scrollId, "Scroll id cannot be null");
+    return new SearchScrollRequest().scroll(DEFAULT_SCROLL_TIMEOUT).scrollId(scrollId);
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollQueryRequest.java

@@ -0,0 +1,89 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.opensearch.request;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.action.search.SearchScrollRequest;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
+import org.opensearch.sql.opensearch.response.OpenSearchResponse;
+
+import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+/**
+ * OpenSearch scroll search request. This has to be stateful because it needs to:
+ *
+ * <p>1) Accumulate search source builder when visiting logical plan to push down operation 2)
+ * Maintain scroll ID between calls to client search method
+ */
+@EqualsAndHashCode
+@Getter
+@ToString
+public class OpenSearchScrollQueryRequest implements OpenSearchRequest {
+
+  /** Default scroll context timeout in minutes. */
+  public static final TimeValue DEFAULT_SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1L);
+
+  /**
+   * {@link IndexName}.
+   */
+  private final IndexName indexName;
+
+  /** Index name. */
+  @EqualsAndHashCode.Exclude
+  @ToString.Exclude
+  private final OpenSearchExprValueFactory exprValueFactory;
+
+  /** Search request source builder. */
+  private final SearchSourceBuilder sourceBuilder;
+
+  public OpenSearchScrollQueryRequest(IndexName indexName, int size, OpenSearchExprValueFactory exprValueFactory) {
+    this.indexName = indexName;
+    this.sourceBuilder = new SearchSourceBuilder();
+    sourceBuilder.from(0);
+    sourceBuilder.size(size);
+    this.exprValueFactory = exprValueFactory;
+  }
+
+  public OpenSearchScrollQueryRequest(String indexName, int size, OpenSearchExprValueFactory exprValueFactory) {
+    this(new IndexName(indexName), size, exprValueFactory);
+  }
+
+  @Override
+  public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
+                                   Function<SearchScrollRequest, SearchResponse> scrollAction) {
+    SearchResponse openSearchResponse;
+    openSearchResponse = searchAction.apply(searchRequest());
+
+    return new OpenSearchResponse(openSearchResponse, exprValueFactory);
+  }
+
+  @Override
+  public void clean(Consumer<String> cleanAction) {
+    // do nothing
+  }
+
+  /**
+   * Generate OpenSearch search request.
+   *
+   * @return search request
+   */
+  public SearchRequest searchRequest() {
+    return new SearchRequest()
+        .indices(indexName.getIndexNames())
+        .scroll(DEFAULT_SCROLL_TIMEOUT)
+        .source(sourceBuilder);
+  }
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java

@@ -11,6 +11,7 @@
 import java.util.Iterator;
 import java.util.Map;
 import lombok.EqualsAndHashCode;
+import lombok.Getter;
 import lombok.ToString;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.search.SearchHits;
@@ -36,6 +37,12 @@ public class OpenSearchResponse implements Iterable<ExprValue> {
    */
   private final Aggregations aggregations;
 
+  /**
+   * Search scrollId result.
+   */
+  @Getter
+  private final String scrollId;
+
   /**
    * ElasticsearchExprValueFactory used to build ExprValue from search result.
    */
@@ -49,6 +56,7 @@ public OpenSearchResponse(SearchResponse searchResponse,
                             OpenSearchExprValueFactory exprValueFactory) {
     this.hits = searchResponse.getHits();
     this.aggregations = searchResponse.getAggregations();
+    this.scrollId = searchResponse.getScrollId();
     this.exprValueFactory = exprValueFactory;
   }
 
@@ -58,6 +66,7 @@ public OpenSearchResponse(SearchResponse searchResponse,
   public OpenSearchResponse(SearchHits hits, OpenSearchExprValueFactory exprValueFactory) {
     this.hits = hits;
     this.aggregations = null;
+    this.scrollId = null;
     this.exprValueFactory = exprValueFactory;
   }