Pagination, phase 1: Add unit tests for :opensearch module with coverage.

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.planner;
 
 import java.util.List;
+import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.data.model.ExprValue;
@@ -15,6 +16,7 @@
 import org.opensearch.sql.planner.physical.ProjectOperator;
 
 @RequiredArgsConstructor
+@EqualsAndHashCode(callSuper = false)
 public class PaginateOperator extends PhysicalPlan {
   @Getter
   private final PhysicalPlan input;

core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java

@@ -108,7 +108,4 @@ public boolean pushDownHighlight(LogicalHighlight highlight) {
   public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
     return visitor.visitTableScanBuilder(this, context);
   }
-
-  public void pushDownOffset(int i) {
-  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/ContinueScrollRequest.java → opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinueScrollRequest.java

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.request;
 
 import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.DEFAULT_SCROLL_TIMEOUT;
 
@@ -17,9 +17,9 @@
 import org.opensearch.action.search.SearchScrollRequest;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
-import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 
+@EqualsAndHashCode
 public class ContinueScrollRequest implements OpenSearchRequest {
   final String initialScrollId;
 
@@ -39,9 +39,7 @@ public ContinueScrollRequest(String scrollId, OpenSearchExprValueFactory exprVal
   @Override
   public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
                                    Function<SearchScrollRequest, SearchResponse> scrollAction) {
-    SearchResponse openSearchResponse;
-
-    openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)
+    SearchResponse openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId)
         .scroll(DEFAULT_SCROLL_TIMEOUT));
 
     // TODO if terminated_early - something went wrong, e.g. no scroll returned.

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/InitialPageRequestBuilder.java → opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java

@@ -3,14 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.request;
 
 import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT;
 
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 import lombok.Getter;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.index.query.QueryBuilder;
@@ -22,8 +21,6 @@
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
-import org.opensearch.sql.opensearch.request.OpenSearchRequest;
-import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest;
 import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser;
 
 public class InitialPageRequestBuilder implements PagedRequestBuilder {
@@ -40,15 +37,16 @@ public class InitialPageRequestBuilder implements PagedRequestBuilder {
    * @param settings other settings
    * @param exprValueFactory value factory
    */
+  // TODO accept indexName as string (same way as `OpenSearchRequestBuilder` does)?
   public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName,
                                    int pageSize,
-                                   Settings settings,
+                                   Settings settings, // TODO: settings are not used - refactor?
                                    OpenSearchExprValueFactory exprValueFactory) {
     this.indexName = indexName;
-    this.sourceBuilder = new SearchSourceBuilder();
     this.exprValueFactory = exprValueFactory;
     this.querySize = pageSize;
-    sourceBuilder.from(0)
+    this.sourceBuilder = new SearchSourceBuilder()
+        .from(0)
         .size(querySize)
         .timeout(DEFAULT_QUERY_TIMEOUT);
   }
@@ -85,9 +83,8 @@ public void pushDownHighlight(String field, Map<String, Literal> arguments) {
    * Push down project expression to OpenSearch.
    */
   public void pushDownProjects(Set<ReferenceExpression> projects) {
-    final Set<String> projectsSet =
-        projects.stream().map(ReferenceExpression::getAttr).collect(Collectors.toSet());
-    sourceBuilder.fetchSource(projectsSet.toArray(new String[0]), new String[0]);
+    sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr)
+        .distinct().toArray(String[]::new), new String[0]);
   }
 
   public void pushTypeMapping(Map<String, ExprType> typeMapping) {

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

@@ -6,6 +6,8 @@
 
 package org.opensearch.sql.opensearch.request;
 
+import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT;
+
 import com.google.common.annotations.VisibleForTesting;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -15,7 +17,6 @@
 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.SearchHits;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
@@ -32,11 +33,6 @@
 @ToString
 public class OpenSearchQueryRequest implements OpenSearchRequest {
 
-  /**
-   * Default query timeout in minutes.
-   */
-  public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L);
-
   /**
    * {@link OpenSearchRequest.IndexName}.
    */

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

@@ -42,6 +42,7 @@
 @EqualsAndHashCode
 @Getter
 @ToString
+// TODO make an interface which defines all pushDown functions?
 public class OpenSearchRequestBuilder {
 
   /**

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/PagedRequestBuilder.java → opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java

@@ -3,10 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
-
-import org.opensearch.sql.opensearch.request.OpenSearchRequest;
-import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest;
+package org.opensearch.sql.opensearch.request;
 
 public interface PagedRequestBuilder {
   OpenSearchRequest build();

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/SubsequentPageRequestBuilder.java → opensearch/src/main/java/org/opensearch/sql/opensearch/request/SubsequentPageRequestBuilder.java

@@ -3,25 +3,22 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.request;
 
+import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
-import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 
 @RequiredArgsConstructor
 public class SubsequentPageRequestBuilder implements PagedRequestBuilder {
-  final OpenSearchRequest.IndexName indexName;
-  final String scrollId;
-  final OpenSearchExprValueFactory exprValueFactory;
+
+  @Getter
+  private final OpenSearchRequest.IndexName indexName;
+  private final String scrollId;
+  private final OpenSearchExprValueFactory exprValueFactory;
 
   @Override
   public OpenSearchRequest build() {
     return new ContinueScrollRequest(scrollId, exprValueFactory);
   }
-
-  @Override
-  public OpenSearchRequest.IndexName getIndexName() {
-    return indexName;
-  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java

@@ -18,10 +18,14 @@
 import org.opensearch.sql.opensearch.planner.physical.ADOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLOperator;
+import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
 import org.opensearch.sql.opensearch.request.system.OpenSearchDescribeIndexRequest;
+import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan;
 import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder;
+import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan;
+import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScanBuilder;
 import org.opensearch.sql.planner.DefaultImplementor;
 import org.opensearch.sql.planner.logical.LogicalAD;
 import org.opensearch.sql.planner.logical.LogicalML;
@@ -132,7 +136,7 @@ public TableScanBuilder createPagedScanBuilder(int pageSize) {
     var requestBuilder = new InitialPageRequestBuilder(indexName, pageSize,
         settings, new OpenSearchExprValueFactory(getFieldTypes()));
     var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder);
-    return new OpenSearchPagedScanBuilder(indexScan);
+    return new OpenSearchPagedIndexScanBuilder(indexScan);
   }
 
   @VisibleForTesting

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchStorageEngine.java

@@ -14,6 +14,8 @@
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
+import org.opensearch.sql.opensearch.request.SubsequentPageRequestBuilder;
+import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan;
 import org.opensearch.sql.opensearch.storage.system.OpenSearchSystemIndex;
 import org.opensearch.sql.storage.StorageEngine;
 import org.opensearch.sql.storage.Table;
@@ -39,6 +41,7 @@ public Table getTable(DataSourceSchemaName dataSourceSchemaName, String name) {
 
   @Override
   public TableScanOperator getTableScan(String indexName, String scrollId) {
+    // TODO call `getTable` here?
     var index = new OpenSearchIndex(client, settings, indexName);
     var requestBuilder = new SubsequentPageRequestBuilder(
         new OpenSearchRequest.IndexName(indexName),

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexScan.java → opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java

@@ -4,7 +4,7 @@
  */
 
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.storage.scan;
 
 import java.util.Collections;
 import java.util.Iterator;

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java

@@ -17,7 +17,6 @@
 import org.opensearch.sql.expression.aggregation.NamedAggregator;
 import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer;
 import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser;
-import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan;
 import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder;
 import org.opensearch.sql.planner.logical.LogicalAggregation;
 import org.opensearch.sql.planner.logical.LogicalSort;

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java

@@ -8,7 +8,6 @@
 import com.google.common.annotations.VisibleForTesting;
 import lombok.EqualsAndHashCode;
 import org.opensearch.sql.expression.ReferenceExpression;
-import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan;
 import org.opensearch.sql.planner.logical.LogicalAggregation;
 import org.opensearch.sql.planner.logical.LogicalFilter;
 import org.opensearch.sql.planner.logical.LogicalHighlight;
@@ -92,11 +91,6 @@ public boolean pushDownProject(LogicalProject project) {
     return delegate.pushDownProject(project);
   }
 
-  @Override
-  public boolean pushDownOffset(int i) {
-    return delegate.pushDownOffset(i);
-  }
-
   @Override
   public boolean pushDownHighlight(LogicalHighlight highlight) {
     return delegate.pushDownHighlight(highlight);

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java

@@ -21,7 +21,6 @@
 import org.opensearch.sql.expression.NamedExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer;
-import org.opensearch.sql.opensearch.storage.OpenSearchIndexScan;
 import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder;
 import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder;
 import org.opensearch.sql.planner.logical.LogicalFilter;
@@ -38,10 +37,6 @@
  */
 @VisibleForTesting
 class OpenSearchIndexScanQueryBuilder extends TableScanBuilder {
-  @Override
-  public void pushDownOffset(int i) {
-    indexScan.getRequestBuilder().getSourceBuilder().from(i);
-  }
 
   /** OpenSearch index scan to be optimized. */
   @EqualsAndHashCode.Include

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedIndexScan.java → opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java

@@ -3,15 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.storage.scan;
 
 import java.util.Collections;
 import java.util.Iterator;
 import lombok.EqualsAndHashCode;
 import lombok.ToString;
+import org.apache.commons.lang3.NotImplementedException;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
+import org.opensearch.sql.opensearch.request.PagedRequestBuilder;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 import org.opensearch.sql.storage.TableScanOperator;
 
@@ -33,7 +35,7 @@ public OpenSearchPagedIndexScan(OpenSearchClient client,
 
   @Override
   public String explain() {
-    throw new RuntimeException("Implement OpenSearchPagedIndexScan.explain");
+    throw new NotImplementedException("Implement OpenSearchPagedIndexScan.explain");
   }
 
   @Override

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchPagedScanBuilder.java → opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java

@@ -3,26 +3,25 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package org.opensearch.sql.opensearch.storage;
+package org.opensearch.sql.opensearch.storage.scan;
 
 import lombok.EqualsAndHashCode;
 import org.opensearch.sql.storage.TableScanOperator;
 import org.opensearch.sql.storage.read.TableScanBuilder;
 
 /**
- * Builder for a paged opensearch request.
- * Override pushDown* methods from TableScaneBuilder as more features
+ * Builder for a paged OpenSearch request.
+ * Override pushDown* methods from TableScanBuilder as more features
  * support pagination.
  */
-public class OpenSearchPagedScanBuilder extends TableScanBuilder {
+public class OpenSearchPagedIndexScanBuilder extends TableScanBuilder {
   @EqualsAndHashCode.Include
   OpenSearchPagedIndexScan indexScan;
 
-  public OpenSearchPagedScanBuilder(OpenSearchPagedIndexScan indexScan) {
+  public OpenSearchPagedIndexScanBuilder(OpenSearchPagedIndexScan indexScan) {
     this.indexScan = indexScan;
   }
 
-
   @Override
   public TableScanOperator build() {
     return indexScan;