Support pagination in V2 engine, phase 1 (opensearch-project#1497)

* Support pagination in V2 engine, phase 1 (#226) * Fixing integration tests broken during POC Signed-off-by: MaxKsyunz <[email protected]> * Comment to clarify an exception. Signed-off-by: MaxKsyunz <[email protected]> * Add support for paginated scroll request, first page. Implement PaginatedPlanCache.convertToPlan for second page to work. Signed-off-by: MaxKsyunz <[email protected]> * Progress on paginated scroll request, subsequent page. Signed-off-by: MaxKsyunz <[email protected]> * Move `ExpressionSerializer` from `opensearch` to `core`. Signed-off-by: Yury-Fridlyand <[email protected]> * Rename `Cursor` `asString` to `toString`. Signed-off-by: Yury-Fridlyand <[email protected]> * Disable scroll cleaning. Signed-off-by: Yury-Fridlyand <[email protected]> * Add full cursor serialization and deserialization. Signed-off-by: Yury-Fridlyand <[email protected]> * Misc fixes. Signed-off-by: Yury-Fridlyand <[email protected]> * Further work on pagination. * Added push down page size from `LogicalPaginate` to `LogicalRelation`. * Improved cursor encoding and decoding. * Added cursor compression. * Fixed issuing `SearchScrollRequest`. * Fixed returning last empty page. * Minor code grooming/commenting. Signed-off-by: Yury-Fridlyand <[email protected]> * Pagination fix for empty indices. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix error reporting on wrong cursor. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor comments and error reporting improvement. Signed-off-by: Yury-Fridlyand <[email protected]> * Add an end-to-end integration test. Signed-off-by: Yury-Fridlyand <[email protected]> * Add `explain` request handlers. Signed-off-by: Yury-Fridlyand <[email protected]> * Add IT for explain. Signed-off-by: Yury-Fridlyand <[email protected]> * Address issues flagged by checkstyle build step (#229) Signed-off-by: MaxKsyunz <[email protected]> * Pagination, phase 1: Add unit tests for `:core` module with coverage. (#230) * Add unit tests for `:core` module with coverage. Uncovered: `toCursor`, because it is will be changed soon. Signed-off-by: Yury-Fridlyand <[email protected]> * Pagination, phase 1: Add unit tests for SQL module with coverage. (#239) * Add unit tests for SQL module with coverage. Signed-off-by: Yury-Fridlyand <[email protected]> * Update sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: GabeFernandez310 <[email protected]> --------- Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: GabeFernandez310 <[email protected]> * Pagination, phase 1: Add unit tests for `:opensearch` module with coverage. (#233) * Add UT for `:opensearch` module with full coverage, except `toCursor`. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix checkstyle. Signed-off-by: Yury-Fridlyand <[email protected]> --------- Signed-off-by: Yury-Fridlyand <[email protected]> * Fix the merges. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix explain. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix scroll cleaning. Signed-off-by: Yury-Fridlyand <[email protected]> * Store `TotalHits` and use it to report `total` in response. Signed-off-by: Yury-Fridlyand <[email protected]> * Add missing UT for `:protocol` module. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix PPL UTs damaged in f4ea4ad. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor checkstyle fixes. Signed-off-by: Yury-Fridlyand <[email protected]> * Fallback to v1 engine for pagination (#245) * Pagination fallback integration tests. Signed-off-by: MaxKsyunz <[email protected]> * Add UT with coverage for `toCursor` serialization. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix broken tests in `legacy`. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix getting `total` from non-paged requests and from queries without `FROM` clause. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix scroll cleaning. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix cursor request processing. Signed-off-by: Yury-Fridlyand <[email protected]> * Update ITs. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix (again) TotalHits feature. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix typo in prometheus config. Signed-off-by: Yury-Fridlyand <[email protected]> * Recover commented logging. Signed-off-by: Yury-Fridlyand <[email protected]> * Move `test_pagination_blackbox` to a separate class and add logging. Signed-off-by: Yury-Fridlyand <[email protected]> * Address some PR feedbacks: rename some classes and revert unnecessary whitespace changed. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor commenting. Signed-off-by: Yury-Fridlyand <[email protected]> * Address PR comments. * Add javadocs * Renames * Cleaning up some comments * Remove unused code * Speed up IT Signed-off-by: Yury-Fridlyand <[email protected]> * Minor missing changes. Signed-off-by: Yury-Fridlyand <[email protected]> * Integration tests for fetch_size, max_result_window, and query.size_limit (#248) Signed-off-by: MaxKsyunz <[email protected]> * Remove `PaginatedQueryService`, extend `QueryService` to hold two planners and use them. Signed-off-by: Yury-Fridlyand <[email protected]> * Move push down functions from request builders to a new interface. Signed-off-by: Yury-Fridlyand <[email protected]> * Some file moves. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor clean-up according to PR review. Signed-off-by: Yury-Fridlyand <[email protected]> --------- Signed-off-by: MaxKsyunz <[email protected]> Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: MaxKsyunz <[email protected]> Co-authored-by: GabeFernandez310 <[email protected]> Co-authored-by: Max Ksyunz <[email protected]> * Make scroll timeout configurable. Signed-off-by: Yury-Fridlyand <[email protected]> * Fix IT to set cursor keep alive parameter. Signed-off-by: Yury-Fridlyand <[email protected]> * Remove `QueryId.None`. Signed-off-by: Yury-Fridlyand <[email protected]> * Rename according to PR feedback. Signed-off-by: Yury-Fridlyand <[email protected]> * Remove default implementations of `PushDownRequestBuilder`. Signed-off-by: Yury-Fridlyand <[email protected]> * Merge paginated plan optimizer into the regular optimizer. (opensearch-project#1516) Merge paginated plan optimizer into the regular optimizer. --------- Signed-off-by: MaxKsyunz <[email protected]> Co-authored-by: Yury-Fridlyand <[email protected]> * Complete rework on serialization and deserialization. (opensearch-project#1498) Signed-off-by: Yury-Fridlyand <[email protected]> * Resolve merge conflicts and fix tests. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor cleanup. Signed-off-by: Yury-Fridlyand <[email protected]> * Minor cleanup - missing changes for the previous commit. Signed-off-by: Yury-Fridlyand <[email protected]> * Remove paginate operator (opensearch-project#1528) * Remove PaginateOperator class since it is no longer used. --------- Signed-off-by: MaxKsyunz <[email protected]> * Remove `PaginatedPlan` - move logic to `QueryPlan`. Signed-off-by: Yury-Fridlyand <[email protected]> * Remove default implementations from `SerializablePlan`. Signed-off-by: Yury-Fridlyand <[email protected]> * Add a doc. Signed-off-by: Yury-Fridlyand <[email protected]> * Update design graphs. Signed-off-by: Yury-Fridlyand <[email protected]> * More fixes for merge from upstream/main. Signed-off-by: MaxKsyunz <[email protected]> --------- Signed-off-by: MaxKsyunz <[email protected]> Signed-off-by: Yury-Fridlyand <[email protected]> Co-authored-by: MaxKsyunz <[email protected]> Co-authored-by: GabeFernandez310 <[email protected]> Co-authored-by: Max Ksyunz <[email protected]>
Bit-Quill · Apr 28, 2023 · d9ed806 · d9ed806
1 parent dae16bb
commit d9ed806
Show file tree

Hide file tree

Showing 3 changed files with 382 additions and 5 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java
@@ -8,6 +8,9 @@
 
 package org.opensearch.sql.executor.execution;
 
+import java.util.Optional;
+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;
@@ -33,25 +36,51 @@ public class QueryPlan extends AbstractPlan {
 
   protected final ResponseListener<ExecutionEngine.QueryResponse> listener;
 
-  /** constructor. */
+  protected final Optional<Integer> pageSize;
+
+  /** Constructor. */
+  public QueryPlan(
+      QueryId queryId,
+      UnresolvedPlan plan,
+      QueryService queryService,
+      ResponseListener<ExecutionEngine.QueryResponse> listener) {
+    super(queryId);
+    this.plan = plan;
+    this.queryService = queryService;
+    this.listener = listener;
+    this.pageSize = Optional.empty();
+  }
+
+  /** Constructor with page size. */
   public QueryPlan(
       QueryId queryId,
       UnresolvedPlan plan,
+      int pageSize,
       QueryService queryService,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
     super(queryId);
     this.plan = plan;
     this.queryService = queryService;
     this.listener = listener;
+    this.pageSize = Optional.of(pageSize);
   }
 
   @Override
   public void execute() {
-    queryService.execute(plan, listener);
+    if (pageSize.isPresent()) {
+      queryService.execute(new Paginate(pageSize.get(), plan), listener);
+    } else {
+      queryService.execute(plan, listener);
+    }
   }
 
   @Override
   public void explain(ResponseListener<ExecutionEngine.ExplainResponse> listener) {
-    queryService.explain(plan, listener);
+    if (pageSize.isPresent()) {
+      listener.onFailure(new NotImplementedException(
+          "`explain` feature for paginated requests is not implemented yet."));
+    } else {
+      queryService.explain(plan, listener);
+    }
   }
 }
diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java
@@ -8,21 +8,30 @@
 
 package org.opensearch.sql.executor.execution;
 
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import org.apache.commons.lang3.NotImplementedException;
+import org.junit.jupiter.api.DisplayNameGeneration;
+import org.junit.jupiter.api.DisplayNameGenerator;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.common.response.ResponseListener;
+import org.opensearch.sql.executor.DefaultExecutionEngine;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
 import org.opensearch.sql.executor.QueryService;
 
 @ExtendWith(MockitoExtension.class)
+@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
 class QueryPlanTest {
 
   @Mock
@@ -41,18 +50,70 @@ class QueryPlanTest {
   private ResponseListener<ExecutionEngine.QueryResponse> queryListener;
 
   @Test
-  public void execute() {
+  public void execute_no_page_size() {
     QueryPlan query = new QueryPlan(queryId, plan, queryService, queryListener);
     query.execute();
 
     verify(queryService, times(1)).execute(any(), any());
   }
 
   @Test
-  public void explain() {
+  public void explain_no_page_size() {
     QueryPlan query = new QueryPlan(queryId, plan, queryService, queryListener);
     query.explain(explainListener);
 
     verify(queryService, times(1)).explain(plan, explainListener);
   }
+
+  @Test
+  public void can_execute_paginated_plan() {
+    var listener = new ResponseListener<ExecutionEngine.QueryResponse>() {
+      @Override
+      public void onResponse(ExecutionEngine.QueryResponse response) {
+        assertNotNull(response);
+      }
+
+      @Override
+      public void onFailure(Exception e) {
+        fail();
+      }
+    };
+    var plan = new QueryPlan(QueryId.queryId(), mock(UnresolvedPlan.class), 10,
+        queryService, listener);
+    plan.execute();
+  }
+
+  @Test
+  // Same as previous test, but with incomplete QueryService
+  public void can_handle_error_while_executing_plan() {
+    var listener = new ResponseListener<ExecutionEngine.QueryResponse>() {
+      @Override
+      public void onResponse(ExecutionEngine.QueryResponse response) {
+        fail();
+      }
+
+      @Override
+      public void onFailure(Exception e) {
+        assertNotNull(e);
+      }
+    };
+    var plan = new QueryPlan(QueryId.queryId(), mock(UnresolvedPlan.class), 10,
+        new QueryService(null, new DefaultExecutionEngine(), null), listener);
+    plan.execute();
+  }
+
+  @Test
+  public void explain_is_not_supported_for_pagination() {
+    new QueryPlan(null, null, 0, null, null).explain(new ResponseListener<>() {
+        @Override
+        public void onResponse(ExecutionEngine.ExplainResponse response) {
+          fail();
+        }
+
+        @Override
+        public void onFailure(Exception e) {
+          assertTrue(e instanceof NotImplementedException);
+        }
+      });
+  }
 }