From 68db7063ea490ce55373a597014d5ddb29915915 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 28 Apr 2023 15:21:18 -0700 Subject: [PATCH 01/38] Remove OpenSearchPagedIndexScan and related classes. Unit tests pass, some integration tests fail. Signed-off-by: MaxKsyunz --- .../optimizer/LogicalPlanOptimizer.java | 3 +- .../planner/optimizer/PushDownPageSize.java | 49 ++++ .../rule/CreatePagingTableScanBuilder.java | 72 ------ .../org/opensearch/sql/storage/Table.java | 4 - .../sql/storage/read/TableScanBuilder.java | 5 +- .../optimizer/LogicalPlanOptimizerTest.java | 111 ++++----- .../CreatePagingTableScanBuilderTest.java | 46 ---- .../org/opensearch/sql/storage/TableTest.java | 25 -- .../request/ContinuePageRequest.java | 8 +- .../request/ContinuePageRequestBuilder.java | 56 +++-- .../request/InitialPageRequestBuilder.java | 114 ---------- .../request/OpenSearchRequestBuilder.java | 135 ++++++----- .../request/OpenSearchScrollRequest.java | 3 +- .../request/PagedRequestBuilder.java | 12 - .../request/PushDownRequestBuilder.java | 18 +- .../opensearch/storage/OpenSearchIndex.java | 34 +-- .../storage/scan/OpenSearchIndexScan.java | 107 ++++++--- ...OpenSearchIndexScanAggregationBuilder.java | 63 +++-- .../scan/OpenSearchIndexScanBuilder.java | 42 ++-- .../scan/OpenSearchIndexScanQueryBuilder.java | 58 ++--- .../scan/OpenSearchPagedIndexScan.java | 115 ---------- .../scan/OpenSearchPagedIndexScanBuilder.java | 29 --- .../storage/scan/PushDownTranslator.java | 35 +++ .../client/OpenSearchNodeClientTest.java | 10 +- .../client/OpenSearchRestClientTest.java | 7 +- .../OpenSearchExecutionEngineTest.java | 80 +++---- .../OpenSearchExecutionProtectorTest.java | 43 ++-- .../ContinuePageRequestBuilderTest.java | 37 +-- .../InitialPageRequestBuilderTest.java | 122 ---------- .../request/OpenSearchRequestBuilderTest.java | 50 ++-- .../storage/OpenSearchIndexTest.java | 103 +++------ .../OpenSearchIndexScanOptimizationTest.java | 69 ++---- .../storage/scan/OpenSearchIndexScanTest.java | 78 +++---- .../scan/OpenSearchPagedIndexScanTest.java | 215 ------------------ 34 files changed, 634 insertions(+), 1324 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java delete mode 100644 core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java delete mode 100644 core/src/test/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilderTest.java delete mode 100644 core/src/test/java/org/opensearch/sql/storage/TableTest.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java delete mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java delete mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java index afe86d0cb1..0f6b3e5195 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java @@ -13,7 +13,6 @@ import java.util.List; import java.util.stream.Collectors; import org.opensearch.sql.planner.logical.LogicalPlan; -import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder; import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter; import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort; import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; @@ -52,7 +51,7 @@ public static LogicalPlanOptimizer create() { * Phase 2: Transformations that rely on data source push down capability */ new CreateTableScanBuilder(), - new CreatePagingTableScanBuilder(), + new PushDownPageSize(), TableScanPushDown.PUSH_DOWN_FILTER, TableScanPushDown.PUSH_DOWN_AGGREGATION, TableScanPushDown.PUSH_DOWN_SORT, diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java new file mode 100644 index 0000000000..47999af37b --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner.optimizer; + +import com.facebook.presto.matching.Captures; +import com.facebook.presto.matching.Pattern; +import java.util.ArrayDeque; +import java.util.Deque; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.storage.read.TableScanBuilder; + +public class PushDownPageSize implements Rule { + TableScanBuilder builder; + @Override + public Pattern pattern() { + return Pattern.typeOf(LogicalPaginate.class).matching(this::findTableScanBuilder); + } + + @Override + public LogicalPlan apply(LogicalPaginate plan, Captures captures) { + + if (!builder.pushDownPageSize(plan)) { + throw new IllegalStateException("Failed to push down LogicalPaginate"); + } + return plan.getChild().get(0); + } + + private boolean findTableScanBuilder(LogicalPaginate logicalPaginate) { + Deque plans = new ArrayDeque<>(); + plans.add(logicalPaginate); + do { + final var plan = plans.removeFirst(); + final var children = plan.getChild(); + if (children.stream().anyMatch(TableScanBuilder.class::isInstance)) { + if (children.size() > 1) { + throw new UnsupportedOperationException( + "Unsupported plan: relation operator cannot have siblings"); + } + builder = (TableScanBuilder) children.get(0); + } + plans.addAll(children); + } while (!plans.isEmpty()); + return builder != null; + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java deleted file mode 100644 index c635400c33..0000000000 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java +++ /dev/null @@ -1,72 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.planner.optimizer.rule; - -import com.facebook.presto.matching.Captures; -import com.facebook.presto.matching.Pattern; -import java.util.ArrayDeque; -import java.util.Deque; -import java.util.List; -import lombok.Getter; -import lombok.experimental.Accessors; -import org.opensearch.sql.planner.logical.LogicalPaginate; -import org.opensearch.sql.planner.logical.LogicalPlan; -import org.opensearch.sql.planner.logical.LogicalRelation; -import org.opensearch.sql.planner.optimizer.Rule; - -/** - * Rule to create a paged TableScanBuilder in pagination request. - */ -public class CreatePagingTableScanBuilder implements Rule { - /** Capture the table inside matched logical paginate operator. */ - private LogicalPlan relationParent = null; - /** Pattern that matches logical relation operator. */ - @Accessors(fluent = true) - @Getter - private final Pattern pattern; - - /** - * Constructor. - */ - public CreatePagingTableScanBuilder() { - this.pattern = Pattern.typeOf(LogicalPaginate.class).matching(this::findLogicalRelation); - } - - /** - * Finds an instance of LogicalRelation and saves a reference in relationParent variable. - * @param logicalPaginate An instance of LogicalPaginate - * @return true if {@link LogicalRelation} node was found among the descendents of - * {@link this.logicalPaginate}, false otherwise. - */ - private boolean findLogicalRelation(LogicalPaginate logicalPaginate) { - Deque plans = new ArrayDeque<>(); - plans.add(logicalPaginate); - do { - final var plan = plans.removeFirst(); - final var children = plan.getChild(); - if (children.stream().anyMatch(LogicalRelation.class::isInstance)) { - if (children.size() > 1) { - throw new UnsupportedOperationException( - "Unsupported plan: relation operator cannot have siblings"); - } - relationParent = plan; - return true; - } - plans.addAll(children); - } while (!plans.isEmpty()); - return false; - } - - - @Override - public LogicalPlan apply(LogicalPaginate plan, Captures captures) { - var logicalRelation = (LogicalRelation) relationParent.getChild().get(0); - var scan = logicalRelation.getTable().createPagedScanBuilder(plan.getPageSize()); - relationParent.replaceChildPlans(List.of(scan)); - - return plan.getChild().get(0); - } -} diff --git a/core/src/main/java/org/opensearch/sql/storage/Table.java b/core/src/main/java/org/opensearch/sql/storage/Table.java index 0194f1d03e..fc1def5a2e 100644 --- a/core/src/main/java/org/opensearch/sql/storage/Table.java +++ b/core/src/main/java/org/opensearch/sql/storage/Table.java @@ -100,8 +100,4 @@ default StreamingSource asStreamingSource() { throw new UnsupportedOperationException(); } - default TableScanBuilder createPagedScanBuilder(int pageSize) { - var error = String.format("'%s' does not support pagination", getClass().toString()); - throw new UnsupportedOperationException(error); - } } diff --git a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java index 9af66e219f..713b428546 100644 --- a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java +++ b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java @@ -11,6 +11,7 @@ import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; import org.opensearch.sql.planner.logical.LogicalProject; @@ -28,7 +29,7 @@ public abstract class TableScanBuilder extends LogicalPlan { /** * Construct and initialize children to empty list. */ - public TableScanBuilder() { + protected TableScanBuilder() { super(Collections.emptyList()); } @@ -116,6 +117,8 @@ public boolean pushDownNested(LogicalNested nested) { return false; } + public boolean pushDownPageSize(LogicalPaginate paginate) { return false; } + @Override public R accept(LogicalPlanNodeVisitor visitor, C context) { return visitor.visitTableScanBuilder(this, context); diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index 543b261d9e..3217252800 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -6,35 +6,7 @@ package org.opensearch.sql.planner.optimizer; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; -import static org.opensearch.sql.data.model.ExprValueUtils.longValue; -import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; -import static org.opensearch.sql.data.type.ExprCoreType.LONG; -import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.limit; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.nested; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.paginate; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.write; - import com.google.common.collect.ImmutableList; -import java.util.Collections; -import java.util.List; -import java.util.Map; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -52,14 +24,27 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; import org.opensearch.sql.planner.logical.LogicalRelation; -import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder; -import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; +import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; +import static org.opensearch.sql.data.model.ExprValueUtils.longValue; +import static org.opensearch.sql.data.type.ExprCoreType.*; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; + @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class LogicalPlanOptimizerTest { @@ -70,13 +55,9 @@ class LogicalPlanOptimizerTest { @Spy private TableScanBuilder tableScanBuilder; - @Spy - private TableScanBuilder pagedTableScanBuilder; - @BeforeEach void setUp() { lenient().when(table.createScanBuilder()).thenReturn(tableScanBuilder); - lenient().when(table.createPagedScanBuilder(anyInt())).thenReturn(pagedTableScanBuilder); } /** @@ -344,45 +325,46 @@ public PhysicalPlan implement(LogicalPlan plan) { @Test void paged_table_scan_builder_support_project_push_down_can_apply_its_rule() { + when(tableScanBuilder.pushDownPageSize(any())).thenReturn(true); var relation = relation("schema", table); + var optimized = LogicalPlanOptimizer.create().optimize(paginate(project(relation), 4)); + verify(tableScanBuilder).pushDownPageSize(any()); - assertEquals( - project(pagedTableScanBuilder), - LogicalPlanOptimizer.create().optimize(paginate(project(relation), 4))); + assertEquals(project(tableScanBuilder), optimized); } - @Test - void push_page_size_noop_if_no_relation() { - var paginate = new LogicalPaginate(42, List.of(project(values()))); - assertEquals(paginate, LogicalPlanOptimizer.create().optimize(paginate)); + void push_down_page_size_multiple_children() { + var relation = relation("schema", table); + var twoChildrenPlan = new LogicalPlan(List.of(relation, relation)) { + @Override + public R accept(LogicalPlanNodeVisitor visitor, C context) { + return null; + } + }; + var queryPlan = paginate(twoChildrenPlan, 4); + var optimizer = LogicalPlanOptimizer.create(); + final var exception = assertThrows(UnsupportedOperationException.class, + () -> optimizer.optimize(queryPlan)); + assertEquals("Unsupported plan: relation operator cannot have siblings", exception.getMessage()); } @Test - void pagination_optimizer_simple_query() { - var projectPlan = project(relation("schema", table), DSL.named(DSL.ref("intV", INTEGER))); - - var optimizer = new LogicalPlanOptimizer( - List.of(new CreateTableScanBuilder(), new CreatePagingTableScanBuilder())); - - { - optimizer.optimize(projectPlan); - verify(table).createScanBuilder(); - verify(table, never()).createPagedScanBuilder(anyInt()); - } + void push_down_page_size_push_failed() { + when(tableScanBuilder.pushDownPageSize(any())).thenReturn(false); + + var queryPlan = paginate(project(relation("schema", table)), 4); + var optimizer = LogicalPlanOptimizer.create(); + final var exception = assertThrows(IllegalStateException.class, + () -> optimizer.optimize(queryPlan)); + assertEquals("Failed to push down LogicalPaginate", exception.getMessage()); } @Test - void pagination_optimizer_paged_query() { - var relation = new LogicalRelation("schema", table); - var projectPlan = project(relation, DSL.named(DSL.ref("intV", INTEGER))); - var pagedPlan = new LogicalPaginate(10, List.of(projectPlan)); - - var optimizer = new LogicalPlanOptimizer( - List.of(new CreateTableScanBuilder(), new CreatePagingTableScanBuilder())); - var optimized = optimizer.optimize(pagedPlan); - verify(table).createPagedScanBuilder(anyInt()); + void push_page_size_noop_if_no_relation() { + var paginate = new LogicalPaginate(42, List.of(project(values()))); + assertEquals(paginate, LogicalPlanOptimizer.create().optimize(paginate)); } @Test @@ -394,19 +376,18 @@ void push_page_size_noop_if_no_sub_plans() { @Test void table_scan_builder_support_offset_push_down_can_apply_its_rule() { - when(table.createPagedScanBuilder(anyInt())).thenReturn(pagedTableScanBuilder); + when(tableScanBuilder.pushDownPageSize(any())).thenReturn(true); var relation = new LogicalRelation("schema", table); var optimized = LogicalPlanOptimizer.create() .optimize(new LogicalPaginate(42, List.of(project(relation)))); // `optimized` structure: LogicalPaginate -> LogicalProject -> TableScanBuilder // LogicalRelation replaced by a TableScanBuilder instance - assertEquals(project(pagedTableScanBuilder), optimized); + assertEquals(project(tableScanBuilder), optimized); } private LogicalPlan optimize(LogicalPlan plan) { final LogicalPlanOptimizer optimizer = LogicalPlanOptimizer.create(); - final LogicalPlan optimize = optimizer.optimize(plan); - return optimize; + return optimizer.optimize(plan); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilderTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilderTest.java deleted file mode 100644 index 79c7b55c60..0000000000 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilderTest.java +++ /dev/null @@ -1,46 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.planner.optimizer.rule; - -import static com.facebook.presto.matching.DefaultMatcher.DEFAULT_MATCHER; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.paginate; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; - -import java.util.List; -import org.junit.jupiter.api.BeforeEach; -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.planner.logical.LogicalPlan; -import org.opensearch.sql.storage.Table; - -@ExtendWith(MockitoExtension.class) -class CreatePagingTableScanBuilderTest { - - @Mock - LogicalPlan multiRelationPaginate; - - @Mock - Table table; - - @BeforeEach - public void setUp() { - when(multiRelationPaginate.getChild()) - .thenReturn( - List.of(relation("t1", table), relation("t2", table))); - } - - @Test - void throws_when_mutliple_children() { - final var pattern = new CreatePagingTableScanBuilder().pattern(); - final var plan = paginate(multiRelationPaginate, 42); - assertThrows(UnsupportedOperationException.class, - () -> DEFAULT_MATCHER.match(pattern, plan)); - } -} diff --git a/core/src/test/java/org/opensearch/sql/storage/TableTest.java b/core/src/test/java/org/opensearch/sql/storage/TableTest.java deleted file mode 100644 index a96ee71af0..0000000000 --- a/core/src/test/java/org/opensearch/sql/storage/TableTest.java +++ /dev/null @@ -1,25 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.storage; - -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.withSettings; - -import org.junit.jupiter.api.DisplayNameGeneration; -import org.junit.jupiter.api.DisplayNameGenerator; -import org.junit.jupiter.api.Test; -import org.mockito.invocation.InvocationOnMock; - -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class TableTest { - - @Test - public void createPagedScanBuilder_throws() { - var table = mock(Table.class, withSettings().defaultAnswer(InvocationOnMock::callRealMethod)); - assertThrows(Throwable.class, () -> table.createPagedScanBuilder(4)); - } -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java index 4789a50896..92ec3dbfe3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java @@ -5,9 +5,6 @@ package org.opensearch.sql.opensearch.request; -import java.util.List; -import java.util.function.Consumer; -import java.util.function.Function; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; @@ -20,11 +17,14 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import java.util.List; +import java.util.function.Consumer; +import java.util.function.Function; + /** * Scroll (cursor) request is used to page the search. This request is not configurable and has * no search query. It just handles paging through responses to the initial request. * It is used on second and next pagination (cursor) requests. - * First (initial) request is handled by {@link InitialPageRequestBuilder}. */ @EqualsAndHashCode @RequiredArgsConstructor diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index b1a6589aca..0df16d73ec 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -5,9 +5,6 @@ package org.opensearch.sql.opensearch.request; -import java.util.List; -import java.util.Map; -import java.util.Set; import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.common.unit.TimeValue; @@ -21,78 +18,91 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; +import java.util.List; +import java.util.Map; +import java.util.Set; + /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. - * Initial search requests is handled by {@link InitialPageRequestBuilder}. */ -public class ContinuePageRequestBuilder extends PagedRequestBuilder { +public class ContinuePageRequestBuilder implements PushDownRequestBuilder { + + public static final String PUSH_DOWN_NOT_SUPPORTED = + "Cursor requests don't support any push down"; - @Getter - private final OpenSearchRequest.IndexName indexName; @Getter private final String scrollId; - private final TimeValue scrollTimeout; private final OpenSearchExprValueFactory exprValueFactory; + private final TimeValue scrollTimeout; /** Constructor. */ - public ContinuePageRequestBuilder(OpenSearchRequest.IndexName indexName, - String scrollId, - Settings settings, + public ContinuePageRequestBuilder(String scrollId, TimeValue scrollTimeout, OpenSearchExprValueFactory exprValueFactory) { - this.indexName = indexName; this.scrollId = scrollId; - this.scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + this.scrollTimeout = scrollTimeout; this.exprValueFactory = exprValueFactory; } @Override - public OpenSearchRequest build() { + public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, + int maxResultWindow, + Settings settings) { return new ContinuePageRequest(scrollId, scrollTimeout, exprValueFactory); } + @Override + public int getQuerySize() { + return Integer.MAX_VALUE; + } + @Override public void pushDownFilter(QueryBuilder query) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownAggregation(Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownSort(List> sortBuilders) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownLimit(Integer limit, Integer offset) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownHighlight(String field, Map arguments) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownProjects(Set projects) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushTypeMapping(Map typeMapping) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownNested(List> nestedArgs) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); } @Override public void pushDownTrackedScore(boolean trackScores) { - throw new UnsupportedOperationException("Cursor requests don't support any push down"); + throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); + } + + @Override + public void pushDownPageSize(int pageSize) { + throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java deleted file mode 100644 index 25b7253eca..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilder.java +++ /dev/null @@ -1,114 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -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 lombok.Getter; -import org.apache.commons.lang3.tuple.Pair; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.search.aggregations.AggregationBuilder; -import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.search.sort.SortBuilder; -import org.opensearch.sql.ast.expression.Literal; -import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; - -/** - * This builder assists creating the initial OpenSearch paging (scrolling) request. - * It is used only on the first page (pagination request). - * Subsequent requests (cursor requests) use {@link ContinuePageRequestBuilder}. - */ -public class InitialPageRequestBuilder extends PagedRequestBuilder { - - @Getter - private final OpenSearchRequest.IndexName indexName; - private final SearchSourceBuilder sourceBuilder; - private final OpenSearchExprValueFactory exprValueFactory; - private final TimeValue scrollTimeout; - - /** - * Constructor. - * @param indexName index being scanned - * @param pageSize page size - * @param exprValueFactory value factory - */ - // TODO accept indexName as string (same way as `OpenSearchRequestBuilder` does)? - public InitialPageRequestBuilder(OpenSearchRequest.IndexName indexName, - int pageSize, - Settings settings, - OpenSearchExprValueFactory exprValueFactory) { - this.indexName = indexName; - this.exprValueFactory = exprValueFactory; - this.scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); - this.sourceBuilder = new SearchSourceBuilder() - .from(0) - .size(pageSize) - .timeout(DEFAULT_QUERY_TIMEOUT); - } - - @Override - public OpenSearchScrollRequest build() { - return new OpenSearchScrollRequest(indexName, scrollTimeout, sourceBuilder, exprValueFactory); - } - - @Override - public void pushDownFilter(QueryBuilder query) { - throw new UnsupportedOperationException("Pagination does not support filter (WHERE clause)"); - } - - @Override - public void pushDownAggregation(Pair, - OpenSearchAggregationResponseParser> aggregationBuilder) { - throw new UnsupportedOperationException("Pagination does not support aggregations"); - } - - @Override - public void pushDownSort(List> sortBuilders) { - throw new UnsupportedOperationException("Pagination does not support sort (ORDER BY clause)"); - } - - @Override - public void pushDownLimit(Integer limit, Integer offset) { - throw new UnsupportedOperationException("Pagination does not support limit (LIMIT clause)"); - } - - @Override - public void pushDownHighlight(String field, Map arguments) { - throw new UnsupportedOperationException("Pagination does not support highlight function"); - } - - /** - * Push down project expression to OpenSearch. - */ - @Override - public void pushDownProjects(Set projects) { - sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr) - .distinct().toArray(String[]::new), new String[0]); - } - - @Override - public void pushTypeMapping(Map typeMapping) { - exprValueFactory.extendTypeMapping(typeMapping); - } - - @Override - public void pushDownNested(List> nestedArgs) { - throw new UnsupportedOperationException("Pagination does not support nested function"); - } - - @Override - public void pushDownTrackedScore(boolean trackScores) { - throw new UnsupportedOperationException("Pagination does not support score function"); - } -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index f8d62ad7ce..7466862eee 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -6,30 +6,13 @@ package org.opensearch.sql.opensearch.request; -import static java.util.stream.Collectors.mapping; -import static java.util.stream.Collectors.toList; -import static org.opensearch.index.query.QueryBuilders.boolQuery; -import static org.opensearch.index.query.QueryBuilders.matchAllQuery; -import static org.opensearch.index.query.QueryBuilders.nestedQuery; -import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; -import static org.opensearch.search.sort.SortOrder.ASC; - -import java.util.Arrays; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.ToString; import org.apache.commons.lang3.tuple.Pair; import org.apache.lucene.search.join.ScoreMode; import org.opensearch.common.unit.TimeValue; -import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.InnerHitBuilder; -import org.opensearch.index.query.NestedQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.*; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.FetchSourceContext; @@ -44,12 +27,24 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -import org.opensearch.sql.planner.logical.LogicalNested; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +import static java.util.stream.Collectors.mapping; +import static java.util.stream.Collectors.toList; +import static org.opensearch.index.query.QueryBuilders.matchAllQuery; +import static org.opensearch.index.query.QueryBuilders.nestedQuery; +import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; +import static org.opensearch.search.sort.SortOrder.ASC; /** * OpenSearch search request builder. */ -@EqualsAndHashCode(callSuper = false) +@EqualsAndHashCode @Getter @ToString public class OpenSearchRequestBuilder implements PushDownRequestBuilder { @@ -60,19 +55,19 @@ public class OpenSearchRequestBuilder implements PushDownRequestBuilder { public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); /** - * {@link OpenSearchRequest.IndexName}. + * Search request source builder. */ - private final OpenSearchRequest.IndexName indexName; + private final SearchSourceBuilder sourceBuilder; /** - * Index max result window. + * Query size of the request -- how many rows will be returned. */ - private final Integer maxResultWindow; + private int querySize; /** - * Search request source builder. + * Size of each page request to return. */ - private final SearchSourceBuilder sourceBuilder; + private Optional pageSize = Optional.empty(); /** * OpenSearchExprValueFactory. @@ -80,62 +75,59 @@ public class OpenSearchRequestBuilder implements PushDownRequestBuilder { @EqualsAndHashCode.Exclude @ToString.Exclude private final OpenSearchExprValueFactory exprValueFactory; - - /** - * Query size of the request -- how many rows will be returned. - */ - private int querySize; - - /** - * Scroll context life time. - */ - private final TimeValue scrollTimeout; - - public OpenSearchRequestBuilder(String indexName, - Integer maxResultWindow, - Settings settings, - OpenSearchExprValueFactory exprValueFactory) { - this(new OpenSearchRequest.IndexName(indexName), maxResultWindow, settings, - exprValueFactory); - } + private int startFrom = 0; /** * Constructor. */ - public OpenSearchRequestBuilder(OpenSearchRequest.IndexName indexName, - Integer maxResultWindow, - Settings settings, - OpenSearchExprValueFactory exprValueFactory) { - this.indexName = indexName; - this.maxResultWindow = maxResultWindow; - this.exprValueFactory = exprValueFactory; - this.scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); - this.querySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprValueFactory) { + this.querySize = querySize; this.sourceBuilder = new SearchSourceBuilder() - .from(0) - .size(querySize) + .from(startFrom) .timeout(DEFAULT_QUERY_TIMEOUT) .trackScores(false); + this.exprValueFactory = exprValueFactory; } + @Override + public int getQuerySize() { + return pageSize.orElse(querySize); + } /** * Build DSL request. * * @return query request or scroll request */ - public OpenSearchRequest build() { - Integer from = sourceBuilder.from(); - Integer size = sourceBuilder.size(); - - if (from + size > maxResultWindow) { - sourceBuilder.size(maxResultWindow - from); - return new OpenSearchScrollRequest( - indexName, scrollTimeout, sourceBuilder, exprValueFactory); + @Override + public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, + int maxResultWindow, + Settings settings) { + int size = querySize; + TimeValue scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + if (pageSize.isEmpty()) { + if (startFrom + size > maxResultWindow) { + sourceBuilder.size(maxResultWindow - startFrom); + return new OpenSearchScrollRequest( + indexName, scrollTimeout, sourceBuilder, exprValueFactory); + } else { + sourceBuilder.from(startFrom); + sourceBuilder.size(size); + return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); + } } else { - return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); + if (startFrom != 0) { + throw new UnsupportedOperationException("Non-zero offset is not supported with pagination"); + } + sourceBuilder.size(pageSize.get()); + return new OpenSearchScrollRequest(indexName, scrollTimeout, + sourceBuilder, exprValueFactory); } } + + boolean isBoolFilterQuery(QueryBuilder current) { + return (current instanceof BoolQueryBuilder); + } /** * Push down query to DSL request. * @@ -170,7 +162,7 @@ public void pushDownFilter(QueryBuilder query) { @Override public void pushDownAggregation( Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { - aggregationBuilder.getLeft().forEach(builder -> sourceBuilder.aggregation(builder)); + aggregationBuilder.getLeft().forEach(sourceBuilder::aggregation); sourceBuilder.size(0); exprValueFactory.setParser(aggregationBuilder.getRight()); } @@ -198,6 +190,7 @@ public void pushDownSort(List> sortBuilders) { @Override public void pushDownLimit(Integer limit, Integer offset) { querySize = limit; + startFrom = offset; sourceBuilder.from(offset).size(limit); } @@ -206,6 +199,11 @@ public void pushDownTrackedScore(boolean trackScores) { sourceBuilder.trackScores(trackScores); } + @Override + public void pushDownPageSize(int pageSize) { + this.pageSize = Optional.of(pageSize); + } + /** * Add highlight to DSL requests. * @param field name of the field to highlight @@ -245,9 +243,8 @@ public void pushDownHighlight(String field, Map arguments) { */ @Override public void pushDownProjects(Set projects) { - final Set 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]); } @Override @@ -258,7 +255,7 @@ public void pushTypeMapping(Map typeMapping) { private boolean isSortByDocOnly() { List> sorts = sourceBuilder.sorts(); if (sorts != null) { - return sorts.equals(Arrays.asList(SortBuilders.fieldSort(DOC_FIELD_NAME))); + return sorts.equals(List.of(SortBuilders.fieldSort(DOC_FIELD_NAME))); } return false; } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 77c6a781fe..03f97cda61 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -88,7 +88,8 @@ public OpenSearchResponse search(Function searchA : List.of(); var response = new OpenSearchResponse(openSearchResponse, exprValueFactory, includes); - if (!(needClean = response.isEmpty())) { + needClean = response.isEmpty(); + if (!needClean) { setScrollId(openSearchResponse.getScrollId()); } return response; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java deleted file mode 100644 index 69309bd7c9..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PagedRequestBuilder.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -public abstract class PagedRequestBuilder implements PushDownRequestBuilder { - public abstract OpenSearchRequest build(); - - public abstract OpenSearchRequest.IndexName getIndexName(); -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java index 59aa1949b6..9d579505a9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java @@ -8,24 +8,19 @@ import java.util.List; import java.util.Map; import java.util.Set; -import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; -import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.sort.SortBuilder; import org.opensearch.sql.ast.expression.Literal; -import org.opensearch.sql.data.type.ExprType; +import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; public interface PushDownRequestBuilder { - default boolean isBoolFilterQuery(QueryBuilder current) { - return (current instanceof BoolQueryBuilder); - } + int getQuerySize(); void pushDownFilter(QueryBuilder query); @@ -45,4 +40,11 @@ void pushDownAggregation(Pair, void pushDownNested(List> nestedArgs); void pushDownTrackedScore(boolean trackScores); -} \ No newline at end of file + + void pushDownPageSize(int pageSize); + + OpenSearchRequest build(OpenSearchRequest.IndexName indexName, + int maxResultWindow, + Settings settings); + +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 949b1e53ec..2cea742eed 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -7,9 +7,6 @@ package org.opensearch.sql.opensearch.storage; import com.google.common.annotations.VisibleForTesting; -import java.util.HashMap; -import java.util.LinkedHashMap; -import java.util.Map; import lombok.RequiredArgsConstructor; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; @@ -20,14 +17,11 @@ 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; @@ -37,6 +31,11 @@ import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.function.Function; + /** OpenSearch table (index) implementation. */ public class OpenSearchIndex implements Table { @@ -168,28 +167,17 @@ public PhysicalPlan implement(LogicalPlan plan) { return plan.accept(new OpenSearchDefaultImplementor(client), null); } - @Override - public LogicalPlan optimize(LogicalPlan plan) { - // No-op because optimization already done in Planner - return plan; - } - @Override public TableScanBuilder createScanBuilder() { Map allFields = new HashMap<>(); getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); allFields.putAll(getFieldOpenSearchTypes()); - OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, indexName, - getMaxResultWindow(), new OpenSearchExprValueFactory(allFields)); - return new OpenSearchIndexScanBuilder(indexScan); - } - - @Override - public TableScanBuilder createPagedScanBuilder(int pageSize) { - var requestBuilder = new InitialPageRequestBuilder(indexName, pageSize, settings, - new OpenSearchExprValueFactory(getFieldOpenSearchTypes())); - var indexScan = new OpenSearchPagedIndexScan(client, requestBuilder); - return new OpenSearchPagedIndexScanBuilder(indexScan); + Function buildScan + = t -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), t); + var builder = new OpenSearchRequestBuilder( + settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), + new OpenSearchExprValueFactory(allFields)); + return new OpenSearchIndexScanBuilder(buildScan, builder); } @VisibleForTesting diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 2171fb564f..33969733e9 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -6,18 +6,28 @@ package org.opensearch.sql.opensearch.storage.scan; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Collections; import java.util.Iterator; import lombok.EqualsAndHashCode; -import lombok.Getter; import lombok.ToString; +import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.exception.NoCursorException; +import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.request.PushDownRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.storage.TableScanOperator; /** @@ -25,21 +35,23 @@ */ @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) @ToString(onlyExplicitlyIncluded = true) -public class OpenSearchIndexScan extends TableScanOperator { +public class OpenSearchIndexScan extends TableScanOperator implements SerializablePlan { /** OpenSearch client. */ - private final OpenSearchClient client; + private transient OpenSearchClient client; + private transient OpenSearchRequest.IndexName indexName; + private transient Settings settings; + private transient Integer maxResultWindow; /** Search request builder. */ @EqualsAndHashCode.Include - @Getter @ToString.Include - private final OpenSearchRequestBuilder requestBuilder; + private transient PushDownRequestBuilder requestBuilder; /** Search request. */ @EqualsAndHashCode.Include @ToString.Include - private OpenSearchRequest request; + private transient OpenSearchRequest request; /** Total query size. */ @EqualsAndHashCode.Include @@ -50,39 +62,42 @@ public class OpenSearchIndexScan extends TableScanOperator { private Integer queryCount; /** Search response for current batch. */ - private Iterator iterator; + private transient Iterator iterator; /** * Constructor. */ - public OpenSearchIndexScan(OpenSearchClient client, Settings settings, - String indexName, Integer maxResultWindow, - OpenSearchExprValueFactory exprValueFactory) { - this( - client, - settings, - new OpenSearchRequest.IndexName(indexName), - maxResultWindow, - exprValueFactory - ); + public static OpenSearchIndexScan create(OpenSearchClient client, + String indexName, + Settings settings, + Integer maxResultWindow, + OpenSearchExprValueFactory exprValueFactory) { + final var name = new OpenSearchRequest.IndexName(indexName); + final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + return new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder); } - /** - * Constructor. - */ - public OpenSearchIndexScan(OpenSearchClient client, Settings settings, - OpenSearchRequest.IndexName indexName, Integer maxResultWindow, - OpenSearchExprValueFactory exprValueFactory) { + + public OpenSearchIndexScan(OpenSearchClient client, + OpenSearchRequest.IndexName indexName, + Settings settings, + Integer maxResultWindow, + PushDownRequestBuilder requestBuilder) { this.client = client; - this.requestBuilder = new OpenSearchRequestBuilder( - indexName, maxResultWindow, settings, exprValueFactory); + this.indexName = indexName; + this.settings = settings; + this.maxResultWindow = maxResultWindow; + this.requestBuilder = requestBuilder; } @Override public void open() { super.open(); + request = requestBuilder.build(indexName, maxResultWindow, settings); + // TODO does this work for paged requests? querySize = requestBuilder.getQuerySize(); - request = requestBuilder.build(); iterator = Collections.emptyIterator(); queryCount = 0; fetchNextBatch(); @@ -126,6 +141,44 @@ public void close() { @Override public String explain() { - return getRequestBuilder().build().toString(); + return requestBuilder.build(indexName, maxResultWindow, settings).toString(); + } + + /** + * @deprecated Exists only to satisfy Java serialization API. + */ + @Deprecated(since="introduction") + public OpenSearchIndexScan() { + } + + @Override + public void readExternal(ObjectInput in) throws IOException { + var serializedName = in.readUTF(); + var scrollId = in.readUTF(); + querySize = in.readInt(); + + var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) + .resolveObject("engine"); + + indexName = new OpenSearchRequest.IndexName(serializedName); + client = engine.getClient(); + settings = engine.getSettings(); + var index = new OpenSearchIndex(client, settings, indexName.toString()); + maxResultWindow = index.getMaxResultWindow(); + + TimeValue scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + requestBuilder = new ContinuePageRequestBuilder(scrollId, scrollTimeout, + new OpenSearchExprValueFactory(index.getFieldOpenSearchTypes())); + } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + if (request.toCursor() == null || request.toCursor().isEmpty()) { + throw new NoCursorException(); + } + + out.writeUTF(indexName.toString()); + out.writeUTF(request.toCursor()); + out.writeInt(querySize); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 74be670dcc..4b9643c4a5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -8,6 +8,7 @@ import java.util.List; import java.util.Set; import java.util.stream.Collectors; +import lombok.EqualsAndHashCode; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.sql.ast.tree.Sort; @@ -15,21 +16,27 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; import org.opensearch.sql.planner.logical.LogicalAggregation; +import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalHighlight; +import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalSort; -import org.opensearch.sql.storage.TableScanOperator; -import org.opensearch.sql.storage.read.TableScanBuilder; /** * Index scan builder for aggregate query used by {@link OpenSearchIndexScanBuilder} internally. */ -class OpenSearchIndexScanAggregationBuilder extends TableScanBuilder { +@EqualsAndHashCode +class OpenSearchIndexScanAggregationBuilder implements PushDownTranslator { /** OpenSearch index scan to be optimized. */ - private final OpenSearchIndexScan indexScan; + private final OpenSearchRequestBuilder requestBuilder; /** Aggregators pushed down. */ private List aggregatorList; @@ -40,25 +47,21 @@ class OpenSearchIndexScanAggregationBuilder extends TableScanBuilder { /** Sorting items pushed down. */ private List> sortList; - /** - * Initialize with given index scan and perform push-down optimization later. - * - * @param indexScan index scan not fully optimized yet - */ - OpenSearchIndexScanAggregationBuilder(OpenSearchIndexScan indexScan) { - this.indexScan = indexScan; + + OpenSearchIndexScanAggregationBuilder(OpenSearchRequestBuilder requestBuilder) { + this.requestBuilder = requestBuilder; } @Override - public TableScanOperator build() { + public OpenSearchRequestBuilder build() { AggregationQueryBuilder builder = new AggregationQueryBuilder(new DefaultExpressionSerializer()); Pair, OpenSearchAggregationResponseParser> aggregationBuilder = builder.buildAggregationBuilder(aggregatorList, groupByList, sortList); - indexScan.getRequestBuilder().pushDownAggregation(aggregationBuilder); - indexScan.getRequestBuilder().pushTypeMapping( + requestBuilder.pushDownAggregation(aggregationBuilder); + requestBuilder.pushTypeMapping( builder.buildTypeMapping(aggregatorList, groupByList)); - return indexScan; + return requestBuilder; } @Override @@ -68,6 +71,11 @@ public boolean pushDownAggregation(LogicalAggregation aggregation) { return true; } + @Override + public boolean pushDownFilter(LogicalFilter filter) { + return false; + } + @Override public boolean pushDownSort(LogicalSort sort) { if (hasAggregatorInSortBy(sort)) { @@ -78,6 +86,31 @@ public boolean pushDownSort(LogicalSort sort) { return true; } + @Override + public boolean pushDownLimit(LogicalLimit limit) { + return false; + } + + @Override + public boolean pushDownProject(LogicalProject project) { + return false; + } + + @Override + public boolean pushDownHighlight(LogicalHighlight highlight) { + return false; + } + + @Override + public boolean pushDownPageSize(LogicalPaginate paginate) { + return false; + } + + @Override + public boolean pushDownNested(LogicalNested nested) { + return false; + } + private boolean hasAggregatorInSortBy(LogicalSort sort) { final Set aggregatorNames = aggregatorList.stream().map(NamedAggregator::getName).collect(Collectors.toSet()); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index 024331d267..ed4259c86a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -5,14 +5,16 @@ package org.opensearch.sql.opensearch.storage.scan; -import com.google.common.annotations.VisibleForTesting; +import java.util.function.Function; import lombok.EqualsAndHashCode; import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalSort; import org.opensearch.sql.storage.TableScanOperator; @@ -25,32 +27,33 @@ */ public class OpenSearchIndexScanBuilder extends TableScanBuilder { + @EqualsAndHashCode.Exclude + private final Function constructor; /** * Delegated index scan builder for non-aggregate or aggregate query. */ @EqualsAndHashCode.Include - private TableScanBuilder delegate; + private PushDownTranslator delegate; /** Is limit operator pushed down. */ private boolean isLimitPushedDown = false; - @VisibleForTesting - OpenSearchIndexScanBuilder(TableScanBuilder delegate) { - this.delegate = delegate; - } + public OpenSearchIndexScanBuilder(Function constructor, + OpenSearchRequestBuilder requestBuilder) { + this.constructor = constructor; + this.delegate + = new OpenSearchIndexScanQueryBuilder(requestBuilder); - /** - * Initialize with given index scan. - * - * @param indexScan index scan to optimize - */ - public OpenSearchIndexScanBuilder(OpenSearchIndexScan indexScan) { - this.delegate = new OpenSearchIndexScanQueryBuilder(indexScan); + } + public OpenSearchIndexScanBuilder(Function constructor, + PushDownTranslator translator) { + this.constructor = constructor; + this.delegate = translator; } @Override public TableScanOperator build() { - return delegate.build(); + return constructor.apply(delegate.build()); } @Override @@ -66,12 +69,19 @@ public boolean pushDownAggregation(LogicalAggregation aggregation) { // Switch to builder for aggregate query which has different push down logic // for later filter, sort and limit operator. - delegate = new OpenSearchIndexScanAggregationBuilder( - (OpenSearchIndexScan) delegate.build()); + delegate = new OpenSearchIndexScanAggregationBuilder(delegate.build()); return delegate.pushDownAggregation(aggregation); } + @Override + public boolean pushDownPageSize(LogicalPaginate paginate) { + if (isLimitPushedDown) { + throw new IllegalStateException("Pagination has to be pushed down before limit."); + } + return delegate.pushDownPageSize(paginate); + } + @Override public boolean pushDownSort(LogicalSort sort) { if (!sortByFieldsOnly(sort)) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index d9b4e6b4e0..2c53257eb7 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -22,41 +22,31 @@ import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.function.OpenSearchFunctions; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; +import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalSort; -import org.opensearch.sql.storage.TableScanOperator; -import org.opensearch.sql.storage.read.TableScanBuilder; /** * Index scan builder for simple non-aggregate query used by * {@link OpenSearchIndexScanBuilder} internally. */ @VisibleForTesting -class OpenSearchIndexScanQueryBuilder extends TableScanBuilder { +@EqualsAndHashCode +class OpenSearchIndexScanQueryBuilder implements PushDownTranslator { - /** OpenSearch index scan to be optimized. */ - @EqualsAndHashCode.Include - private final OpenSearchIndexScan indexScan; + OpenSearchRequestBuilder requestBuilder; - /** - * Initialize with given index scan and perform push-down optimization later. - * - * @param indexScan index scan not optimized yet - */ - OpenSearchIndexScanQueryBuilder(OpenSearchIndexScan indexScan) { - this.indexScan = indexScan; - } - - @Override - public TableScanOperator build() { - return indexScan; + public OpenSearchIndexScanQueryBuilder(OpenSearchRequestBuilder requestBuilder) { + this.requestBuilder = requestBuilder; } @Override @@ -65,8 +55,8 @@ public boolean pushDownFilter(LogicalFilter filter) { new DefaultExpressionSerializer()); Expression queryCondition = filter.getCondition(); QueryBuilder query = queryBuilder.build(queryCondition); - indexScan.getRequestBuilder().pushDownFilter(query); - indexScan.getRequestBuilder().pushDownTrackedScore( + requestBuilder.pushDownFilter(query); + requestBuilder.pushDownTrackedScore( trackScoresFromOpenSearchFunction(queryCondition)); return true; } @@ -75,7 +65,7 @@ public boolean pushDownFilter(LogicalFilter filter) { public boolean pushDownSort(LogicalSort sort) { List> sortList = sort.getSortList(); final SortQueryBuilder builder = new SortQueryBuilder(); - indexScan.getRequestBuilder().pushDownSort(sortList.stream() + requestBuilder.pushDownSort(sortList.stream() .map(sortItem -> builder.build(sortItem.getValue(), sortItem.getKey())) .collect(Collectors.toList())); return true; @@ -83,13 +73,13 @@ public boolean pushDownSort(LogicalSort sort) { @Override public boolean pushDownLimit(LogicalLimit limit) { - indexScan.getRequestBuilder().pushDownLimit(limit.getLimit(), limit.getOffset()); + requestBuilder.pushDownLimit(limit.getLimit(), limit.getOffset()); return true; } @Override public boolean pushDownProject(LogicalProject project) { - indexScan.getRequestBuilder().pushDownProjects( + requestBuilder.pushDownProjects( findReferenceExpressions(project.getProjectList())); // Return false intentionally to keep the original project operator @@ -98,12 +88,18 @@ public boolean pushDownProject(LogicalProject project) { @Override public boolean pushDownHighlight(LogicalHighlight highlight) { - indexScan.getRequestBuilder().pushDownHighlight( + requestBuilder.pushDownHighlight( StringUtils.unquoteText(highlight.getHighlightField().toString()), highlight.getArguments()); return true; } + @Override + public boolean pushDownPageSize(LogicalPaginate paginate) { + requestBuilder.pushDownPageSize(paginate.getPageSize()); + return true; + } + private boolean trackScoresFromOpenSearchFunction(Expression condition) { if (condition instanceof OpenSearchFunctions.OpenSearchFunction && ((OpenSearchFunctions.OpenSearchFunction) condition).isScoreTracked()) { @@ -118,8 +114,8 @@ private boolean trackScoresFromOpenSearchFunction(Expression condition) { @Override public boolean pushDownNested(LogicalNested nested) { - indexScan.getRequestBuilder().pushDownNested(nested.getFields()); - indexScan.getRequestBuilder().pushDownProjects( + requestBuilder.pushDownNested(nested.getFields()); + requestBuilder.pushDownProjects( findReferenceExpressions(nested.getProjectList())); // Return false intentionally to keep the original nested operator // Since we return false we need to pushDownProject here as it won't be @@ -128,6 +124,16 @@ public boolean pushDownNested(LogicalNested nested) { return false; } + @Override + public boolean pushDownAggregation(LogicalAggregation aggregation) { + return false; + } + + @Override + public OpenSearchRequestBuilder build() { + return requestBuilder; + } + /** * Find reference expression from expression. * @param expressions a list of expression. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java deleted file mode 100644 index 3667a3ffdf..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScan.java +++ /dev/null @@ -1,115 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.storage.scan; - -import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; -import java.util.Collections; -import java.util.Iterator; -import lombok.EqualsAndHashCode; -import lombok.Getter; -import lombok.ToString; -import org.apache.commons.lang3.NotImplementedException; -import org.opensearch.sql.data.model.ExprValue; -import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.executor.pagination.PlanSerializer; -import org.opensearch.sql.opensearch.client.OpenSearchClient; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; -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.opensearch.storage.OpenSearchIndex; -import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; -import org.opensearch.sql.planner.SerializablePlan; -import org.opensearch.sql.storage.TableScanOperator; - -@EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) -@ToString(onlyExplicitlyIncluded = true) -public class OpenSearchPagedIndexScan extends TableScanOperator implements SerializablePlan { - private OpenSearchClient client; - @Getter - private PagedRequestBuilder requestBuilder; - @EqualsAndHashCode.Include - @ToString.Include - private OpenSearchRequest request; - private Iterator iterator; - private long totalHits = 0; - - public OpenSearchPagedIndexScan(OpenSearchClient client, PagedRequestBuilder requestBuilder) { - this.client = client; - this.requestBuilder = requestBuilder; - } - - @Override - public String explain() { - throw new NotImplementedException("Implement OpenSearchPagedIndexScan.explain"); - } - - @Override - public boolean hasNext() { - return iterator.hasNext(); - } - - @Override - public ExprValue next() { - return iterator.next(); - } - - @Override - public void open() { - super.open(); - request = requestBuilder.build(); - OpenSearchResponse response = client.search(request); - if (!response.isEmpty()) { - iterator = response.iterator(); - totalHits = response.getTotalHits(); - } else { - iterator = Collections.emptyIterator(); - } - } - - @Override - public void close() { - super.close(); - client.cleanup(request); - } - - @Override - public long getTotalHits() { - return totalHits; - } - - /** Don't use, it is for deserialization needs only. */ - @Deprecated - public OpenSearchPagedIndexScan() { - } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) - .resolveObject("engine"); - var indexName = (String) in.readUTF(); - var scrollId = (String) in.readUTF(); - client = engine.getClient(); - var index = new OpenSearchIndex(client, engine.getSettings(), indexName); - requestBuilder = new ContinuePageRequestBuilder( - new OpenSearchRequest.IndexName(indexName), - scrollId, engine.getSettings(), - new OpenSearchExprValueFactory(index.getFieldOpenSearchTypes())); - } - - @Override - public void writeExternal(ObjectOutput out) throws IOException { - if (request.toCursor() == null || request.toCursor().isEmpty()) { - throw new NoCursorException(); - } - - out.writeUTF(requestBuilder.getIndexName().toString()); - out.writeUTF(request.toCursor()); - } -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java deleted file mode 100644 index 779df4ebec..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanBuilder.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -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 TableScanBuilder as more features - * support pagination. - */ -public class OpenSearchPagedIndexScanBuilder extends TableScanBuilder { - @EqualsAndHashCode.Include - OpenSearchPagedIndexScan indexScan; - - public OpenSearchPagedIndexScanBuilder(OpenSearchPagedIndexScan indexScan) { - this.indexScan = indexScan; - } - - @Override - public TableScanOperator build() { - return indexScan; - } -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java new file mode 100644 index 0000000000..f178681706 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan; + +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.planner.logical.LogicalAggregation; +import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalHighlight; +import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalProject; +import org.opensearch.sql.planner.logical.LogicalSort; + +public interface PushDownTranslator { + boolean pushDownFilter(LogicalFilter filter); + + boolean pushDownSort(LogicalSort sort); + + boolean pushDownLimit(LogicalLimit limit); + + boolean pushDownProject(LogicalProject project); + + boolean pushDownHighlight(LogicalHighlight highlight); + + boolean pushDownPageSize(LogicalPaginate paginate); + + boolean pushDownNested(LogicalNested nested); + boolean pushDownAggregation(LogicalAggregation aggregation); + + OpenSearchRequestBuilder build(); +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index e3c9291ced..caafa40d57 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -29,7 +29,6 @@ import com.google.common.io.Resources; import java.io.IOException; import java.net.URL; -import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -97,13 +96,10 @@ class OpenSearchNodeClientTest { @Mock private SearchHit searchHit; - @Mock - private ThreadContext threadContext; - @Mock private GetIndexResponse indexResponse; - private ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", + private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", new ExprIntegerValue(1))); private OpenSearchClient client; @@ -455,8 +451,8 @@ private void mockNodeClientSettings(String indexName, String indexMetadata) GetSettingsResponse mockResponse = mock(GetSettingsResponse.class); when(nodeClient.admin().indices().prepareGetSettings(any()).setLocal(anyBoolean()).get()) .thenReturn(mockResponse); - Map metadata = Map.of(indexName, - IndexMetadata.fromXContent(createParser(indexMetadata)).getSettings()); + Map metadata =Map.of(indexName, + IndexMetadata.fromXContent(createParser(indexMetadata)).getSettings()); when(mockResponse.getIndexToSettings()).thenReturn(metadata); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index 22d02d1ab5..e019a919fc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -78,7 +78,6 @@ class OpenSearchRestClientTest { private static final String TEST_MAPPING_FILE = "mappings/accounts.json"; - @Mock(answer = RETURNS_DEEP_STUBS) private RestHighLevelClient restClient; @@ -93,7 +92,7 @@ class OpenSearchRestClientTest { @Mock private GetIndexResponse getIndexResponse; - private ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", + private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", new ExprIntegerValue(1))); @BeforeEach @@ -362,9 +361,7 @@ void scroll_with_IOException() throws IOException { void schedule() { AtomicBoolean isRun = new AtomicBoolean(false); client.schedule( - () -> { - isRun.set(true); - }); + () -> isRun.set(true)); assertTrue(isRun.get()); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index c96782abea..245067e51e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -96,17 +96,17 @@ void execute_successfully() { List actual = new ArrayList<>(); executor.execute( plan, - new ResponseListener() { - @Override - public void onResponse(QueryResponse response) { - actual.addAll(response.getResults()); - } - - @Override - public void onFailure(Exception e) { - fail("Error occurred during execution", e); - } - }); + new ResponseListener<>() { + @Override + public void onResponse(QueryResponse response) { + actual.addAll(response.getResults()); + } + + @Override + public void onFailure(Exception e) { + fail("Error occurred during execution", e); + } + }); assertTrue(plan.hasOpen); assertEquals(expected, actual); @@ -126,18 +126,18 @@ void execute_with_cursor() { List actual = new ArrayList<>(); executor.execute( plan, - new ResponseListener() { - @Override - public void onResponse(QueryResponse response) { - actual.addAll(response.getResults()); - assertTrue(response.getCursor().toString().startsWith("n:")); - } - - @Override - public void onFailure(Exception e) { - fail("Error occurred during execution", e); - } - }); + new ResponseListener<>() { + @Override + public void onResponse(QueryResponse response) { + actual.addAll(response.getResults()); + assertTrue(response.getCursor().toString().startsWith("n:")); + } + + @Override + public void onFailure(Exception e) { + fail("Error occurred during execution", e); + } + }); assertEquals(expected, actual); } @@ -154,17 +154,17 @@ void execute_with_failure() { AtomicReference actual = new AtomicReference<>(); executor.execute( plan, - new ResponseListener() { - @Override - public void onResponse(QueryResponse response) { - fail("Expected error didn't happen"); - } - - @Override - public void onFailure(Exception e) { - actual.set(e); - } - }); + new ResponseListener<>() { + @Override + public void onResponse(QueryResponse response) { + fail("Expected error didn't happen"); + } + + @Override + public void onFailure(Exception e) { + actual.set(e); + } + }); assertEquals(expected, actual.get()); verify(plan).close(); } @@ -178,11 +178,11 @@ void explain_successfully() { when(settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE)) .thenReturn(TimeValue.timeValueMinutes(1)); - PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), settings, - "test", 10000, mock(OpenSearchExprValueFactory.class)); + PhysicalPlan plan = OpenSearchIndexScan.create(mock(OpenSearchClient.class), "test", settings, + 10000, mock(OpenSearchExprValueFactory.class)); AtomicReference result = new AtomicReference<>(); - executor.explain(plan, new ResponseListener() { + executor.explain(plan, new ResponseListener<>() { @Override public void onResponse(ExplainResponse response) { result.set(response); @@ -205,7 +205,7 @@ void explain_with_failure() { when(plan.accept(any(), any())).thenThrow(IllegalStateException.class); AtomicReference result = new AtomicReference<>(); - executor.explain(plan, new ResponseListener() { + executor.explain(plan, new ResponseListener<>() { @Override public void onResponse(ExplainResponse response) { fail("Should fail as expected"); @@ -261,11 +261,11 @@ private static class FakePhysicalPlan extends TableScanOperator implements Seria private boolean hasSplit; @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + public void readExternal(ObjectInput in) { } @Override - public void writeExternal(ObjectOutput out) throws IOException { + public void writeExternal(ObjectOutput out) { } @Override diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index fe0077914e..d2aba45f0b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -11,8 +11,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC; -import static org.opensearch.sql.common.setting.Settings.Key.QUERY_SIZE_LIMIT; -import static org.opensearch.sql.common.setting.Settings.Key.SQL_CURSOR_KEEP_ALIVE; import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.STRING; @@ -39,11 +37,11 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.client.node.NodeClient; -import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; @@ -89,22 +87,21 @@ public void setup() { } @Test - public void testProtectIndexScan() { - when(settings.getSettingValue(QUERY_SIZE_LIMIT)).thenReturn(200); - when(settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); + void testProtectIndexScan() { + when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)) + .thenReturn(200); String indexName = "test"; Integer maxResultWindow = 10000; NamedExpression include = named("age", ref("age", INTEGER)); ReferenceExpression exclude = ref("name", STRING); ReferenceExpression dedupeField = ref("name", STRING); ReferenceExpression topField = ref("name", STRING); - List topExprs = Arrays.asList(ref("age", INTEGER)); + List topExprs = List.of(ref("age", INTEGER)); Expression filterExpr = literal(ExprBooleanValue.of(true)); - List groupByExprs = Arrays.asList(named("age", ref("age", INTEGER))); + List groupByExprs = List.of(named("age", ref("age", INTEGER))); List aggregators = - Arrays.asList(named("avg(age)", new AvgAggregator(Arrays.asList(ref("age", INTEGER)), - DOUBLE))); + List.of(named("avg(age)", new AvgAggregator(List.of(ref("age", INTEGER)), + DOUBLE))); Map mappings = ImmutableMap.of(ref("name", STRING), ref("lastname", STRING)); Pair newEvalField = @@ -127,8 +124,8 @@ public void testProtectIndexScan() { PhysicalPlanDSL.agg( filter( resourceMonitor( - new OpenSearchIndexScan(client, settings, - indexName, + OpenSearchIndexScan.create(client, + indexName, settings, maxResultWindow, exprValueFactory)), filterExpr), @@ -156,8 +153,8 @@ public void testProtectIndexScan() { PhysicalPlanDSL.rename( PhysicalPlanDSL.agg( filter( - new OpenSearchIndexScan(client, settings, - indexName, + OpenSearchIndexScan.create(client, indexName, + settings, maxResultWindow, exprValueFactory), filterExpr), @@ -178,7 +175,7 @@ public void testProtectIndexScan() { @SuppressWarnings("unchecked") @Test - public void testProtectSortForWindowOperator() { + void testProtectSortForWindowOperator() { NamedExpression rank = named(mock(RankFunction.class)); Pair sortItem = ImmutablePair.of(DEFAULT_ASC, DSL.ref("age", INTEGER)); @@ -204,7 +201,7 @@ public void testProtectSortForWindowOperator() { } @Test - public void testProtectWindowOperatorInput() { + void testProtectWindowOperatorInput() { NamedExpression avg = named(mock(AggregateWindowFunction.class)); WindowDefinition windowDefinition = mock(WindowDefinition.class); @@ -223,7 +220,7 @@ public void testProtectWindowOperatorInput() { @SuppressWarnings("unchecked") @Test - public void testNotProtectWindowOperatorInputIfAlreadyProtected() { + void testNotProtectWindowOperatorInputIfAlreadyProtected() { NamedExpression avg = named(mock(AggregateWindowFunction.class)); Pair sortItem = ImmutablePair.of(DEFAULT_ASC, DSL.ref("age", INTEGER)); @@ -248,7 +245,7 @@ public void testNotProtectWindowOperatorInputIfAlreadyProtected() { } @Test - public void testWithoutProtection() { + void testWithoutProtection() { Expression filterExpr = literal(ExprBooleanValue.of(true)); assertEquals( @@ -264,7 +261,7 @@ public void testWithoutProtection() { } @Test - public void testVisitMlCommons() { + void testVisitMlCommons() { NodeClient nodeClient = mock(NodeClient.class); MLCommonsOperator mlCommonsOperator = new MLCommonsOperator( @@ -282,7 +279,7 @@ public void testVisitMlCommons() { } @Test - public void testVisitAD() { + void testVisitAD() { NodeClient nodeClient = mock(NodeClient.class); ADOperator adOperator = new ADOperator( @@ -300,7 +297,7 @@ public void testVisitAD() { } @Test - public void testVisitML() { + void testVisitML() { NodeClient nodeClient = mock(NodeClient.class); MLOperator mlOperator = new MLOperator( @@ -320,7 +317,7 @@ public void testVisitML() { } @Test - public void testVisitNested() { + void testVisitNested() { Set args = Set.of("message.info"); Map> groupedFieldsByPath = Map.of("message", List.of("message.info")); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java index 5cabe1930d..f149e0d3b6 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java @@ -5,14 +5,6 @@ package org.opensearch.sql.opensearch.request; -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import java.util.List; -import java.util.Map; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -24,44 +16,39 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.mock; + @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) -public class ContinuePageRequestBuilderTest { +class ContinuePageRequestBuilderTest { @Mock private OpenSearchExprValueFactory exprValueFactory; - @Mock - private Settings settings; - - private final OpenSearchRequest.IndexName indexName = new OpenSearchRequest.IndexName("test"); private final String scrollId = "scroll"; private ContinuePageRequestBuilder requestBuilder; @BeforeEach void setup() { - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); - requestBuilder = new ContinuePageRequestBuilder( - indexName, scrollId, settings, exprValueFactory); + var timeout = TimeValue.timeValueMinutes(1); + requestBuilder = new ContinuePageRequestBuilder(scrollId, timeout, exprValueFactory); } @Test - public void build() { + void build() { assertEquals( new ContinuePageRequest(scrollId, TimeValue.timeValueMinutes(1), exprValueFactory), - requestBuilder.build() + requestBuilder.build(null, 0, null) ); } @Test - public void getIndexName() { - assertEquals(indexName, requestBuilder.getIndexName()); - } - - @Test - public void pushDown_not_supported() { + void pushDown_not_supported() { assertAll( () -> assertThrows(UnsupportedOperationException.class, () -> requestBuilder.pushDownFilter(mock())), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java deleted file mode 100644 index ef850380d4..0000000000 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/InitialPageRequestBuilderTest.java +++ /dev/null @@ -1,122 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; -import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.junit.jupiter.api.BeforeEach; -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.common.unit.TimeValue; -import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.data.type.ExprType; -import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; - -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -@ExtendWith(MockitoExtension.class) -public class InitialPageRequestBuilderTest { - - @Mock - private OpenSearchExprValueFactory exprValueFactory; - - @Mock - private Settings settings; - - private final int pageSize = 42; - - private final OpenSearchRequest.IndexName indexName = new OpenSearchRequest.IndexName("test"); - - private InitialPageRequestBuilder requestBuilder; - - @BeforeEach - void setup() { - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); - requestBuilder = new InitialPageRequestBuilder( - indexName, pageSize, settings, exprValueFactory); - } - - @Test - public void build() { - assertEquals( - new OpenSearchScrollRequest(indexName, TimeValue.timeValueMinutes(1), - new SearchSourceBuilder() - .from(0) - .size(pageSize) - .timeout(DEFAULT_QUERY_TIMEOUT), - exprValueFactory), - requestBuilder.build() - ); - } - - @Test - public void pushDown_not_supported() { - assertAll( - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownFilter(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownAggregation(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownSort(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownLimit(1, 2)), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownHighlight("", Map.of())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownNested(List.of())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownTrackedScore(true)) - ); - } - - @Test - public void pushTypeMapping() { - Map typeMapping = Map.of("intA", OpenSearchDataType.of(INTEGER)); - requestBuilder.pushTypeMapping(typeMapping); - - verify(exprValueFactory).extendTypeMapping(typeMapping); - } - - @Test - public void pushDownProject() { - Set references = Set.of(DSL.ref("intA", INTEGER)); - requestBuilder.pushDownProjects(references); - - assertEquals( - new OpenSearchScrollRequest(indexName, TimeValue.timeValueMinutes(1), - new SearchSourceBuilder() - .from(0) - .size(pageSize) - .timeout(DEFAULT_QUERY_TIMEOUT) - .fetchSource(new String[]{"intA"}, new String[0]), - exprValueFactory), - requestBuilder.build() - ); - } - - @Test - public void getIndexName() { - assertEquals(indexName, requestBuilder.getIndexName()); - } -} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 94433c29b9..d8e3dd5d5b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -6,9 +6,10 @@ package org.opensearch.sql.opensearch.request; +import static org.junit.Assert.assertThrows; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; import static org.opensearch.index.query.QueryBuilders.matchAllQuery; import static org.opensearch.index.query.QueryBuilders.nestedQuery; import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; @@ -43,6 +44,7 @@ import org.opensearch.search.sort.ScoreSortBuilder; import org.opensearch.search.sort.SortBuilders; import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; @@ -55,13 +57,15 @@ @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class OpenSearchRequestBuilderTest { +class OpenSearchRequestBuilderTest { private static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); private static final Integer DEFAULT_OFFSET = 0; private static final Integer DEFAULT_LIMIT = 200; private static final Integer MAX_RESULT_WINDOW = 500; + private static final OpenSearchRequest.IndexName indexName + = new OpenSearchRequest.IndexName("test"); @Mock private Settings settings; @@ -72,18 +76,18 @@ public class OpenSearchRequestBuilderTest { @BeforeEach void setup() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) + lenient().when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)) + .thenReturn(DEFAULT_LIMIT); + lenient().when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) .thenReturn(TimeValue.timeValueMinutes(1)); - requestBuilder = new OpenSearchRequestBuilder( - "test", MAX_RESULT_WINDOW, settings, exprValueFactory); + requestBuilder = new OpenSearchRequestBuilder(DEFAULT_LIMIT, exprValueFactory); } @Test void build_query_request() { - Integer limit = 200; - Integer offset = 0; + int limit = 200; + int offset = 0; requestBuilder.pushDownLimit(limit, offset); requestBuilder.pushDownTrackedScore(true); @@ -96,7 +100,7 @@ void build_query_request() { .timeout(DEFAULT_QUERY_TIMEOUT) .trackScores(true), exprValueFactory), - requestBuilder.build()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); } @Test @@ -113,7 +117,7 @@ void build_scroll_request_with_correct_size() { .size(MAX_RESULT_WINDOW - offset) .timeout(DEFAULT_QUERY_TIMEOUT), exprValueFactory), - requestBuilder.build()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); } @Test @@ -128,7 +132,7 @@ void test_push_down_query() { .timeout(DEFAULT_QUERY_TIMEOUT) .query(query) .sort(DOC_FIELD_NAME, ASC), - requestBuilder.getSourceBuilder() + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder() ); } @@ -168,7 +172,7 @@ void test_push_down_query_and_sort() { .timeout(DEFAULT_QUERY_TIMEOUT) .query(query) .sort(sortBuilder), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -182,7 +186,7 @@ void test_push_down_sort() { .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(sortBuilder), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -196,7 +200,7 @@ void test_push_down_non_field_sort() { .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(sortBuilder), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -212,7 +216,7 @@ void test_push_down_multiple_sort() { .timeout(DEFAULT_QUERY_TIMEOUT) .sort(SortBuilders.fieldSort("intA")) .sort(SortBuilders.fieldSort("intB")), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -226,7 +230,7 @@ void test_push_down_project() { .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .fetchSource(new String[]{"intA"}, new String[0]), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -256,7 +260,7 @@ void test_push_down_nested() { .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -289,7 +293,7 @@ void test_push_down_multiple_nested_with_same_path() { .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -326,7 +330,7 @@ void test_push_down_nested_with_filter() { .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -336,4 +340,12 @@ void test_push_type_mapping() { verify(exprValueFactory).extendTypeMapping(typeMapping); } + + @Test + void push_down_highlight_with_repeating_fields() { + requestBuilder.pushDownHighlight("name", Map.of()); + var exception = assertThrows(SemanticCheckException.class, () -> + requestBuilder.pushDownHighlight("name", Map.of())); + assertEquals("Duplicate field name in highlight", exception.getMessage()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 2ff1de862b..12f12dc03f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -6,35 +6,7 @@ package org.opensearch.sql.opensearch.storage; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.aMapWithSize; -import static org.hamcrest.Matchers.allOf; -import static org.hamcrest.Matchers.hasEntry; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; -import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; -import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; -import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.expression.DSL.literal; -import static org.opensearch.sql.expression.DSL.named; -import static org.opensearch.sql.expression.DSL.ref; -import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.eval; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.remove; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.rename; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; - import com.google.common.collect.ImmutableMap; -import java.util.Arrays; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; @@ -42,32 +14,39 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.ast.tree.Sort; import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.expression.aggregation.AvgAggregator; -import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.type.OpenSearchTextType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.mapping.IndexMapping; -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.PagedRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; -import org.opensearch.sql.opensearch.storage.scan.OpenSearchPagedIndexScan; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; -import org.opensearch.sql.storage.Table; + +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.*; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.expression.DSL.named; +import static org.opensearch.sql.expression.DSL.ref; +import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; @ExtendWith(MockitoExtension.class) class OpenSearchIndexTest { @@ -83,9 +62,6 @@ class OpenSearchIndexTest { @Mock private Settings settings; - @Mock - private Table table; - @Mock private IndexMapping mapping; @@ -117,6 +93,7 @@ void createIndex() { OpenSearchDataType.of(MappingType.Keyword)))); schema.put("age", INTEGER); index.create(schema); + verify(client).createIndex(any(), any()); } @Test @@ -137,7 +114,7 @@ void getFieldTypes() { .put("id2", MappingType.Short) .put("blob", MappingType.Binary) .build().entrySet().stream().collect(Collectors.toMap( - e -> e.getKey(), e -> OpenSearchDataType.of(e.getValue()) + Map.Entry::getKey, e -> OpenSearchDataType.of(e.getValue()) ))); when(client.getIndexMappings("test")).thenReturn(ImmutableMap.of("test", mapping)); @@ -208,59 +185,31 @@ void getReservedFieldTypes() { @Test void implementRelationOperatorOnly() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); when(client.getIndexMaxResultWindows("test")).thenReturn(Map.of("test", 10000)); - + when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - assertEquals(new OpenSearchIndexScan(client, settings, indexName, + assertEquals(OpenSearchIndexScan.create(client, indexName, settings, maxResultWindow, exprValueFactory), index.implement(index.optimize(plan))); } - @Test - void implementPagedRelationOperatorOnly() { - when(client.getIndexMaxResultWindows("test")).thenReturn(Map.of("test", 10000)); - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); - - LogicalPlan plan = index.createPagedScanBuilder(42); - Integer maxResultWindow = index.getMaxResultWindow(); - PagedRequestBuilder builder = new InitialPageRequestBuilder( - new OpenSearchRequest.IndexName(indexName), - maxResultWindow, mock(), exprValueFactory); - assertEquals(new OpenSearchPagedIndexScan(client, builder), index.implement(plan)); - } - @Test void implementRelationOperatorWithOptimization() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); when(client.getIndexMaxResultWindows("test")).thenReturn(Map.of("test", 10000)); - + when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - assertEquals(new OpenSearchIndexScan(client, settings, indexName, - maxResultWindow, exprValueFactory), index.implement(plan)); + assertEquals(OpenSearchIndexScan.create(client, indexName, settings, + maxResultWindow, exprValueFactory), index.implement(plan)); } @Test void implementOtherLogicalOperators() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); when(client.getIndexMaxResultWindows("test")).thenReturn(Map.of("test", 10000)); - + when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); NamedExpression include = named("age", ref("age", INTEGER)); ReferenceExpression exclude = ref("name", STRING); ReferenceExpression dedupeField = ref("name", STRING); - Expression filterExpr = literal(ExprBooleanValue.of(true)); - List groupByExprs = Arrays.asList(named("age", ref("age", INTEGER))); - List aggregators = - Arrays.asList(named("avg(age)", new AvgAggregator(Arrays.asList(ref("age", INTEGER)), - DOUBLE))); Map mappings = ImmutableMap.of(ref("name", STRING), ref("lastname", STRING)); Pair newEvalField = @@ -292,7 +241,7 @@ void implementOtherLogicalOperators() { PhysicalPlanDSL.eval( PhysicalPlanDSL.remove( PhysicalPlanDSL.rename( - new OpenSearchIndexScan(client, settings, indexName, + OpenSearchIndexScan.create(client, indexName, settings, maxResultWindow, exprValueFactory), mappings), exclude), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index bde940a939..73f1bde7a5 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -6,44 +6,8 @@ package org.opensearch.sql.opensearch.storage.scan; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.reset; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; -import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; -import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; -import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; -import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; -import static org.opensearch.sql.data.type.ExprCoreType.LONG; -import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.limit; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.nested; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_AGGREGATION; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_FILTER; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_HIGHLIGHT; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_LIMIT; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_NESTED; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_PROJECT; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_SORT; - import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashSet; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; import lombok.Builder; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; @@ -53,7 +17,6 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.query.SpanOrQueryBuilder; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregationBuilders; import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder; @@ -65,13 +28,8 @@ import org.opensearch.sql.ast.tree.Sort.SortOption; import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValueUtils; -import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; -import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.expression.FunctionExpression; -import org.opensearch.sql.expression.HighlightExpression; -import org.opensearch.sql.expression.NamedExpression; -import org.opensearch.sql.expression.ReferenceExpression; +import org.opensearch.sql.expression.*; import org.opensearch.sql.expression.function.OpenSearchFunctions; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; @@ -79,13 +37,24 @@ import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; -import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalNested; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.storage.Table; +import java.util.*; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.*; +import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; +import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; +import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; +import static org.opensearch.sql.data.type.ExprCoreType.*; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.*; + @ExtendWith(MockitoExtension.class) class OpenSearchIndexScanOptimizationTest { @@ -105,16 +74,15 @@ class OpenSearchIndexScanOptimizationTest { @BeforeEach void setUp() { - indexScanBuilder = new OpenSearchIndexScanBuilder(indexScan); + indexScanBuilder = new OpenSearchIndexScanBuilder(t -> indexScan, requestBuilder); when(table.createScanBuilder()).thenReturn(indexScanBuilder); - when(indexScan.getRequestBuilder()).thenReturn(requestBuilder); } @Test void test_project_push_down() { assertEqualsAfterOptimization( project( - indexScanAggBuilder( + indexScanBuilder( withProjectPushedDown(DSL.ref("intV", INTEGER))), DSL.named("i", DSL.ref("intV", INTEGER)) ), @@ -678,16 +646,17 @@ void project_literal_should_not_be_pushed_down() { private OpenSearchIndexScanBuilder indexScanBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(new OpenSearchIndexScanQueryBuilder(indexScan)); + return new OpenSearchIndexScanBuilder(t -> indexScan, new OpenSearchIndexScanQueryBuilder(requestBuilder)); } private OpenSearchIndexScanBuilder indexScanAggBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(new OpenSearchIndexScanAggregationBuilder(indexScan)); + return new OpenSearchIndexScanBuilder(t -> indexScan, new OpenSearchIndexScanAggregationBuilder(requestBuilder)); } private void assertEqualsAfterOptimization(LogicalPlan expected, LogicalPlan actual) { - assertEquals(expected, optimize(actual)); + final var optimized = optimize(actual); + assertEquals(expected, optimized); // Trigger build to make sure all push down actually happened in scan builder indexScanBuilder.build(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index c788e78f1a..45c6fe1bcc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -42,7 +42,6 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; -import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; @@ -55,19 +54,20 @@ @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchIndexScanTest { + public static final int QUERY_SIZE = 200; @Mock private OpenSearchClient client; @Mock private Settings settings; - private OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( + private final OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( Map.of("name", OpenSearchDataType.of(STRING), "department", OpenSearchDataType.of(STRING))); @BeforeEach void setup() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); + lenient().when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(QUERY_SIZE); when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) .thenReturn(TimeValue.timeValueMinutes(1)); } @@ -75,8 +75,8 @@ void setup() { @Test void query_empty_result() { mockResponse(client); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "test", 3, exprValueFactory)) { + try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "test", settings, + 3, exprValueFactory)) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -93,8 +93,8 @@ void query_all_results_with_query() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "employees", 10, exprValueFactory)) { + try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, + 10, exprValueFactory)) { indexScan.open(); assertAll( @@ -114,14 +114,16 @@ void query_all_results_with_query() { verify(client).cleanup(any()); } + final static OpenSearchRequest.IndexName EMPLOYEES_INDEX + = new OpenSearchRequest.IndexName("employees"); @Test void query_all_results_with_scroll() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "employees", 10, exprValueFactory)) { + try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, + 10, exprValueFactory)) { indexScan.open(); assertAll( @@ -149,9 +151,10 @@ void query_some_results_with_query() { employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "employees", 10, exprValueFactory)) { - indexScan.getRequestBuilder().pushDownLimit(3, 0); + OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(0, exprValueFactory); + builder.pushDownLimit(3, 0); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, + 10, builder)) { indexScan.open(); assertAll( @@ -176,10 +179,10 @@ void query_some_results_with_scroll() { mockResponse(client, new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "employees", 2, exprValueFactory)) { - indexScan.getRequestBuilder().pushDownLimit(3, 0); + final var requestuilder = new OpenSearchRequestBuilder(10, exprValueFactory); + requestuilder.pushDownLimit(3, 0); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, + 2, requestuilder)) { indexScan.open(); assertAll( @@ -208,8 +211,8 @@ void query_results_limited_by_query_size() { employee(4, "Bob", "HR")}); when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(2); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "employees", 10, exprValueFactory)) { + try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, + 10, exprValueFactory)) { indexScan.open(); assertAll( @@ -270,32 +273,14 @@ void push_down_highlight_with_arguments() { highlightBuilder); } - @Test - void push_down_highlight_with_repeating_fields() { - mockResponse(client, - new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, - new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, settings, - "test", 2, exprValueFactory)) { - indexScan.getRequestBuilder().pushDownLimit(3, 0); - indexScan.open(); - Map args = new HashMap<>(); - indexScan.getRequestBuilder().pushDownHighlight("name", args); - indexScan.getRequestBuilder().pushDownHighlight("name", args); - } catch (SemanticCheckException e) { - assertTrue(e.getClass().equals(SemanticCheckException.class)); - } - verify(client).cleanup(any()); - } - private PushDownAssertion assertThat() { return new PushDownAssertion(client, exprValueFactory, settings); } private static class PushDownAssertion { private final OpenSearchClient client; - private final OpenSearchIndexScan indexScan; + private final OpenSearchRequestBuilder requestBuilder; + private final Settings settings; private final OpenSearchResponse response; private final OpenSearchExprValueFactory factory; @@ -303,40 +288,45 @@ public PushDownAssertion(OpenSearchClient client, OpenSearchExprValueFactory valueFactory, Settings settings) { this.client = client; - this.indexScan = new OpenSearchIndexScan(client, settings, - "test", 10000, valueFactory); + this.requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, valueFactory); + this.settings = settings; + this.response = mock(OpenSearchResponse.class); this.factory = valueFactory; when(response.isEmpty()).thenReturn(true); } PushDownAssertion pushDown(QueryBuilder query) { - indexScan.getRequestBuilder().pushDownFilter(query); + requestBuilder.pushDownFilter(query); return this; } PushDownAssertion pushDownHighlight(String query, Map arguments) { - indexScan.getRequestBuilder().pushDownHighlight(query, arguments); + requestBuilder.pushDownHighlight(query, arguments); return this; } PushDownAssertion shouldQueryHighlight(QueryBuilder query, HighlightBuilder highlight) { - OpenSearchRequest request = new OpenSearchQueryRequest("test", 200, factory); + OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, QUERY_SIZE, factory); request.getSourceBuilder() .query(query) .highlighter(highlight) .sort(DOC_FIELD_NAME, ASC); when(client.search(request)).thenReturn(response); + var indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, + 10000, requestBuilder); indexScan.open(); return this; } PushDownAssertion shouldQuery(QueryBuilder expected) { - OpenSearchRequest request = new OpenSearchQueryRequest("test", 200, factory); + OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, QUERY_SIZE, factory); request.getSourceBuilder() .query(expected) .sort(DOC_FIELD_NAME, ASC); when(client.search(request)).thenReturn(response); + var indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, + 10000, requestBuilder); indexScan.open(); return this; } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java deleted file mode 100644 index cd94154012..0000000000 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchPagedIndexScanTest.java +++ /dev/null @@ -1,215 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.storage.scan; - -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.CALLS_REAL_METHODS; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.mockito.Mockito.withSettings; -import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.employee; -import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; - -import com.google.common.collect.ImmutableMap; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; -import java.io.ObjectInputStream; -import java.io.ObjectOutputStream; -import java.util.Map; -import lombok.SneakyThrows; -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.data.model.ExprValue; -import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.executor.pagination.PlanSerializer; -import org.opensearch.sql.opensearch.client.OpenSearchClient; -import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; -import org.opensearch.sql.opensearch.request.InitialPageRequestBuilder; -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.opensearch.storage.OpenSearchStorageEngine; - -@ExtendWith(MockitoExtension.class) -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class OpenSearchPagedIndexScanTest { - @Mock - private OpenSearchClient client; - - private final OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( - ImmutableMap.of( - "name", OpenSearchDataType.of(STRING), - "department", OpenSearchDataType.of(STRING))); - - @Test - void query_empty_result() { - mockResponse(client); - InitialPageRequestBuilder builder = new InitialPageRequestBuilder( - new OpenSearchRequest.IndexName("test"), 3, mock(), exprValueFactory); - try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { - indexScan.open(); - assertFalse(indexScan.hasNext()); - } - verify(client).cleanup(any()); - } - - @Test - void query_all_results_initial_scroll_request() { - mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); - - PagedRequestBuilder builder = new InitialPageRequestBuilder( - new OpenSearchRequest.IndexName("test"), 3, mock(), exprValueFactory); - try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { - indexScan.open(); - - assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) - ); - } - verify(client).cleanup(any()); - - builder = new ContinuePageRequestBuilder( - new OpenSearchRequest.IndexName("test"), "scroll", mock(), exprValueFactory); - try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { - indexScan.open(); - - assertFalse(indexScan.hasNext()); - } - verify(client, times(2)).cleanup(any()); - } - - @Test - void query_all_results_continuation_scroll_request() { - mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); - - ContinuePageRequestBuilder builder = new ContinuePageRequestBuilder( - new OpenSearchRequest.IndexName("test"), "scroll", mock(), exprValueFactory); - try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { - indexScan.open(); - - assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) - ); - } - verify(client).cleanup(any()); - - builder = new ContinuePageRequestBuilder( - new OpenSearchRequest.IndexName("test"), "scroll", mock(), exprValueFactory); - try (OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder)) { - indexScan.open(); - - assertFalse(indexScan.hasNext()); - } - verify(client, times(2)).cleanup(any()); - } - - @Test - void explain_not_implemented() { - assertThrows(Throwable.class, () -> mock(OpenSearchPagedIndexScan.class, - withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); - } - - @Test - @SneakyThrows - void serialization() { - PagedRequestBuilder builder = mock(); - OpenSearchRequest request = mock(); - OpenSearchResponse response = mock(); - when(request.toCursor()).thenReturn("cu-cursor"); - when(builder.build()).thenReturn(request); - var indexName = new OpenSearchRequest.IndexName("index"); - when(builder.getIndexName()).thenReturn(indexName); - when(client.search(any())).thenReturn(response); - OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder); - indexScan.open(); - - ByteArrayOutputStream output = new ByteArrayOutputStream(); - ObjectOutputStream objectOutput = new ObjectOutputStream(output); - objectOutput.writeObject(indexScan); - objectOutput.flush(); - - when(client.getIndexMappings(any())).thenReturn(Map.of()); - OpenSearchStorageEngine engine = mock(); - when(engine.getClient()).thenReturn(client); - when(engine.getSettings()).thenReturn(mock()); - ObjectInputStream objectInput = new PlanSerializer(engine) - .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); - var roundTripScan = (OpenSearchPagedIndexScan) objectInput.readObject(); - roundTripScan.open(); - - // indexScan's request could be a OpenSearchScrollRequest or a ContinuePageRequest, but - // roundTripScan's request is always a ContinuePageRequest - // Thus, we can't compare those scans - //assertEquals(indexScan, roundTripScan); - // But we can validate that index name and scroll was serialized-deserialized correctly - assertEquals(indexName, roundTripScan.getRequestBuilder().getIndexName()); - assertTrue(roundTripScan.getRequestBuilder() instanceof ContinuePageRequestBuilder); - assertEquals("cu-cursor", - ((ContinuePageRequestBuilder) roundTripScan.getRequestBuilder()).getScrollId()); - } - - @Test - @SneakyThrows - void dont_serialize_if_no_cursor() { - PagedRequestBuilder builder = mock(); - OpenSearchRequest request = mock(); - OpenSearchResponse response = mock(); - when(builder.build()).thenReturn(request); - when(client.search(any())).thenReturn(response); - OpenSearchPagedIndexScan indexScan = new OpenSearchPagedIndexScan(client, builder); - indexScan.open(); - - when(request.toCursor()).thenReturn(null, ""); - for (int i = 0; i < 2; i++) { - assertThrows(NoCursorException.class, () -> { - ByteArrayOutputStream output = new ByteArrayOutputStream(); - ObjectOutputStream objectOutput = new ObjectOutputStream(output); - objectOutput.writeObject(indexScan); - objectOutput.flush(); - }); - } - } -} From ef9337532401a1ecb1efc39f446aba0d460a70a2 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 1 May 2023 22:33:37 -0700 Subject: [PATCH 02/38] Some bug fixes Signed-off-by: MaxKsyunz --- .../sql/planner/optimizer/PushDownPageSize.java | 14 ++++++++------ .../request/OpenSearchRequestBuilder.java | 2 +- .../storage/scan/OpenSearchIndexScan.java | 2 +- .../storage/scan/OpenSearchIndexScanTest.java | 1 - 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java index 47999af37b..066f908330 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java @@ -9,27 +9,29 @@ import com.facebook.presto.matching.Pattern; import java.util.ArrayDeque; import java.util.Deque; +import java.util.Optional; + import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.storage.read.TableScanBuilder; public class PushDownPageSize implements Rule { - TableScanBuilder builder; @Override public Pattern pattern() { - return Pattern.typeOf(LogicalPaginate.class).matching(this::findTableScanBuilder); + return Pattern.typeOf(LogicalPaginate.class).matching(lp -> findTableScanBuilder(lp).isPresent()); } @Override public LogicalPlan apply(LogicalPaginate plan, Captures captures) { + var builder = findTableScanBuilder(plan).orElseThrow(); if (!builder.pushDownPageSize(plan)) { throw new IllegalStateException("Failed to push down LogicalPaginate"); } return plan.getChild().get(0); } - private boolean findTableScanBuilder(LogicalPaginate logicalPaginate) { + private Optional findTableScanBuilder(LogicalPaginate logicalPaginate) { Deque plans = new ArrayDeque<>(); plans.add(logicalPaginate); do { @@ -38,12 +40,12 @@ private boolean findTableScanBuilder(LogicalPaginate logicalPaginate) { if (children.stream().anyMatch(TableScanBuilder.class::isInstance)) { if (children.size() > 1) { throw new UnsupportedOperationException( - "Unsupported plan: relation operator cannot have siblings"); + "Unsupported plan: relation operator cannot have siblings"); } - builder = (TableScanBuilder) children.get(0); + return Optional.of((TableScanBuilder) children.get(0)); } plans.addAll(children); } while (!plans.isEmpty()); - return builder != null; + return Optional.empty(); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 7466862eee..8c233e34d6 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -111,7 +111,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, indexName, scrollTimeout, sourceBuilder, exprValueFactory); } else { sourceBuilder.from(startFrom); - sourceBuilder.size(size); + sourceBuilder.size(querySize); return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); } } else { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 33969733e9..4c5d6a75e8 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -65,7 +65,7 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab private transient Iterator iterator; /** - * Constructor. + * Factory method used in tests. */ public static OpenSearchIndexScan create(OpenSearchClient client, String indexName, diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 45c6fe1bcc..c5c7bce586 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -346,7 +346,6 @@ public OpenSearchResponse answer(InvocationOnMock invocation) { when(response.isEmpty()).thenReturn(false); ExprValue[] searchHit = searchHitBatches[batchNum]; when(response.iterator()).thenReturn(Arrays.asList(searchHit).iterator()); - // used in OpenSearchPagedIndexScanTest lenient().when(response.getTotalHits()) .thenReturn((long) searchHitBatches[batchNum].length); } else { From 306a7fe8bf2d9247f7ce7d2d3ddd2e9e94294b50 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 2 May 2023 00:30:39 -0700 Subject: [PATCH 03/38] Updating tests 10000 is default query size set in `integ-test/build.gradle` Previous request builder wasn't setting size correctly for explain requests so output didn't match what was executed. Signed-off-by: MaxKsyunz --- .../resources/expectedOutput/ppl/explain_filter_agg_push.json | 2 +- .../src/test/resources/expectedOutput/ppl/explain_output.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json index 2d7f5f8c08..568b397f07 100644 --- a/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json +++ b/integ-test/src/test/resources/expectedOutput/ppl/explain_filter_agg_push.json @@ -8,7 +8,7 @@ { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":0,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)" + "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)" }, "children": [] } diff --git a/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json b/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json index 45988e35c7..8d45714283 100644 --- a/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json +++ b/integ-test/src/test/resources/expectedOutput/ppl/explain_output.json @@ -31,7 +31,7 @@ { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":0,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)" + "request": "OpenSearchQueryRequest(indexName\u003dopensearch-sql_test_index_account, sourceBuilder\u003d{\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":30,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"state\":{\"terms\":{\"field\":\"state.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}},{\"city\":{\"terms\":{\"field\":\"city.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg_age\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone\u003dfalse)" }, "children": [] } From 7d5b126ed59e6fd79af07d6497130c9ee7029c14 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 2 May 2023 00:35:24 -0700 Subject: [PATCH 04/38] Updating PaginationWindowIT Paginated response no longer sends an empty last page. Last page of the response now lacks cursor property. This matches legacy behaviour. Signed-off-by: MaxKsyunz --- .../test/java/org/opensearch/sql/sql/PaginationWindowIT.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java index 724451ef65..3f111feb28 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java @@ -39,6 +39,7 @@ public void testFetchSizeLessThanMaxResultWindow() throws IOException { numRows += response.getJSONArray("datarows").length(); response = executeCursorQuery(cursor); } while (response.has("cursor")); + numRows += response.getJSONArray("datarows").length(); var countRows = executeJdbcRequest("SELECT COUNT(*) FROM " + TEST_INDEX_PHRASE) .getJSONArray("datarows") @@ -62,7 +63,7 @@ public void testQuerySizeLimitDoesNotEffectTotalRowsReturned() throws IOExceptio numRows += response.getJSONArray("datarows").length(); response = executeCursorQuery(cursor); } while (response.has("cursor")); - + numRows += response.getJSONArray("datarows").length(); var countRows = executeJdbcRequest("SELECT COUNT(*) FROM " + TEST_INDEX_PHRASE) .getJSONArray("datarows") .getJSONArray(0) From dc2c47199dcf8a31d22f67fef04c41978c57bb15 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 2 May 2023 00:46:16 -0700 Subject: [PATCH 05/38] Addressing Code Analysis warnings Signed-off-by: MaxKsyunz --- .../opensearch/sql/sql/PaginationWindowIT.java | 6 +++--- .../request/OpenSearchRequestBuilder.java | 17 ++++++++++------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java index 3f111feb28..6a978a1116 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java @@ -21,7 +21,7 @@ public void init() throws IOException { } @After - void resetParams() throws IOException { + public void resetParams() throws IOException { resetMaxResultWindow(TEST_INDEX_PHRASE); resetQuerySizeLimit(); } @@ -31,7 +31,7 @@ public void testFetchSizeLessThanMaxResultWindow() throws IOException { setMaxResultWindow(TEST_INDEX_PHRASE, 6); JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5); - String cursor = ""; + String cursor; int numRows = 0; do { // Process response @@ -55,7 +55,7 @@ public void testQuerySizeLimitDoesNotEffectTotalRowsReturned() throws IOExceptio JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5); assertTrue(response.getInt("size") > querySizeLimit); - String cursor = ""; + String cursor; int numRows = 0; do { // Process response diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 8c233e34d6..405253e741 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -12,7 +12,11 @@ import org.apache.commons.lang3.tuple.Pair; import org.apache.lucene.search.join.ScoreMode; import org.opensearch.common.unit.TimeValue; -import org.opensearch.index.query.*; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.InnerHitBuilder; +import org.opensearch.index.query.NestedQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.FetchSourceContext; @@ -30,7 +34,6 @@ import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; @@ -67,7 +70,7 @@ public class OpenSearchRequestBuilder implements PushDownRequestBuilder { /** * Size of each page request to return. */ - private Optional pageSize = Optional.empty(); + private Integer pageSize = null; /** * OpenSearchExprValueFactory. @@ -91,7 +94,7 @@ public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprVa @Override public int getQuerySize() { - return pageSize.orElse(querySize); + return pageSize == null? querySize : pageSize; } /** * Build DSL request. @@ -104,7 +107,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, Settings settings) { int size = querySize; TimeValue scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); - if (pageSize.isEmpty()) { + if (pageSize == null) { if (startFrom + size > maxResultWindow) { sourceBuilder.size(maxResultWindow - startFrom); return new OpenSearchScrollRequest( @@ -118,7 +121,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, if (startFrom != 0) { throw new UnsupportedOperationException("Non-zero offset is not supported with pagination"); } - sourceBuilder.size(pageSize.get()); + sourceBuilder.size(pageSize); return new OpenSearchScrollRequest(indexName, scrollTimeout, sourceBuilder, exprValueFactory); } @@ -201,7 +204,7 @@ public void pushDownTrackedScore(boolean trackScores) { @Override public void pushDownPageSize(int pageSize) { - this.pageSize = Optional.of(pageSize); + this.pageSize = pageSize; } /** From 91cea6124ffbad28a1fb104e92c298af4b89d7c7 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 2 May 2023 11:26:51 -0700 Subject: [PATCH 06/38] Do not serialize OpenSearchScrollRequest when needClean is true. When OpenSearchScrollRequest.needClean is true, the object represents the last request. Signed-off-by: MaxKsyunz --- .../sql/sql/PaginationBlackboxIT.java | 83 ++++++++----------- .../request/OpenSearchScrollRequest.java | 2 +- 2 files changed, 37 insertions(+), 48 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java index d8213b1fe4..285fc1f989 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java @@ -6,14 +6,9 @@ package org.opensearch.sql.sql; -import static org.opensearch.sql.legacy.TestUtils.getResponseBody; -import static org.opensearch.sql.legacy.TestUtils.isIndexExist; -import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ONLINE; - +import java.io.IOException; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; -import java.util.stream.Collectors; import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; @@ -23,25 +18,30 @@ import org.junit.Test; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; -import org.opensearch.client.Request; import org.opensearch.sql.legacy.SQLIntegTestCase; +import org.opensearch.sql.util.TestUtils; // This class has only one test case, because it is parametrized and takes significant time @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class PaginationBlackboxIT extends SQLIntegTestCase { - private final String index; + private final Index index; private final Integer pageSize; - public PaginationBlackboxIT(@Name("index") String index, + public PaginationBlackboxIT(@Name("index") Index index, @Name("pageSize") Integer pageSize) { this.index = index; this.pageSize = pageSize; } + @Override + protected void init() throws IOException { + loadIndex(index); + } + @ParametersFactory(argumentFormatting = "index = %1$s, page_size = %2$d") public static Iterable compareTwoDates() { - var indices = new PaginationBlackboxHelper().getIndices(); + var indices = List.of(Index.ACCOUNT, Index.BEER, Index.BANK); var pageSizes = List.of(5, 10, 100, 1000); var testData = new ArrayList(); for (var index : indices) { @@ -55,63 +55,52 @@ public static Iterable compareTwoDates() { @Test @SneakyThrows public void test_pagination_blackbox() { - var response = executeJdbcRequest(String.format("select * from %s", index)); + var response = executeJdbcRequest(String.format("select * from %s", index.getName())); var indexSize = response.getInt("total"); var rows = response.getJSONArray("datarows"); var schema = response.getJSONArray("schema"); - var testReportPrefix = String.format("index: %s, page size: %d || ", index, pageSize); + var testReportPrefix = String.format("index: %s, page size: %d || ", index.getName(), pageSize); var rowsPaged = new JSONArray(); var rowsReturned = 0; - response = new JSONObject(executeFetchQuery( - String.format("select * from %s", index), pageSize, "jdbc")); + var responseCounter = 1; this.logger.info(testReportPrefix + "first response"); - while (response.has("cursor")) { - assertEquals(indexSize, response.getInt("total")); + response = new JSONObject(executeFetchQuery( + String.format("select * from %s", index.getName()), pageSize, "jdbc")); + + var cursor = response.has("cursor")? response.getString("cursor") : ""; + do { + this.logger.info(testReportPrefix + + String.format("subsequent response %d/%d", responseCounter++, (indexSize / pageSize) + 1)); assertTrue("Paged response schema doesn't match to non-paged", schema.similar(response.getJSONArray("schema"))); - var cursor = response.getString("cursor"); - assertTrue(testReportPrefix + "Cursor returned from legacy engine", - cursor.startsWith("n:")); + rowsReturned += response.getInt("size"); var datarows = response.getJSONArray("datarows"); for (int i = 0; i < datarows.length(); i++) { rowsPaged.put(datarows.get(i)); } - response = executeCursorQuery(cursor); - this.logger.info(testReportPrefix - + String.format("subsequent response %d/%d", responseCounter++, (indexSize / pageSize) + 1)); - } + + if (response.has("cursor")) { + TestUtils.verifyIsV2Cursor(response); + cursor = response.getString("cursor"); + response = executeCursorQuery(cursor); + } else { + cursor = ""; + } + + } while(!cursor.isEmpty()); assertTrue("Paged response schema doesn't match to non-paged", schema.similar(response.getJSONArray("schema"))); - assertEquals(0, response.getInt("total")); +// assertEquals(0, response.getInt("total")); - assertEquals(testReportPrefix + "Last page is not empty", - 0, response.getInt("size")); - assertEquals(testReportPrefix + "Last page is not empty", - 0, response.getJSONArray("datarows").length()); +// assertEquals(testReportPrefix + "Last page is not empty", +// 0, response.getInt("size")); +// assertEquals(testReportPrefix + "Last page is not empty", +// 0, response.getJSONArray("datarows").length()); assertEquals(testReportPrefix + "Paged responses return another row count that non-paged", indexSize, rowsReturned); assertTrue(testReportPrefix + "Paged accumulated result has other rows than non-paged", rows.similar(rowsPaged)); } - - // A dummy class created, because accessing to `client()` isn't available from a static context, - // but it is needed before an instance of `PaginationBlackboxIT` is created. - private static class PaginationBlackboxHelper extends SQLIntegTestCase { - - @SneakyThrows - private List getIndices() { - initClient(); - loadIndex(Index.ACCOUNT); - loadIndex(Index.BEER); - loadIndex(Index.BANK); - if (!isIndexExist(client(), "empty")) { - executeRequest(new Request("PUT", "/empty")); - } - return Arrays.stream(getResponseBody(client().performRequest(new Request("GET", "_cat/indices?h=i")), true).split("\n")) - // exclude this index, because it is too big and extends test time too long (almost 10k docs) - .map(String::trim).filter(i -> !i.equals(TEST_INDEX_ONLINE)).collect(Collectors.toList()); - } - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 03f97cda61..46663e1fdc 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -153,6 +153,6 @@ public void reset() { */ @Override public String toCursor() { - return scrollId; + return needClean? "" : scrollId; } } From 2ee810c004ea86ac33a4c5a510793c3aecfbcc94 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 2 May 2023 14:02:08 -0700 Subject: [PATCH 07/38] Fix checkstyle errors. Signed-off-by: MaxKsyunz --- .../planner/optimizer/PushDownPageSize.java | 4 +- .../sql/storage/read/TableScanBuilder.java | 4 +- .../optimizer/LogicalPlanOptimizerTest.java | 139 +++++++++--------- .../request/ContinuePageRequest.java | 7 +- .../request/ContinuePageRequestBuilder.java | 9 +- .../request/OpenSearchRequestBuilder.java | 31 ++-- .../request/OpenSearchScrollRequest.java | 2 +- .../opensearch/storage/OpenSearchIndex.java | 9 +- .../storage/scan/OpenSearchIndexScan.java | 12 +- .../scan/OpenSearchIndexScanBuilder.java | 31 ++-- .../storage/scan/PushDownTranslator.java | 1 + .../client/OpenSearchNodeClientTest.java | 8 +- .../client/OpenSearchRestClientTest.java | 4 +- .../OpenSearchExecutionProtectorTest.java | 2 +- .../ContinuePageRequestBuilderTest.java | 14 +- .../request/OpenSearchRequestBuilderTest.java | 2 +- .../storage/OpenSearchIndexTest.java | 46 +++--- .../OpenSearchIndexScanOptimizationTest.java | 60 ++++++-- .../storage/scan/OpenSearchIndexScanTest.java | 3 +- 19 files changed, 225 insertions(+), 163 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java index 066f908330..8d88c7ef43 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java @@ -10,7 +10,6 @@ import java.util.ArrayDeque; import java.util.Deque; import java.util.Optional; - import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.storage.read.TableScanBuilder; @@ -18,7 +17,8 @@ public class PushDownPageSize implements Rule { @Override public Pattern pattern() { - return Pattern.typeOf(LogicalPaginate.class).matching(lp -> findTableScanBuilder(lp).isPresent()); + return Pattern.typeOf(LogicalPaginate.class) + .matching(lp -> findTableScanBuilder(lp).isPresent()); } @Override diff --git a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java index 713b428546..f0158c52b8 100644 --- a/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java +++ b/core/src/main/java/org/opensearch/sql/storage/read/TableScanBuilder.java @@ -117,7 +117,9 @@ public boolean pushDownNested(LogicalNested nested) { return false; } - public boolean pushDownPageSize(LogicalPaginate paginate) { return false; } + public boolean pushDownPageSize(LogicalPaginate paginate) { + return false; + } @Override public R accept(LogicalPlanNodeVisitor visitor, C context) { diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index 3217252800..ce61380d1d 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -6,7 +6,23 @@ package org.opensearch.sql.planner.optimizer; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; +import static org.opensearch.sql.data.model.ExprValueUtils.longValue; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.data.type.ExprCoreType.LONG; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values; + import com.google.common.collect.ImmutableList; +import java.util.Collections; +import java.util.List; +import java.util.Map; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -24,6 +40,7 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; import org.opensearch.sql.planner.logical.LogicalRelation; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -31,20 +48,6 @@ import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; -import java.util.Collections; -import java.util.List; -import java.util.Map; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.*; -import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; -import static org.opensearch.sql.data.model.ExprValueUtils.longValue; -import static org.opensearch.sql.data.type.ExprCoreType.*; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; - @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class LogicalPlanOptimizerTest { @@ -66,15 +69,15 @@ void setUp() { @Test void filter_merge_filter() { assertEquals( - filter( + LogicalPlanDSL.filter( tableScanBuilder, DSL.and(DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(2))), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(1)))) ), optimize( - filter( - filter( - relation("schema", table), + LogicalPlanDSL.filter( + LogicalPlanDSL.filter( + LogicalPlanDSL.relation("schema", table), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(1))) ), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(2))) @@ -89,17 +92,17 @@ void filter_merge_filter() { @Test void push_filter_under_sort() { assertEquals( - sort( - filter( + LogicalPlanDSL.sort( + LogicalPlanDSL.filter( tableScanBuilder, DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) ), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), optimize( - filter( - sort( - relation("schema", table), + LogicalPlanDSL.filter( + LogicalPlanDSL.sort( + LogicalPlanDSL.relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) @@ -114,8 +117,8 @@ void push_filter_under_sort() { @Test void multiple_filter_should_eventually_be_merged() { assertEquals( - sort( - filter( + LogicalPlanDSL.sort( + LogicalPlanDSL.filter( tableScanBuilder, DSL.and(DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))), DSL.less(DSL.ref("longV", INTEGER), DSL.literal(longValue(1L)))) @@ -123,10 +126,10 @@ void multiple_filter_should_eventually_be_merged() { Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), optimize( - filter( - sort( - filter( - relation("schema", table), + LogicalPlanDSL.filter( + LogicalPlanDSL.sort( + LogicalPlanDSL.filter( + LogicalPlanDSL.relation("schema", table), DSL.less(DSL.ref("longV", INTEGER), DSL.literal(longValue(1L))) ), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) @@ -140,25 +143,25 @@ void multiple_filter_should_eventually_be_merged() { @Test void default_table_scan_builder_should_not_push_down_anything() { LogicalPlan[] plans = { - project( - relation("schema", table), + LogicalPlanDSL.project( + LogicalPlanDSL.relation("schema", table), DSL.named("i", DSL.ref("intV", INTEGER)) ), - filter( - relation("schema", table), + LogicalPlanDSL.filter( + LogicalPlanDSL.relation("schema", table), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) ), - aggregation( - relation("schema", table), + LogicalPlanDSL.aggregation( + LogicalPlanDSL.relation("schema", table), ImmutableList .of(DSL.named("AVG(intV)", DSL.avg(DSL.ref("intV", INTEGER)))), ImmutableList.of(DSL.named("longV", DSL.ref("longV", LONG)))), - sort( - relation("schema", table), + LogicalPlanDSL.sort( + LogicalPlanDSL.relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("intV", INTEGER))), - limit( - relation("schema", table), + LogicalPlanDSL.limit( + LogicalPlanDSL.relation("schema", table), 1, 1) }; @@ -174,8 +177,8 @@ void table_scan_builder_support_project_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - project( - relation("schema", table), + LogicalPlanDSL.project( + LogicalPlanDSL.relation("schema", table), DSL.named("i", DSL.ref("intV", INTEGER))) ) ); @@ -188,8 +191,8 @@ void table_scan_builder_support_filter_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - filter( - relation("schema", table), + LogicalPlanDSL.filter( + LogicalPlanDSL.relation("schema", table), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1)))) ) ); @@ -202,8 +205,8 @@ void table_scan_builder_support_aggregation_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - aggregation( - relation("schema", table), + LogicalPlanDSL.aggregation( + LogicalPlanDSL.relation("schema", table), ImmutableList .of(DSL.named("AVG(intV)", DSL.avg(DSL.ref("intV", INTEGER)))), @@ -219,8 +222,8 @@ void table_scan_builder_support_sort_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - sort( - relation("schema", table), + LogicalPlanDSL.sort( + LogicalPlanDSL.relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("intV", INTEGER))) ) ); @@ -233,8 +236,8 @@ void table_scan_builder_support_limit_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - limit( - relation("schema", table), + LogicalPlanDSL.limit( + LogicalPlanDSL.relation("schema", table), 1, 1) ) ); @@ -247,8 +250,8 @@ void table_scan_builder_support_highlight_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - highlight( - relation("schema", table), + LogicalPlanDSL.highlight( + LogicalPlanDSL.relation("schema", table), DSL.literal("*"), Collections.emptyMap()) ) @@ -262,8 +265,8 @@ void table_scan_builder_support_nested_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - nested( - relation("schema", table), + LogicalPlanDSL.nested( + LogicalPlanDSL.relation("schema", table), List.of(Map.of("field", new ReferenceExpression("message.info", STRING))), List.of(new NamedExpression( "message.info", @@ -289,8 +292,8 @@ public PhysicalPlan implement(LogicalPlan plan) { }; assertEquals( - relation("schema", table), - optimize(relation("schema", table)) + LogicalPlanDSL.relation("schema", table), + optimize(LogicalPlanDSL.relation("schema", table)) ); } @@ -301,7 +304,7 @@ void table_support_write_builder_should_be_replaced() { assertEquals( writeBuilder, - optimize(write(values(), table, Collections.emptyList())) + optimize(LogicalPlanDSL.write(values(), table, Collections.emptyList())) ); } @@ -327,34 +330,38 @@ public PhysicalPlan implement(LogicalPlan plan) { void paged_table_scan_builder_support_project_push_down_can_apply_its_rule() { when(tableScanBuilder.pushDownPageSize(any())).thenReturn(true); - var relation = relation("schema", table); - var optimized = LogicalPlanOptimizer.create().optimize(paginate(project(relation), 4)); + var relation = LogicalPlanDSL.relation("schema", table); + var optimized = LogicalPlanOptimizer.create() + .optimize(LogicalPlanDSL.paginate(LogicalPlanDSL.project(relation), 4)); verify(tableScanBuilder).pushDownPageSize(any()); - assertEquals(project(tableScanBuilder), optimized); + assertEquals(LogicalPlanDSL.project(tableScanBuilder), optimized); } @Test void push_down_page_size_multiple_children() { - var relation = relation("schema", table); + var relation = LogicalPlanDSL.relation("schema", table); var twoChildrenPlan = new LogicalPlan(List.of(relation, relation)) { @Override public R accept(LogicalPlanNodeVisitor visitor, C context) { return null; } }; - var queryPlan = paginate(twoChildrenPlan, 4); + var queryPlan = LogicalPlanDSL.paginate(twoChildrenPlan, 4); var optimizer = LogicalPlanOptimizer.create(); final var exception = assertThrows(UnsupportedOperationException.class, () -> optimizer.optimize(queryPlan)); - assertEquals("Unsupported plan: relation operator cannot have siblings", exception.getMessage()); + assertEquals("Unsupported plan: relation operator cannot have siblings", + exception.getMessage()); } @Test void push_down_page_size_push_failed() { when(tableScanBuilder.pushDownPageSize(any())).thenReturn(false); - var queryPlan = paginate(project(relation("schema", table)), 4); + var queryPlan = LogicalPlanDSL.paginate( + LogicalPlanDSL.project( + LogicalPlanDSL.relation("schema", table)), 4); var optimizer = LogicalPlanOptimizer.create(); final var exception = assertThrows(IllegalStateException.class, () -> optimizer.optimize(queryPlan)); @@ -363,7 +370,7 @@ void push_down_page_size_push_failed() { @Test void push_page_size_noop_if_no_relation() { - var paginate = new LogicalPaginate(42, List.of(project(values()))); + var paginate = new LogicalPaginate(42, List.of(LogicalPlanDSL.project(values()))); assertEquals(paginate, LogicalPlanOptimizer.create().optimize(paginate)); } @@ -380,10 +387,10 @@ void table_scan_builder_support_offset_push_down_can_apply_its_rule() { var relation = new LogicalRelation("schema", table); var optimized = LogicalPlanOptimizer.create() - .optimize(new LogicalPaginate(42, List.of(project(relation)))); + .optimize(new LogicalPaginate(42, List.of(LogicalPlanDSL.project(relation)))); // `optimized` structure: LogicalPaginate -> LogicalProject -> TableScanBuilder // LogicalRelation replaced by a TableScanBuilder instance - assertEquals(project(tableScanBuilder), optimized); + assertEquals(LogicalPlanDSL.project(tableScanBuilder), optimized); } private LogicalPlan optimize(LogicalPlan plan) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java index 92ec3dbfe3..0e06283429 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java @@ -5,6 +5,9 @@ package org.opensearch.sql.opensearch.request; +import java.util.List; +import java.util.function.Consumer; +import java.util.function.Function; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; @@ -17,10 +20,6 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; -import java.util.List; -import java.util.function.Consumer; -import java.util.function.Function; - /** * Scroll (cursor) request is used to page the search. This request is not configurable and has * no search query. It just handles paging through responses to the initial request. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index 0df16d73ec..e51ec5db2c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -5,6 +5,9 @@ package org.opensearch.sql.opensearch.request; +import java.util.List; +import java.util.Map; +import java.util.Set; import lombok.Getter; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.common.unit.TimeValue; @@ -18,10 +21,6 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -import java.util.List; -import java.util.Map; -import java.util.Set; - /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. */ @@ -62,7 +61,7 @@ public void pushDownFilter(QueryBuilder query) { @Override public void pushDownAggregation(Pair, - OpenSearchAggregationResponseParser> aggregationBuilder) { + OpenSearchAggregationResponseParser> aggregationBuilder) { throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 405253e741..71172e3d2f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -6,6 +6,17 @@ package org.opensearch.sql.opensearch.request; +import static java.util.stream.Collectors.mapping; +import static java.util.stream.Collectors.toList; +import static org.opensearch.index.query.QueryBuilders.matchAllQuery; +import static org.opensearch.index.query.QueryBuilders.nestedQuery; +import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; +import static org.opensearch.search.sort.SortOrder.ASC; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.ToString; @@ -32,17 +43,6 @@ import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; - -import static java.util.stream.Collectors.mapping; -import static java.util.stream.Collectors.toList; -import static org.opensearch.index.query.QueryBuilders.matchAllQuery; -import static org.opensearch.index.query.QueryBuilders.nestedQuery; -import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; -import static org.opensearch.search.sort.SortOrder.ASC; /** * OpenSearch search request builder. @@ -94,8 +94,9 @@ public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprVa @Override public int getQuerySize() { - return pageSize == null? querySize : pageSize; + return pageSize == null ? querySize : pageSize; } + /** * Build DSL request. * @@ -131,6 +132,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, boolean isBoolFilterQuery(QueryBuilder current) { return (current instanceof BoolQueryBuilder); } + /** * Push down query to DSL request. * @@ -246,8 +248,9 @@ public void pushDownHighlight(String field, Map arguments) { */ @Override public void pushDownProjects(Set projects) { - sourceBuilder.fetchSource(projects.stream().map(ReferenceExpression::getAttr).distinct().toArray(String[]::new), - new String[0]); + sourceBuilder.fetchSource( + projects.stream().map(ReferenceExpression::getAttr).distinct().toArray(String[]::new), + new String[0]); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 46663e1fdc..b85c69a60e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -153,6 +153,6 @@ public void reset() { */ @Override public String toCursor() { - return needClean? "" : scrollId; + return needClean ? "" : scrollId; } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 2cea742eed..2a2c71022b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -7,6 +7,10 @@ package org.opensearch.sql.opensearch.storage; import com.google.common.annotations.VisibleForTesting; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.function.Function; import lombok.RequiredArgsConstructor; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; @@ -31,11 +35,6 @@ import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.read.TableScanBuilder; -import java.util.HashMap; -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.function.Function; - /** OpenSearch table (index) implementation. */ public class OpenSearchIndex implements Table { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 4c5d6a75e8..9c78f36543 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -79,7 +79,9 @@ public static OpenSearchIndexScan create(OpenSearchClient client, settings, maxResultWindow, requestBuilder); } - + /** + * Creates index scan based on a provided OpenSearchRequestBuilder. + */ public OpenSearchIndexScan(OpenSearchClient client, OpenSearchRequest.IndexName indexName, Settings settings, @@ -144,17 +146,17 @@ public String explain() { return requestBuilder.build(indexName, maxResultWindow, settings).toString(); } - /** + /** No-args constructor. * @deprecated Exists only to satisfy Java serialization API. */ - @Deprecated(since="introduction") + @Deprecated(since = "introduction") public OpenSearchIndexScan() { } @Override public void readExternal(ObjectInput in) throws IOException { - var serializedName = in.readUTF(); - var scrollId = in.readUTF(); + final var serializedName = in.readUTF(); + final var scrollId = in.readUTF(); querySize = in.readInt(); var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index ed4259c86a..b0b742723a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -27,8 +27,12 @@ */ public class OpenSearchIndexScanBuilder extends TableScanBuilder { + /** + * A function to create an index scan based on a OpenSearchRequestBuilder. + */ @EqualsAndHashCode.Exclude - private final Function constructor; + private final Function scanConstructor; + /** * Delegated index scan builder for non-aggregate or aggregate query. */ @@ -38,22 +42,31 @@ public class OpenSearchIndexScanBuilder extends TableScanBuilder { /** Is limit operator pushed down. */ private boolean isLimitPushedDown = false; - public OpenSearchIndexScanBuilder(Function constructor, + + /** + * Constructor used during query execution. + */ + public OpenSearchIndexScanBuilder( + Function scanConstructor, OpenSearchRequestBuilder requestBuilder) { - this.constructor = constructor; - this.delegate - = new OpenSearchIndexScanQueryBuilder(requestBuilder); + this.scanConstructor = scanConstructor; + this.delegate = new OpenSearchIndexScanQueryBuilder(requestBuilder); } - public OpenSearchIndexScanBuilder(Function constructor, - PushDownTranslator translator) { - this.constructor = constructor; + + /** + * Constructor used for unit tests. + */ + public OpenSearchIndexScanBuilder(Function constructor, + PushDownTranslator translator) { + this.scanConstructor = constructor; this.delegate = translator; } @Override public TableScanOperator build() { - return constructor.apply(delegate.build()); + return scanConstructor.apply(delegate.build()); } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java index f178681706..bb67c472dd 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java @@ -29,6 +29,7 @@ public interface PushDownTranslator { boolean pushDownPageSize(LogicalPaginate paginate); boolean pushDownNested(LogicalNested nested); + boolean pushDownAggregation(LogicalAggregation aggregation); OpenSearchRequestBuilder build(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java index caafa40d57..b378fae297 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClientTest.java @@ -99,8 +99,8 @@ class OpenSearchNodeClientTest { @Mock private GetIndexResponse indexResponse; - private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", - new ExprIntegerValue(1))); + private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap( + Map.of("id", new ExprIntegerValue(1))); private OpenSearchClient client; @@ -451,8 +451,8 @@ private void mockNodeClientSettings(String indexName, String indexMetadata) GetSettingsResponse mockResponse = mock(GetSettingsResponse.class); when(nodeClient.admin().indices().prepareGetSettings(any()).setLocal(anyBoolean()).get()) .thenReturn(mockResponse); - Map metadata =Map.of(indexName, - IndexMetadata.fromXContent(createParser(indexMetadata)).getSettings()); + Map metadata = Map.of(indexName, + IndexMetadata.fromXContent(createParser(indexMetadata)).getSettings()); when(mockResponse.getIndexToSettings()).thenReturn(metadata); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java index e019a919fc..2958fa1100 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/client/OpenSearchRestClientTest.java @@ -92,8 +92,8 @@ class OpenSearchRestClientTest { @Mock private GetIndexResponse getIndexResponse; - private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap(ImmutableMap.of("id", - new ExprIntegerValue(1))); + private final ExprTupleValue exprTupleValue = ExprTupleValue.fromExprValueMap( + Map.of("id", new ExprIntegerValue(1))); @BeforeEach void setUp() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index d2aba45f0b..b5fff14614 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -100,7 +100,7 @@ void testProtectIndexScan() { Expression filterExpr = literal(ExprBooleanValue.of(true)); List groupByExprs = List.of(named("age", ref("age", INTEGER))); List aggregators = - List.of(named("avg(age)", new AvgAggregator(List.of(ref("age", INTEGER)), + List.of(named("avg(age)", new AvgAggregator(List.of(ref("age", INTEGER)), DOUBLE))); Map mappings = ImmutableMap.of(ref("name", STRING), ref("lastname", STRING)); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java index f149e0d3b6..3384e35f3d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java @@ -5,6 +5,13 @@ package org.opensearch.sql.opensearch.request; +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; + +import java.util.List; +import java.util.Map; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -13,15 +20,8 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.common.unit.TimeValue; -import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import java.util.List; -import java.util.Map; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.mock; - @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @ExtendWith(MockitoExtension.class) class ContinuePageRequestBuilderTest { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index d8e3dd5d5b..265bbc7142 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -345,7 +345,7 @@ void test_push_type_mapping() { void push_down_highlight_with_repeating_fields() { requestBuilder.pushDownHighlight("name", Map.of()); var exception = assertThrows(SemanticCheckException.class, () -> - requestBuilder.pushDownHighlight("name", Map.of())); + requestBuilder.pushDownHighlight("name", Map.of())); assertEquals("Duplicate field name in highlight", exception.getMessage()); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 12f12dc03f..0afaaa3e7c 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -6,7 +6,32 @@ package org.opensearch.sql.opensearch.storage; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.aMapWithSize; +import static org.hamcrest.Matchers.allOf; +import static org.hamcrest.Matchers.hasEntry; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.expression.DSL.named; +import static org.opensearch.sql.expression.DSL.ref; +import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.eval; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.remove; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.rename; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; + import com.google.common.collect.ImmutableMap; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Collectors; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; @@ -31,23 +56,6 @@ import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; -import java.util.HashMap; -import java.util.Map; -import java.util.stream.Collectors; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.*; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.*; -import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; -import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.expression.DSL.named; -import static org.opensearch.sql.expression.DSL.ref; -import static org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; - @ExtendWith(MockitoExtension.class) class OpenSearchIndexTest { @@ -81,9 +89,9 @@ void isExist() { @Test void createIndex() { - Map mappings = ImmutableMap.of( + Map mappings = Map.of( "properties", - ImmutableMap.of( + Map.of( "name", "text", "age", "integer")); doNothing().when(client).createIndex(indexName, mappings); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index 73f1bde7a5..12fa2552fe 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -6,8 +6,43 @@ package org.opensearch.sql.opensearch.storage.scan; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; +import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; +import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; +import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; +import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; +import static org.opensearch.sql.data.type.ExprCoreType.LONG; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.limit; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.nested; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_AGGREGATION; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_FILTER; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_HIGHLIGHT; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_LIMIT; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_NESTED; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_PROJECT; +import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.PUSH_DOWN_SORT; + import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; import lombok.Builder; import org.apache.commons.lang3.tuple.Pair; import org.junit.jupiter.api.BeforeEach; @@ -29,7 +64,11 @@ import org.opensearch.sql.data.model.ExprTupleValue; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.data.type.ExprType; -import org.opensearch.sql.expression.*; +import org.opensearch.sql.expression.DSL; +import org.opensearch.sql.expression.FunctionExpression; +import org.opensearch.sql.expression.HighlightExpression; +import org.opensearch.sql.expression.NamedExpression; +import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.function.OpenSearchFunctions; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; @@ -43,19 +82,6 @@ import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.storage.Table; -import java.util.*; -import java.util.stream.Collectors; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.*; -import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; -import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; -import static org.opensearch.sql.data.model.ExprValueUtils.integerValue; -import static org.opensearch.sql.data.type.ExprCoreType.*; -import static org.opensearch.sql.planner.logical.LogicalPlanDSL.*; -import static org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown.*; - - @ExtendWith(MockitoExtension.class) class OpenSearchIndexScanOptimizationTest { @@ -646,12 +672,14 @@ void project_literal_should_not_be_pushed_down() { private OpenSearchIndexScanBuilder indexScanBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(t -> indexScan, new OpenSearchIndexScanQueryBuilder(requestBuilder)); + return new OpenSearchIndexScanBuilder(t -> indexScan, + new OpenSearchIndexScanQueryBuilder(requestBuilder)); } private OpenSearchIndexScanBuilder indexScanAggBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(t -> indexScan, new OpenSearchIndexScanAggregationBuilder(requestBuilder)); + return new OpenSearchIndexScanBuilder(t -> indexScan, + new OpenSearchIndexScanAggregationBuilder(requestBuilder)); } private void assertEqualsAfterOptimization(LogicalPlan expected, LogicalPlan actual) { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index c5c7bce586..f1a8d6714b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -114,8 +114,9 @@ void query_all_results_with_query() { verify(client).cleanup(any()); } - final static OpenSearchRequest.IndexName EMPLOYEES_INDEX + static final OpenSearchRequest.IndexName EMPLOYEES_INDEX = new OpenSearchRequest.IndexName("employees"); + @Test void query_all_results_with_scroll() { mockResponse(client, From 39ce9028c9d8e003498675ff488e0bdbc394c3ab Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 3 May 2023 22:26:24 -0700 Subject: [PATCH 08/38] Complete unit test coverage. Signed-off-by: MaxKsyunz --- .../storage/scan/OpenSearchIndexScan.java | 1 - ...OpenSearchIndexScanAggregationBuilder.java | 15 +- .../scan/OpenSearchIndexScanBuilder.java | 5 +- .../scan/OpenSearchIndexScanQueryBuilder.java | 8 +- .../storage/scan/PushDownTranslator.java | 3 - .../ContinuePageRequestBuilderTest.java | 2 + .../request/OpenSearchRequestBuilderTest.java | 26 ++ .../request/OpenSearchScrollRequestTest.java | 15 ++ ...SearchIndexScanAggregationBuilderTest.java | 75 ++++++ .../OpenSearchIndexScanOptimizationTest.java | 39 ++- .../OpenSearchIndexScanPaginationTest.java | 228 ++++++++++++++++++ .../storage/scan/OpenSearchIndexScanTest.java | 10 +- 12 files changed, 399 insertions(+), 28 deletions(-) create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilderTest.java create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 9c78f36543..6f279a8a1f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -98,7 +98,6 @@ public OpenSearchIndexScan(OpenSearchClient client, public void open() { super.open(); request = requestBuilder.build(indexName, maxResultWindow, settings); - // TODO does this work for paged requests? querySize = requestBuilder.getQuerySize(); iterator = Collections.emptyIterator(); queryCount = 0; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 4b9643c4a5..c72c3e1064 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -39,17 +39,19 @@ class OpenSearchIndexScanAggregationBuilder implements PushDownTranslator { private final OpenSearchRequestBuilder requestBuilder; /** Aggregators pushed down. */ - private List aggregatorList; + private final List aggregatorList; /** Grouping items pushed down. */ - private List groupByList; + private final List groupByList; /** Sorting items pushed down. */ private List> sortList; - OpenSearchIndexScanAggregationBuilder(OpenSearchRequestBuilder requestBuilder) { + OpenSearchIndexScanAggregationBuilder(OpenSearchRequestBuilder requestBuilder, LogicalAggregation aggregation) { this.requestBuilder = requestBuilder; + aggregatorList = aggregation.getAggregatorList(); + groupByList = aggregation.getGroupByList(); } @Override @@ -64,13 +66,6 @@ public OpenSearchRequestBuilder build() { return requestBuilder; } - @Override - public boolean pushDownAggregation(LogicalAggregation aggregation) { - aggregatorList = aggregation.getAggregatorList(); - groupByList = aggregation.getGroupByList(); - return true; - } - @Override public boolean pushDownFilter(LogicalFilter filter) { return false; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index b0b742723a..1d0036fa74 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -82,9 +82,8 @@ public boolean pushDownAggregation(LogicalAggregation aggregation) { // Switch to builder for aggregate query which has different push down logic // for later filter, sort and limit operator. - delegate = new OpenSearchIndexScanAggregationBuilder(delegate.build()); - - return delegate.pushDownAggregation(aggregation); + delegate = new OpenSearchIndexScanAggregationBuilder(delegate.build(), aggregation); + return true; } @Override diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index 2c53257eb7..6df5b6802f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -26,7 +26,6 @@ import org.opensearch.sql.opensearch.storage.script.filter.FilterQueryBuilder; import org.opensearch.sql.opensearch.storage.script.sort.SortQueryBuilder; import org.opensearch.sql.opensearch.storage.serialization.DefaultExpressionSerializer; -import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; @@ -124,11 +123,6 @@ public boolean pushDownNested(LogicalNested nested) { return false; } - @Override - public boolean pushDownAggregation(LogicalAggregation aggregation) { - return false; - } - @Override public OpenSearchRequestBuilder build() { return requestBuilder; @@ -138,7 +132,7 @@ public OpenSearchRequestBuilder build() { * Find reference expression from expression. * @param expressions a list of expression. * - * @return a list of ReferenceExpression + * @return a set of ReferenceExpression */ public static Set findReferenceExpressions( List expressions) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java index bb67c472dd..79d5896fd0 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java @@ -6,7 +6,6 @@ package org.opensearch.sql.opensearch.storage.scan; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; -import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalHighlight; import org.opensearch.sql.planner.logical.LogicalLimit; @@ -30,7 +29,5 @@ public interface PushDownTranslator { boolean pushDownNested(LogicalNested nested); - boolean pushDownAggregation(LogicalAggregation aggregation); - OpenSearchRequestBuilder build(); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java index 3384e35f3d..b3c5c6ac05 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java @@ -66,6 +66,8 @@ void pushDown_not_supported() { () -> requestBuilder.pushTypeMapping(mock())), () -> assertThrows(UnsupportedOperationException.class, () -> requestBuilder.pushDownNested(List.of())), + () -> assertThrows(UnsupportedOperationException.class, + () -> requestBuilder.pushDownPageSize(3)), () -> assertThrows(UnsupportedOperationException.class, () -> requestBuilder.pushDownTrackedScore(true)) ); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 265bbc7142..bf7a9284a2 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -348,4 +348,30 @@ void push_down_highlight_with_repeating_fields() { requestBuilder.pushDownHighlight("name", Map.of())); assertEquals("Duplicate field name in highlight", exception.getMessage()); } + + @Test + void push_down_page_size() { + requestBuilder.pushDownPageSize(3); + assertEquals( + new SearchSourceBuilder() + .from(DEFAULT_OFFSET) + .size(3) + .timeout(DEFAULT_QUERY_TIMEOUT), + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + } + + @Test + void exception_when_non_zero_offset_and_page_size() { + requestBuilder.pushDownPageSize(3); + requestBuilder.pushDownLimit(300, 2); + assertThrows(UnsupportedOperationException.class, + () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); + } + + @Test + void query_size_is_page_size() { + requestBuilder.pushDownPageSize(3); + requestBuilder.pushDownLimit(4, 0); + assertEquals(3, requestBuilder.getQuerySize()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 461184e6d5..addacf04ee 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -206,6 +206,16 @@ void no_clean_on_non_empty_response() { assertNull(request.getScrollId()); } + @Test + void no_cursor_on_empty_response() { + SearchResponse searchResponse = mock(); + when(searchResponse.getHits()).thenReturn( + new SearchHits(new SearchHit[0], null, 1f)); + + request.search((x) -> searchResponse, (x) -> searchResponse); + assertEquals("", request.toCursor()); + } + @Test void no_clean_if_no_scroll_in_response() { SearchResponse searchResponse = mock(); @@ -217,4 +227,9 @@ void no_clean_if_no_scroll_in_response() { request.clean((s) -> fail()); } + + @Test + void noCursor() { + assertNull(request.toCursor()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilderTest.java new file mode 100644 index 0000000000..5a510fefec --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilderTest.java @@ -0,0 +1,75 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; + +import org.junit.jupiter.api.BeforeEach; +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.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.planner.logical.LogicalAggregation; +import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalHighlight; +import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalProject; +import org.opensearch.sql.planner.logical.LogicalSort; + +@ExtendWith(MockitoExtension.class) +class OpenSearchIndexScanAggregationBuilderTest { + @Mock + OpenSearchRequestBuilder requestBuilder; + @Mock + LogicalAggregation logicalAggregation; + OpenSearchIndexScanAggregationBuilder builder; + + @BeforeEach + void setup() { + builder = new OpenSearchIndexScanAggregationBuilder(requestBuilder, logicalAggregation); + } + + @Test + void pushDownFilter() { + assertFalse(builder.pushDownFilter(mock(LogicalFilter.class))); + } + + @Test + void pushDownSort() { + assertTrue(builder.pushDownSort(mock(LogicalSort.class))); + } + + @Test + void pushDownLimit() { + assertFalse(builder.pushDownLimit(mock(LogicalLimit.class))); + } + + @Test + void pushDownProject() { + assertFalse(builder.pushDownProject(mock(LogicalProject.class))); + } + + @Test + void pushDownHighlight() { + assertFalse(builder.pushDownHighlight(mock(LogicalHighlight.class))); + } + + @Test + void pushDownPageSize() { + assertFalse(builder.pushDownPageSize(mock(LogicalPaginate.class))); + } + + @Test + void pushDownNested() { + assertFalse(builder.pushDownNested(mock(LogicalNested.class))); + } + +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index 12fa2552fe..ac67394582 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -7,6 +7,8 @@ package org.opensearch.sql.opensearch.storage.scan; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -23,6 +25,7 @@ import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.limit; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.nested; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.paginate; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; @@ -76,9 +79,11 @@ import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; import org.opensearch.sql.opensearch.response.agg.SingleValueParser; import org.opensearch.sql.opensearch.storage.script.aggregation.AggregationQueryBuilder; +import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalNested; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.optimizer.LogicalPlanOptimizer; +import org.opensearch.sql.planner.optimizer.PushDownPageSize; import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.storage.Table; @@ -330,6 +335,33 @@ void test_sort_push_down() { ); } + @Test + void test_page_push_down() { + assertEqualsAfterOptimization( + project( + indexScanBuilder( + withPageSizePushDown(5)), + DSL.named("intV", DSL.ref("intV", INTEGER)) + ), + paginate(project( + relation("schema", table), + DSL.named("intV", DSL.ref("intV", INTEGER)) + ), 5 + )); + } + + @Test + void exception_when_page_size_after_limit() { + LogicalPlan plan = paginate( + project( + limit( + relation("schema", table), + 1, 1), + DSL.named("intV", DSL.ref("intV", INTEGER)) + ), 4); + assertThrows(IllegalStateException.class, () -> optimize(plan)); + } + @Test void test_score_sort_push_down() { assertEqualsAfterOptimization( @@ -679,7 +711,7 @@ private OpenSearchIndexScanBuilder indexScanBuilder(Runnable... verifyPushDownCa private OpenSearchIndexScanBuilder indexScanAggBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; return new OpenSearchIndexScanBuilder(t -> indexScan, - new OpenSearchIndexScanAggregationBuilder(requestBuilder)); + new OpenSearchIndexScanAggregationBuilder(requestBuilder, mock(LogicalAggregation.class))); } private void assertEqualsAfterOptimization(LogicalPlan expected, LogicalPlan actual) { @@ -756,6 +788,10 @@ private Runnable withTrackedScoresPushedDown(boolean trackScores) { return () -> verify(requestBuilder, times(1)).pushDownTrackedScore(trackScores); } + private Runnable withPageSizePushDown(int pageSize) { + return () -> verify(requestBuilder, times(1)).pushDownPageSize(pageSize); + } + private static AggregationAssertHelper.AggregationAssertHelperBuilder aggregate(String aggName) { var aggBuilder = new AggregationAssertHelper.AggregationAssertHelperBuilder(); aggBuilder.aggregateName = aggName; @@ -781,6 +817,7 @@ private static class AggregationAssertHelper { private LogicalPlan optimize(LogicalPlan plan) { LogicalPlanOptimizer optimizer = new LogicalPlanOptimizer(List.of( new CreateTableScanBuilder(), + new PushDownPageSize(), PUSH_DOWN_FILTER, PUSH_DOWN_AGGREGATION, PUSH_DOWN_SORT, diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java new file mode 100644 index 0000000000..db0e43db7c --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -0,0 +1,228 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.opensearch.storage.scan; + +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; +import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.QUERY_SIZE; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.employee; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; +import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockTwoPageResponse; + +import com.google.common.collect.ImmutableMap; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.util.Map; +import lombok.SneakyThrows; +import org.junit.jupiter.api.BeforeEach; +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.common.unit.TimeValue; +import org.opensearch.sql.common.setting.Settings; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.exception.NoCursorException; +import org.opensearch.sql.executor.pagination.PlanSerializer; +import org.opensearch.sql.opensearch.client.OpenSearchClient; +import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; +import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; +import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.request.PushDownRequestBuilder; +import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; + +@ExtendWith(MockitoExtension.class) +@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) +public class OpenSearchIndexScanPaginationTest { + + public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); + public static final int MAX_RESULT_WINDOW = 3; + public static final String SCROLL_ID = "0xbadbeef"; + public static final TimeValue SCROLL_TIMEOUT = TimeValue.timeValueMinutes(4); + @Mock + private Settings settings; + + @BeforeEach + void setup() { + lenient().when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(QUERY_SIZE); + lenient().when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) + .thenReturn(TimeValue.timeValueMinutes(1)); + } + + @Mock + private OpenSearchClient client; + + private final OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( + ImmutableMap.of( + "name", OpenSearchDataType.of(STRING), + "department", OpenSearchDataType.of(STRING))); + + @Test + void query_empty_result() { + mockResponse(client); + var builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + try (var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + assertFalse(indexScan.hasNext()); + } + verify(client).cleanup(any()); + } + + @Test + void query_all_results_initial_scroll_request() { + mockResponse(client, new ExprValue[]{ + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); + + PushDownRequestBuilder builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); + } + verify(client).cleanup(any()); + + builder = new ContinuePageRequestBuilder(SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + + assertFalse(indexScan.hasNext()); + } + verify(client, times(2)).cleanup(any()); + } + + @Test + void query_all_results_continuation_scroll_request() { + mockResponse(client, new ExprValue[]{ + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); + + ContinuePageRequestBuilder builder = new ContinuePageRequestBuilder( + SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + + assertAll( + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), + + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) + ); + } + verify(client).cleanup(any()); + + builder = new ContinuePageRequestBuilder(SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + + assertFalse(indexScan.hasNext()); + } + verify(client, times(2)).cleanup(any()); + } + + @Test + void explain_not_implemented() { + assertThrows(Throwable.class, () -> mock(OpenSearchIndexScan.class, + withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); + } + + @Test + @SneakyThrows + void serialization() { + mockTwoPageResponse(client); + OpenSearchRequestBuilder builder = mock(); + OpenSearchRequest request = mock(); + when(request.toCursor()).thenReturn("cu-cursor"); + when(builder.build(any(), eq(MAX_RESULT_WINDOW), any())).thenReturn(request); + var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, + MAX_RESULT_WINDOW, builder); + indexScan.open(); + + ByteArrayOutputStream output = new ByteArrayOutputStream(); + ObjectOutputStream objectOutput = new ObjectOutputStream(output); + objectOutput.writeObject(indexScan); + objectOutput.flush(); + + when(client.getIndexMappings(any())).thenReturn(Map.of()); + OpenSearchStorageEngine engine = mock(); + when(engine.getClient()).thenReturn(client); + when(engine.getSettings()).thenReturn(mock()); + when(client.getIndexMaxResultWindows(any())) + .thenReturn((Map.of(INDEX_NAME.getIndexNames()[0], MAX_RESULT_WINDOW))); + ObjectInputStream objectInput = new PlanSerializer(engine) + .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); + var roundTripScan = (OpenSearchIndexScan) objectInput.readObject(); + roundTripScan.open(); + assertTrue(roundTripScan.hasNext()); + } + + @Test + @SneakyThrows + void dont_serialize_if_no_cursor() { + OpenSearchRequestBuilder builder = mock(); + OpenSearchRequest request = mock(); + OpenSearchResponse response = mock(); + when(builder.build(any(), anyInt(), any())).thenReturn(request); + when(client.search(any())).thenReturn(response); + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + indexScan.open(); + + when(request.toCursor()).thenReturn(null, ""); + for (int i = 0; i < 2; i++) { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + ObjectOutputStream objectOutput = new ObjectOutputStream(output); + assertThrows(NoCursorException.class, () -> objectOutput.writeObject(indexScan)); + } + } + } +} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index f1a8d6714b..47e6e7ba2d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -177,9 +177,7 @@ void query_some_results_with_query() { @Test void query_some_results_with_scroll() { - mockResponse(client, - new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, - new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); + mockTwoPageResponse(client); final var requestuilder = new OpenSearchRequestBuilder(10, exprValueFactory); requestuilder.pushDownLimit(3, 0); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, @@ -203,6 +201,12 @@ void query_some_results_with_scroll() { verify(client).cleanup(any()); } + static void mockTwoPageResponse(OpenSearchClient client) { + mockResponse(client, + new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, + new ExprValue[]{employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); + } + @Test void query_results_limited_by_query_size() { mockResponse(client, new ExprValue[]{ From fabd4eeac521933f7c0dfd3161f93b7decb99374 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 3 May 2023 22:46:25 -0700 Subject: [PATCH 09/38] Code improvements - Rename `getQuerySize` to `getMaxResponseSize` to reflect what the value means. - Create ExecutableRequestBuilder interface -- what OpenSearchIndexScan requires from a request builder. - Remove `PushDownRequestBuilder` interface as unnecessary Signed-off-by: MaxKsyunz --- .../request/ContinuePageRequestBuilder.java | 55 +------------------ .../request/ExecutableRequestBuilder.java | 11 ++++ .../request/OpenSearchRequestBuilder.java | 16 +----- .../request/PushDownRequestBuilder.java | 50 ----------------- .../opensearch/storage/OpenSearchIndex.java | 5 +- .../storage/scan/OpenSearchIndexScan.java | 18 +++--- .../storage/scan/PushDownTranslator.java | 3 + .../ContinuePageRequestBuilderTest.java | 26 --------- .../request/OpenSearchRequestBuilderTest.java | 2 +- .../OpenSearchIndexScanPaginationTest.java | 4 +- 10 files changed, 34 insertions(+), 156 deletions(-) create mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index e51ec5db2c..5a08bd155e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -24,7 +24,7 @@ /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. */ -public class ContinuePageRequestBuilder implements PushDownRequestBuilder { +public class ContinuePageRequestBuilder implements ExecutableRequestBuilder { public static final String PUSH_DOWN_NOT_SUPPORTED = "Cursor requests don't support any push down"; @@ -50,58 +50,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, } @Override - public int getQuerySize() { + public int getMaxResponseSize() { return Integer.MAX_VALUE; } - - @Override - public void pushDownFilter(QueryBuilder query) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownAggregation(Pair, - OpenSearchAggregationResponseParser> aggregationBuilder) { - throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownSort(List> sortBuilders) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownLimit(Integer limit, Integer offset) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownHighlight(String field, Map arguments) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownProjects(Set projects) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushTypeMapping(Map typeMapping) { - throw new UnsupportedOperationException(PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownNested(List> nestedArgs) { - throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownTrackedScore(boolean trackScores) { - throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); - } - - @Override - public void pushDownPageSize(int pageSize) { - throw new UnsupportedOperationException(ContinuePageRequestBuilder.PUSH_DOWN_NOT_SUPPORTED); - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java new file mode 100644 index 0000000000..6ca3dbebe7 --- /dev/null +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java @@ -0,0 +1,11 @@ +package org.opensearch.sql.opensearch.request; + +import org.opensearch.sql.common.setting.Settings; + +public interface ExecutableRequestBuilder { + int getMaxResponseSize(); + + OpenSearchRequest build(OpenSearchRequest.IndexName indexName, + int maxResultWindow, + Settings settings); +} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 71172e3d2f..c519932c64 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -50,7 +50,7 @@ @EqualsAndHashCode @Getter @ToString -public class OpenSearchRequestBuilder implements PushDownRequestBuilder { +public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { /** * Default query timeout in minutes. @@ -93,7 +93,7 @@ public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprVa } @Override - public int getQuerySize() { + public int getMaxResponseSize() { return pageSize == null ? querySize : pageSize; } @@ -138,7 +138,6 @@ boolean isBoolFilterQuery(QueryBuilder current) { * * @param query query request */ - @Override public void pushDownFilter(QueryBuilder query) { QueryBuilder current = sourceBuilder.query(); @@ -164,7 +163,6 @@ public void pushDownFilter(QueryBuilder query) { * * @param aggregationBuilder pair of aggregation query and aggregation parser. */ - @Override public void pushDownAggregation( Pair, OpenSearchAggregationResponseParser> aggregationBuilder) { aggregationBuilder.getLeft().forEach(sourceBuilder::aggregation); @@ -177,7 +175,6 @@ public void pushDownAggregation( * * @param sortBuilders sortBuilders. */ - @Override public void pushDownSort(List> sortBuilders) { // TODO: Sort by _doc is added when filter push down. Remove both logic once doctest fixed. if (isSortByDocOnly()) { @@ -190,21 +187,18 @@ public void pushDownSort(List> sortBuilders) { } /** - * Push down size (limit) and from (offset) to DSL request. + * Pushdown size (limit) and from (offset) to DSL request. */ - @Override public void pushDownLimit(Integer limit, Integer offset) { querySize = limit; startFrom = offset; sourceBuilder.from(offset).size(limit); } - @Override public void pushDownTrackedScore(boolean trackScores) { sourceBuilder.trackScores(trackScores); } - @Override public void pushDownPageSize(int pageSize) { this.pageSize = pageSize; } @@ -213,7 +207,6 @@ public void pushDownPageSize(int pageSize) { * Add highlight to DSL requests. * @param field name of the field to highlight */ - @Override public void pushDownHighlight(String field, Map arguments) { String unquotedField = StringUtils.unquoteText(field); if (sourceBuilder.highlighter() != null) { @@ -246,14 +239,12 @@ public void pushDownHighlight(String field, Map arguments) { /** * Push down project list to DSL requests. */ - @Override public void pushDownProjects(Set projects) { sourceBuilder.fetchSource( projects.stream().map(ReferenceExpression::getAttr).distinct().toArray(String[]::new), new String[0]); } - @Override public void pushTypeMapping(Map typeMapping) { exprValueFactory.extendTypeMapping(typeMapping); } @@ -270,7 +261,6 @@ private boolean isSortByDocOnly() { * Push down nested to sourceBuilder. * @param nestedArgs : Nested arguments to push down. */ - @Override public void pushDownNested(List> nestedArgs) { initBoolQueryFilter(); groupFieldNamesByPath(nestedArgs).forEach( diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java deleted file mode 100644 index 9d579505a9..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PushDownRequestBuilder.java +++ /dev/null @@ -1,50 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.apache.commons.lang3.tuple.Pair; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.search.aggregations.AggregationBuilder; -import org.opensearch.search.sort.SortBuilder; -import org.opensearch.sql.ast.expression.Literal; -import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; -import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; - -public interface PushDownRequestBuilder { - - int getQuerySize(); - - void pushDownFilter(QueryBuilder query); - - void pushDownAggregation(Pair, - OpenSearchAggregationResponseParser> aggregationBuilder); - - void pushDownSort(List> sortBuilders); - - void pushDownLimit(Integer limit, Integer offset); - - void pushDownHighlight(String field, Map arguments); - - void pushDownProjects(Set projects); - - void pushTypeMapping(Map typeMapping); - - void pushDownNested(List> nestedArgs); - - void pushDownTrackedScore(boolean trackScores); - - void pushDownPageSize(int pageSize); - - OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, - Settings settings); - -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 2a2c71022b..699886cfa2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -171,8 +171,9 @@ public TableScanBuilder createScanBuilder() { Map allFields = new HashMap<>(); getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); allFields.putAll(getFieldOpenSearchTypes()); - Function buildScan - = t -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), t); + Function buildScan = requestBuilder + -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), + requestBuilder); var builder = new OpenSearchRequestBuilder( settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), new OpenSearchExprValueFactory(allFields)); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 6f279a8a1f..bdd6e38548 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -21,9 +21,9 @@ import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; +import org.opensearch.sql.opensearch.request.ExecutableRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; -import org.opensearch.sql.opensearch.request.PushDownRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; @@ -46,17 +46,17 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab /** Search request builder. */ @EqualsAndHashCode.Include @ToString.Include - private transient PushDownRequestBuilder requestBuilder; + private transient ExecutableRequestBuilder requestBuilder; /** Search request. */ @EqualsAndHashCode.Include @ToString.Include private transient OpenSearchRequest request; - /** Total query size. */ + /** Largest number of rows allowed in the response. */ @EqualsAndHashCode.Include @ToString.Include - private Integer querySize; + private Integer maxResponseSize; /** Number of rows returned. */ private Integer queryCount; @@ -86,7 +86,7 @@ public OpenSearchIndexScan(OpenSearchClient client, OpenSearchRequest.IndexName indexName, Settings settings, Integer maxResultWindow, - PushDownRequestBuilder requestBuilder) { + ExecutableRequestBuilder requestBuilder) { this.client = client; this.indexName = indexName; this.settings = settings; @@ -98,7 +98,7 @@ public OpenSearchIndexScan(OpenSearchClient client, public void open() { super.open(); request = requestBuilder.build(indexName, maxResultWindow, settings); - querySize = requestBuilder.getQuerySize(); + maxResponseSize = requestBuilder.getMaxResponseSize(); iterator = Collections.emptyIterator(); queryCount = 0; fetchNextBatch(); @@ -106,7 +106,7 @@ public void open() { @Override public boolean hasNext() { - if (queryCount >= querySize) { + if (queryCount >= maxResponseSize) { iterator = Collections.emptyIterator(); } else if (!iterator.hasNext()) { fetchNextBatch(); @@ -156,7 +156,7 @@ public OpenSearchIndexScan() { public void readExternal(ObjectInput in) throws IOException { final var serializedName = in.readUTF(); final var scrollId = in.readUTF(); - querySize = in.readInt(); + maxResponseSize = in.readInt(); var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) .resolveObject("engine"); @@ -180,6 +180,6 @@ public void writeExternal(ObjectOutput out) throws IOException { out.writeUTF(indexName.toString()); out.writeUTF(request.toCursor()); - out.writeInt(querySize); + out.writeInt(maxResponseSize); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java index 79d5896fd0..49c0583b4b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java @@ -14,6 +14,9 @@ import org.opensearch.sql.planner.logical.LogicalProject; import org.opensearch.sql.planner.logical.LogicalSort; +/** + * Translates a logical query plan into OpenSearch DSL and an appropriate request. + */ public interface PushDownTranslator { boolean pushDownFilter(LogicalFilter filter); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java index b3c5c6ac05..bf03b332d2 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java @@ -46,30 +46,4 @@ void build() { requestBuilder.build(null, 0, null) ); } - - @Test - void pushDown_not_supported() { - assertAll( - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownFilter(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownAggregation(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownSort(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownLimit(1, 2)), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownHighlight("", Map.of())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownProjects(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushTypeMapping(mock())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownNested(List.of())), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownPageSize(3)), - () -> assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.pushDownTrackedScore(true)) - ); - } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index bf7a9284a2..f8c40c47fd 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -372,6 +372,6 @@ void exception_when_non_zero_offset_and_page_size() { void query_size_is_page_size() { requestBuilder.pushDownPageSize(3); requestBuilder.pushDownLimit(4, 0); - assertEquals(3, requestBuilder.getQuerySize()); + assertEquals(3, requestBuilder.getMaxResponseSize()); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java index db0e43db7c..de7e807891 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -49,9 +49,9 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; +import org.opensearch.sql.opensearch.request.ExecutableRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; -import org.opensearch.sql.opensearch.request.PushDownRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; @@ -99,7 +99,7 @@ void query_all_results_initial_scroll_request() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - PushDownRequestBuilder builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + ExecutableRequestBuilder builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { indexScan.open(); From 6eb0498a76c3391e8d10134b8f3d635419efbdcc Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 4 May 2023 13:12:45 -0700 Subject: [PATCH 10/38] Checkstyle fixes Signed-off-by: MaxKsyunz --- .../request/ContinuePageRequestBuilder.java | 14 ---- .../opensearch/storage/OpenSearchIndex.java | 4 +- ...OpenSearchIndexScanAggregationBuilder.java | 3 +- .../request/OpenSearchRequestBuilderTest.java | 4 +- .../OpenSearchIndexScanOptimizationTest.java | 22 +++--- .../OpenSearchIndexScanPaginationTest.java | 67 ++++++++++--------- 6 files changed, 51 insertions(+), 63 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index 5a08bd155e..dfb9708a41 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -5,30 +5,16 @@ package org.opensearch.sql.opensearch.request; -import java.util.List; -import java.util.Map; -import java.util.Set; import lombok.Getter; -import org.apache.commons.lang3.tuple.Pair; import org.opensearch.common.unit.TimeValue; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.search.aggregations.AggregationBuilder; -import org.opensearch.search.sort.SortBuilder; -import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.expression.ReferenceExpression; -import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser; /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. */ public class ContinuePageRequestBuilder implements ExecutableRequestBuilder { - public static final String PUSH_DOWN_NOT_SUPPORTED = - "Cursor requests don't support any push down"; - @Getter private final String scrollId; private final OpenSearchExprValueFactory exprValueFactory; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 699886cfa2..ebf6dd41e5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -172,8 +172,8 @@ public TableScanBuilder createScanBuilder() { getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); allFields.putAll(getFieldOpenSearchTypes()); Function buildScan = requestBuilder - -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), - requestBuilder); + -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), + requestBuilder); var builder = new OpenSearchRequestBuilder( settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), new OpenSearchExprValueFactory(allFields)); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index c72c3e1064..5518399f86 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -48,7 +48,8 @@ class OpenSearchIndexScanAggregationBuilder implements PushDownTranslator { private List> sortList; - OpenSearchIndexScanAggregationBuilder(OpenSearchRequestBuilder requestBuilder, LogicalAggregation aggregation) { + OpenSearchIndexScanAggregationBuilder(OpenSearchRequestBuilder requestBuilder, + LogicalAggregation aggregation) { this.requestBuilder = requestBuilder; aggregatorList = aggregation.getAggregatorList(); groupByList = aggregation.getGroupByList(); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index f8c40c47fd..35dbc9dedb 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -357,7 +357,7 @@ void push_down_page_size() { .from(DEFAULT_OFFSET) .size(3) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); } @Test @@ -365,7 +365,7 @@ void exception_when_non_zero_offset_and_page_size() { requestBuilder.pushDownPageSize(3); requestBuilder.pushDownLimit(300, 2); assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); + () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index ac67394582..8edcc467cb 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -338,22 +338,22 @@ void test_sort_push_down() { @Test void test_page_push_down() { assertEqualsAfterOptimization( - project( - indexScanBuilder( - withPageSizePushDown(5)), - DSL.named("intV", DSL.ref("intV", INTEGER)) - ), - paginate(project( - relation("schema", table), - DSL.named("intV", DSL.ref("intV", INTEGER)) - ), 5 - )); + project( + indexScanBuilder( + withPageSizePushDown(5)), + DSL.named("intV", DSL.ref("intV", INTEGER)) + ), + paginate(project( + relation("schema", table), + DSL.named("intV", DSL.ref("intV", INTEGER)) + ), 5 + )); } @Test void exception_when_page_size_after_limit() { LogicalPlan plan = paginate( - project( + project( limit( relation("schema", table), 1, 1), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java index de7e807891..7ab8900e59 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -26,7 +26,6 @@ import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockTwoPageResponse; -import com.google.common.collect.ImmutableMap; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.ObjectInputStream; @@ -59,7 +58,8 @@ @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class OpenSearchIndexScanPaginationTest { - public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); + public static final OpenSearchRequest.IndexName INDEX_NAME + = new OpenSearchRequest.IndexName("test"); public static final int MAX_RESULT_WINDOW = 3; public static final String SCROLL_ID = "0xbadbeef"; public static final TimeValue SCROLL_TIMEOUT = TimeValue.timeValueMinutes(4); @@ -76,8 +76,8 @@ void setup() { @Mock private OpenSearchClient client; - private final OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( - ImmutableMap.of( + private final OpenSearchExprValueFactory exprValueFactory + = new OpenSearchExprValueFactory(Map.of( "name", OpenSearchDataType.of(STRING), "department", OpenSearchDataType.of(STRING))); @@ -85,7 +85,8 @@ void setup() { void query_empty_result() { mockResponse(client); var builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + try (var indexScan + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { indexScan.open(); assertFalse(indexScan.hasNext()); } @@ -95,9 +96,9 @@ void query_empty_result() { @Test void query_all_results_initial_scroll_request() { mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); ExecutableRequestBuilder builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (var indexScan @@ -105,17 +106,17 @@ void query_all_results_initial_scroll_request() { indexScan.open(); assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); @@ -133,28 +134,28 @@ void query_all_results_initial_scroll_request() { @Test void query_all_results_continuation_scroll_request() { mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); + employee(1, "John", "IT"), + employee(2, "Smith", "HR"), + employee(3, "Allen", "IT")}); ContinuePageRequestBuilder builder = new ContinuePageRequestBuilder( - SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); + SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); try (var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { indexScan.open(); assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), + () -> assertTrue(indexScan.hasNext()), + () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) + () -> assertFalse(indexScan.hasNext()), + () -> assertEquals(3, indexScan.getTotalHits()) ); } verify(client).cleanup(any()); @@ -172,7 +173,7 @@ void query_all_results_continuation_scroll_request() { @Test void explain_not_implemented() { assertThrows(Throwable.class, () -> mock(OpenSearchIndexScan.class, - withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); + withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); } @Test @@ -184,7 +185,7 @@ void serialization() { when(request.toCursor()).thenReturn("cu-cursor"); when(builder.build(any(), eq(MAX_RESULT_WINDOW), any())).thenReturn(request); var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, - MAX_RESULT_WINDOW, builder); + MAX_RESULT_WINDOW, builder); indexScan.open(); ByteArrayOutputStream output = new ByteArrayOutputStream(); @@ -197,9 +198,9 @@ void serialization() { when(engine.getClient()).thenReturn(client); when(engine.getSettings()).thenReturn(mock()); when(client.getIndexMaxResultWindows(any())) - .thenReturn((Map.of(INDEX_NAME.getIndexNames()[0], MAX_RESULT_WINDOW))); + .thenReturn((Map.of(INDEX_NAME.getIndexNames()[0], MAX_RESULT_WINDOW))); ObjectInputStream objectInput = new PlanSerializer(engine) - .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); + .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); var roundTripScan = (OpenSearchIndexScan) objectInput.readObject(); roundTripScan.open(); assertTrue(roundTripScan.hasNext()); @@ -214,7 +215,7 @@ void dont_serialize_if_no_cursor() { when(builder.build(any(), anyInt(), any())).thenReturn(request); when(client.search(any())).thenReturn(response); try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { indexScan.open(); when(request.toCursor()).thenReturn(null, ""); From bdfe563673a993deb9384b14913e5fd5a7714516 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 4 May 2023 14:59:25 -0700 Subject: [PATCH 11/38] Updating expected out in doctest Refactoring fixed a small bug in explain. Previously, explain output did not show the correct `SearchSourceBuilder.size` value . It now does and this affected some tests. Signed-off-by: MaxKsyunz --- docs/user/optimization/optimization.rst | 6 +++--- docs/user/ppl/interfaces/endpoint.rst | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/user/optimization/optimization.rst b/docs/user/optimization/optimization.rst index e0fe943560..8ab998309d 100644 --- a/docs/user/optimization/optimization.rst +++ b/docs/user/optimization/optimization.rst @@ -287,7 +287,7 @@ The Aggregation operator will merge into OpenSearch Aggregation:: { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" + "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":200,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" }, "children": [] } @@ -313,7 +313,7 @@ The Sort operator will merge into OpenSearch Aggregation.:: { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"last\",\"order\":\"desc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" + "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":200,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"last\",\"order\":\"desc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" }, "children": [] } @@ -348,7 +348,7 @@ Because the OpenSearch Composite Aggregation doesn't support order by metrics fi { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" + "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":200,\"timeout\":\"1m\",\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"gender\":{\"terms\":{\"field\":\"gender.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}}}, searchDone=false)" }, "children": [] } diff --git a/docs/user/ppl/interfaces/endpoint.rst b/docs/user/ppl/interfaces/endpoint.rst index fb64eff688..793b94eb8d 100644 --- a/docs/user/ppl/interfaces/endpoint.rst +++ b/docs/user/ppl/interfaces/endpoint.rst @@ -91,7 +91,7 @@ The following PPL query demonstrated that where and stats command were pushed do { "name": "OpenSearchIndexScan", "description": { - "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":10,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}, searchDone=false)" + "request": "OpenSearchQueryRequest(indexName=accounts, sourceBuilder={\"from\":0,\"size\":200,\"timeout\":\"1m\",\"query\":{\"range\":{\"age\":{\"from\":10,\"to\":null,\"include_lower\":false,\"include_upper\":true,\"boost\":1.0}}},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"avg(age)\":{\"avg\":{\"field\":\"age\"}}}}, searchDone=false)" }, "children": [] } From 1f6b6f1d27efcd013b09b72213402d927dd0516a Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 5 May 2023 00:58:22 -0700 Subject: [PATCH 12/38] Refactoring - Inlined `OpenSearchIndexScan.create` that was only used in tests. - Changed Function parameter of `OpenSearchIndexScanBuilder` to an abstract method -- better describes intent. - Renamed PushDownTranslator to PushDownQueryBuilder for better consistency. Signed-off-by: MaxKsyunz --- .../opensearch/storage/OpenSearchIndex.java | 13 +++++---- .../storage/scan/OpenSearchIndexScan.java | 16 ----------- ...OpenSearchIndexScanAggregationBuilder.java | 2 +- .../scan/OpenSearchIndexScanBuilder.java | 26 +++++------------- .../scan/OpenSearchIndexScanQueryBuilder.java | 2 +- ...nslator.java => PushDownQueryBuilder.java} | 2 +- .../OpenSearchExecutionEngineTest.java | 14 ++++++---- .../OpenSearchExecutionProtectorTest.java | 24 +++++++---------- .../storage/OpenSearchIndexTest.java | 23 +++++++++++----- .../OpenSearchIndexScanOptimizationTest.java | 25 +++++++++++++---- .../storage/scan/OpenSearchIndexScanTest.java | 27 ++++++++++++------- 11 files changed, 91 insertions(+), 83 deletions(-) rename opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/{PushDownTranslator.java => PushDownQueryBuilder.java} (96%) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index ebf6dd41e5..90939f799d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -10,7 +10,6 @@ import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; -import java.util.function.Function; import lombok.RequiredArgsConstructor; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; @@ -33,6 +32,7 @@ import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; /** OpenSearch table (index) implementation. */ @@ -171,13 +171,16 @@ public TableScanBuilder createScanBuilder() { Map allFields = new HashMap<>(); getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); allFields.putAll(getFieldOpenSearchTypes()); - Function buildScan = requestBuilder - -> new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), - requestBuilder); var builder = new OpenSearchRequestBuilder( settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), new OpenSearchExprValueFactory(allFields)); - return new OpenSearchIndexScanBuilder(buildScan, builder); + return new OpenSearchIndexScanBuilder(builder) { + @Override + protected TableScanOperator createScan(OpenSearchRequestBuilder requestBuilder) { + return new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), + requestBuilder); + } + }; } @VisibleForTesting diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index bdd6e38548..de2c2e3e12 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -23,7 +23,6 @@ import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; import org.opensearch.sql.opensearch.request.ExecutableRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; -import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; @@ -64,21 +63,6 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab /** Search response for current batch. */ private transient Iterator iterator; - /** - * Factory method used in tests. - */ - public static OpenSearchIndexScan create(OpenSearchClient client, - String indexName, - Settings settings, - Integer maxResultWindow, - OpenSearchExprValueFactory exprValueFactory) { - final var name = new OpenSearchRequest.IndexName(indexName); - final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); - final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); - return new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder); - } - /** * Creates index scan based on a provided OpenSearchRequestBuilder. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 5518399f86..58946801fb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -33,7 +33,7 @@ * Index scan builder for aggregate query used by {@link OpenSearchIndexScanBuilder} internally. */ @EqualsAndHashCode -class OpenSearchIndexScanAggregationBuilder implements PushDownTranslator { +class OpenSearchIndexScanAggregationBuilder implements PushDownQueryBuilder { /** OpenSearch index scan to be optimized. */ private final OpenSearchRequestBuilder requestBuilder; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index 1d0036fa74..cc429698de 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -5,7 +5,6 @@ package org.opensearch.sql.opensearch.storage.scan; -import java.util.function.Function; import lombok.EqualsAndHashCode; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; @@ -25,31 +24,21 @@ * by delegated builder internally. This is to avoid conditional check of different push down logic * for non-aggregate and aggregate query everywhere. */ -public class OpenSearchIndexScanBuilder extends TableScanBuilder { - - /** - * A function to create an index scan based on a OpenSearchRequestBuilder. - */ - @EqualsAndHashCode.Exclude - private final Function scanConstructor; +public abstract class OpenSearchIndexScanBuilder extends TableScanBuilder { /** * Delegated index scan builder for non-aggregate or aggregate query. */ @EqualsAndHashCode.Include - private PushDownTranslator delegate; + private PushDownQueryBuilder delegate; /** Is limit operator pushed down. */ private boolean isLimitPushedDown = false; - /** * Constructor used during query execution. */ - public OpenSearchIndexScanBuilder( - Function scanConstructor, - OpenSearchRequestBuilder requestBuilder) { - this.scanConstructor = scanConstructor; + protected OpenSearchIndexScanBuilder(OpenSearchRequestBuilder requestBuilder) { this.delegate = new OpenSearchIndexScanQueryBuilder(requestBuilder); } @@ -57,18 +46,17 @@ public OpenSearchIndexScanBuilder( /** * Constructor used for unit tests. */ - public OpenSearchIndexScanBuilder(Function constructor, - PushDownTranslator translator) { - this.scanConstructor = constructor; + protected OpenSearchIndexScanBuilder(PushDownQueryBuilder translator) { this.delegate = translator; } @Override public TableScanOperator build() { - return scanConstructor.apply(delegate.build()); + return createScan(delegate.build()); } + protected abstract TableScanOperator createScan(OpenSearchRequestBuilder requestBuilder); + @Override public boolean pushDownFilter(LogicalFilter filter) { return delegate.pushDownFilter(filter); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java index 6df5b6802f..590272a9f1 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanQueryBuilder.java @@ -40,7 +40,7 @@ */ @VisibleForTesting @EqualsAndHashCode -class OpenSearchIndexScanQueryBuilder implements PushDownTranslator { +class OpenSearchIndexScanQueryBuilder implements PushDownQueryBuilder { OpenSearchRequestBuilder requestBuilder; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java similarity index 96% rename from opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java rename to opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java index 49c0583b4b..563405049f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownTranslator.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java @@ -17,7 +17,7 @@ /** * Translates a logical query plan into OpenSearch DSL and an appropriate request. */ -public interface PushDownTranslator { +public interface PushDownQueryBuilder { boolean pushDownFilter(LogicalFilter filter); boolean pushDownSort(LogicalSort sort); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index 245067e51e..f72e863656 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -17,12 +17,10 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.sql.common.setting.Settings.Key.QUERY_SIZE_LIMIT; import static org.opensearch.sql.common.setting.Settings.Key.SQL_CURSOR_KEEP_ALIVE; import static org.opensearch.sql.data.model.ExprValueUtils.tupleValue; import static org.opensearch.sql.executor.ExecutionEngine.QueryResponse; -import java.io.IOException; import java.io.ObjectInput; import java.io.ObjectOutput; import java.util.ArrayList; @@ -50,6 +48,8 @@ import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.executor.protector.OpenSearchExecutionProtector; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -174,12 +174,16 @@ void explain_successfully() { OpenSearchExecutionEngine executor = new OpenSearchExecutionEngine(client, protector, new PlanSerializer(null)); Settings settings = mock(Settings.class); - when(settings.getSettingValue(QUERY_SIZE_LIMIT)).thenReturn(100); when(settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE)) .thenReturn(TimeValue.timeValueMinutes(1)); - PhysicalPlan plan = OpenSearchIndexScan.create(mock(OpenSearchClient.class), "test", settings, - 10000, mock(OpenSearchExprValueFactory.class)); + OpenSearchExprValueFactory exprValueFactory = mock(OpenSearchExprValueFactory.class); + final var name = new OpenSearchRequest.IndexName("test"); + final int defaultQuerySize = 100; + final int maxResultWindow = 10000; + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), name, + settings, maxResultWindow, requestBuilder); AtomicReference result = new AtomicReference<>(); executor.explain(plan, new ResponseListener<>() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index b5fff14614..3a8386e75a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -9,7 +9,6 @@ import static java.util.Collections.emptyList; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; import static org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC; import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; @@ -24,7 +23,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -41,7 +39,6 @@ import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.ast.tree.Sort; -import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; @@ -58,6 +55,8 @@ 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.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.setting.OpenSearchSettings; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.planner.physical.NestedOperator; @@ -88,10 +87,9 @@ public void setup() { @Test void testProtectIndexScan() { - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)) - .thenReturn(200); String indexName = "test"; - Integer maxResultWindow = 10000; + final int maxResultWindow = 10000; + final int querySizeLimit = 200; NamedExpression include = named("age", ref("age", INTEGER)); ReferenceExpression exclude = ref("name", STRING); ReferenceExpression dedupeField = ref("name", STRING); @@ -111,6 +109,8 @@ void testProtectIndexScan() { Integer limit = 10; Integer offset = 10; + final var name = new OpenSearchRequest.IndexName(indexName); + final var requestBuilder = new OpenSearchRequestBuilder(querySizeLimit, exprValueFactory); assertEquals( PhysicalPlanDSL.project( PhysicalPlanDSL.limit( @@ -124,10 +124,8 @@ void testProtectIndexScan() { PhysicalPlanDSL.agg( filter( resourceMonitor( - OpenSearchIndexScan.create(client, - indexName, settings, - maxResultWindow, - exprValueFactory)), + new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder)), filterExpr), aggregators, groupByExprs), @@ -153,10 +151,8 @@ void testProtectIndexScan() { PhysicalPlanDSL.rename( PhysicalPlanDSL.agg( filter( - OpenSearchIndexScan.create(client, indexName, - settings, - maxResultWindow, - exprValueFactory), + new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder), filterExpr), aggregators, groupByExprs), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 0afaaa3e7c..dea99222a8 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -51,6 +51,8 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchTextType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.mapping.IndexMapping; +import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; @@ -197,8 +199,11 @@ void implementRelationOperatorOnly() { when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - assertEquals(OpenSearchIndexScan.create(client, indexName, settings, - maxResultWindow, exprValueFactory), index.implement(index.optimize(plan))); + final var name = new OpenSearchRequest.IndexName(indexName); + final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + assertEquals(new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder), index.implement(index.optimize(plan))); } @Test @@ -207,8 +212,11 @@ void implementRelationOperatorWithOptimization() { when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - assertEquals(OpenSearchIndexScan.create(client, indexName, settings, - maxResultWindow, exprValueFactory), index.implement(plan)); + final var name = new OpenSearchRequest.IndexName(indexName); + final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + assertEquals(new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder), index.implement(plan)); } @Test @@ -242,6 +250,9 @@ void implementOtherLogicalOperators() { include); Integer maxResultWindow = index.getMaxResultWindow(); + final var name = new OpenSearchRequest.IndexName(indexName); + final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); assertEquals( PhysicalPlanDSL.project( PhysicalPlanDSL.dedupe( @@ -249,8 +260,8 @@ void implementOtherLogicalOperators() { PhysicalPlanDSL.eval( PhysicalPlanDSL.remove( PhysicalPlanDSL.rename( - OpenSearchIndexScan.create(client, indexName, settings, - maxResultWindow, exprValueFactory), + new OpenSearchIndexScan(client, name, + settings, maxResultWindow, requestBuilder), mappings), exclude), newEvalField), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index 8edcc467cb..8007f28b85 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -86,6 +86,7 @@ import org.opensearch.sql.planner.optimizer.PushDownPageSize; import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder; import org.opensearch.sql.storage.Table; +import org.opensearch.sql.storage.TableScanOperator; @ExtendWith(MockitoExtension.class) class OpenSearchIndexScanOptimizationTest { @@ -105,7 +106,12 @@ class OpenSearchIndexScanOptimizationTest { @BeforeEach void setUp() { - indexScanBuilder = new OpenSearchIndexScanBuilder(t -> indexScan, requestBuilder); + indexScanBuilder = new OpenSearchIndexScanBuilder(requestBuilder) { + @Override + protected TableScanOperator createScan(OpenSearchRequestBuilder build) { + return indexScan; + } + }; when(table.createScanBuilder()).thenReturn(indexScanBuilder); } @@ -704,14 +710,23 @@ void project_literal_should_not_be_pushed_down() { private OpenSearchIndexScanBuilder indexScanBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(t -> indexScan, - new OpenSearchIndexScanQueryBuilder(requestBuilder)); + return new OpenSearchIndexScanBuilder(new OpenSearchIndexScanQueryBuilder(requestBuilder)) { + @Override + protected TableScanOperator createScan(OpenSearchRequestBuilder build) { + return indexScan; + } + }; } private OpenSearchIndexScanBuilder indexScanAggBuilder(Runnable... verifyPushDownCalls) { this.verifyPushDownCalls = verifyPushDownCalls; - return new OpenSearchIndexScanBuilder(t -> indexScan, - new OpenSearchIndexScanAggregationBuilder(requestBuilder, mock(LogicalAggregation.class))); + return new OpenSearchIndexScanBuilder(new OpenSearchIndexScanAggregationBuilder( + requestBuilder, mock(LogicalAggregation.class))) { + @Override + protected TableScanOperator createScan(OpenSearchRequestBuilder build) { + return indexScan; + } + }; } private void assertEqualsAfterOptimization(LogicalPlan expected, LogicalPlan actual) { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 47e6e7ba2d..bc3ee8f952 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -67,7 +67,6 @@ class OpenSearchIndexScanTest { @BeforeEach void setup() { - lenient().when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(QUERY_SIZE); when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) .thenReturn(TimeValue.timeValueMinutes(1)); } @@ -75,8 +74,10 @@ void setup() { @Test void query_empty_result() { mockResponse(client); - try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "test", settings, - 3, exprValueFactory)) { + final var name = new OpenSearchRequest.IndexName("test"); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, + settings, 3, requestBuilder)) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -93,8 +94,10 @@ void query_all_results_with_query() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, - 10, exprValueFactory)) { + final var name = new OpenSearchRequest.IndexName("employees"); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, + settings, 10, requestBuilder)) { indexScan.open(); assertAll( @@ -123,8 +126,10 @@ void query_all_results_with_scroll() { new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); - try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, - 10, exprValueFactory)) { + final var name = new OpenSearchRequest.IndexName("employees"); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, + settings, 10, requestBuilder)) { indexScan.open(); assertAll( @@ -214,10 +219,12 @@ void query_results_limited_by_query_size() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(2); - try (OpenSearchIndexScan indexScan = OpenSearchIndexScan.create(client, "employees", settings, - 10, exprValueFactory)) { + final var name = new OpenSearchRequest.IndexName("employees"); + final int defaultQuerySize = 2; + final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, + settings, 10, requestBuilder)) { indexScan.open(); assertAll( From a8295e4e507e5096cd676f869e449e87aa9539b5 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 5 May 2023 01:30:50 -0700 Subject: [PATCH 13/38] Checkstyle fixes Signed-off-by: MaxKsyunz --- .../executor/OpenSearchExecutionEngineTest.java | 2 +- .../sql/opensearch/storage/OpenSearchIndexTest.java | 4 ++-- .../opensearch/storage/scan/OpenSearchIndexScanTest.java | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index f72e863656..8e7f0570bc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -183,7 +183,7 @@ void explain_successfully() { final int maxResultWindow = 10000; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), name, - settings, maxResultWindow, requestBuilder); + settings, maxResultWindow, requestBuilder); AtomicReference result = new AtomicReference<>(); executor.explain(plan, new ResponseListener<>() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index dea99222a8..04bbe62d03 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -203,7 +203,7 @@ void implementRelationOperatorOnly() { final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); assertEquals(new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), index.implement(index.optimize(plan))); + settings, maxResultWindow, requestBuilder), index.implement(index.optimize(plan))); } @Test @@ -216,7 +216,7 @@ void implementRelationOperatorWithOptimization() { final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); assertEquals(new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), index.implement(plan)); + settings, maxResultWindow, requestBuilder), index.implement(plan)); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index bc3ee8f952..8ddb2ad45d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -77,7 +77,7 @@ void query_empty_result() { final var name = new OpenSearchRequest.IndexName("test"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 3, requestBuilder)) { + settings, 3, requestBuilder)) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -97,7 +97,7 @@ void query_all_results_with_query() { final var name = new OpenSearchRequest.IndexName("employees"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + settings, 10, requestBuilder)) { indexScan.open(); assertAll( @@ -129,7 +129,7 @@ void query_all_results_with_scroll() { final var name = new OpenSearchRequest.IndexName("employees"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + settings, 10, requestBuilder)) { indexScan.open(); assertAll( @@ -224,7 +224,7 @@ void query_results_limited_by_query_size() { final int defaultQuerySize = 2; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + settings, 10, requestBuilder)) { indexScan.open(); assertAll( From 86429f8ed3c8a0a59fdb5872c9ce41d4aefad946 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 5 May 2023 14:06:03 -0700 Subject: [PATCH 14/38] Rename createContinuePaginatedPlan in QueryPlanFactory to create Old name doesn't quite make sense. Signed-off-by: MaxKsyunz --- .../sql/executor/execution/QueryPlanFactory.java | 10 +++++----- .../executor/execution/QueryPlanFactoryTest.java | 16 ++++++++-------- .../java/org/opensearch/sql/ppl/PPLService.java | 2 +- .../java/org/opensearch/sql/sql/SQLService.java | 4 ++-- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 18455c2a02..6facfff847 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -65,7 +65,7 @@ public void onFailure(Exception e) { /** * Create QueryExecution from Statement. */ - public AbstractPlan createContinuePaginatedPlan( + public AbstractPlan create( Statement statement, Optional> queryListener, Optional> explainListener) { @@ -75,9 +75,9 @@ public AbstractPlan createContinuePaginatedPlan( /** * Creates a ContinuePaginatedPlan from a cursor. */ - public AbstractPlan createContinuePaginatedPlan(String cursor, boolean isExplain, - ResponseListener queryResponseListener, - ResponseListener explainListener) { + public AbstractPlan create(String cursor, boolean isExplain, + ResponseListener queryResponseListener, + ResponseListener explainListener) { QueryId queryId = QueryId.queryId(); var plan = new ContinuePaginatedPlan(queryId, cursor, queryService, planSerializer, queryResponseListener); @@ -119,7 +119,7 @@ public AbstractPlan visitExplain( return new ExplainPlan( QueryId.queryId(), - createContinuePaginatedPlan(node.getStatement(), + create(node.getStatement(), Optional.of(NO_CONSUMER_RESPONSE_LISTENER), Optional.empty()), context.getRight().get()); } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index 6bdbf1c4c9..8daf1edd43 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -62,7 +62,7 @@ void init() { public void createFromQueryShouldSuccess() { Statement query = new Query(plan, 0); AbstractPlan queryExecution = - factory.createContinuePaginatedPlan(query, Optional.of(queryListener), Optional.empty()); + factory.create(query, Optional.of(queryListener), Optional.empty()); assertTrue(queryExecution instanceof QueryPlan); } @@ -70,15 +70,15 @@ public void createFromQueryShouldSuccess() { public void createFromExplainShouldSuccess() { Statement query = new Explain(new Query(plan, 0)); AbstractPlan queryExecution = - factory.createContinuePaginatedPlan(query, Optional.empty(), Optional.of(explainListener)); + factory.create(query, Optional.empty(), Optional.of(explainListener)); assertTrue(queryExecution instanceof ExplainPlan); } @Test public void createFromCursorShouldSuccess() { - AbstractPlan queryExecution = factory.createContinuePaginatedPlan("", false, + AbstractPlan queryExecution = factory.create("", false, queryListener, explainListener); - AbstractPlan explainExecution = factory.createContinuePaginatedPlan("", true, + AbstractPlan explainExecution = factory.create("", true, queryListener, explainListener); assertAll( () -> assertTrue(queryExecution instanceof ContinuePaginatedPlan), @@ -91,7 +91,7 @@ public void createFromQueryWithoutQueryListenerShouldThrowException() { Statement query = new Query(plan, 0); IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> factory.createContinuePaginatedPlan( + assertThrows(IllegalArgumentException.class, () -> factory.create( query, Optional.empty(), Optional.empty())); assertEquals("[BUG] query listener must be not null", exception.getMessage()); } @@ -101,7 +101,7 @@ public void createFromExplainWithoutExplainListenerShouldThrowException() { Statement query = new Explain(new Query(plan, 0)); IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> factory.createContinuePaginatedPlan( + assertThrows(IllegalArgumentException.class, () -> factory.create( query, Optional.empty(), Optional.empty())); assertEquals("[BUG] explain listener must be not null", exception.getMessage()); } @@ -129,7 +129,7 @@ public void createQueryWithFetchSizeWhichCanBePaged() { factory = new QueryPlanFactory(queryService, planSerializer); Statement query = new Query(plan, 10); AbstractPlan queryExecution = - factory.createContinuePaginatedPlan(query, Optional.of(queryListener), Optional.empty()); + factory.create(query, Optional.of(queryListener), Optional.empty()); assertTrue(queryExecution instanceof QueryPlan); } @@ -139,7 +139,7 @@ public void createQueryWithFetchSizeWhichCannotBePaged() { factory = new QueryPlanFactory(queryService, planSerializer); Statement query = new Query(plan, 10); assertThrows(UnsupportedCursorRequestException.class, - () -> factory.createContinuePaginatedPlan(query, + () -> factory.create(query, Optional.of(queryListener), Optional.empty())); } } diff --git a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java index f91ac7222f..40a7a85f78 100644 --- a/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java +++ b/ppl/src/main/java/org/opensearch/sql/ppl/PPLService.java @@ -90,7 +90,7 @@ private AbstractPlan plan( QueryContext.getRequestId(), anonymizer.anonymizeStatement(statement)); - return queryExecutionFactory.createContinuePaginatedPlan( + return queryExecutionFactory.create( statement, queryListener, explainListener); } } diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index 4ecf9e699b..a3ed89e9d7 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -67,7 +67,7 @@ private AbstractPlan plan( Optional> explainListener) { if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. - return queryExecutionFactory.createContinuePaginatedPlan(request.getCursor().get(), + return queryExecutionFactory.create(request.getCursor().get(), request.isExplainRequest(), queryListener.orElse(null), explainListener.orElse(null)); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) @@ -81,7 +81,7 @@ private AbstractPlan plan( .fetchSize(request.getFetchSize()) .build())); - return queryExecutionFactory.createContinuePaginatedPlan( + return queryExecutionFactory.create( statement, queryListener, explainListener); } } From 4f5b69d49c4189d52c257b319a79e0807c2ac0c9 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 9 May 2023 00:15:47 -0700 Subject: [PATCH 15/38] Address PR comments. Signed-off-by: MaxKsyunz --- .../optimizer/LogicalPlanOptimizerTest.java | 117 ++++++++++-------- .../request/ContinuePageRequestBuilder.java | 1 + .../scan/OpenSearchIndexScanBuilder.java | 3 - .../OpenSearchIndexScanOptimizationTest.java | 12 -- 4 files changed, 64 insertions(+), 69 deletions(-) diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index ce61380d1d..5370f0c0be 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -17,7 +17,17 @@ import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.LONG; import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.aggregation; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.filter; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.highlight; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.limit; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.nested; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.paginate; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.project; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.relation; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.sort; import static org.opensearch.sql.planner.logical.LogicalPlanDSL.values; +import static org.opensearch.sql.planner.logical.LogicalPlanDSL.write; import com.google.common.collect.ImmutableList; import java.util.Collections; @@ -40,7 +50,6 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; -import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; import org.opensearch.sql.planner.logical.LogicalRelation; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -69,15 +78,15 @@ void setUp() { @Test void filter_merge_filter() { assertEquals( - LogicalPlanDSL.filter( + filter( tableScanBuilder, DSL.and(DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(2))), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(1)))) ), optimize( - LogicalPlanDSL.filter( - LogicalPlanDSL.filter( - LogicalPlanDSL.relation("schema", table), + filter( + filter( + relation("schema", table), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(1))) ), DSL.equal(DSL.ref("integer_value", INTEGER), DSL.literal(integerValue(2))) @@ -92,17 +101,17 @@ void filter_merge_filter() { @Test void push_filter_under_sort() { assertEquals( - LogicalPlanDSL.sort( - LogicalPlanDSL.filter( + sort( + filter( tableScanBuilder, DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) ), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), optimize( - LogicalPlanDSL.filter( - LogicalPlanDSL.sort( - LogicalPlanDSL.relation("schema", table), + filter( + sort( + relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) @@ -117,8 +126,8 @@ void push_filter_under_sort() { @Test void multiple_filter_should_eventually_be_merged() { assertEquals( - LogicalPlanDSL.sort( - LogicalPlanDSL.filter( + sort( + filter( tableScanBuilder, DSL.and(DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))), DSL.less(DSL.ref("longV", INTEGER), DSL.literal(longValue(1L)))) @@ -126,10 +135,10 @@ void multiple_filter_should_eventually_be_merged() { Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) ), optimize( - LogicalPlanDSL.filter( - LogicalPlanDSL.sort( - LogicalPlanDSL.filter( - LogicalPlanDSL.relation("schema", table), + filter( + sort( + filter( + relation("schema", table), DSL.less(DSL.ref("longV", INTEGER), DSL.literal(longValue(1L))) ), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("longV", LONG)) @@ -143,25 +152,25 @@ void multiple_filter_should_eventually_be_merged() { @Test void default_table_scan_builder_should_not_push_down_anything() { LogicalPlan[] plans = { - LogicalPlanDSL.project( - LogicalPlanDSL.relation("schema", table), + project( + relation("schema", table), DSL.named("i", DSL.ref("intV", INTEGER)) ), - LogicalPlanDSL.filter( - LogicalPlanDSL.relation("schema", table), + filter( + relation("schema", table), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1))) ), - LogicalPlanDSL.aggregation( - LogicalPlanDSL.relation("schema", table), + aggregation( + relation("schema", table), ImmutableList .of(DSL.named("AVG(intV)", DSL.avg(DSL.ref("intV", INTEGER)))), ImmutableList.of(DSL.named("longV", DSL.ref("longV", LONG)))), - LogicalPlanDSL.sort( - LogicalPlanDSL.relation("schema", table), + sort( + relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("intV", INTEGER))), - LogicalPlanDSL.limit( - LogicalPlanDSL.relation("schema", table), + limit( + relation("schema", table), 1, 1) }; @@ -177,8 +186,8 @@ void table_scan_builder_support_project_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.project( - LogicalPlanDSL.relation("schema", table), + project( + relation("schema", table), DSL.named("i", DSL.ref("intV", INTEGER))) ) ); @@ -191,8 +200,8 @@ void table_scan_builder_support_filter_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.filter( - LogicalPlanDSL.relation("schema", table), + filter( + relation("schema", table), DSL.equal(DSL.ref("intV", INTEGER), DSL.literal(integerValue(1)))) ) ); @@ -205,8 +214,8 @@ void table_scan_builder_support_aggregation_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.aggregation( - LogicalPlanDSL.relation("schema", table), + aggregation( + relation("schema", table), ImmutableList .of(DSL.named("AVG(intV)", DSL.avg(DSL.ref("intV", INTEGER)))), @@ -222,8 +231,8 @@ void table_scan_builder_support_sort_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.sort( - LogicalPlanDSL.relation("schema", table), + sort( + relation("schema", table), Pair.of(Sort.SortOption.DEFAULT_ASC, DSL.ref("intV", INTEGER))) ) ); @@ -236,8 +245,8 @@ void table_scan_builder_support_limit_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.limit( - LogicalPlanDSL.relation("schema", table), + limit( + relation("schema", table), 1, 1) ) ); @@ -250,8 +259,8 @@ void table_scan_builder_support_highlight_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.highlight( - LogicalPlanDSL.relation("schema", table), + highlight( + relation("schema", table), DSL.literal("*"), Collections.emptyMap()) ) @@ -265,8 +274,8 @@ void table_scan_builder_support_nested_push_down_can_apply_its_rule() { assertEquals( tableScanBuilder, optimize( - LogicalPlanDSL.nested( - LogicalPlanDSL.relation("schema", table), + nested( + relation("schema", table), List.of(Map.of("field", new ReferenceExpression("message.info", STRING))), List.of(new NamedExpression( "message.info", @@ -292,8 +301,8 @@ public PhysicalPlan implement(LogicalPlan plan) { }; assertEquals( - LogicalPlanDSL.relation("schema", table), - optimize(LogicalPlanDSL.relation("schema", table)) + relation("schema", table), + optimize(relation("schema", table)) ); } @@ -304,7 +313,7 @@ void table_support_write_builder_should_be_replaced() { assertEquals( writeBuilder, - optimize(LogicalPlanDSL.write(values(), table, Collections.emptyList())) + optimize(write(values(), table, Collections.emptyList())) ); } @@ -330,24 +339,24 @@ public PhysicalPlan implement(LogicalPlan plan) { void paged_table_scan_builder_support_project_push_down_can_apply_its_rule() { when(tableScanBuilder.pushDownPageSize(any())).thenReturn(true); - var relation = LogicalPlanDSL.relation("schema", table); + var relation = relation("schema", table); var optimized = LogicalPlanOptimizer.create() - .optimize(LogicalPlanDSL.paginate(LogicalPlanDSL.project(relation), 4)); + .optimize(paginate(project(relation), 4)); verify(tableScanBuilder).pushDownPageSize(any()); - assertEquals(LogicalPlanDSL.project(tableScanBuilder), optimized); + assertEquals(project(tableScanBuilder), optimized); } @Test void push_down_page_size_multiple_children() { - var relation = LogicalPlanDSL.relation("schema", table); + var relation = relation("schema", table); var twoChildrenPlan = new LogicalPlan(List.of(relation, relation)) { @Override public R accept(LogicalPlanNodeVisitor visitor, C context) { return null; } }; - var queryPlan = LogicalPlanDSL.paginate(twoChildrenPlan, 4); + var queryPlan = paginate(twoChildrenPlan, 4); var optimizer = LogicalPlanOptimizer.create(); final var exception = assertThrows(UnsupportedOperationException.class, () -> optimizer.optimize(queryPlan)); @@ -359,9 +368,9 @@ public R accept(LogicalPlanNodeVisitor visitor, C context) { void push_down_page_size_push_failed() { when(tableScanBuilder.pushDownPageSize(any())).thenReturn(false); - var queryPlan = LogicalPlanDSL.paginate( - LogicalPlanDSL.project( - LogicalPlanDSL.relation("schema", table)), 4); + var queryPlan = paginate( + project( + relation("schema", table)), 4); var optimizer = LogicalPlanOptimizer.create(); final var exception = assertThrows(IllegalStateException.class, () -> optimizer.optimize(queryPlan)); @@ -370,7 +379,7 @@ void push_down_page_size_push_failed() { @Test void push_page_size_noop_if_no_relation() { - var paginate = new LogicalPaginate(42, List.of(LogicalPlanDSL.project(values()))); + var paginate = new LogicalPaginate(42, List.of(project(values()))); assertEquals(paginate, LogicalPlanOptimizer.create().optimize(paginate)); } @@ -387,10 +396,10 @@ void table_scan_builder_support_offset_push_down_can_apply_its_rule() { var relation = new LogicalRelation("schema", table); var optimized = LogicalPlanOptimizer.create() - .optimize(new LogicalPaginate(42, List.of(LogicalPlanDSL.project(relation)))); + .optimize(new LogicalPaginate(42, List.of(project(relation)))); // `optimized` structure: LogicalPaginate -> LogicalProject -> TableScanBuilder // LogicalRelation replaced by a TableScanBuilder instance - assertEquals(LogicalPlanDSL.project(tableScanBuilder), optimized); + assertEquals(project(tableScanBuilder), optimized); } private LogicalPlan optimize(LogicalPlan plan) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index dfb9708a41..a255d0495b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -12,6 +12,7 @@ /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. + * Initial search requests are handled by {@link OpenSearchRequestBuilder} */ public class ContinuePageRequestBuilder implements ExecutableRequestBuilder { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java index cc429698de..c6df692095 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanBuilder.java @@ -76,9 +76,6 @@ public boolean pushDownAggregation(LogicalAggregation aggregation) { @Override public boolean pushDownPageSize(LogicalPaginate paginate) { - if (isLimitPushedDown) { - throw new IllegalStateException("Pagination has to be pushed down before limit."); - } return delegate.pushDownPageSize(paginate); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java index 8007f28b85..6bf9002a67 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanOptimizationTest.java @@ -356,18 +356,6 @@ void test_page_push_down() { )); } - @Test - void exception_when_page_size_after_limit() { - LogicalPlan plan = paginate( - project( - limit( - relation("schema", table), - 1, 1), - DSL.named("intV", DSL.ref("intV", INTEGER)) - ), 4); - assertThrows(IllegalStateException.class, () -> optimize(plan)); - } - @Test void test_score_sort_push_down() { assertEqualsAfterOptimization( From b1050dddd13dddb23df0968b2326ca378ab3a175 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 10 May 2023 13:22:45 -0700 Subject: [PATCH 16/38] Address PR comments. Remove transient keyword as it applies to serialization only, not exetnalization. Renamed SerializablePlan to ExternalizablePlan since Serializable has specific implications in Java which are not true for classes that implement ExeternalizablePlan. Signed-off-by: MaxKsyunz --- .../executor/pagination/PlanSerializer.java | 4 ++-- ...izablePlan.java => ExternalizablePlan.java} | 6 +++--- .../sql/planner/physical/ProjectOperator.java | 6 +++--- .../pagination/PlanSerializerTest.java | 6 +++--- ...anTest.java => ExternalizablePlanTest.java} | 4 ++-- .../planner/physical/ProjectOperatorTest.java | 4 ++-- .../protector/ResourceMonitorPlan.java | 8 ++++---- .../storage/scan/OpenSearchIndexScan.java | 18 +++++++++--------- .../OpenSearchExecutionEngineTest.java | 4 ++-- .../executor/ResourceMonitorPlanTest.java | 4 ++-- 10 files changed, 32 insertions(+), 32 deletions(-) rename core/src/main/java/org/opensearch/sql/planner/{SerializablePlan.java => ExternalizablePlan.java} (90%) rename core/src/test/java/org/opensearch/sql/planner/{SerializablePlanTest.java => ExternalizablePlanTest.java} (92%) diff --git a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java index d6d10ee89c..fe4d576d1a 100644 --- a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java +++ b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java @@ -20,7 +20,7 @@ import lombok.RequiredArgsConstructor; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; @@ -44,7 +44,7 @@ public boolean canConvertToCursor(UnresolvedPlan plan) { public Cursor convertToCursor(PhysicalPlan plan) { try { return new Cursor(CURSOR_PREFIX - + serialize(((SerializablePlan) plan).getPlanForSerialization())); + + serialize(((ExternalizablePlan) plan).getPlanForSerialization())); // ClassCastException thrown when a plan in the tree doesn't implement SerializablePlan } catch (NotSerializableException | ClassCastException | NoCursorException e) { return Cursor.None; diff --git a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java b/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java similarity index 90% rename from core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java rename to core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java index 487b1da6bd..dec64ffdc7 100644 --- a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java @@ -23,11 +23,11 @@ * *
  • * Overwrite {@link #getPlanForSerialization} to return - * another instance of {@link SerializablePlan}. + * another instance of {@link ExternalizablePlan}. *
    • * */ -public interface SerializablePlan extends Externalizable { +public interface ExternalizablePlan extends Externalizable { /** * Argument is an instance of {@link PlanSerializer.CursorDeserializationStream}. @@ -57,7 +57,7 @@ public interface SerializablePlan extends Externalizable { * It is needed to skip a `ResourceMonitorPlan` instance only, actually. * @return Next plan for serialization. */ - default SerializablePlan getPlanForSerialization() { + default ExternalizablePlan getPlanForSerialization() { return this; } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 1699c97c15..3a0da4f0dc 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -25,7 +25,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; /** * Project the fields specified in {@link ProjectOperator#projectList} from input. @@ -33,7 +33,7 @@ @ToString @EqualsAndHashCode(callSuper = false) @AllArgsConstructor -public class ProjectOperator extends PhysicalPlan implements SerializablePlan { +public class ProjectOperator extends PhysicalPlan implements ExternalizablePlan { @Getter private PhysicalPlan input; @Getter @@ -116,6 +116,6 @@ public void readExternal(ObjectInput in) throws IOException, ClassNotFoundExcept @Override public void writeExternal(ObjectOutput out) throws IOException { out.writeObject(projectList); - out.writeObject(((SerializablePlan) input).getPlanForSerialization()); + out.writeObject(((ExternalizablePlan) input).getPlanForSerialization()); } } diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java index b1e97920c8..622f10b3ba 100644 --- a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java @@ -32,7 +32,7 @@ import org.opensearch.sql.ast.dsl.AstDSL; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; import org.opensearch.sql.storage.StorageEngine; @@ -168,7 +168,7 @@ void serialize_and_deserialize() { @Test void convertToCursor_and_convertToPlan() { var plan = new TestOperator(100500); - var roundTripPlan = (SerializablePlan) + var roundTripPlan = (ExternalizablePlan) planCache.convertToPlan(planCache.convertToCursor(plan).toString()); assertEquals(plan, roundTripPlan); assertNotSame(plan, roundTripPlan); @@ -191,7 +191,7 @@ void resolveObject() { // Helpers and auxiliary classes section below - public static class TestOperator extends PhysicalPlan implements SerializablePlan { + public static class TestOperator extends PhysicalPlan implements ExternalizablePlan { private int field; private boolean throwNoCursorOnWrite = false; private boolean throwIoOnWrite = false; diff --git a/core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java b/core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java similarity index 92% rename from core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java rename to core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java index 8073445dc0..9055fff605 100644 --- a/core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java @@ -18,9 +18,9 @@ @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class SerializablePlanTest { +public class ExternalizablePlanTest { @Mock(answer = CALLS_REAL_METHODS) - SerializablePlan plan; + ExternalizablePlan plan; @Test void getPlanForSerialization_defaults_to_self() { diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index 77fcb7a505..ba4b1f2cbd 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -40,7 +40,7 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; @ExtendWith(MockitoExtension.class) class ProjectOperatorTest extends PhysicalPlanTestBase { @@ -236,7 +236,7 @@ public void serializable() { } @EqualsAndHashCode(callSuper = false) - public static class TestOperator extends PhysicalPlan implements SerializablePlan { + public static class TestOperator extends PhysicalPlan implements ExternalizablePlan { @Override public R accept(PhysicalPlanNodeVisitor visitor, C context) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index 0ec4d743b3..a65c72ff7e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -15,7 +15,7 @@ import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.monitor.ResourceMonitor; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -25,7 +25,7 @@ @ToString @RequiredArgsConstructor @EqualsAndHashCode(callSuper = false) -public class ResourceMonitorPlan extends PhysicalPlan implements SerializablePlan { +public class ResourceMonitorPlan extends PhysicalPlan implements ExternalizablePlan { /** * How many method calls to delegate's next() to perform resource check once. @@ -93,8 +93,8 @@ public long getTotalHits() { } @Override - public SerializablePlan getPlanForSerialization() { - return (SerializablePlan) delegate; + public ExternalizablePlan getPlanForSerialization() { + return (ExternalizablePlan) delegate; } /** diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index de2c2e3e12..da84e6db07 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -26,7 +26,7 @@ import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.storage.TableScanOperator; /** @@ -34,23 +34,23 @@ */ @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) @ToString(onlyExplicitlyIncluded = true) -public class OpenSearchIndexScan extends TableScanOperator implements SerializablePlan { +public class OpenSearchIndexScan extends TableScanOperator implements ExternalizablePlan { /** OpenSearch client. */ - private transient OpenSearchClient client; + private OpenSearchClient client; - private transient OpenSearchRequest.IndexName indexName; - private transient Settings settings; - private transient Integer maxResultWindow; + private OpenSearchRequest.IndexName indexName; + private Settings settings; + private Integer maxResultWindow; /** Search request builder. */ @EqualsAndHashCode.Include @ToString.Include - private transient ExecutableRequestBuilder requestBuilder; + private ExecutableRequestBuilder requestBuilder; /** Search request. */ @EqualsAndHashCode.Include @ToString.Include - private transient OpenSearchRequest request; + private OpenSearchRequest request; /** Largest number of rows allowed in the response. */ @EqualsAndHashCode.Include @@ -61,7 +61,7 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab private Integer queryCount; /** Search response for current batch. */ - private transient Iterator iterator; + private Iterator iterator; /** * Creates index scan based on a provided OpenSearchRequestBuilder. diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index 8e7f0570bc..aa08ba7e00 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -51,7 +51,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; @@ -258,7 +258,7 @@ public void onFailure(Exception e) { } @RequiredArgsConstructor - private static class FakePhysicalPlan extends TableScanOperator implements SerializablePlan { + private static class FakePhysicalPlan extends TableScanOperator implements ExternalizablePlan { private final Iterator it; private boolean hasOpen; private boolean hasClosed; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java index 0b9f302ceb..53ddc40cb3 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java @@ -21,7 +21,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.sql.monitor.ResourceMonitor; import org.opensearch.sql.opensearch.executor.protector.ResourceMonitorPlan; -import org.opensearch.sql.planner.SerializablePlan; +import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -119,7 +119,7 @@ void getTotalHitsSuccess() { @Test void getPlanForSerialization() { - plan = mock(PhysicalPlan.class, withSettings().extraInterfaces(SerializablePlan.class)); + plan = mock(PhysicalPlan.class, withSettings().extraInterfaces(ExternalizablePlan.class)); monitorPlan = new ResourceMonitorPlan(plan, resourceMonitor); assertEquals(plan, monitorPlan.getPlanForSerialization()); } From a16f2fa3bcd0afdd3142789fdac379d7f761edc7 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 11 May 2023 23:34:21 -0700 Subject: [PATCH 17/38] WIP OpenSearchIndexScan refactor. Unit tests pass. Signed-off-by: MaxKsyunz --- .../planner/optimizer/PushDownPageSize.java | 4 + .../sql/sql/PaginationWindowIT.java | 6 +- .../request/ContinuePageRequest.java | 14 +-- .../request/ContinuePageRequestBuilder.java | 26 +---- .../request/ExecutableRequestBuilder.java | 6 +- .../request/OpenSearchQueryRequest.java | 27 +++-- .../opensearch/request/OpenSearchRequest.java | 18 +-- .../request/OpenSearchRequestBuilder.java | 25 ++-- .../request/OpenSearchScrollRequest.java | 77 ++++++++----- .../opensearch/storage/OpenSearchIndex.java | 23 ++-- .../storage/scan/OpenSearchIndexScan.java | 50 ++------ .../OpenSearchExecutionEngineTest.java | 4 +- .../OpenSearchExecutionProtectorTest.java | 12 +- .../ContinuePageRequestBuilderTest.java | 49 -------- .../request/ContinuePageRequestTest.java | 25 +--- .../request/OpenSearchQueryRequestTest.java | 27 +++-- .../request/OpenSearchRequestBuilderTest.java | 93 ++++++++------- .../request/OpenSearchRequestTest.java | 23 ---- .../request/OpenSearchScrollRequestTest.java | 50 ++++---- .../storage/OpenSearchIndexTest.java | 38 +++--- .../OpenSearchIndexScanPaginationTest.java | 109 ++---------------- .../storage/scan/OpenSearchIndexScanTest.java | 71 ++++++------ 22 files changed, 286 insertions(+), 491 deletions(-) delete mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java delete mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java index 8d88c7ef43..831428001b 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java @@ -14,6 +14,10 @@ import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.storage.read.TableScanBuilder; +/** + * A {@link LogicalPlanOptimizer} rule that pushes down page size + * to table scan builder. + */ public class PushDownPageSize implements Rule { @Override public Pattern pattern() { diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java index 6a978a1116..be208cd137 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationWindowIT.java @@ -31,11 +31,10 @@ public void testFetchSizeLessThanMaxResultWindow() throws IOException { setMaxResultWindow(TEST_INDEX_PHRASE, 6); JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5); - String cursor; int numRows = 0; do { // Process response - cursor = response.getString("cursor"); + String cursor = response.getString("cursor"); numRows += response.getJSONArray("datarows").length(); response = executeCursorQuery(cursor); } while (response.has("cursor")); @@ -55,11 +54,10 @@ public void testQuerySizeLimitDoesNotEffectTotalRowsReturned() throws IOExceptio JSONObject response = executeQueryTemplate("SELECT * FROM %s", TEST_INDEX_PHRASE, 5); assertTrue(response.getInt("size") > querySizeLimit); - String cursor; int numRows = 0; do { // Process response - cursor = response.getString("cursor"); + String cursor = response.getString("cursor"); numRows += response.getJSONArray("datarows").length(); response = executeCursorQuery(cursor); } while (response.has("cursor")); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java index 0e06283429..ff8f945b1e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java @@ -5,6 +5,9 @@ package org.opensearch.sql.opensearch.request; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.List; import java.util.function.Consumer; import java.util.function.Function; @@ -16,7 +19,6 @@ 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; @@ -63,14 +65,12 @@ public void clean(Consumer cleanAction) { } @Override - public SearchSourceBuilder getSourceBuilder() { - throw new UnsupportedOperationException( - "SearchSourceBuilder is unavailable for ContinueScrollRequest"); + public void writeExternal(ObjectOutput out) { + throw new UnsupportedOperationException("ContinuePageRequest.writeExternal"); } @Override - public String toCursor() { - // on the last page, we shouldn't return the scroll to user, it is kept for closing (clean) - return scrollFinished ? null : responseScrollId; + public void readExternal(ObjectInput in) { + throw new UnsupportedOperationException("ContinuePageRequest.readExternal"); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java index a255d0495b..c480136ac3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java @@ -5,10 +5,7 @@ package org.opensearch.sql.opensearch.request; -import lombok.Getter; import org.opensearch.common.unit.TimeValue; -import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; /** * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. @@ -16,28 +13,11 @@ */ public class ContinuePageRequestBuilder implements ExecutableRequestBuilder { - @Getter - private final String scrollId; - private final OpenSearchExprValueFactory exprValueFactory; - private final TimeValue scrollTimeout; - - /** Constructor. */ - public ContinuePageRequestBuilder(String scrollId, TimeValue scrollTimeout, - OpenSearchExprValueFactory exprValueFactory) { - this.scrollId = scrollId; - this.scrollTimeout = scrollTimeout; - this.exprValueFactory = exprValueFactory; - } - @Override public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, - Settings settings) { - return new ContinuePageRequest(scrollId, scrollTimeout, exprValueFactory); + int maxResultWindow, TimeValue scrollTimeout) { + throw new UnsupportedOperationException("Implement me"); + //return new OpenSearchScrollRequest(scrollId, scrollTimeout, exprValueFactory); } - @Override - public int getMaxResponseSize() { - return Integer.MAX_VALUE; - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java index 6ca3dbebe7..9b1ec039e2 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java @@ -1,11 +1,9 @@ package org.opensearch.sql.opensearch.request; -import org.opensearch.sql.common.setting.Settings; +import org.opensearch.common.unit.TimeValue; public interface ExecutableRequestBuilder { - int getMaxResponseSize(); OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, - Settings settings); + int maxResultWindow, TimeValue scrollTimeout); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 63aeed02f0..2e54febb8e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -8,7 +8,9 @@ import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; -import com.google.common.annotations.VisibleForTesting; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Arrays; import java.util.List; import java.util.function.Consumer; @@ -47,6 +49,7 @@ public class OpenSearchQueryRequest implements OpenSearchRequest { private final SearchSourceBuilder sourceBuilder; + /** * OpenSearchExprValueFactory. */ @@ -102,7 +105,9 @@ public OpenSearchResponse search(Function searchA } else { searchDone = true; return new OpenSearchResponse( - searchAction.apply(searchRequest()), exprValueFactory, includes); + searchAction.apply(new SearchRequest() + .indices(indexName.getIndexNames()) + .source(sourceBuilder)), exprValueFactory, includes); } } @@ -111,15 +116,13 @@ public void clean(Consumer cleanAction) { //do nothing. } - /** - * Generate OpenSearch search request. - * - * @return search request - */ - @VisibleForTesting - protected SearchRequest searchRequest() { - return new SearchRequest() - .indices(indexName.getIndexNames()) - .source(sourceBuilder); + @Override + public void writeExternal(ObjectOutput out) throws IOException { + throw new UnsupportedOperationException("OpenSearchQueryRequest.writeExternal"); + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + throw new UnsupportedOperationException("OpenSearchQueryRequest.readExternal"); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index c5b6d60af3..c0704ea030 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -6,20 +6,21 @@ package org.opensearch.sql.opensearch.request; +import java.io.Externalizable; +import java.io.Serializable; import java.util.function.Consumer; import java.util.function.Function; import lombok.EqualsAndHashCode; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; 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.response.OpenSearchResponse; /** * OpenSearch search request. */ -public interface OpenSearchRequest { +public interface OpenSearchRequest extends Externalizable { /** * Apply the search action or scroll action on request based on context. * @@ -37,21 +38,14 @@ OpenSearchResponse search(Function searchAction, */ void clean(Consumer cleanAction); - /** - * Get the SearchSourceBuilder. - * - * @return SearchSourceBuilder. - */ - SearchSourceBuilder getSourceBuilder(); - /** * Get the ElasticsearchExprValueFactory. * @return ElasticsearchExprValueFactory. */ OpenSearchExprValueFactory getExprValueFactory(); - default String toCursor() { - return ""; + default boolean hasAnotherBatch() { + return false; } /** @@ -59,7 +53,7 @@ default String toCursor() { * Indices are separated by ",". */ @EqualsAndHashCode - class IndexName { + class IndexName implements Serializable { private static final String COMMA = ","; private final String[] indexNames; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index c519932c64..2e91ed4935 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -35,7 +35,6 @@ import org.opensearch.search.sort.SortBuilder; import org.opensearch.search.sort.SortBuilders; import org.opensearch.sql.ast.expression.Literal; -import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.common.utils.StringUtils; import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.ReferenceExpression; @@ -65,7 +64,7 @@ public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { /** * Query size of the request -- how many rows will be returned. */ - private int querySize; + private int requestedTotalSize; /** * Size of each page request to return. @@ -83,8 +82,8 @@ public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { /** * Constructor. */ - public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprValueFactory) { - this.querySize = querySize; + public OpenSearchRequestBuilder(int requestedTotalSize, OpenSearchExprValueFactory exprValueFactory) { + this.requestedTotalSize = requestedTotalSize; this.sourceBuilder = new SearchSourceBuilder() .from(startFrom) .timeout(DEFAULT_QUERY_TIMEOUT) @@ -92,11 +91,6 @@ public OpenSearchRequestBuilder(int querySize, OpenSearchExprValueFactory exprVa this.exprValueFactory = exprValueFactory; } - @Override - public int getMaxResponseSize() { - return pageSize == null ? querySize : pageSize; - } - /** * Build DSL request. * @@ -104,10 +98,8 @@ public int getMaxResponseSize() { */ @Override public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, - Settings settings) { - int size = querySize; - TimeValue scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + int maxResultWindow, TimeValue scrollTimeout) { + int size = requestedTotalSize; if (pageSize == null) { if (startFrom + size > maxResultWindow) { sourceBuilder.size(maxResultWindow - startFrom); @@ -115,7 +107,7 @@ public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, indexName, scrollTimeout, sourceBuilder, exprValueFactory); } else { sourceBuilder.from(startFrom); - sourceBuilder.size(querySize); + sourceBuilder.size(requestedTotalSize); return new OpenSearchQueryRequest(indexName, sourceBuilder, exprValueFactory); } } else { @@ -190,7 +182,7 @@ public void pushDownSort(List> sortBuilders) { * Pushdown size (limit) and from (offset) to DSL request. */ public void pushDownLimit(Integer limit, Integer offset) { - querySize = limit; + requestedTotalSize = limit; startFrom = offset; sourceBuilder.from(offset).size(limit); } @@ -270,6 +262,9 @@ fieldNames, createEmptyNestedQuery(path) ); } + public int getMaxResponseSize() { + return pageSize == null? requestedTotalSize : pageSize; + } /** * Initialize bool query for push down. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index b85c69a60e..7789c5d3f0 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -6,6 +6,9 @@ package org.opensearch.sql.opensearch.request; +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; import java.util.Arrays; import java.util.List; import java.util.Objects; @@ -20,7 +23,6 @@ import org.opensearch.action.search.SearchScrollRequest; import org.opensearch.common.unit.TimeValue; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.search.fetch.subphase.FetchSourceContext; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -34,20 +36,19 @@ @Getter @ToString public class OpenSearchScrollRequest implements OpenSearchRequest { - + private SearchRequest initialSearchRequest; /** Scroll context timeout. */ - private final TimeValue scrollTimeout; + private TimeValue scrollTimeout; /** * {@link OpenSearchRequest.IndexName}. */ - private final IndexName indexName; + private IndexName indexName; /** Index name. */ @EqualsAndHashCode.Exclude @ToString.Exclude - private final OpenSearchExprValueFactory exprValueFactory; - + private OpenSearchExprValueFactory exprValueFactory; /** * Scroll id which is set after first request issued. Because ElasticsearchClient is shared by * multi-thread so this state has to be maintained here. @@ -58,8 +59,12 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { private boolean needClean = false; - /** Search request source builder. */ - private final SearchSourceBuilder sourceBuilder; + private List includes; + + /** Default constructor for Externalizable only. + */ + public OpenSearchScrollRequest() { + } /** Constructor. */ public OpenSearchScrollRequest(IndexName indexName, @@ -68,9 +73,17 @@ public OpenSearchScrollRequest(IndexName indexName, OpenSearchExprValueFactory exprValueFactory) { this.indexName = indexName; this.scrollTimeout = scrollTimeout; - this.sourceBuilder = sourceBuilder; this.exprValueFactory = exprValueFactory; - } + this.initialSearchRequest = new SearchRequest() + .indices(indexName.getIndexNames()) + .scroll(scrollTimeout) + .source(sourceBuilder); + + includes = sourceBuilder.fetchSource() != null && sourceBuilder.fetchSource().includes() != null + ? Arrays.asList(sourceBuilder.fetchSource().includes()) + : List.of(); + } + /** Constructor. */ @Override @@ -80,12 +93,8 @@ public OpenSearchResponse search(Function searchA if (isScroll()) { openSearchResponse = scrollAction.apply(scrollRequest()); } else { - openSearchResponse = searchAction.apply(searchRequest()); + openSearchResponse = searchAction.apply(initialSearchRequest); } - FetchSourceContext fetchSource = this.sourceBuilder.fetchSource(); - List includes = fetchSource != null && fetchSource.includes() != null - ? Arrays.asList(this.sourceBuilder.fetchSource().includes()) - : List.of(); var response = new OpenSearchResponse(openSearchResponse, exprValueFactory, includes); needClean = response.isEmpty(); @@ -108,18 +117,6 @@ public void clean(Consumer cleanAction) { } } - /** - * Generate OpenSearch search request. - * - * @return search request - */ - public SearchRequest searchRequest() { - return new SearchRequest() - .indices(indexName.getIndexNames()) - .scroll(scrollTimeout) - .source(sourceBuilder); - } - /** * Is scroll started which means pages after first is being requested. * @@ -152,7 +149,29 @@ public void reset() { * @return a string representing the scroll request. */ @Override - public String toCursor() { - return needClean ? "" : scrollId; + public boolean hasAnotherBatch() { + return scrollId != null && !scrollId.equals(""); + } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + out.writeObject(initialSearchRequest); + out.writeLong(scrollTimeout.getMillis()); + out.writeObject(indexName); + out.writeObject(exprValueFactory); + out.writeUTF(scrollId); + out.writeBoolean(needClean); + out.writeObject(includes); + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + initialSearchRequest = (SearchRequest) in.readObject(); + scrollTimeout = TimeValue.timeValueMillis(in.readLong()); + indexName = (IndexName) in.readObject(); + exprValueFactory = (OpenSearchExprValueFactory) in.readObject(); + scrollId = in.readUTF(); + needClean = in.readBoolean(); + includes = (List) in.readObject(); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 90939f799d..b62dd162bf 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -11,6 +11,7 @@ import java.util.LinkedHashMap; import java.util.Map; import lombok.RequiredArgsConstructor; +import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; @@ -168,21 +169,29 @@ public PhysicalPlan implement(LogicalPlan plan) { @Override public TableScanBuilder createScanBuilder() { - Map allFields = new HashMap<>(); - getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); - allFields.putAll(getFieldOpenSearchTypes()); + final int querySizeLimit = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); + var builder = new OpenSearchRequestBuilder( - settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT), - new OpenSearchExprValueFactory(allFields)); + querySizeLimit, + createExprValueFactory()); + return new OpenSearchIndexScanBuilder(builder) { @Override protected TableScanOperator createScan(OpenSearchRequestBuilder requestBuilder) { - return new OpenSearchIndexScan(client, indexName, settings, getMaxResultWindow(), - requestBuilder); + final TimeValue cursorKeepAlive = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + return new OpenSearchIndexScan(client, requestBuilder.getMaxResponseSize(), + requestBuilder.build(indexName, getMaxResultWindow(), cursorKeepAlive)); } }; } + private OpenSearchExprValueFactory createExprValueFactory() { + Map allFields = new HashMap<>(); + getReservedFieldTypes().forEach((k, v) -> allFields.put(k, OpenSearchDataType.of(v))); + allFields.putAll(getFieldOpenSearchTypes()); + return new OpenSearchExprValueFactory(allFields); + } + @VisibleForTesting @RequiredArgsConstructor public static class OpenSearchDefaultImplementor diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index da84e6db07..d2e61bb378 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -13,18 +13,12 @@ import java.util.Iterator; import lombok.EqualsAndHashCode; import lombok.ToString; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; -import org.opensearch.sql.opensearch.request.ExecutableRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; -import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.storage.TableScanOperator; @@ -39,14 +33,6 @@ public class OpenSearchIndexScan extends TableScanOperator implements Externaliz /** OpenSearch client. */ private OpenSearchClient client; - private OpenSearchRequest.IndexName indexName; - private Settings settings; - private Integer maxResultWindow; - /** Search request builder. */ - @EqualsAndHashCode.Include - @ToString.Include - private ExecutableRequestBuilder requestBuilder; - /** Search request. */ @EqualsAndHashCode.Include @ToString.Include @@ -55,7 +41,7 @@ public class OpenSearchIndexScan extends TableScanOperator implements Externaliz /** Largest number of rows allowed in the response. */ @EqualsAndHashCode.Include @ToString.Include - private Integer maxResponseSize; + private int maxResponseSize; /** Number of rows returned. */ private Integer queryCount; @@ -67,22 +53,16 @@ public class OpenSearchIndexScan extends TableScanOperator implements Externaliz * Creates index scan based on a provided OpenSearchRequestBuilder. */ public OpenSearchIndexScan(OpenSearchClient client, - OpenSearchRequest.IndexName indexName, - Settings settings, - Integer maxResultWindow, - ExecutableRequestBuilder requestBuilder) { + int maxResponseSize, + OpenSearchRequest request) { this.client = client; - this.indexName = indexName; - this.settings = settings; - this.maxResultWindow = maxResultWindow; - this.requestBuilder = requestBuilder; + this.maxResponseSize = maxResponseSize; + this.request = request; } @Override public void open() { super.open(); - request = requestBuilder.build(indexName, maxResultWindow, settings); - maxResponseSize = requestBuilder.getMaxResponseSize(); iterator = Collections.emptyIterator(); queryCount = 0; fetchNextBatch(); @@ -126,7 +106,7 @@ public void close() { @Override public String explain() { - return requestBuilder.build(indexName, maxResultWindow, settings).toString(); + return request.toString(); } /** No-args constructor. @@ -137,33 +117,23 @@ public OpenSearchIndexScan() { } @Override - public void readExternal(ObjectInput in) throws IOException { - final var serializedName = in.readUTF(); - final var scrollId = in.readUTF(); + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + request = (OpenSearchRequest) in.readObject(); maxResponseSize = in.readInt(); var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) .resolveObject("engine"); - indexName = new OpenSearchRequest.IndexName(serializedName); client = engine.getClient(); - settings = engine.getSettings(); - var index = new OpenSearchIndex(client, settings, indexName.toString()); - maxResultWindow = index.getMaxResultWindow(); - - TimeValue scrollTimeout = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); - requestBuilder = new ContinuePageRequestBuilder(scrollId, scrollTimeout, - new OpenSearchExprValueFactory(index.getFieldOpenSearchTypes())); } @Override public void writeExternal(ObjectOutput out) throws IOException { - if (request.toCursor() == null || request.toCursor().isEmpty()) { + if (!request.hasAnotherBatch()) { throw new NoCursorException(); } - out.writeUTF(indexName.toString()); - out.writeUTF(request.toCursor()); + out.writeObject(request); out.writeInt(maxResponseSize); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index aa08ba7e00..684e59af5d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -182,8 +182,8 @@ void explain_successfully() { final int defaultQuerySize = 100; final int maxResultWindow = 10000; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); - PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), name, - settings, maxResultWindow, requestBuilder); + PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), + maxResultWindow, requestBuilder.build(name, maxResultWindow, settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE))); AtomicReference result = new AtomicReference<>(); executor.explain(plan, new ResponseListener<>() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index 3a8386e75a..bde2e0b156 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -39,6 +39,7 @@ import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.ast.tree.Sort; +import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; @@ -110,7 +111,8 @@ void testProtectIndexScan() { Integer offset = 10; final var name = new OpenSearchRequest.IndexName(indexName); - final var requestBuilder = new OpenSearchRequestBuilder(querySizeLimit, exprValueFactory); + final var request = new OpenSearchRequestBuilder(querySizeLimit, exprValueFactory) + .build(name, maxResultWindow, settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)); assertEquals( PhysicalPlanDSL.project( PhysicalPlanDSL.limit( @@ -124,8 +126,8 @@ void testProtectIndexScan() { PhysicalPlanDSL.agg( filter( resourceMonitor( - new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder)), + new OpenSearchIndexScan(client, + maxResultWindow, request)), filterExpr), aggregators, groupByExprs), @@ -151,8 +153,8 @@ void testProtectIndexScan() { PhysicalPlanDSL.rename( PhysicalPlanDSL.agg( filter( - new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), + new OpenSearchIndexScan(client, + maxResultWindow, request), filterExpr), aggregators, groupByExprs), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java deleted file mode 100644 index bf03b332d2..0000000000 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilderTest.java +++ /dev/null @@ -1,49 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.mock; - -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.BeforeEach; -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.common.unit.TimeValue; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; - -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -@ExtendWith(MockitoExtension.class) -class ContinuePageRequestBuilderTest { - - @Mock - private OpenSearchExprValueFactory exprValueFactory; - - private final String scrollId = "scroll"; - - private ContinuePageRequestBuilder requestBuilder; - - @BeforeEach - void setup() { - var timeout = TimeValue.timeValueMinutes(1); - requestBuilder = new ContinuePageRequestBuilder(scrollId, timeout, exprValueFactory); - } - - @Test - void build() { - assertEquals( - new ContinuePageRequest(scrollId, TimeValue.timeValueMinutes(1), exprValueFactory), - requestBuilder.build(null, 0, null) - ); - } -} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java index e991fc5787..40512eb46d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java @@ -6,15 +6,10 @@ package org.opensearch.sql.opensearch.request; import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -41,7 +36,7 @@ @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class ContinuePageRequestTest { +class ContinuePageRequestTest { @Mock private Function searchAction; @@ -71,7 +66,7 @@ public class ContinuePageRequestTest { scroll, TimeValue.timeValueMinutes(1), factory); @Test - public void search_with_non_empty_response() { + void search_with_non_empty_response() { when(scrollAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); @@ -80,7 +75,6 @@ public void search_with_non_empty_response() { OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); assertAll( () -> assertFalse(searchResponse.isEmpty()), - () -> assertEquals(nextScroll, request.toCursor()), () -> verify(scrollAction, times(1)).apply(any()), () -> verify(searchAction, never()).apply(any()) ); @@ -88,7 +82,7 @@ public void search_with_non_empty_response() { @Test // Empty response means scroll search is done and no cursor/scroll should be set - public void search_with_empty_response() { + void search_with_empty_response() { when(scrollAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(null); @@ -97,7 +91,6 @@ public void search_with_empty_response() { OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); assertAll( () -> assertTrue(searchResponse.isEmpty()), - () -> assertNull(request.toCursor()), () -> verify(scrollAction, times(1)).apply(any()), () -> verify(searchAction, never()).apply(any()) ); @@ -105,7 +98,7 @@ public void search_with_empty_response() { @Test @SneakyThrows - public void clean() { + void clean() { request.clean(cleanAction); verify(cleanAction, never()).accept(any()); // Enforce cleaning by setting a private field. @@ -113,14 +106,4 @@ public void clean() { request.clean(cleanAction); verify(cleanAction, times(1)).accept(any()); } - - @Test - // Added for coverage only - public void getters() { - factory = mock(); - assertAll( - () -> assertThrows(Throwable.class, request::getSourceBuilder), - () -> assertSame(factory, new ContinuePageRequest("", null, factory).getExprValueFactory()) - ); - } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index adb2a16a84..955e4899bc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -10,15 +10,16 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; -import java.util.Iterator; import java.util.function.Consumer; import java.util.function.Function; +import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; @@ -31,7 +32,6 @@ import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.FetchSourceContext; -import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -143,14 +143,21 @@ void clean() { void searchRequest() { request.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - assertEquals( + request.search(searchRequest ->{ + assertEquals( new SearchRequest() - .indices("test") - .source(new SearchSourceBuilder() - .timeout(DEFAULT_QUERY_TIMEOUT) - .from(0) - .size(200) - .query(QueryBuilders.termQuery("name", "John"))), - request.searchRequest()); + .indices("test") + .source(new SearchSourceBuilder() + .timeout(DEFAULT_QUERY_TIMEOUT) + .from(0) + .size(200) + .query(QueryBuilders.termQuery("name", "John"))), + searchRequest); + return when(mock(SearchResponse.class).getHits()) + .thenReturn(new SearchHits(new SearchHit[0], + new TotalHits(0, TotalHits.Relation.EQUAL_TO),0.0f)) + .getMock(); + }, searchScrollRequest -> null); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 35dbc9dedb..9b03f5dcab 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -8,8 +8,9 @@ import static org.junit.Assert.assertThrows; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.opensearch.index.query.QueryBuilders.matchAllQuery; import static org.opensearch.index.query.QueryBuilders.nestedQuery; import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME; @@ -22,6 +23,7 @@ import java.util.Map; import java.util.Set; import org.apache.commons.lang3.tuple.Pair; +import org.apache.lucene.search.TotalHits; import org.apache.lucene.search.join.ScoreMode; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -30,11 +32,14 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.action.search.SearchResponse; import org.opensearch.common.unit.TimeValue; import org.opensearch.index.query.InnerHitBuilder; import org.opensearch.index.query.NestedQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregationBuilders; import org.opensearch.search.aggregations.bucket.composite.TermsValuesSourceBuilder; @@ -66,8 +71,6 @@ class OpenSearchRequestBuilderTest { private static final OpenSearchRequest.IndexName indexName = new OpenSearchRequest.IndexName("test"); - @Mock - private Settings settings; @Mock private OpenSearchExprValueFactory exprValueFactory; @@ -76,11 +79,6 @@ class OpenSearchRequestBuilderTest { @BeforeEach void setup() { - lenient().when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)) - .thenReturn(DEFAULT_LIMIT); - lenient().when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); - requestBuilder = new OpenSearchRequestBuilder(DEFAULT_LIMIT, exprValueFactory); } @@ -100,7 +98,7 @@ void build_query_request() { .timeout(DEFAULT_QUERY_TIMEOUT) .trackScores(true), exprValueFactory), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } @Test @@ -117,7 +115,7 @@ void build_scroll_request_with_correct_size() { .size(MAX_RESULT_WINDOW - offset) .timeout(DEFAULT_QUERY_TIMEOUT), exprValueFactory), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } @Test @@ -125,15 +123,20 @@ void test_push_down_query() { QueryBuilder query = QueryBuilders.termQuery("intA", 1); requestBuilder.pushDownFilter(query); - assertEquals( + var r = requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT); + r.search(searchRequest -> { + assertEquals( new SearchSourceBuilder() - .from(DEFAULT_OFFSET) - .size(DEFAULT_LIMIT) - .timeout(DEFAULT_QUERY_TIMEOUT) - .query(query) - .sort(DOC_FIELD_NAME, ASC), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder() - ); + .from(DEFAULT_OFFSET) + .size(DEFAULT_LIMIT) + .timeout(DEFAULT_QUERY_TIMEOUT) + .query(query) + .sort(DOC_FIELD_NAME, ASC), + searchRequest.source() + ); + return mock(); + }, searchScrollRequest -> {throw new UnsupportedOperationException();}); + } @Test @@ -165,14 +168,23 @@ void test_push_down_query_and_sort() { FieldSortBuilder sortBuilder = SortBuilders.fieldSort("intA"); requestBuilder.pushDownSort(List.of(sortBuilder)); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .query(query) .sort(sortBuilder), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); + } + + void assertSearchSourceBuilder(SearchSourceBuilder expected, OpenSearchRequestBuilder requestBuilder) + throws UnsupportedOperationException { + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT).search(searchRequest -> { + assertEquals(expected, searchRequest.source()); + return when(mock(SearchResponse.class).getHits()).thenReturn(new SearchHits(new SearchHit[0], + new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)).getMock(); + }, searchScrollRequest ->{ throw new UnsupportedOperationException();}); } @Test @@ -180,13 +192,13 @@ void test_push_down_sort() { FieldSortBuilder sortBuilder = SortBuilders.fieldSort("intA"); requestBuilder.pushDownSort(List.of(sortBuilder)); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(sortBuilder), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -194,13 +206,13 @@ void test_push_down_non_field_sort() { ScoreSortBuilder sortBuilder = SortBuilders.scoreSort(); requestBuilder.pushDownSort(List.of(sortBuilder)); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(sortBuilder), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -209,14 +221,14 @@ void test_push_down_multiple_sort() { SortBuilders.fieldSort("intA"), SortBuilders.fieldSort("intB"))); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .sort(SortBuilders.fieldSort("intA")) .sort(SortBuilders.fieldSort("intB")), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -224,13 +236,13 @@ void test_push_down_project() { Set references = Set.of(DSL.ref("intA", INTEGER)); requestBuilder.pushDownProjects(references); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT) .fetchSource(new String[]{"intA"}, new String[0]), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -254,13 +266,13 @@ void test_push_down_nested() { .innerHit(new InnerHitBuilder().setFetchSourceContext( new FetchSourceContext(true, new String[]{"message.info"}, null))); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .query(QueryBuilders.boolQuery().filter(QueryBuilders.boolQuery().must(nestedQuery))) .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -287,13 +299,13 @@ void test_push_down_multiple_nested_with_same_path() { NestedQueryBuilder nestedQuery = nestedQuery("message", matchAllQuery(), ScoreMode.None) .innerHit(new InnerHitBuilder().setFetchSourceContext( new FetchSourceContext(true, new String[]{"message.info", "message.from"}, null))); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .query(QueryBuilders.boolQuery().filter(QueryBuilders.boolQuery().must(nestedQuery))) .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -318,7 +330,7 @@ void test_push_down_nested_with_filter() { .innerHit(new InnerHitBuilder().setFetchSourceContext( new FetchSourceContext(true, new String[]{"message.info"}, null))); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .query( QueryBuilders.boolQuery().filter( @@ -330,7 +342,7 @@ void test_push_down_nested_with_filter() { .from(DEFAULT_OFFSET) .size(DEFAULT_LIMIT) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -352,12 +364,12 @@ void push_down_highlight_with_repeating_fields() { @Test void push_down_page_size() { requestBuilder.pushDownPageSize(3); - assertEquals( + assertSearchSourceBuilder( new SearchSourceBuilder() .from(DEFAULT_OFFSET) .size(3) .timeout(DEFAULT_QUERY_TIMEOUT), - requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings).getSourceBuilder()); + requestBuilder); } @Test @@ -365,13 +377,6 @@ void exception_when_non_zero_offset_and_page_size() { requestBuilder.pushDownPageSize(3); requestBuilder.pushDownLimit(300, 2); assertThrows(UnsupportedOperationException.class, - () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, settings)); - } - - @Test - void query_size_is_page_size() { - requestBuilder.pushDownPageSize(3); - requestBuilder.pushDownLimit(4, 0); - assertEquals(3, requestBuilder.getMaxResponseSize()); + () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java deleted file mode 100644 index d0a274ce2a..0000000000 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestTest.java +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - - -package org.opensearch.sql.opensearch.request; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.CALLS_REAL_METHODS; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.withSettings; - -import org.junit.jupiter.api.Test; - -public class OpenSearchRequestTest { - - @Test - void toCursor() { - var request = mock(OpenSearchRequest.class, withSettings().defaultAnswer(CALLS_REAL_METHODS)); - assertEquals("", request.toCursor()); - } -} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index addacf04ee..8ebfea3b03 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -60,25 +60,27 @@ class OpenSearchScrollRequestTest { @Mock private SearchSourceBuilder sourceBuilder; - @Mock - private FetchSourceContext fetchSourceContext; @Mock private OpenSearchExprValueFactory factory; + private final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); private final OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - new SearchSourceBuilder(), factory); + searchSourceBuilder, factory); @Test void searchRequest() { - request.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - - assertEquals( + searchSourceBuilder.query(QueryBuilders.termQuery("name", "John")); + request.search(searchRequest -> { + assertEquals( new SearchRequest() - .indices("test") - .scroll(TimeValue.timeValueMinutes(1)) - .source(new SearchSourceBuilder().query(QueryBuilders.termQuery("name", "John"))), - request.searchRequest()); + .indices("test") + .scroll(TimeValue.timeValueMinutes(1)) + .source(new SearchSourceBuilder().query(QueryBuilders.termQuery("name", "John"))), + searchRequest); + SearchHits searchHitsMock = when(mock(SearchHits.class).getHits()).thenReturn(new SearchHit[0]).getMock(); + return when(mock(SearchResponse.class).getHits()).thenReturn(searchHitsMock).getMock(); + }, searchScrollRequest -> null); } @Test @@ -111,16 +113,13 @@ void search() { factory ); - String[] includes = {"_id", "_index"}; - when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext); - when(fetchSourceContext.includes()).thenReturn(includes); - when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); - OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - verify(fetchSourceContext, times(2)).includes(); - assertFalse(searchResponse.isEmpty()); + OpenSearchResponse openSearchResponse = request.search(searchRequest -> searchResponse, + searchScrollRequest -> {throw new AssertionError();}); + + assertFalse(openSearchResponse.isEmpty()); } @Test @@ -132,7 +131,6 @@ void search_withoutContext() { factory ); - when(sourceBuilder.fetchSource()).thenReturn(null); when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); @@ -151,24 +149,21 @@ void search_withoutIncludes() { factory ); - when(sourceBuilder.fetchSource()).thenReturn(fetchSourceContext); - when(fetchSourceContext.includes()).thenReturn(null); when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - verify(fetchSourceContext, times(1)).includes(); assertFalse(searchResponse.isEmpty()); } @Test - void toCursor() { + void hasAnotherBatch() { request.setScrollId("scroll123"); - assertEquals("scroll123", request.toCursor()); + assertTrue(request.hasAnotherBatch()); request.reset(); - assertNull(request.toCursor()); + assertFalse(request.hasAnotherBatch()); } @Test @@ -213,7 +208,7 @@ void no_cursor_on_empty_response() { new SearchHits(new SearchHit[0], null, 1f)); request.search((x) -> searchResponse, (x) -> searchResponse); - assertEquals("", request.toCursor()); + assertFalse(request.hasAnotherBatch()); } @Test @@ -227,9 +222,4 @@ void no_clean_if_no_scroll_in_response() { request.clean((s) -> fail()); } - - @Test - void noCursor() { - assertNull(request.toCursor()); - } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 04bbe62d03..326e5dcb9b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -3,7 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ - package org.opensearch.sql.opensearch.storage; import static org.hamcrest.MatcherAssert.assertThat; @@ -39,6 +38,7 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.ast.tree.Sort; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; @@ -61,7 +61,9 @@ @ExtendWith(MockitoExtension.class) class OpenSearchIndexTest { - private final String indexName = "test"; + public static final int QUERY_SIZE_LIMIT = 200; + public static final TimeValue SCROLL_TIMEOUT = new TimeValue(1); + public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); @Mock private OpenSearchClient client; @@ -79,12 +81,12 @@ class OpenSearchIndexTest { @BeforeEach void setUp() { - this.index = new OpenSearchIndex(client, settings, indexName); + this.index = new OpenSearchIndex(client, settings, "test"); } @Test void isExist() { - when(client.exists(indexName)).thenReturn(true); + when(client.exists("test")).thenReturn(true); assertTrue(index.exists()); } @@ -96,7 +98,7 @@ void createIndex() { Map.of( "name", "text", "age", "integer")); - doNothing().when(client).createIndex(indexName, mappings); + doNothing().when(client).createIndex("test", mappings); Map schema = new HashMap<>(); schema.put("name", OpenSearchTextType.of(Map.of("keyword", @@ -199,11 +201,10 @@ void implementRelationOperatorOnly() { when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - final var name = new OpenSearchRequest.IndexName(indexName); - final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); - final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); - assertEquals(new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), index.implement(index.optimize(plan))); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE_LIMIT, exprValueFactory); + assertEquals(new OpenSearchIndexScan(client, + 200, requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), + index.implement(index.optimize(plan))); } @Test @@ -212,11 +213,9 @@ void implementRelationOperatorWithOptimization() { when(settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT)).thenReturn(200); LogicalPlan plan = index.createScanBuilder(); Integer maxResultWindow = index.getMaxResultWindow(); - final var name = new OpenSearchRequest.IndexName(indexName); - final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); - final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); - assertEquals(new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), index.implement(plan)); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE_LIMIT, exprValueFactory); + assertEquals(new OpenSearchIndexScan(client, 200, + requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), index.implement(plan)); } @Test @@ -230,7 +229,6 @@ void implementOtherLogicalOperators() { ImmutableMap.of(ref("name", STRING), ref("lastname", STRING)); Pair newEvalField = ImmutablePair.of(ref("name1", STRING), ref("name", STRING)); - Integer sortCount = 100; Pair sortField = ImmutablePair.of(Sort.SortOption.DEFAULT_ASC, ref("name1", STRING)); @@ -250,9 +248,7 @@ void implementOtherLogicalOperators() { include); Integer maxResultWindow = index.getMaxResultWindow(); - final var name = new OpenSearchRequest.IndexName(indexName); - final int defaultQuerySize = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); - final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); + final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE_LIMIT, exprValueFactory); assertEquals( PhysicalPlanDSL.project( PhysicalPlanDSL.dedupe( @@ -260,8 +256,8 @@ void implementOtherLogicalOperators() { PhysicalPlanDSL.eval( PhysicalPlanDSL.remove( PhysicalPlanDSL.rename( - new OpenSearchIndexScan(client, name, - settings, maxResultWindow, requestBuilder), + new OpenSearchIndexScan(client, + QUERY_SIZE_LIMIT, requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), mappings), exclude), newEvalField), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java index 7ab8900e59..cc22b315c7 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -5,8 +5,6 @@ package org.opensearch.sql.opensearch.storage.scan; -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -16,13 +14,11 @@ import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; import static org.opensearch.sql.data.type.ExprCoreType.STRING; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.QUERY_SIZE; -import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.employee; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockTwoPageResponse; @@ -41,14 +37,11 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.common.setting.Settings; -import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.request.ContinuePageRequestBuilder; -import org.opensearch.sql.opensearch.request.ExecutableRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -86,90 +79,13 @@ void query_empty_result() { mockResponse(client); var builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + = new OpenSearchIndexScan(client, MAX_RESULT_WINDOW, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)))) { indexScan.open(); assertFalse(indexScan.hasNext()); } verify(client).cleanup(any()); } - @Test - void query_all_results_initial_scroll_request() { - mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); - - ExecutableRequestBuilder builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { - indexScan.open(); - - assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) - ); - } - verify(client).cleanup(any()); - - builder = new ContinuePageRequestBuilder(SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); - try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { - indexScan.open(); - - assertFalse(indexScan.hasNext()); - } - verify(client, times(2)).cleanup(any()); - } - - @Test - void query_all_results_continuation_scroll_request() { - mockResponse(client, new ExprValue[]{ - employee(1, "John", "IT"), - employee(2, "Smith", "HR"), - employee(3, "Allen", "IT")}); - - ContinuePageRequestBuilder builder = new ContinuePageRequestBuilder( - SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); - try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { - indexScan.open(); - - assertAll( - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(1, "John", "IT"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(2, "Smith", "HR"), indexScan.next()), - - () -> assertTrue(indexScan.hasNext()), - () -> assertEquals(employee(3, "Allen", "IT"), indexScan.next()), - - () -> assertFalse(indexScan.hasNext()), - () -> assertEquals(3, indexScan.getTotalHits()) - ); - } - verify(client).cleanup(any()); - - builder = new ContinuePageRequestBuilder(SCROLL_ID, SCROLL_TIMEOUT, exprValueFactory); - try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { - indexScan.open(); - - assertFalse(indexScan.hasNext()); - } - verify(client, times(2)).cleanup(any()); - } - @Test void explain_not_implemented() { assertThrows(Throwable.class, () -> mock(OpenSearchIndexScan.class, @@ -182,10 +98,10 @@ void serialization() { mockTwoPageResponse(client); OpenSearchRequestBuilder builder = mock(); OpenSearchRequest request = mock(); - when(request.toCursor()).thenReturn("cu-cursor"); + when(request.hasAnotherBatch()).thenReturn(true); when(builder.build(any(), eq(MAX_RESULT_WINDOW), any())).thenReturn(request); - var indexScan = new OpenSearchIndexScan(client, INDEX_NAME, settings, - MAX_RESULT_WINDOW, builder); + var indexScan = new OpenSearchIndexScan(client, + MAX_RESULT_WINDOW, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, SCROLL_TIMEOUT)); indexScan.open(); ByteArrayOutputStream output = new ByteArrayOutputStream(); @@ -193,12 +109,8 @@ void serialization() { objectOutput.writeObject(indexScan); objectOutput.flush(); - when(client.getIndexMappings(any())).thenReturn(Map.of()); OpenSearchStorageEngine engine = mock(); when(engine.getClient()).thenReturn(client); - when(engine.getSettings()).thenReturn(mock()); - when(client.getIndexMaxResultWindows(any())) - .thenReturn((Map.of(INDEX_NAME.getIndexNames()[0], MAX_RESULT_WINDOW))); ObjectInputStream objectInput = new PlanSerializer(engine) .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); var roundTripScan = (OpenSearchIndexScan) objectInput.readObject(); @@ -215,15 +127,14 @@ void dont_serialize_if_no_cursor() { when(builder.build(any(), anyInt(), any())).thenReturn(request); when(client.search(any())).thenReturn(response); try (var indexScan - = new OpenSearchIndexScan(client, INDEX_NAME, settings, MAX_RESULT_WINDOW, builder)) { + = new OpenSearchIndexScan(client, MAX_RESULT_WINDOW, + builder.build(INDEX_NAME, MAX_RESULT_WINDOW, SCROLL_TIMEOUT))) { indexScan.open(); - when(request.toCursor()).thenReturn(null, ""); - for (int i = 0; i < 2; i++) { - ByteArrayOutputStream output = new ByteArrayOutputStream(); - ObjectOutputStream objectOutput = new ObjectOutputStream(output); - assertThrows(NoCursorException.class, () -> objectOutput.writeObject(indexScan)); - } + when(request.hasAnotherBatch()).thenReturn(false); + ByteArrayOutputStream output = new ByteArrayOutputStream(); + ObjectOutputStream objectOutput = new ObjectOutputStream(output); + assertThrows(NoCursorException.class, () -> objectOutput.writeObject(indexScan)); } } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 8ddb2ad45d..80033e1e25 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -36,6 +36,7 @@ import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.SearchHit; +import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder; import org.opensearch.sql.ast.expression.DataType; import org.opensearch.sql.ast.expression.Literal; @@ -55,6 +56,9 @@ class OpenSearchIndexScanTest { public static final int QUERY_SIZE = 200; + public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("employees"); + public static final int MAX_RESULT_WINDOW = 10000; + public static final TimeValue CURSOR_KEEP_ALIVE = TimeValue.timeValueMinutes(1); @Mock private OpenSearchClient client; @@ -67,8 +71,6 @@ class OpenSearchIndexScanTest { @BeforeEach void setup() { - when(settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)) - .thenReturn(TimeValue.timeValueMinutes(1)); } @Test @@ -76,8 +78,8 @@ void query_empty_result() { mockResponse(client); final var name = new OpenSearchRequest.IndexName("test"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 3, requestBuilder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + QUERY_SIZE, requestBuilder.build(name, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -94,10 +96,9 @@ void query_all_results_with_query() { employee(2, "Smith", "HR"), employee(3, "Allen", "IT")}); - final var name = new OpenSearchRequest.IndexName("employees"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -126,10 +127,9 @@ void query_all_results_with_scroll() { new ExprValue[]{employee(1, "John", "IT"), employee(2, "Smith", "HR")}, new ExprValue[]{employee(3, "Allen", "IT")}); - final var name = new OpenSearchRequest.IndexName("employees"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -157,10 +157,10 @@ void query_some_results_with_query() { employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); + final int limit = 3; OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(0, exprValueFactory); - builder.pushDownLimit(3, 0); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, - 10, builder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + limit, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -184,9 +184,8 @@ void query_some_results_with_query() { void query_some_results_with_scroll() { mockTwoPageResponse(client); final var requestuilder = new OpenSearchRequestBuilder(10, exprValueFactory); - requestuilder.pushDownLimit(3, 0); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, - 2, requestuilder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + 3, requestuilder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -220,11 +219,10 @@ void query_results_limited_by_query_size() { employee(3, "Allen", "IT"), employee(4, "Bob", "HR")}); - final var name = new OpenSearchRequest.IndexName("employees"); final int defaultQuerySize = 2; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); - try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, name, - settings, 10, requestBuilder)) { + try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, + defaultQuerySize, requestBuilder.build(INDEX_NAME, QUERY_SIZE, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -292,7 +290,6 @@ private PushDownAssertion assertThat() { private static class PushDownAssertion { private final OpenSearchClient client; private final OpenSearchRequestBuilder requestBuilder; - private final Settings settings; private final OpenSearchResponse response; private final OpenSearchExprValueFactory factory; @@ -301,7 +298,6 @@ public PushDownAssertion(OpenSearchClient client, Settings settings) { this.client = client; this.requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, valueFactory); - this.settings = settings; this.response = mock(OpenSearchResponse.class); this.factory = valueFactory; @@ -319,26 +315,33 @@ PushDownAssertion pushDownHighlight(String query, Map arguments } PushDownAssertion shouldQueryHighlight(QueryBuilder query, HighlightBuilder highlight) { - OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, QUERY_SIZE, factory); - request.getSourceBuilder() - .query(query) - .highlighter(highlight) - .sort(DOC_FIELD_NAME, ASC); + var sourceBuilder = new SearchSourceBuilder() + .from(0) + .timeout(CURSOR_KEEP_ALIVE) + .query(query) + .size(QUERY_SIZE) + .highlighter(highlight) + .sort(DOC_FIELD_NAME, ASC); + OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, sourceBuilder, factory); + when(client.search(request)).thenReturn(response); - var indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, - 10000, requestBuilder); + var indexScan = new OpenSearchIndexScan(client, + QUERY_SIZE, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); indexScan.open(); return this; } PushDownAssertion shouldQuery(QueryBuilder expected) { - OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, QUERY_SIZE, factory); - request.getSourceBuilder() - .query(expected) - .sort(DOC_FIELD_NAME, ASC); + var builder = new SearchSourceBuilder() + .from(0) + .query(expected) + .size(QUERY_SIZE) + .timeout(CURSOR_KEEP_ALIVE) + .sort(DOC_FIELD_NAME, ASC); + OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, builder, factory); when(client.search(request)).thenReturn(response); - var indexScan = new OpenSearchIndexScan(client, EMPLOYEES_INDEX, settings, - 10000, requestBuilder); + var indexScan = new OpenSearchIndexScan(client, + 10000, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); indexScan.open(); return this; } From eff76b2c04c6192d037fbe694a2bb9b203f013ab Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 16 May 2023 11:11:26 -0700 Subject: [PATCH 18/38] Refactor OpenSearchIndexScan and OpenSearchRequest. Signed-off-by: MaxKsyunz --- .../request/ContinuePageRequest.java | 76 ------------ .../request/ContinuePageRequestBuilder.java | 23 ---- .../request/OpenSearchQueryRequest.java | 10 +- .../opensearch/request/OpenSearchRequest.java | 18 ++- .../request/OpenSearchScrollRequest.java | 40 ++++--- .../storage/scan/OpenSearchIndexScan.java | 33 +++++- .../request/ContinuePageRequestTest.java | 109 ------------------ .../OpenSearchIndexScanPaginationTest.java | 33 ------ .../storage/scan/OpenSearchIndexScanTest.java | 4 + 9 files changed, 71 insertions(+), 275 deletions(-) delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java delete mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java deleted file mode 100644 index ff8f945b1e..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequest.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; -import java.util.List; -import java.util.function.Consumer; -import java.util.function.Function; -import lombok.EqualsAndHashCode; -import lombok.Getter; -import lombok.RequiredArgsConstructor; -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.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.response.OpenSearchResponse; - -/** - * Scroll (cursor) request is used to page the search. This request is not configurable and has - * no search query. It just handles paging through responses to the initial request. - * It is used on second and next pagination (cursor) requests. - */ -@EqualsAndHashCode -@RequiredArgsConstructor -public class ContinuePageRequest implements OpenSearchRequest { - private final String initialScrollId; - private final TimeValue scrollTimeout; - // ScrollId that OpenSearch returns after search. - private String responseScrollId; - - @EqualsAndHashCode.Exclude - @ToString.Exclude - @Getter - private final OpenSearchExprValueFactory exprValueFactory; - - @EqualsAndHashCode.Exclude - private boolean scrollFinished = false; - - @Override - public OpenSearchResponse search(Function searchAction, - Function scrollAction) { - SearchResponse openSearchResponse = scrollAction.apply(new SearchScrollRequest(initialScrollId) - .scroll(scrollTimeout)); - - // TODO if terminated_early - something went wrong, e.g. no scroll returned. - var response = new OpenSearchResponse(openSearchResponse, exprValueFactory, List.of()); - // on the last empty page, we should close the scroll - scrollFinished = response.isEmpty(); - responseScrollId = openSearchResponse.getScrollId(); - return response; - } - - @Override - public void clean(Consumer cleanAction) { - if (scrollFinished) { - cleanAction.accept(responseScrollId); - } - } - - @Override - public void writeExternal(ObjectOutput out) { - throw new UnsupportedOperationException("ContinuePageRequest.writeExternal"); - } - - @Override - public void readExternal(ObjectInput in) { - throw new UnsupportedOperationException("ContinuePageRequest.readExternal"); - } -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java deleted file mode 100644 index c480136ac3..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ContinuePageRequestBuilder.java +++ /dev/null @@ -1,23 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import org.opensearch.common.unit.TimeValue; - -/** - * Builds a {@link ContinuePageRequest} to handle subsequent pagination/scroll/cursor requests. - * Initial search requests are handled by {@link OpenSearchRequestBuilder} - */ -public class ContinuePageRequestBuilder implements ExecutableRequestBuilder { - - @Override - public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, TimeValue scrollTimeout) { - throw new UnsupportedOperationException("Implement me"); - //return new OpenSearchScrollRequest(scrollId, scrollTimeout, exprValueFactory); - } - -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 2e54febb8e..880977d669 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -21,6 +21,7 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.common.io.stream.StreamOutput; import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.fetch.subphase.FetchSourceContext; @@ -117,12 +118,7 @@ public void clean(Consumer cleanAction) { } @Override - public void writeExternal(ObjectOutput out) throws IOException { - throw new UnsupportedOperationException("OpenSearchQueryRequest.writeExternal"); - } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - throw new UnsupportedOperationException("OpenSearchQueryRequest.readExternal"); + public void writeTo(StreamOutput out) throws IOException { + throw new UnsupportedOperationException("Not necessary"); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index c0704ea030..29e8544ab5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -6,21 +6,23 @@ package org.opensearch.sql.opensearch.request; -import java.io.Externalizable; -import java.io.Serializable; +import java.io.IOException; import java.util.function.Consumer; import java.util.function.Function; import lombok.EqualsAndHashCode; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; +import org.opensearch.common.io.stream.Writeable; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; /** * OpenSearch search request. */ -public interface OpenSearchRequest extends Externalizable { +public interface OpenSearchRequest extends Writeable { /** * Apply the search action or scroll action on request based on context. * @@ -53,11 +55,14 @@ default boolean hasAnotherBatch() { * Indices are separated by ",". */ @EqualsAndHashCode - class IndexName implements Serializable { + class IndexName implements Writeable { private static final String COMMA = ","; private final String[] indexNames; + public IndexName(StreamInput si) throws IOException { + indexNames = si.readStringArray(); + } public IndexName(String indexName) { this.indexNames = indexName.split(COMMA); } @@ -70,5 +75,10 @@ public String[] getIndexNames() { public String toString() { return String.join(COMMA, indexNames); } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeStringArray(indexNames); + } } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 7789c5d3f0..17a0580813 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -7,8 +7,6 @@ package org.opensearch.sql.opensearch.request; import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; import java.util.Arrays; import java.util.List; import java.util.Objects; @@ -21,10 +19,14 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; 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 org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; /** * OpenSearch scroll search request. This has to be stateful because it needs to: @@ -150,28 +152,30 @@ public void reset() { */ @Override public boolean hasAnotherBatch() { - return scrollId != null && !scrollId.equals(""); + return !needClean && scrollId != null && !scrollId.equals(""); } @Override - public void writeExternal(ObjectOutput out) throws IOException { - out.writeObject(initialSearchRequest); - out.writeLong(scrollTimeout.getMillis()); - out.writeObject(indexName); - out.writeObject(exprValueFactory); - out.writeUTF(scrollId); + public void writeTo(StreamOutput out) throws IOException { + initialSearchRequest.writeTo(out); + out.writeTimeValue(scrollTimeout); + if (!needClean) { + // If needClean is true, there is no more data to get from OpenSearch and scrollId is + // used only to clean up OpenSearch context. + out.writeString(scrollId); + } out.writeBoolean(needClean); - out.writeObject(includes); + out.writeStringCollection(includes); + indexName.writeTo(out); } - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - initialSearchRequest = (SearchRequest) in.readObject(); - scrollTimeout = TimeValue.timeValueMillis(in.readLong()); - indexName = (IndexName) in.readObject(); - exprValueFactory = (OpenSearchExprValueFactory) in.readObject(); - scrollId = in.readUTF(); + public OpenSearchScrollRequest(StreamInput in, OpenSearchStorageEngine engine) throws IOException { + initialSearchRequest = new SearchRequest(in); + scrollTimeout = in.readTimeValue(); + scrollId = in.readString(); needClean = in.readBoolean(); - includes = (List) in.readObject(); + includes = in.readStringList(); + indexName = new IndexName(in); + exprValueFactory = new OpenSearchExprValueFactory(((OpenSearchIndex) engine.getTable(null, indexName.toString())).getFieldOpenSearchTypes()); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index d2e61bb378..a1da02914c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -6,6 +6,7 @@ package org.opensearch.sql.opensearch.storage.scan; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.ObjectInput; import java.io.ObjectOutput; @@ -13,11 +14,14 @@ import java.util.Iterator; import lombok.EqualsAndHashCode; import lombok.ToString; +import org.opensearch.common.io.stream.BytesStreamInput; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.request.OpenSearchRequest; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; import org.opensearch.sql.planner.ExternalizablePlan; @@ -117,14 +121,22 @@ public OpenSearchIndexScan() { } @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - request = (OpenSearchRequest) in.readObject(); - maxResponseSize = in.readInt(); + public void readExternal(ObjectInput in) throws IOException { + int reqSize = in.readInt(); + byte[] requestStream = new byte[reqSize]; + in.read(requestStream); var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) - .resolveObject("engine"); + .resolveObject("engine"); + + try (BytesStreamInput bsi = new BytesStreamInput(requestStream)) { + + request = new OpenSearchScrollRequest(bsi, engine); + } + maxResponseSize = in.readInt(); client = engine.getClient(); + } @Override @@ -133,7 +145,18 @@ public void writeExternal(ObjectOutput out) throws IOException { throw new NoCursorException(); } - out.writeObject(request); + BytesStreamOutput reqOut = new BytesStreamOutput(); + request.writeTo(reqOut); + reqOut.flush(); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + reqOut.bytes().writeTo(baos); + baos.flush(); + + var reqAsBytes = baos.toByteArray(); + out.writeInt(reqAsBytes.length); + out.write(reqAsBytes); + out.writeInt(maxResponseSize); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java deleted file mode 100644 index 40512eb46d..0000000000 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/ContinuePageRequestTest.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.opensearch.request; - -import static org.junit.jupiter.api.Assertions.assertAll; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.lenient; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.util.function.Consumer; -import java.util.function.Function; -import lombok.SneakyThrows; -import org.apache.commons.lang3.reflect.FieldUtils; -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.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.SearchHit; -import org.opensearch.search.SearchHits; -import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; -import org.opensearch.sql.opensearch.response.OpenSearchResponse; - -@ExtendWith(MockitoExtension.class) -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -class ContinuePageRequestTest { - - @Mock - private Function searchAction; - - @Mock - private Function scrollAction; - - @Mock - private Consumer cleanAction; - - @Mock - private SearchResponse searchResponse; - - @Mock - private SearchHits searchHits; - - @Mock - private SearchHit searchHit; - - @Mock - private OpenSearchExprValueFactory factory; - - private final String scroll = "scroll"; - private final String nextScroll = "nextScroll"; - - private final ContinuePageRequest request = new ContinuePageRequest( - scroll, TimeValue.timeValueMinutes(1), factory); - - @Test - void search_with_non_empty_response() { - when(scrollAction.apply(any())).thenReturn(searchResponse); - when(searchResponse.getHits()).thenReturn(searchHits); - when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); - when(searchResponse.getScrollId()).thenReturn(nextScroll); - - OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - assertAll( - () -> assertFalse(searchResponse.isEmpty()), - () -> verify(scrollAction, times(1)).apply(any()), - () -> verify(searchAction, never()).apply(any()) - ); - } - - @Test - // Empty response means scroll search is done and no cursor/scroll should be set - void search_with_empty_response() { - when(scrollAction.apply(any())).thenReturn(searchResponse); - when(searchResponse.getHits()).thenReturn(searchHits); - when(searchHits.getHits()).thenReturn(null); - lenient().when(searchResponse.getScrollId()).thenReturn(nextScroll); - - OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); - assertAll( - () -> assertTrue(searchResponse.isEmpty()), - () -> verify(scrollAction, times(1)).apply(any()), - () -> verify(searchAction, never()).apply(any()) - ); - } - - @Test - @SneakyThrows - void clean() { - request.clean(cleanAction); - verify(cleanAction, never()).accept(any()); - // Enforce cleaning by setting a private field. - FieldUtils.writeField(request, "scrollFinished", true, true); - request.clean(cleanAction); - verify(cleanAction, times(1)).accept(any()); - } -} diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java index cc22b315c7..1b890a286d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -7,10 +7,8 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.CALLS_REAL_METHODS; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; @@ -20,11 +18,8 @@ import static org.opensearch.sql.data.type.ExprCoreType.STRING; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.QUERY_SIZE; import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockResponse; -import static org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanTest.mockTwoPageResponse; -import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.util.Map; import lombok.SneakyThrows; @@ -38,14 +33,12 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.response.OpenSearchResponse; -import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -92,32 +85,6 @@ void explain_not_implemented() { withSettings().defaultAnswer(CALLS_REAL_METHODS)).explain()); } - @Test - @SneakyThrows - void serialization() { - mockTwoPageResponse(client); - OpenSearchRequestBuilder builder = mock(); - OpenSearchRequest request = mock(); - when(request.hasAnotherBatch()).thenReturn(true); - when(builder.build(any(), eq(MAX_RESULT_WINDOW), any())).thenReturn(request); - var indexScan = new OpenSearchIndexScan(client, - MAX_RESULT_WINDOW, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, SCROLL_TIMEOUT)); - indexScan.open(); - - ByteArrayOutputStream output = new ByteArrayOutputStream(); - ObjectOutputStream objectOutput = new ObjectOutputStream(output); - objectOutput.writeObject(indexScan); - objectOutput.flush(); - - OpenSearchStorageEngine engine = mock(); - when(engine.getClient()).thenReturn(client); - ObjectInputStream objectInput = new PlanSerializer(engine) - .getCursorDeserializationStream(new ByteArrayInputStream(output.toByteArray())); - var roundTripScan = (OpenSearchIndexScan) objectInput.readObject(); - roundTripScan.open(); - assertTrue(roundTripScan.hasNext()); - } - @Test @SneakyThrows void dont_serialize_if_no_cursor() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 80033e1e25..bd5740178b 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -19,9 +19,12 @@ import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import java.io.ByteArrayOutputStream; +import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; import java.util.Map; +import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -49,6 +52,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) From 8548d904dfd0c34d7e3339ea9e353b2cd717260b Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Tue, 16 May 2023 11:38:45 -0700 Subject: [PATCH 19/38] Refactor OpenSearchIndexScan and OpenSearchRequest. Signed-off-by: MaxKsyunz --- .../request/OpenSearchQueryRequest.java | 2 - .../opensearch/request/OpenSearchRequest.java | 1 + .../request/OpenSearchRequestBuilder.java | 6 ++- .../request/OpenSearchScrollRequest.java | 14 ++++-- .../opensearch/storage/OpenSearchIndex.java | 9 ++-- .../storage/scan/OpenSearchIndexScan.java | 2 +- .../OpenSearchExecutionEngineTest.java | 3 +- .../OpenSearchExecutionProtectorTest.java | 3 +- .../request/OpenSearchQueryRequestTest.java | 12 +++--- .../request/OpenSearchRequestBuilderTest.java | 31 +++++++++---- .../request/OpenSearchScrollRequestTest.java | 12 ++++-- .../storage/OpenSearchIndexTest.java | 12 +++--- .../OpenSearchIndexScanPaginationTest.java | 5 +-- .../storage/scan/OpenSearchIndexScanTest.java | 43 +++++++++---------- 14 files changed, 92 insertions(+), 63 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 880977d669..b20da41b84 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -9,8 +9,6 @@ import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; import java.util.Arrays; import java.util.List; import java.util.function.Consumer; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index 29e8544ab5..184f729bfb 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -63,6 +63,7 @@ class IndexName implements Writeable { public IndexName(StreamInput si) throws IOException { indexNames = si.readStringArray(); } + public IndexName(String indexName) { this.indexNames = indexName.split(COMMA); } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 2e91ed4935..72ef85d8c5 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -82,7 +82,8 @@ public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { /** * Constructor. */ - public OpenSearchRequestBuilder(int requestedTotalSize, OpenSearchExprValueFactory exprValueFactory) { + public OpenSearchRequestBuilder(int requestedTotalSize, + OpenSearchExprValueFactory exprValueFactory) { this.requestedTotalSize = requestedTotalSize; this.sourceBuilder = new SearchSourceBuilder() .from(startFrom) @@ -263,8 +264,9 @@ fieldNames, createEmptyNestedQuery(path) } public int getMaxResponseSize() { - return pageSize == null? requestedTotalSize : pageSize; + return pageSize == null ? requestedTotalSize : pageSize; } + /** * Initialize bool query for push down. */ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 17a0580813..c9ac24141e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -84,7 +84,7 @@ public OpenSearchScrollRequest(IndexName indexName, includes = sourceBuilder.fetchSource() != null && sourceBuilder.fetchSource().includes() != null ? Arrays.asList(sourceBuilder.fetchSource().includes()) : List.of(); - } + } /** Constructor. */ @@ -169,13 +169,21 @@ public void writeTo(StreamOutput out) throws IOException { indexName.writeTo(out); } - public OpenSearchScrollRequest(StreamInput in, OpenSearchStorageEngine engine) throws IOException { + /** + * Constructs OpenSearchScrollRequest from serialized representation. + * @param in stream to read data from. + * @param engine OpenSearchSqlEngine to get node-specific context. + * @throws IOException thrown if reading from input {@param in} fails. + */ + public OpenSearchScrollRequest(StreamInput in, OpenSearchStorageEngine engine) + throws IOException { initialSearchRequest = new SearchRequest(in); scrollTimeout = in.readTimeValue(); scrollId = in.readString(); needClean = in.readBoolean(); includes = in.readStringList(); indexName = new IndexName(in); - exprValueFactory = new OpenSearchExprValueFactory(((OpenSearchIndex) engine.getTable(null, indexName.toString())).getFieldOpenSearchTypes()); + OpenSearchIndex index = (OpenSearchIndex) engine.getTable(null, indexName.toString()); + exprValueFactory = new OpenSearchExprValueFactory(index.getFieldOpenSearchTypes()); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index b62dd162bf..532d62333d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -172,15 +172,16 @@ public TableScanBuilder createScanBuilder() { final int querySizeLimit = settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT); var builder = new OpenSearchRequestBuilder( - querySizeLimit, - createExprValueFactory()); + querySizeLimit, + createExprValueFactory()); return new OpenSearchIndexScanBuilder(builder) { @Override protected TableScanOperator createScan(OpenSearchRequestBuilder requestBuilder) { - final TimeValue cursorKeepAlive = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); + final TimeValue cursorKeepAlive = + settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE); return new OpenSearchIndexScan(client, requestBuilder.getMaxResponseSize(), - requestBuilder.build(indexName, getMaxResultWindow(), cursorKeepAlive)); + requestBuilder.build(indexName, getMaxResultWindow(), cursorKeepAlive)); } }; } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index a1da02914c..049d67b92b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -127,7 +127,7 @@ public void readExternal(ObjectInput in) throws IOException { in.read(requestStream); var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) - .resolveObject("engine"); + .resolveObject("engine"); try (BytesStreamInput bsi = new BytesStreamInput(requestStream)) { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index 684e59af5d..63e0b8e08f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -183,7 +183,8 @@ void explain_successfully() { final int maxResultWindow = 10000; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); PhysicalPlan plan = new OpenSearchIndexScan(mock(OpenSearchClient.class), - maxResultWindow, requestBuilder.build(name, maxResultWindow, settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE))); + maxResultWindow, requestBuilder.build(name, maxResultWindow, + settings.getSettingValue(SQL_CURSOR_KEEP_ALIVE))); AtomicReference result = new AtomicReference<>(); executor.explain(plan, new ResponseListener<>() { diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java index bde2e0b156..1c978c849e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java @@ -112,7 +112,8 @@ void testProtectIndexScan() { final var name = new OpenSearchRequest.IndexName(indexName); final var request = new OpenSearchRequestBuilder(querySizeLimit, exprValueFactory) - .build(name, maxResultWindow, settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)); + .build(name, maxResultWindow, + settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)); assertEquals( PhysicalPlanDSL.project( PhysicalPlanDSL.limit( diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index 955e4899bc..1a2855e5ef 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -143,21 +143,21 @@ void clean() { void searchRequest() { request.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - request.search(searchRequest ->{ - assertEquals( - new SearchRequest() + Function querySearch = searchRequest -> { + assertEquals(new SearchRequest() .indices("test") .source(new SearchSourceBuilder() .timeout(DEFAULT_QUERY_TIMEOUT) .from(0) .size(200) .query(QueryBuilders.termQuery("name", "John"))), - searchRequest); + searchRequest); return when(mock(SearchResponse.class).getHits()) .thenReturn(new SearchHits(new SearchHit[0], - new TotalHits(0, TotalHits.Relation.EQUAL_TO),0.0f)) + new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)) .getMock(); - }, searchScrollRequest -> null); + }; + request.search(querySearch, searchScrollRequest -> null); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 9b03f5dcab..ef6c879183 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.function.Function; import org.apache.commons.lang3.tuple.Pair; import org.apache.lucene.search.TotalHits; import org.apache.lucene.search.join.ScoreMode; @@ -32,7 +33,9 @@ import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; +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.index.query.InnerHitBuilder; import org.opensearch.index.query.NestedQueryBuilder; @@ -124,7 +127,7 @@ void test_push_down_query() { requestBuilder.pushDownFilter(query); var r = requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT); - r.search(searchRequest -> { + Function querySearch = searchRequest -> { assertEquals( new SearchSourceBuilder() .from(DEFAULT_OFFSET) @@ -132,10 +135,14 @@ void test_push_down_query() { .timeout(DEFAULT_QUERY_TIMEOUT) .query(query) .sort(DOC_FIELD_NAME, ASC), - searchRequest.source() + searchRequest.source() ); return mock(); - }, searchScrollRequest -> {throw new UnsupportedOperationException();}); + }; + Function scrollSearch = searchScrollRequest -> { + throw new UnsupportedOperationException(); + }; + r.search(querySearch, scrollSearch); } @@ -178,13 +185,21 @@ void test_push_down_query_and_sort() { requestBuilder); } - void assertSearchSourceBuilder(SearchSourceBuilder expected, OpenSearchRequestBuilder requestBuilder) + void assertSearchSourceBuilder(SearchSourceBuilder expected, + OpenSearchRequestBuilder requestBuilder) throws UnsupportedOperationException { - requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT).search(searchRequest -> { + Function querySearch = searchRequest -> { assertEquals(expected, searchRequest.source()); - return when(mock(SearchResponse.class).getHits()).thenReturn(new SearchHits(new SearchHit[0], - new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)).getMock(); - }, searchScrollRequest ->{ throw new UnsupportedOperationException();}); + return when(mock(SearchResponse.class).getHits()) + .thenReturn(new SearchHits(new SearchHit[0], new TotalHits(0, + TotalHits.Relation.EQUAL_TO), 0.0f)) + .getMock(); + }; + Function scrollSearch = searchScrollRequest -> { + throw new UnsupportedOperationException(); + }; + requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT).search( + querySearch, scrollSearch); } @Test diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 8ebfea3b03..78285a5275 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -66,7 +66,7 @@ class OpenSearchScrollRequestTest { private final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); private final OpenSearchScrollRequest request = new OpenSearchScrollRequest( new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), - searchSourceBuilder, factory); + searchSourceBuilder, factory); @Test void searchRequest() { @@ -77,8 +77,9 @@ void searchRequest() { .indices("test") .scroll(TimeValue.timeValueMinutes(1)) .source(new SearchSourceBuilder().query(QueryBuilders.termQuery("name", "John"))), - searchRequest); - SearchHits searchHitsMock = when(mock(SearchHits.class).getHits()).thenReturn(new SearchHit[0]).getMock(); + searchRequest); + SearchHits searchHitsMock = when(mock(SearchHits.class).getHits()) + .thenReturn(new SearchHit[0]).getMock(); return when(mock(SearchResponse.class).getHits()).thenReturn(searchHitsMock).getMock(); }, searchScrollRequest -> null); } @@ -116,8 +117,11 @@ void search() { when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); + Function scrollSearch = searchScrollRequest -> { + throw new AssertionError(); + }; OpenSearchResponse openSearchResponse = request.search(searchRequest -> searchResponse, - searchScrollRequest -> {throw new AssertionError();}); + scrollSearch); assertFalse(openSearchResponse.isEmpty()); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java index 326e5dcb9b..11694813cc 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java @@ -63,7 +63,8 @@ class OpenSearchIndexTest { public static final int QUERY_SIZE_LIMIT = 200; public static final TimeValue SCROLL_TIMEOUT = new TimeValue(1); - public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); + public static final OpenSearchRequest.IndexName INDEX_NAME + = new OpenSearchRequest.IndexName("test"); @Mock private OpenSearchClient client; @@ -203,8 +204,8 @@ void implementRelationOperatorOnly() { Integer maxResultWindow = index.getMaxResultWindow(); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE_LIMIT, exprValueFactory); assertEquals(new OpenSearchIndexScan(client, - 200, requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), - index.implement(index.optimize(plan))); + 200, requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), + index.implement(index.optimize(plan))); } @Test @@ -215,7 +216,7 @@ void implementRelationOperatorWithOptimization() { Integer maxResultWindow = index.getMaxResultWindow(); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE_LIMIT, exprValueFactory); assertEquals(new OpenSearchIndexScan(client, 200, - requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), index.implement(plan)); + requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), index.implement(plan)); } @Test @@ -257,7 +258,8 @@ void implementOtherLogicalOperators() { PhysicalPlanDSL.remove( PhysicalPlanDSL.rename( new OpenSearchIndexScan(client, - QUERY_SIZE_LIMIT, requestBuilder.build(INDEX_NAME, maxResultWindow, SCROLL_TIMEOUT)), + QUERY_SIZE_LIMIT, requestBuilder.build(INDEX_NAME, maxResultWindow, + SCROLL_TIMEOUT)), mappings), exclude), newEvalField), diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java index 1b890a286d..67f0869d6e 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanPaginationTest.java @@ -47,7 +47,6 @@ public class OpenSearchIndexScanPaginationTest { public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); public static final int MAX_RESULT_WINDOW = 3; - public static final String SCROLL_ID = "0xbadbeef"; public static final TimeValue SCROLL_TIMEOUT = TimeValue.timeValueMinutes(4); @Mock private Settings settings; @@ -71,8 +70,8 @@ void setup() { void query_empty_result() { mockResponse(client); var builder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); - try (var indexScan - = new OpenSearchIndexScan(client, MAX_RESULT_WINDOW, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE)))) { + try (var indexScan = new OpenSearchIndexScan(client, MAX_RESULT_WINDOW, + builder.build(INDEX_NAME, MAX_RESULT_WINDOW, SCROLL_TIMEOUT))) { indexScan.open(); assertFalse(indexScan.hasNext()); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index bd5740178b..00878ed0a4 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -19,12 +19,9 @@ import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import java.io.ByteArrayOutputStream; -import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; import java.util.Map; -import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -52,7 +49,6 @@ import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; -import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) @@ -60,7 +56,8 @@ class OpenSearchIndexScanTest { public static final int QUERY_SIZE = 200; - public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("employees"); + public static final OpenSearchRequest.IndexName INDEX_NAME + = new OpenSearchRequest.IndexName("employees"); public static final int MAX_RESULT_WINDOW = 10000; public static final TimeValue CURSOR_KEEP_ALIVE = TimeValue.timeValueMinutes(1); @Mock @@ -83,7 +80,7 @@ void query_empty_result() { final var name = new OpenSearchRequest.IndexName("test"); final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - QUERY_SIZE, requestBuilder.build(name, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { + QUERY_SIZE, requestBuilder.build(name, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( () -> assertFalse(indexScan.hasNext()), @@ -102,7 +99,7 @@ void query_all_results_with_query() { final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { + 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -133,7 +130,7 @@ void query_all_results_with_scroll() { final var requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { + 10, requestBuilder.build(INDEX_NAME, 10000, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -164,7 +161,7 @@ void query_some_results_with_query() { final int limit = 3; OpenSearchRequestBuilder builder = new OpenSearchRequestBuilder(0, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - limit, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { + limit, builder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -189,7 +186,7 @@ void query_some_results_with_scroll() { mockTwoPageResponse(client); final var requestuilder = new OpenSearchRequestBuilder(10, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - 3, requestuilder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { + 3, requestuilder.build(INDEX_NAME, MAX_RESULT_WINDOW, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -226,7 +223,7 @@ void query_results_limited_by_query_size() { final int defaultQuerySize = 2; final var requestBuilder = new OpenSearchRequestBuilder(defaultQuerySize, exprValueFactory); try (OpenSearchIndexScan indexScan = new OpenSearchIndexScan(client, - defaultQuerySize, requestBuilder.build(INDEX_NAME, QUERY_SIZE, CURSOR_KEEP_ALIVE))) { + defaultQuerySize, requestBuilder.build(INDEX_NAME, QUERY_SIZE, CURSOR_KEEP_ALIVE))) { indexScan.open(); assertAll( @@ -288,7 +285,7 @@ void push_down_highlight_with_arguments() { } private PushDownAssertion assertThat() { - return new PushDownAssertion(client, exprValueFactory, settings); + return new PushDownAssertion(client, exprValueFactory); } private static class PushDownAssertion { @@ -298,8 +295,7 @@ private static class PushDownAssertion { private final OpenSearchExprValueFactory factory; public PushDownAssertion(OpenSearchClient client, - OpenSearchExprValueFactory valueFactory, - Settings settings) { + OpenSearchExprValueFactory valueFactory) { this.client = client; this.requestBuilder = new OpenSearchRequestBuilder(QUERY_SIZE, valueFactory); @@ -320,17 +316,18 @@ PushDownAssertion pushDownHighlight(String query, Map arguments PushDownAssertion shouldQueryHighlight(QueryBuilder query, HighlightBuilder highlight) { var sourceBuilder = new SearchSourceBuilder() - .from(0) - .timeout(CURSOR_KEEP_ALIVE) - .query(query) - .size(QUERY_SIZE) - .highlighter(highlight) - .sort(DOC_FIELD_NAME, ASC); - OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, sourceBuilder, factory); + .from(0) + .timeout(CURSOR_KEEP_ALIVE) + .query(query) + .size(QUERY_SIZE) + .highlighter(highlight) + .sort(DOC_FIELD_NAME, ASC); + OpenSearchRequest request = + new OpenSearchQueryRequest(EMPLOYEES_INDEX, sourceBuilder, factory); when(client.search(request)).thenReturn(response); var indexScan = new OpenSearchIndexScan(client, - QUERY_SIZE, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); + QUERY_SIZE, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); indexScan.open(); return this; } @@ -345,7 +342,7 @@ PushDownAssertion shouldQuery(QueryBuilder expected) { OpenSearchRequest request = new OpenSearchQueryRequest(EMPLOYEES_INDEX, builder, factory); when(client.search(request)).thenReturn(response); var indexScan = new OpenSearchIndexScan(client, - 10000, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); + 10000, requestBuilder.build(EMPLOYEES_INDEX, 10000, CURSOR_KEEP_ALIVE)); indexScan.open(); return this; } From f4770a8a15734b9ea842a7583629fe969bb54338 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 17 May 2023 15:31:33 -0700 Subject: [PATCH 20/38] WIP Cleanup and start of refactor to remove ContinuePaginatedPlan. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/analysis/Analyzer.java | 7 +- .../sql/ast/AbstractNodeVisitor.java | 5 + .../org/opensearch/sql/ast/tree/Cursor.java | 21 +++++ .../execution/ContinuePaginatedPlan.java | 58 ------------ .../executor/execution/QueryPlanFactory.java | 12 ++- .../executor/pagination/PlanSerializer.java | 4 - .../opensearch/sql/planner/LogicalCursor.java | 21 +++++ .../logical/LogicalPlanNodeVisitor.java | 5 + .../execution/ContinuePaginatedPlanTest.java | 93 ------------------- .../execution/QueryPlanFactoryTest.java | 5 +- .../pagination/PlanSerializerTest.java | 7 +- .../executor/OpenSearchExecutionEngine.java | 3 +- .../request/OpenSearchScrollRequest.java | 3 +- .../opensearch/storage/OpenSearchIndex.java | 1 + .../storage/scan/OpenSearchIndexScan.java | 1 + .../request/OpenSearchScrollRequestTest.java | 45 ++++++++- .../storage/scan/OpenSearchIndexScanTest.java | 60 ++++++++++++ 17 files changed, 180 insertions(+), 171 deletions(-) create mode 100644 core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java delete mode 100644 core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java create mode 100644 core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java delete mode 100644 core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index aef7de69a8..dcafa73dcb 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -42,6 +42,7 @@ import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; +import org.opensearch.sql.ast.tree.Cursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; @@ -64,7 +65,6 @@ import org.opensearch.sql.common.antlr.SyntaxCheckException; import org.opensearch.sql.data.model.ExprMissingValue; import org.opensearch.sql.data.type.ExprCoreType; -import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.datasource.DataSourceService; import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.DSL; @@ -80,6 +80,7 @@ import org.opensearch.sql.expression.function.FunctionName; import org.opensearch.sql.expression.function.TableFunctionImplementation; import org.opensearch.sql.expression.parse.ParseExpression; +import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalDedupe; @@ -211,6 +212,10 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex tableFunctionImplementation.applyArguments()); } + @Override + public LogicalPlan visitCursor(Cursor cursor, AnalysisContext context) { + return new LogicalCursor(cursor.getCursor()); + } @Override public LogicalPlan visitLimit(Limit node, AnalysisContext context) { diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index 9c283d95f6..5665cdaa47 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -41,6 +41,7 @@ import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; +import org.opensearch.sql.ast.tree.Cursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; @@ -299,4 +300,8 @@ public T visitExplain(Explain node, C context) { public T visitPaginate(Paginate paginate, C context) { return visitChildren(paginate, context); } + + public T visitCursor(Cursor cursor, C context) { + return visit(cursor, context); + } } diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java b/core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java new file mode 100644 index 0000000000..eb898d7727 --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java @@ -0,0 +1,21 @@ +package org.opensearch.sql.ast.tree; + +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import org.opensearch.sql.ast.AbstractNodeVisitor; + +@RequiredArgsConstructor +public class Cursor extends UnresolvedPlan { + @Getter + final String cursor; + + @Override + public T accept(AbstractNodeVisitor nodeVisitor, C context) { + return nodeVisitor.visitCursor(this, context); + } + + @Override + public UnresolvedPlan attach(UnresolvedPlan child) { + throw new UnsupportedOperationException("Cursor unresolved plan does not support children"); + } +} diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java b/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java deleted file mode 100644 index eda65aba2d..0000000000 --- a/core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.executor.execution; - -import org.opensearch.sql.common.response.ResponseListener; -import org.opensearch.sql.executor.ExecutionEngine; -import org.opensearch.sql.executor.QueryId; -import org.opensearch.sql.executor.QueryService; -import org.opensearch.sql.executor.pagination.PlanSerializer; -import org.opensearch.sql.planner.physical.PhysicalPlan; - -/** - * ContinuePaginatedPlan represents cursor a request. - * It returns subsequent pages to the user (2nd page and all next). - */ -public class ContinuePaginatedPlan extends AbstractPlan { - - private final String cursor; - private final QueryService queryService; - private final PlanSerializer planSerializer; - - private final ResponseListener queryResponseListener; - - - /** - * Create an abstract plan that can continue paginating a given cursor. - */ - public ContinuePaginatedPlan(QueryId queryId, String cursor, QueryService queryService, - PlanSerializer planCache, - ResponseListener - queryResponseListener) { - super(queryId); - this.cursor = cursor; - this.planSerializer = planCache; - this.queryService = queryService; - this.queryResponseListener = queryResponseListener; - } - - @Override - public void execute() { - try { - PhysicalPlan plan = planSerializer.convertToPlan(cursor); - queryService.executePlan(plan, queryResponseListener); - } catch (Exception e) { - queryResponseListener.onFailure(e); - } - } - - @Override - public void explain(ResponseListener listener) { - listener.onFailure(new UnsupportedOperationException( - "Explain of a paged query continuation is not supported. " - + "Use `explain` for the initial query request.")); - } -} diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 6facfff847..7fa0950e77 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -17,11 +17,14 @@ import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; +import org.opensearch.sql.ast.tree.Cursor; +import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.exception.UnsupportedCursorRequestException; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.pagination.CanPaginateVisitor; import org.opensearch.sql.executor.pagination.PlanSerializer; /** @@ -79,11 +82,14 @@ public AbstractPlan create(String cursor, boolean isExplain, ResponseListener queryResponseListener, ResponseListener explainListener) { QueryId queryId = QueryId.queryId(); - var plan = new ContinuePaginatedPlan(queryId, cursor, queryService, - planSerializer, queryResponseListener); + var plan = new QueryPlan(queryId, new Cursor(cursor), queryService, queryResponseListener); return isExplain ? new ExplainPlan(queryId, plan, explainListener) : plan; } + boolean canConvertToCursor(UnresolvedPlan plan) { + return plan.accept(new CanPaginateVisitor(), null); + } + @Override public AbstractPlan visitQuery( Query node, @@ -94,7 +100,7 @@ public AbstractPlan visitQuery( context.getLeft().isPresent(), "[BUG] query listener must be not null"); if (node.getFetchSize() > 0) { - if (planSerializer.canConvertToCursor(node.getPlan())) { + if (canConvertToCursor(node.getPlan())) { return new QueryPlan(QueryId.queryId(), node.getPlan(), node.getFetchSize(), queryService, context.getLeft().get()); diff --git a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java index fe4d576d1a..276d0515dc 100644 --- a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java +++ b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java @@ -18,7 +18,6 @@ import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; import lombok.RequiredArgsConstructor; -import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; @@ -34,9 +33,6 @@ public class PlanSerializer { private final StorageEngine engine; - public boolean canConvertToCursor(UnresolvedPlan plan) { - return plan.accept(new CanPaginateVisitor(), null); - } /** * Converts a physical plan tree to a cursor. diff --git a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java new file mode 100644 index 0000000000..352427327d --- /dev/null +++ b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java @@ -0,0 +1,21 @@ +package org.opensearch.sql.planner; + +import java.util.List; +import lombok.Getter; +import org.opensearch.sql.planner.logical.LogicalPlan; +import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; + +public class LogicalCursor extends LogicalPlan { + @Getter + final String cursor; + + public LogicalCursor(String cursor) { + super(List.of()); + this.cursor = cursor; + } + + @Override + public R accept(LogicalPlanNodeVisitor visitor, C context) { + return visitor.visitCursor(this, context); + } +} diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java index b3d63e843f..8aec99988d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java @@ -6,6 +6,7 @@ package org.opensearch.sql.planner.logical; +import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; @@ -108,4 +109,8 @@ public R visitAD(LogicalAD plan, C context) { public R visitPaginate(LogicalPaginate plan, C context) { return visitNode(plan, context); } + + public R visitCursor(LogicalCursor plan, C context) { + return visitNode(plan, context); + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java deleted file mode 100644 index 3e08280acb..0000000000 --- a/core/src/test/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlanTest.java +++ /dev/null @@ -1,93 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.executor.execution; - -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.fail; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.CALLS_REAL_METHODS; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static org.mockito.Mockito.withSettings; - -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.DisplayNameGeneration; -import org.junit.jupiter.api.DisplayNameGenerator; -import org.junit.jupiter.api.Test; -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; -import org.opensearch.sql.executor.pagination.PlanSerializer; -import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.storage.StorageEngine; - -@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class ContinuePaginatedPlanTest { - - private static PlanSerializer planSerializer; - - private static QueryService queryService; - - /** - * Initialize the mocks. - */ - @BeforeAll - public static void setUp() { - var storageEngine = mock(StorageEngine.class); - planSerializer = new PlanSerializer(storageEngine); - queryService = new QueryService(null, new DefaultExecutionEngine(), null); - } - - @Test - public void can_execute_plan() { - var planSerializer = mock(PlanSerializer.class); - when(planSerializer.convertToPlan(anyString())).thenReturn(mock(PhysicalPlan.class)); - var listener = new ResponseListener() { - @Override - public void onResponse(ExecutionEngine.QueryResponse response) { - assertNotNull(response); - } - - @Override - public void onFailure(Exception e) { - fail(e); - } - }; - var plan = new ContinuePaginatedPlan(QueryId.queryId(), "", - queryService, planSerializer, listener); - plan.execute(); - } - - @Test - public void can_handle_error_while_executing_plan() { - var listener = new ResponseListener() { - @Override - public void onResponse(ExecutionEngine.QueryResponse response) { - fail(); - } - - @Override - public void onFailure(Exception e) { - assertNotNull(e); - } - }; - var plan = new ContinuePaginatedPlan(QueryId.queryId(), "", queryService, - planSerializer, listener); - plan.execute(); - } - - @Test - public void explain_is_not_supported() { - var listener = mock(ResponseListener.class); - mock(ContinuePaginatedPlan.class, withSettings().defaultAnswer(CALLS_REAL_METHODS)) - .explain(listener); - verify(listener).onFailure(any(UnsupportedOperationException.class)); - } -} diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index 8daf1edd43..50c4b2e761 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -12,7 +12,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.Mockito.when; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; import java.util.Optional; @@ -81,7 +80,7 @@ public void createFromCursorShouldSuccess() { AbstractPlan explainExecution = factory.create("", true, queryListener, explainListener); assertAll( - () -> assertTrue(queryExecution instanceof ContinuePaginatedPlan), + () -> assertTrue(queryExecution instanceof QueryPlan), () -> assertTrue(explainExecution instanceof ExplainPlan) ); } @@ -125,7 +124,6 @@ public void noConsumerResponseChannel() { @Test public void createQueryWithFetchSizeWhichCanBePaged() { - when(planSerializer.canConvertToCursor(plan)).thenReturn(true); factory = new QueryPlanFactory(queryService, planSerializer); Statement query = new Query(plan, 10); AbstractPlan queryExecution = @@ -135,7 +133,6 @@ public void createQueryWithFetchSizeWhichCanBePaged() { @Test public void createQueryWithFetchSizeWhichCannotBePaged() { - when(planSerializer.canConvertToCursor(plan)).thenReturn(false); factory = new QueryPlanFactory(queryService, planSerializer); Statement query = new Query(plan, 10); assertThrows(UnsupportedCursorRequestException.class, diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java index 622f10b3ba..4eb4e28446 100644 --- a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java @@ -22,7 +22,6 @@ import java.io.Serializable; import java.util.List; import lombok.SneakyThrows; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -52,19 +51,19 @@ void setUp() { @Test void canConvertToCursor_relation() { - assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); +// assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); } @Test void canConvertToCursor_project_allFields_relation() { var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.allFields()); - assertTrue(planCache.canConvertToCursor(unresolvedPlan)); +// assertTrue(planCache.canConvertToCursor(unresolvedPlan)); } @Test void canConvertToCursor_project_some_fields_relation() { var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.field("rando")); - Assertions.assertFalse(planCache.canConvertToCursor(unresolvedPlan)); +// Assertions.assertFalse(planCache.canConvertToCursor(unresolvedPlan)); } @ParameterizedTest diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index bfc29b02d2..b2eee3850f 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -9,6 +9,7 @@ import com.google.common.collect.ImmutableMap; import java.util.ArrayList; import java.util.List; +import java.util.Map; import lombok.RequiredArgsConstructor; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.data.model.ExprValue; @@ -70,7 +71,7 @@ public void explain(PhysicalPlan plan, ResponseListener listene @Override public ExplainResponseNode visitTableScan(TableScanOperator node, Object context) { return explain(node, context, explainNode -> { - explainNode.setDescription(ImmutableMap.of("request", node.explain())); + explainNode.setDescription(Map.of("request", node.explain())); }); } }; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index c9ac24141e..9aa8225a4e 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -162,7 +162,8 @@ public void writeTo(StreamOutput out) throws IOException { if (!needClean) { // If needClean is true, there is no more data to get from OpenSearch and scrollId is // used only to clean up OpenSearch context. - out.writeString(scrollId); + + out.writeString(scrollId == null? "" : scrollId); } out.writeBoolean(needClean); out.writeStringCollection(includes); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 532d62333d..4d34fcdaff 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -27,6 +27,7 @@ import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder; import org.opensearch.sql.planner.DefaultImplementor; +import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalML; import org.opensearch.sql.planner.logical.LogicalMLCommons; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 049d67b92b..050f0ae71b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -126,6 +126,7 @@ public void readExternal(ObjectInput in) throws IOException { byte[] requestStream = new byte[reqSize]; in.read(requestStream); + var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) .resolveObject("engine"); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 78285a5275..5e31b505b3 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -8,6 +8,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -17,8 +18,10 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; +import lombok.SneakyThrows; import org.apache.lucene.search.TotalHits; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -29,19 +32,24 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.common.io.stream.BytesStreamInput; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.unit.TimeValue; import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.search.fetch.subphase.FetchSourceContext; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchScrollRequestTest { + public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); + public static final TimeValue SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1); @Mock private Function searchAction; @@ -65,9 +73,24 @@ class OpenSearchScrollRequestTest { private final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); private final OpenSearchScrollRequest request = new OpenSearchScrollRequest( - new OpenSearchRequest.IndexName("test"), TimeValue.timeValueMinutes(1), + INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, factory); + @Test + void constructor() { + searchSourceBuilder.fetchSource(new String[] {"test"}, null); + var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, factory); + assertNotEquals(List.of(), request.getIncludes()); + } + + @Test + void constructor2() { + searchSourceBuilder.fetchSource(new String[]{"test"}, null); + var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, + factory); + assertNotEquals(List.of(), request.getIncludes()); + } + @Test void searchRequest() { searchSourceBuilder.query(QueryBuilders.termQuery("name", "John")); @@ -226,4 +249,22 @@ void no_clean_if_no_scroll_in_response() { request.clean((s) -> fail()); } + + @Test + @SneakyThrows + void serialize_deserialize_no_needClean() { + var stream = new BytesStreamOutput(); + request.writeTo(stream); + stream.flush(); + assertTrue(stream.size() > 0); + + // deserialize + var inStream = new BytesStreamInput(stream.bytes().toBytesRef().bytes); + var indexMock = mock(OpenSearchIndex.class); + var engine = mock(OpenSearchStorageEngine.class); + when(engine.getTable(any(), any())).thenReturn(indexMock); + var newRequest = new OpenSearchScrollRequest(inStream, engine); + assertEquals(request.getInitialSearchRequest(), newRequest.getInitialSearchRequest()); + assertEquals("", newRequest.getScrollId()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 00878ed0a4..d4de3a892f 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -9,6 +9,7 @@ import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.lenient; @@ -19,9 +20,13 @@ import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; import java.util.Map; +import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -43,12 +48,14 @@ import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; +import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @ExtendWith(MockitoExtension.class) @@ -74,6 +81,59 @@ class OpenSearchIndexScanTest { void setup() { } + @Test + void explain() { + var request = mock(OpenSearchRequest.class); + when(request.toString()).thenReturn("explain works!"); + try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request)) { + assertEquals("explain works!", indexScan.explain()); + } + } + + @Test + @SneakyThrows + void throws_no_cursor_exception() { + var request = mock(OpenSearchRequest.class); + when(request.hasAnotherBatch()).thenReturn(false); + try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request); + var byteStream = new ByteArrayOutputStream(); + var objectStream = new ObjectOutputStream(byteStream)) { + assertThrows(NoCursorException.class, () -> objectStream.writeObject(indexScan)); + } + } + + @Test + @SneakyThrows + void serialize() { + var searchSourceBuilder = new SearchSourceBuilder().size(4); + + var factory = mock(OpenSearchExprValueFactory.class); + + var request = new OpenSearchScrollRequest( + INDEX_NAME, CURSOR_KEEP_ALIVE,searchSourceBuilder, factory); + request.setScrollId("valid-id"); + try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request); + var byteStream = new ByteArrayOutputStream(); + var objectStream = new ObjectOutputStream(byteStream)) { + try { + indexScan.writeExternal(objectStream); + objectStream.flush(); + byteStream.flush(); + assertTrue(byteStream.size() > 0); + } catch (IOException e) { + throw new AssertionError("Expected to serialize the indexScan"); + } + } + } + + @Test + void plan_for_serialization() { + var request = mock(OpenSearchRequest.class); + try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request)) { + assertEquals(indexScan, indexScan.getPlanForSerialization()); + } + } + @Test void query_empty_result() { mockResponse(client); From e350662279263ed491031265a0adda2c45041a79 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Tue, 23 May 2023 11:24:24 -0700 Subject: [PATCH 21/38] Updating imports to reflect changes in opensearch core. (#1645) Signed-off-by: MaxKsyunz (cherry picked from commit 8e5d766b0fbea11b449eeada9fb1ccab375cec1f) --- .../test/java/org/opensearch/sql/legacy/RestIntegTestCase.java | 3 ++- .../main/java/org/opensearch/sql/legacy/domain/Paramer.java | 2 +- .../sql/legacy/executor/format/PrettyFormatRestExecutor.java | 2 +- .../opensearch/sql/legacy/executor/format/SelectResultSet.java | 2 +- .../java/org/opensearch/sql/legacy/query/maker/AggMaker.java | 2 +- 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java index 166e21ec82..555d0af0d2 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java @@ -6,6 +6,7 @@ package org.opensearch.sql.legacy; +import static org.opensearch.core.common.Strings.isNullOrEmpty; import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient; import static org.opensearch.sql.legacy.TestUtils.getAccountIndexMapping; import static org.opensearch.sql.legacy.TestUtils.getBankIndexMapping; @@ -98,7 +99,7 @@ public static void dumpCoverage() { // jacoco.dir is set in sqlplugin-coverage.gradle, if it doesn't exist we don't // want to collect coverage so we can return early String jacocoBuildPath = System.getProperty("jacoco.dir"); - if (Strings.isNullOrEmpty(jacocoBuildPath)) { + if (isNullOrEmpty(jacocoBuildPath)) { return; } diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/domain/Paramer.java b/legacy/src/main/java/org/opensearch/sql/legacy/domain/Paramer.java index 5c2c84f9b8..6cdf0148a8 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/domain/Paramer.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/domain/Paramer.java @@ -14,7 +14,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import org.opensearch.common.Strings; +import org.opensearch.core.common.Strings; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.index.query.MatchPhraseQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutor.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutor.java index ceccd7ca4a..cf034a3d0b 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutor.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/PrettyFormatRestExecutor.java @@ -12,7 +12,7 @@ import org.opensearch.OpenSearchException; import org.opensearch.action.search.SearchResponse; import org.opensearch.client.Client; -import org.opensearch.common.Strings; +import org.opensearch.core.common.Strings; import org.opensearch.rest.BytesRestResponse; import org.opensearch.rest.RestChannel; import org.opensearch.rest.RestStatus; diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java index 8e9d48ce06..a6f4cf815a 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/executor/format/SelectResultSet.java @@ -31,8 +31,8 @@ import org.opensearch.action.admin.indices.mapping.get.GetFieldMappingsResponse; import org.opensearch.action.search.ClearScrollResponse; import org.opensearch.client.Client; -import org.opensearch.common.Strings; import org.opensearch.common.document.DocumentField; +import org.opensearch.core.common.Strings; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.search.aggregations.Aggregation; diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/query/maker/AggMaker.java b/legacy/src/main/java/org/opensearch/sql/legacy/query/maker/AggMaker.java index 7960d80669..c3cb6c8016 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/query/maker/AggMaker.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/query/maker/AggMaker.java @@ -19,10 +19,10 @@ import java.util.stream.Collectors; import org.apache.commons.lang3.StringUtils; import org.opensearch.common.ParsingException; -import org.opensearch.common.Strings; import org.opensearch.common.xcontent.LoggingDeprecationHandler; import org.opensearch.common.xcontent.json.JsonXContent; import org.opensearch.common.xcontent.json.JsonXContentParser; +import org.opensearch.core.common.Strings; import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoHashGridAggregationBuilder; From 7e812e550584310388f7610e085d3865cc3cda2d Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 24 May 2023 13:41:48 -0700 Subject: [PATCH 22/38] Integrating with main. Signed-off-by: MaxKsyunz --- .../request/OpenSearchQueryRequest.java | 2 -- .../opensearch/request/OpenSearchRequest.java | 6 ++++ .../request/OpenSearchRequestBuilder.java | 7 +---- .../request/OpenSearchQueryRequestTest.java | 28 +++++++++++-------- 4 files changed, 23 insertions(+), 20 deletions(-) diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index b20da41b84..a5efac3422 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -6,8 +6,6 @@ package org.opensearch.sql.opensearch.request; -import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; - import java.io.IOException; import java.util.Arrays; import java.util.List; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index 184f729bfb..8a5eb6ec4c 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -16,6 +16,7 @@ import org.opensearch.common.io.stream.StreamInput; import org.opensearch.common.io.stream.StreamOutput; import org.opensearch.common.io.stream.Writeable; +import org.opensearch.common.unit.TimeValue; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; import org.opensearch.sql.opensearch.response.OpenSearchResponse; @@ -23,6 +24,11 @@ * OpenSearch search request. */ public interface OpenSearchRequest extends Writeable { + /** + * Default query timeout in minutes. + */ + TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); + /** * Apply the search action or scroll action on request based on context. * diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 72ef85d8c5..276b3b0220 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -51,11 +51,6 @@ @ToString public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { - /** - * Default query timeout in minutes. - */ - public static final TimeValue DEFAULT_QUERY_TIMEOUT = TimeValue.timeValueMinutes(1L); - /** * Search request source builder. */ @@ -87,7 +82,7 @@ public OpenSearchRequestBuilder(int requestedTotalSize, this.requestedTotalSize = requestedTotalSize; this.sourceBuilder = new SearchSourceBuilder() .from(startFrom) - .timeout(DEFAULT_QUERY_TIMEOUT) + .timeout(OpenSearchRequest.DEFAULT_QUERY_TIMEOUT) .trackScores(false); this.exprValueFactory = exprValueFactory; } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index c79e0e1576..5162aecfa2 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -15,7 +15,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder.DEFAULT_QUERY_TIMEOUT; +import static org.opensearch.sql.opensearch.request.OpenSearchRequest.DEFAULT_QUERY_TIMEOUT; import java.util.function.Consumer; import java.util.function.Function; @@ -146,21 +146,14 @@ void clean() { void searchRequest() { request.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - Function querySearch = searchRequest -> { - assertEquals(new SearchRequest() + assertSearchRequest(new SearchRequest() .indices("test") .source(new SearchSourceBuilder() .timeout(DEFAULT_QUERY_TIMEOUT) .from(0) .size(200) .query(QueryBuilders.termQuery("name", "John"))), - searchRequest); - return when(mock(SearchResponse.class).getHits()) - .thenReturn(new SearchHits(new SearchHit[0], - new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)) - .getMock(); - }; - request.search(querySearch, searchScrollRequest -> null); + request); } @@ -168,7 +161,7 @@ void searchRequest() { void searchCrossClusterRequest() { remoteRequest.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - assertEquals( + assertSearchRequest( new SearchRequest() .indices("ccs:test") .source(new SearchSourceBuilder() @@ -176,6 +169,17 @@ void searchCrossClusterRequest() { .from(0) .size(200) .query(QueryBuilders.termQuery("name", "John"))), - remoteRequest.searchRequest()); + remoteRequest); + } + + private void assertSearchRequest(SearchRequest expected, OpenSearchQueryRequest request) { + Function querySearch = searchRequest -> { + assertEquals(expected, searchRequest); + return when(mock(SearchResponse.class).getHits()) + .thenReturn(new SearchHits(new SearchHit[0], + new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)) + .getMock(); + }; + request.search(querySearch, searchScrollRequest -> null); } } From f1e1342ad3f66530b007b36a7cc5c275699bacb5 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Wed, 24 May 2023 23:28:50 -0700 Subject: [PATCH 23/38] Address refactoring comments WIP Signed-off-by: MaxKsyunz --- .../java/org/opensearch/sql/analysis/Analyzer.java | 4 +++- .../sql/executor/execution/QueryPlanFactory.java | 2 -- .../opensearch/sql/planner/DefaultImplementor.java | 7 +++++++ .../org/opensearch/sql/planner/LogicalCursor.java | 13 +++++++++++-- .../executor/execution/QueryPlanFactoryTest.java | 13 ++++++++----- docs/dev/Pagination-v2.md | 3 ++- .../java/org/opensearch/sql/ppl/StandaloneIT.java | 2 +- .../org/opensearch/sql/util/StandaloneModule.java | 6 ++---- .../sql/plugin/config/OpenSearchPluginModule.java | 2 +- .../java/org/opensearch/sql/ppl/PPLServiceTest.java | 6 +----- .../java/org/opensearch/sql/sql/SQLServiceTest.java | 5 +---- 11 files changed, 37 insertions(+), 26 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index dcafa73dcb..f05c66aaa8 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -6,6 +6,7 @@ package org.opensearch.sql.analysis; +import static org.opensearch.sql.analysis.DataSourceSchemaIdentifierNameResolver.DEFAULT_DATASOURCE_NAME; import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_FIRST; import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_LAST; import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC; @@ -214,7 +215,8 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex @Override public LogicalPlan visitCursor(Cursor cursor, AnalysisContext context) { - return new LogicalCursor(cursor.getCursor()); + return new LogicalCursor(cursor.getCursor(), + dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); } @Override diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 7fa0950e77..582e19888c 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -25,7 +25,6 @@ import org.opensearch.sql.executor.QueryId; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.pagination.CanPaginateVisitor; -import org.opensearch.sql.executor.pagination.PlanSerializer; /** * QueryExecution Factory. @@ -42,7 +41,6 @@ public class QueryPlanFactory * Query Service. */ private final QueryService queryService; - private final PlanSerializer planSerializer; /** * NO_CONSUMER_RESPONSE_LISTENER should never be called. It is only used as constructor diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index 9bde4ab647..5a10ddd913 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -6,6 +6,7 @@ package org.opensearch.sql.planner; +import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalDedupe; import org.opensearch.sql.planner.logical.LogicalEval; @@ -148,6 +149,12 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { + "implementing and optimizing logical plan with relation involved"); } + + @Override + public PhysicalPlan visitCursor(LogicalCursor plan, C context) { + return new PlanSerializer(plan.getEngine()).convertToPlan(plan.getCursor()); + } + protected PhysicalPlan visitChild(LogicalPlan node, C context) { // Logical operators visited here must have a single child return node.getChild().get(0).accept(this, context); diff --git a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java index 352427327d..703f4bf0fb 100644 --- a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java +++ b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java @@ -1,17 +1,26 @@ package org.opensearch.sql.planner; import java.util.List; +import lombok.EqualsAndHashCode; import lombok.Getter; +import lombok.ToString; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor; +import org.opensearch.sql.storage.StorageEngine; +@EqualsAndHashCode(callSuper = false) +@ToString public class LogicalCursor extends LogicalPlan { @Getter - final String cursor; + private final String cursor; - public LogicalCursor(String cursor) { + @Getter + private final StorageEngine engine; + + public LogicalCursor(String cursor, StorageEngine engine) { super(List.of()); this.cursor = cursor; + this.engine = engine; } @Override diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index 50c4b2e761..e2e1f3684d 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -12,6 +12,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; import static org.opensearch.sql.executor.execution.QueryPlanFactory.NO_CONSUMER_RESPONSE_LISTENER; import java.util.Optional; @@ -28,6 +30,7 @@ import org.opensearch.sql.exception.UnsupportedCursorRequestException; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryService; +import org.opensearch.sql.executor.pagination.CanPaginateVisitor; import org.opensearch.sql.executor.pagination.PlanSerializer; @ExtendWith(MockitoExtension.class) @@ -48,13 +51,11 @@ class QueryPlanFactoryTest { @Mock private ExecutionEngine.QueryResponse queryResponse; - @Mock - private PlanSerializer planSerializer; private QueryPlanFactory factory; @BeforeEach void init() { - factory = new QueryPlanFactory(queryService, planSerializer); + factory = new QueryPlanFactory(queryService); } @Test @@ -124,7 +125,8 @@ public void noConsumerResponseChannel() { @Test public void createQueryWithFetchSizeWhichCanBePaged() { - factory = new QueryPlanFactory(queryService, planSerializer); + when(plan.accept(any(CanPaginateVisitor.class), any())).thenReturn(Boolean.TRUE); + factory = new QueryPlanFactory(queryService); Statement query = new Query(plan, 10); AbstractPlan queryExecution = factory.create(query, Optional.of(queryListener), Optional.empty()); @@ -133,7 +135,8 @@ public void createQueryWithFetchSizeWhichCanBePaged() { @Test public void createQueryWithFetchSizeWhichCannotBePaged() { - factory = new QueryPlanFactory(queryService, planSerializer); + when(plan.accept(any(CanPaginateVisitor.class), any())).thenReturn(Boolean.FALSE); + factory = new QueryPlanFactory(queryService); Statement query = new Query(plan, 10); assertThrows(UnsupportedCursorRequestException.class, () -> factory.create(query, diff --git a/docs/dev/Pagination-v2.md b/docs/dev/Pagination-v2.md index 2416ae51de..60f8e1e708 100644 --- a/docs/dev/Pagination-v2.md +++ b/docs/dev/Pagination-v2.md @@ -504,6 +504,8 @@ sequenceDiagram participant QueryPlanFactory participant QueryService participant OpenSearchExecutionEngine + participant LogicalCursor + participant DefaultImplementor participant PlanSerializer SQLService ->>+ QueryPlanFactory : execute @@ -614,7 +616,6 @@ sequenceDiagram participant ResourceMonitorPlan participant OpenSearchIndexScan participant OpenSearchScrollRequest - participant OpenSearchScrollRequest PlanSerializer ->>+ ProjectOperator : getPlanForSerialization ProjectOperator -->>- PlanSerializer : this diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java index 595fd8acd5..5a4e38b6b5 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java @@ -241,7 +241,7 @@ public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); Planner planner = new Planner(LogicalPlanOptimizer.create()); QueryService queryService = new QueryService(analyzer, executionEngine, planner); - return new QueryPlanFactory(queryService, planSerializer); + return new QueryPlanFactory(queryService); } } diff --git a/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java index a86f251377..e38f408514 100644 --- a/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java +++ b/integ-test/src/test/java/org/opensearch/sql/util/StandaloneModule.java @@ -104,11 +104,9 @@ public PlanSerializer paginatedPlanCache(StorageEngine storageEngine) { } @Provides - public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, - PlanSerializer planSerializer, - QueryService qs) { + public QueryPlanFactory queryPlanFactory(QueryService qs) { - return new QueryPlanFactory(qs, planSerializer); + return new QueryPlanFactory(qs); } @Provides diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index b80cb3faab..7415c78a8c 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -108,6 +108,6 @@ public QueryPlanFactory queryPlanFactory(DataSourceService dataSourceService, Planner planner = new Planner(LogicalPlanOptimizer.create()); QueryService queryService = new QueryService( analyzer, executionEngine, planner); - return new QueryPlanFactory(queryService, planSerializer); + return new QueryPlanFactory(queryService); } } diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java index 117aca50bf..74e5b0f82e 100644 --- a/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java +++ b/ppl/src/test/java/org/opensearch/sql/ppl/PPLServiceTest.java @@ -27,7 +27,6 @@ import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; import org.opensearch.sql.executor.pagination.Cursor; -import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.ppl.antlr.PPLSyntaxParser; import org.opensearch.sql.ppl.domain.PPLQueryRequest; @@ -48,9 +47,6 @@ public class PPLServiceTest { @Mock private ExecutionEngine.Schema schema; - @Mock - private PlanSerializer planSerializer; - /** * Setup the test context. */ @@ -59,7 +55,7 @@ public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); pplService = new PPLService(new PPLSyntaxParser(), queryManager, - new QueryPlanFactory(queryService, planSerializer)); + new QueryPlanFactory(queryService)); } @After diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java index 39c27c5e06..ca4939d507 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -49,14 +49,11 @@ class SQLServiceTest { @Mock private QueryService queryService; - @Mock - private PlanSerializer planSerializer; - @BeforeEach public void setUp() { queryManager = DefaultQueryManager.defaultQueryManager(); sqlService = new SQLService(new SQLSyntaxParser(), queryManager, - new QueryPlanFactory(queryService, planSerializer)); + new QueryPlanFactory(queryService)); } @AfterEach From 73f18df8c79653790381e4203652c74d8e274826 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 25 May 2023 00:56:53 -0700 Subject: [PATCH 24/38] Restore error to explain requests containing only a cursor. Signed-off-by: MaxKsyunz --- sql/src/main/java/org/opensearch/sql/sql/SQLService.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index a3ed89e9d7..5f9701b187 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -65,10 +65,15 @@ private AbstractPlan plan( SQLQueryRequest request, Optional> queryListener, Optional> explainListener) { + boolean isExplainRequest = request.isExplainRequest(); if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. + if (isExplainRequest) { + throw new UnsupportedOperationException("Explain of a paged query continuation " + + "is not supported. Use `explain` for the initial query request."); + } return queryExecutionFactory.create(request.getCursor().get(), - request.isExplainRequest(), queryListener.orElse(null), explainListener.orElse(null)); + isExplainRequest, queryListener.orElse(null), explainListener.orElse(null)); } else { // 1.Parse query and convert parse tree (CST) to abstract syntax tree (AST) ParseTree cst = parser.parse(request.getQuery()); @@ -77,7 +82,7 @@ private AbstractPlan plan( new AstStatementBuilder( new AstBuilder(request.getQuery()), AstStatementBuilder.StatementBuilderContext.builder() - .isExplain(request.isExplainRequest()) + .isExplain(isExplainRequest) .fetchSize(request.getFetchSize()) .build())); From 60226beb366999842e24011e4d332938700304db Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 25 May 2023 22:13:36 -0700 Subject: [PATCH 25/38] Complete test coverage. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/analysis/Analyzer.java | 4 +- .../sql/ast/AbstractNodeVisitor.java | 4 +- .../tree/{Cursor.java => FetchCursor.java} | 4 +- .../opensearch/sql/executor/QueryService.java | 8 -- .../executor/execution/QueryPlanFactory.java | 4 +- .../sql/planner/DefaultImplementor.java | 1 - .../sql/planner/logical/LogicalPlanDSL.java | 8 +- .../opensearch/sql/analysis/AnalyzerTest.java | 6 ++ .../pagination/PlanSerializerTest.java | 82 +------------------ .../sql/planner/DefaultImplementorTest.java | 14 +++- .../logical/LogicalPlanNodeVisitorTest.java | 6 +- .../planner/physical/ProjectOperatorTest.java | 3 +- .../opensearch/sql/utils/TestOperator.java | 69 ++++++++++++++++ .../sql/sql/StandalonePaginationIT.java | 3 +- .../sql/legacy/plugin/RestSQLQueryAction.java | 4 +- .../request/OpenSearchQueryRequest.java | 5 ++ .../opensearch/request/OpenSearchRequest.java | 4 +- .../request/OpenSearchScrollRequest.java | 27 +++--- .../request/OpenSearchQueryRequestTest.java | 9 +- .../request/OpenSearchRequestBuilderTest.java | 12 +++ .../request/OpenSearchScrollRequestTest.java | 67 ++++++++++++++- .../storage/scan/OpenSearchIndexScanTest.java | 33 ++++---- .../opensearch/sql/sql/SQLServiceTest.java | 5 +- 23 files changed, 246 insertions(+), 136 deletions(-) rename core/src/main/java/org/opensearch/sql/ast/tree/{Cursor.java => FetchCursor.java} (81%) create mode 100644 core/src/test/java/org/opensearch/sql/utils/TestOperator.java diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index f05c66aaa8..92317f0be4 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -43,7 +43,7 @@ import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.Cursor; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; @@ -214,7 +214,7 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex } @Override - public LogicalPlan visitCursor(Cursor cursor, AnalysisContext context) { + public LogicalPlan visitCursor(FetchCursor cursor, AnalysisContext context) { return new LogicalCursor(cursor.getCursor(), dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); } diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index 5665cdaa47..db9340f4cc 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -41,7 +41,7 @@ import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.Cursor; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; import org.opensearch.sql.ast.tree.Filter; @@ -301,7 +301,7 @@ public T visitPaginate(Paginate paginate, C context) { return visitChildren(paginate, context); } - public T visitCursor(Cursor cursor, C context) { + public T visitCursor(FetchCursor cursor, C context) { return visit(cursor, context); } } diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java similarity index 81% rename from core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java rename to core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java index eb898d7727..64655d1100 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/Cursor.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java @@ -1,11 +1,13 @@ package org.opensearch.sql.ast.tree; +import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.RequiredArgsConstructor; import org.opensearch.sql.ast.AbstractNodeVisitor; @RequiredArgsConstructor -public class Cursor extends UnresolvedPlan { +@EqualsAndHashCode(callSuper = false) +public class FetchCursor extends UnresolvedPlan { @Getter final String cursor; diff --git a/core/src/main/java/org/opensearch/sql/executor/QueryService.java b/core/src/main/java/org/opensearch/sql/executor/QueryService.java index a4cd1982cd..94e7081920 100644 --- a/core/src/main/java/org/opensearch/sql/executor/QueryService.java +++ b/core/src/main/java/org/opensearch/sql/executor/QueryService.java @@ -46,14 +46,6 @@ public void execute(UnresolvedPlan plan, } } - /** - * Execute a physical plan without analyzing or planning anything. - */ - public void executePlan(PhysicalPlan plan, - ResponseListener listener) { - executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener); - } - /** * Execute the {@link UnresolvedPlan}, with {@link PlanContext} and using {@link ResponseListener} * to get response. diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 582e19888c..2231d5a7c1 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -17,7 +17,7 @@ import org.opensearch.sql.ast.statement.Explain; import org.opensearch.sql.ast.statement.Query; import org.opensearch.sql.ast.statement.Statement; -import org.opensearch.sql.ast.tree.Cursor; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.UnresolvedPlan; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.exception.UnsupportedCursorRequestException; @@ -80,7 +80,7 @@ public AbstractPlan create(String cursor, boolean isExplain, ResponseListener queryResponseListener, ResponseListener explainListener) { QueryId queryId = QueryId.queryId(); - var plan = new QueryPlan(queryId, new Cursor(cursor), queryService, queryResponseListener); + var plan = new QueryPlan(queryId, new FetchCursor(cursor), queryService, queryResponseListener); return isExplain ? new ExplainPlan(queryId, plan, explainListener) : plan; } diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index 5a10ddd913..d82ea138b5 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -149,7 +149,6 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { + "implementing and optimizing logical plan with relation involved"); } - @Override public PhysicalPlan visitCursor(LogicalCursor plan, C context) { return new PlanSerializer(plan.getEngine()).convertToPlan(plan.getCursor()); diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java index e95e47a013..937f346078 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java @@ -11,19 +11,19 @@ import java.util.Arrays; import java.util.List; import java.util.Map; -import java.util.stream.Collectors; import lombok.experimental.UtilityClass; import org.apache.commons.lang3.tuple.Pair; import org.opensearch.sql.ast.expression.Literal; import org.opensearch.sql.ast.tree.RareTopN.CommandType; import org.opensearch.sql.ast.tree.Sort.SortOption; -import org.opensearch.sql.data.model.ExprCollectionValue; import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.LiteralExpression; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.window.WindowDefinition; +import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; /** @@ -32,6 +32,10 @@ @UtilityClass public class LogicalPlanDSL { + public static LogicalPlan cursor(String cursor, StorageEngine engine) { + return new LogicalCursor(cursor, engine); + } + public static LogicalPlan write(LogicalPlan input, Table table, List columns) { return new LogicalWrite(input, table, columns); } diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index 20927f262c..b9ede4eea7 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -75,6 +75,7 @@ import org.opensearch.sql.ast.expression.ScoreFunction; import org.opensearch.sql.ast.expression.SpanUnit; import org.opensearch.sql.ast.tree.AD; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Kmeans; import org.opensearch.sql.ast.tree.ML; import org.opensearch.sql.ast.tree.Paginate; @@ -1641,4 +1642,9 @@ public void visit_paginate() { assertTrue(actual instanceof LogicalPaginate); assertEquals(10, ((LogicalPaginate) actual).getPageSize()); } + + @Test + void visit_cursor() { + LogicalPlan actual = analyze((new FetchCursor("test"))); + } } diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java index 4eb4e28446..ce3c35c1c0 100644 --- a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java @@ -15,12 +15,8 @@ import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; import java.io.ObjectOutputStream; import java.io.Serializable; -import java.util.List; import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayNameGeneration; @@ -29,12 +25,11 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import org.opensearch.sql.ast.dsl.AstDSL; -import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; -import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; import org.opensearch.sql.storage.StorageEngine; +import org.opensearch.sql.utils.TestOperator; @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) public class PlanSerializerTest { @@ -49,23 +44,6 @@ void setUp() { planCache = new PlanSerializer(storageEngine); } - @Test - void canConvertToCursor_relation() { -// assertTrue(planCache.canConvertToCursor(AstDSL.relation("Table"))); - } - - @Test - void canConvertToCursor_project_allFields_relation() { - var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.allFields()); -// assertTrue(planCache.canConvertToCursor(unresolvedPlan)); - } - - @Test - void canConvertToCursor_project_some_fields_relation() { - var unresolvedPlan = AstDSL.project(AstDSL.relation("table"), AstDSL.field("rando")); -// Assertions.assertFalse(planCache.canConvertToCursor(unresolvedPlan)); - } - @ParameterizedTest @ValueSource(strings = {"pewpew", "asdkfhashdfjkgakgfwuigfaijkb", "ajdhfgajklghadfjkhgjkadhgad" + "kadfhgadhjgfjklahdgqheygvskjfbvgsdklgfuirehiluANUIfgauighbahfuasdlhfnhaughsdlfhaughaggf" @@ -111,7 +89,7 @@ void serialize_deserialize_obj() { void serialize_throws() { assertThrows(Throwable.class, () -> serialize(new NotSerializableTestClass())); var testObj = new TestOperator(); - testObj.throwIoOnWrite = true; + testObj.setThrowIoOnWrite(true); assertThrows(Throwable.class, () -> serialize(testObj)); } @@ -129,7 +107,7 @@ void deserialize_throws() { @SneakyThrows void convertToCursor_returns_no_cursor_if_cant_serialize() { var plan = new TestOperator(42); - plan.throwNoCursorOnWrite = true; + plan.setThrowNoCursorOnWrite(true); assertAll( () -> assertThrows(NoCursorException.class, () -> serialize(plan)), () -> assertEquals(Cursor.None, planCache.convertToCursor(plan)) @@ -190,60 +168,6 @@ void resolveObject() { // Helpers and auxiliary classes section below - public static class TestOperator extends PhysicalPlan implements ExternalizablePlan { - private int field; - private boolean throwNoCursorOnWrite = false; - private boolean throwIoOnWrite = false; - - public TestOperator() { - } - - public TestOperator(int value) { - field = value; - } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - field = in.readInt(); - } - - @Override - public void writeExternal(ObjectOutput out) throws IOException { - if (throwNoCursorOnWrite) { - throw new NoCursorException(); - } - if (throwIoOnWrite) { - throw new IOException(); - } - out.writeInt(field); - } - - @Override - public boolean equals(Object o) { - return field == ((TestOperator) o).field; - } - - @Override - public R accept(PhysicalPlanNodeVisitor visitor, C context) { - return null; - } - - @Override - public boolean hasNext() { - return false; - } - - @Override - public ExprValue next() { - return null; - } - - @Override - public List getChild() { - return null; - } - } - @SneakyThrows private String serialize(Serializable input) { return new PlanSerializer(null).serialize(input); diff --git a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java index bf1464f5f6..c88fd20d40 100644 --- a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java @@ -46,6 +46,7 @@ import org.opensearch.sql.ast.tree.Sort; import org.opensearch.sql.data.model.ExprBooleanValue; import org.opensearch.sql.data.type.ExprCoreType; +import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.Expression; import org.opensearch.sql.expression.NamedExpression; @@ -54,17 +55,18 @@ import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.window.WindowDefinition; import org.opensearch.sql.expression.window.ranking.RowNumberFunction; -import org.opensearch.sql.planner.logical.LogicalPaginate; import org.opensearch.sql.planner.logical.LogicalPlan; import org.opensearch.sql.planner.logical.LogicalPlanDSL; import org.opensearch.sql.planner.logical.LogicalRelation; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanDSL; +import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; import org.opensearch.sql.storage.write.TableWriteOperator; +import org.opensearch.sql.utils.TestOperator; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -222,6 +224,16 @@ public void visitWindowOperator_should_return_PhysicalWindowOperator() { assertEquals(physicalPlan, logicalPlan.accept(implementor, null)); } + @Test + void visitLogicalCursor_deserializes_it() { + var engine = Mockito.mock(StorageEngine.class); + + var physicalPlan = new TestOperator(); + var logicalPlan = LogicalPlanDSL.cursor(new PlanSerializer(engine) + .convertToCursor(physicalPlan).toString(), engine); + assertEquals(physicalPlan, logicalPlan.accept(implementor, null)); + } + @Test public void visitTableScanBuilder_should_build_TableScanOperator() { TableScanOperator tableScanOperator = Mockito.mock(TableScanOperator.class); diff --git a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java index 34e0e39d87..dc064abe8d 100644 --- a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java @@ -36,7 +36,9 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.Aggregator; import org.opensearch.sql.expression.window.WindowDefinition; +import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.read.TableScanBuilder; @@ -130,9 +132,11 @@ public TableWriteOperator build(PhysicalPlan child) { LogicalNested nested = new LogicalNested(null, nestedArgs, projectList); + LogicalCursor cursor = new LogicalCursor("n:test", mock(StorageEngine.class)); return Stream.of( relation, tableScanBuilder, write, tableWriteBuilder, filter, aggregation, rename, project, - remove, eval, sort, dedup, window, rareTopN, highlight, mlCommons, ad, ml, paginate, nested + remove, eval, sort, dedup, window, rareTopN, highlight, mlCommons, ad, ml, paginate, nested, + cursor ).map(Arguments::of); } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index ba4b1f2cbd..b0422be142 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -41,6 +41,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.utils.TestOperator; @ExtendWith(MockitoExtension.class) class ProjectOperatorTest extends PhysicalPlanTestBase { @@ -236,7 +237,7 @@ public void serializable() { } @EqualsAndHashCode(callSuper = false) - public static class TestOperator extends PhysicalPlan implements ExternalizablePlan { + public static class TestOperator2 extends PhysicalPlan implements ExternalizablePlan { @Override public R accept(PhysicalPlanNodeVisitor visitor, C context) { diff --git a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java new file mode 100644 index 0000000000..1a5ebcc1ed --- /dev/null +++ b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java @@ -0,0 +1,69 @@ +package org.opensearch.sql.utils; + +import java.io.IOException; +import java.io.ObjectInput; +import java.io.ObjectOutput; +import java.util.List; +import lombok.Setter; +import org.opensearch.sql.data.model.ExprValue; +import org.opensearch.sql.exception.NoCursorException; +import org.opensearch.sql.executor.pagination.PlanSerializerTest; +import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.physical.PhysicalPlan; +import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; + +public class TestOperator extends PhysicalPlan implements ExternalizablePlan { + private int field; + @Setter + private boolean throwNoCursorOnWrite = false; + @Setter + private boolean throwIoOnWrite = false; + + public TestOperator() { + } + + public TestOperator(int value) { + field = value; + } + + @Override + public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + field = in.readInt(); + } + + @Override + public void writeExternal(ObjectOutput out) throws IOException { + if (throwNoCursorOnWrite) { + throw new NoCursorException(); + } + if (throwIoOnWrite) { + throw new IOException(); + } + out.writeInt(field); + } + + @Override + public boolean equals(Object o) { + return field == ((TestOperator) o).field; + } + + @Override + public R accept(PhysicalPlanNodeVisitor visitor, C context) { + return null; + } + + @Override + public boolean hasNext() { + return false; + } + + @Override + public ExprValue next() { + return null; + } + + @Override + public List getChild() { + return null; + } +} diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 0095bec7ca..3488110ac9 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -26,6 +26,7 @@ import org.opensearch.common.inject.Injector; import org.opensearch.common.inject.ModulesBuilder; import org.opensearch.common.unit.TimeValue; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.common.response.ResponseListener; import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.data.type.ExprCoreType; @@ -126,7 +127,7 @@ public void onFailure(Exception e) { PhysicalPlan plan = planSerializer.convertToPlan(firstResponder.getCursor().toString()); var secondResponder = new TestResponder(); - queryService.executePlan(plan, secondResponder); + queryService.execute(new FetchCursor(firstResponder.getCursor().toString()), secondResponder); // act 3: confirm that there's no cursor. } diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java index cbbc8c7b9c..c48b18a609 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java @@ -102,7 +102,9 @@ public RestChannelConsumer prepareRequest( channel, createExplainResponseListener(channel, executionErrorHandler), fallbackHandler)); - } else { + } + // If close request, sqlService.closeCursor + else { return channel -> sqlService.execute( request, diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index a5efac3422..473c30347b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -113,6 +113,11 @@ public void clean(Consumer cleanAction) { //do nothing. } + @Override + public boolean hasAnotherBatch() { + return false; + } + @Override public void writeTo(StreamOutput out) throws IOException { throw new UnsupportedOperationException("Not necessary"); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java index 8a5eb6ec4c..e6fe9f32e4 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequest.java @@ -52,9 +52,7 @@ OpenSearchResponse search(Function searchAction, */ OpenSearchExprValueFactory getExprValueFactory(); - default boolean hasAnotherBatch() { - return false; - } + boolean hasAnotherBatch(); /** * OpenSearch Index Name. diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 9aa8225a4e..8c90542289 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -57,10 +57,13 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { */ @Setter @Getter - private String scrollId; + private String scrollId = NO_SCROLL_ID; + + public static final String NO_SCROLL_ID = ""; private boolean needClean = false; + @Getter private List includes; /** Default constructor for Externalizable only. @@ -81,9 +84,9 @@ public OpenSearchScrollRequest(IndexName indexName, .scroll(scrollTimeout) .source(sourceBuilder); - includes = sourceBuilder.fetchSource() != null && sourceBuilder.fetchSource().includes() != null - ? Arrays.asList(sourceBuilder.fetchSource().includes()) - : List.of(); + includes = sourceBuilder.fetchSource() == null + ? List.of() + : Arrays.asList(sourceBuilder.fetchSource().includes()); } @@ -112,7 +115,7 @@ public void clean(Consumer cleanAction) { // clean on the last page only, to prevent closing the scroll/cursor in the middle of paging. if (needClean && isScroll()) { cleanAction.accept(getScrollId()); - setScrollId(null); + setScrollId(NO_SCROLL_ID); } } finally { reset(); @@ -125,7 +128,7 @@ public void clean(Consumer cleanAction) { * @return true if scroll started */ public boolean isScroll() { - return scrollId != null; + return scrollId != NO_SCROLL_ID; } /** @@ -143,7 +146,7 @@ public SearchScrollRequest scrollRequest() { * to be reused across different physical plan. */ public void reset() { - scrollId = null; + scrollId = NO_SCROLL_ID; } /** @@ -152,20 +155,20 @@ public void reset() { */ @Override public boolean hasAnotherBatch() { - return !needClean && scrollId != null && !scrollId.equals(""); + return !needClean && !scrollId.equals(NO_SCROLL_ID); } @Override public void writeTo(StreamOutput out) throws IOException { initialSearchRequest.writeTo(out); out.writeTimeValue(scrollTimeout); + out.writeBoolean(needClean); if (!needClean) { // If needClean is true, there is no more data to get from OpenSearch and scrollId is // used only to clean up OpenSearch context. - out.writeString(scrollId == null? "" : scrollId); + out.writeString(scrollId); } - out.writeBoolean(needClean); out.writeStringCollection(includes); indexName.writeTo(out); } @@ -180,8 +183,10 @@ public OpenSearchScrollRequest(StreamInput in, OpenSearchStorageEngine engine) throws IOException { initialSearchRequest = new SearchRequest(in); scrollTimeout = in.readTimeValue(); - scrollId = in.readString(); needClean = in.readBoolean(); + if (!needClean) { + scrollId = in.readString(); + } includes = in.readStringList(); indexName = new IndexName(in); OpenSearchIndex index = (OpenSearchIndex) engine.getTable(null, indexName.toString()); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index 5162aecfa2..530ca57779 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -8,6 +8,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; @@ -27,6 +28,7 @@ import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchResponse; import org.opensearch.action.search.SearchScrollRequest; +import org.opensearch.common.io.stream.StreamOutput; import org.opensearch.index.query.QueryBuilders; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; @@ -105,10 +107,10 @@ void search_withoutContext() { when(searchAction.apply(any())).thenReturn(searchResponse); when(searchResponse.getHits()).thenReturn(searchHits); when(searchHits.getHits()).thenReturn(new SearchHit[] {searchHit}); - OpenSearchResponse searchResponse = request.search(searchAction, scrollAction); verify(sourceBuilder, times(1)).fetchSource(); assertFalse(searchResponse.isEmpty()); + assertFalse(request.hasAnotherBatch()); } @Test @@ -172,6 +174,11 @@ void searchCrossClusterRequest() { remoteRequest); } + @Test + void writeTo_unsupported() { + assertThrows(UnsupportedOperationException.class, () -> request.writeTo(mock(StreamOutput.class))); + } + private void assertSearchRequest(SearchRequest expected, OpenSearchQueryRequest request) { Function querySearch = searchRequest -> { assertEquals(expected, searchRequest); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index ef6c879183..c38ba36057 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -394,4 +394,16 @@ void exception_when_non_zero_offset_and_page_size() { assertThrows(UnsupportedOperationException.class, () -> requestBuilder.build(indexName, MAX_RESULT_WINDOW, DEFAULT_QUERY_TIMEOUT)); } + + @Test + void maxResponseSize_is_page_size() { + requestBuilder.pushDownPageSize(4); + assertEquals(4, requestBuilder.getMaxResponseSize()); + } + + @Test + void maxResponseSize_is_limit() { + requestBuilder.pushDownLimit(100, 0); + assertEquals(100, requestBuilder.getMaxResponseSize()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 5e31b505b3..34aff19cd6 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -6,20 +6,26 @@ package org.opensearch.sql.opensearch.request; +import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; 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.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.opensearch.sql.opensearch.request.OpenSearchScrollRequest.NO_SCROLL_ID; import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiConsumer; +import java.util.function.Consumer; import java.util.function.Function; import lombok.SneakyThrows; import org.apache.lucene.search.TotalHits; @@ -191,6 +197,9 @@ void hasAnotherBatch() { request.reset(); assertFalse(request.hasAnotherBatch()); + + request.setScrollId(""); + assertFalse(request.hasAnotherBatch()); } @Test @@ -210,7 +219,7 @@ void clean_on_empty_response() { AtomicBoolean cleanCalled = new AtomicBoolean(false); request.clean((s) -> cleanCalled.set(true)); - assertNull(request.getScrollId()); + assertEquals(NO_SCROLL_ID, request.getScrollId()); assertTrue(cleanCalled.get()); } @@ -225,7 +234,7 @@ void no_clean_on_non_empty_response() { assertEquals("scroll", request.getScrollId()); request.clean((s) -> fail()); - assertNull(request.getScrollId()); + assertEquals(NO_SCROLL_ID, request.getScrollId()); } @Test @@ -245,7 +254,7 @@ void no_clean_if_no_scroll_in_response() { new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1F)); request.search((x) -> searchResponse, (x) -> searchResponse); - assertNull(request.getScrollId()); + assertEquals(NO_SCROLL_ID, request.getScrollId()); request.clean((s) -> fail()); } @@ -267,4 +276,56 @@ void serialize_deserialize_no_needClean() { assertEquals(request.getInitialSearchRequest(), newRequest.getInitialSearchRequest()); assertEquals("", newRequest.getScrollId()); } + + @Test + @SneakyThrows + void serialize_deserialize_needClean() { + var stream = new BytesStreamOutput(); + lenient().when(searchResponse.getHits()).thenReturn( + new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1F)); + lenient().when(searchResponse.getScrollId()).thenReturn(""); + + request.search(searchRequest -> searchResponse, null); + request.writeTo(stream); + stream.flush(); + assertTrue(stream.size() > 0); + + // deserialize + var inStream = new BytesStreamInput(stream.bytes().toBytesRef().bytes); + var indexMock = mock(OpenSearchIndex.class); + var engine = mock(OpenSearchStorageEngine.class); + when(engine.getTable(any(), any())).thenReturn(indexMock); + var newRequest = new OpenSearchScrollRequest(inStream, engine); + assertEquals(request.getInitialSearchRequest(), newRequest.getInitialSearchRequest()); + assertEquals("", newRequest.getScrollId()); + } + + @Test + void default_constructor() { + assertEquals(NO_SCROLL_ID, new OpenSearchScrollRequest().getScrollId()); + } + + @Test + void setScrollId() { + request.setScrollId("test"); + assertEquals("test", request.getScrollId()); + } + + @Test + void includes() { + + assertIncludes(List.of(), searchSourceBuilder); + + searchSourceBuilder.fetchSource((String[])null, (String[])null); + assertIncludes(List.of(), searchSourceBuilder); + + searchSourceBuilder.fetchSource(new String[] {"test"}, null); + assertIncludes(List.of("test"), searchSourceBuilder); + + } + + void assertIncludes(List expected, SearchSourceBuilder sourceBuilder) { + assertEquals(expected, new OpenSearchScrollRequest( + INDEX_NAME, SCROLL_TIMEOUT, sourceBuilder, factory).getIncludes()); + } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index d4de3a892f..297c0115c1 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -20,8 +20,10 @@ import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.data.type.ExprCoreType.STRING; +import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; @@ -49,6 +51,7 @@ import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.exception.NoCursorException; +import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.opensearch.client.OpenSearchClient; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory; @@ -57,6 +60,9 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; +import org.opensearch.sql.opensearch.storage.OpenSearchIndex; +import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; +import org.opensearch.sql.storage.StorageEngine; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -108,22 +114,21 @@ void serialize() { var searchSourceBuilder = new SearchSourceBuilder().size(4); var factory = mock(OpenSearchExprValueFactory.class); - + var engine = mock(OpenSearchStorageEngine.class); + var index = mock(OpenSearchIndex.class); + when(engine.getClient()).thenReturn(client); + when(engine.getTable(any(), any())).thenReturn(index); var request = new OpenSearchScrollRequest( - INDEX_NAME, CURSOR_KEEP_ALIVE,searchSourceBuilder, factory); + INDEX_NAME, CURSOR_KEEP_ALIVE, searchSourceBuilder, factory); request.setScrollId("valid-id"); - try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request); - var byteStream = new ByteArrayOutputStream(); - var objectStream = new ObjectOutputStream(byteStream)) { - try { - indexScan.writeExternal(objectStream); - objectStream.flush(); - byteStream.flush(); - assertTrue(byteStream.size() > 0); - } catch (IOException e) { - throw new AssertionError("Expected to serialize the indexScan"); - } - } + + try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request)) { + var planSerializer = new PlanSerializer(engine); + var cursor = planSerializer.convertToCursor(indexScan); + var newPlan = planSerializer.convertToPlan(cursor.toString()); + assertEquals(indexScan, newPlan); + } + } @Test diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java index ca4939d507..7eb07b0bd0 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -6,6 +6,7 @@ package org.opensearch.sql.sql; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; @@ -146,8 +147,8 @@ public void onResponse(ExplainResponse response) { @Override public void onFailure(Exception e) { - assertTrue(e.getMessage() - .contains("`explain` request for cursor requests is not supported.")); + assertEquals("Explain of a paged query continuation is not supported." + + " Use `explain` for the initial query request.", e.getMessage()); } }); } From fabb1793efc37739675192669fd7c0419164d89b Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 25 May 2023 22:33:42 -0700 Subject: [PATCH 26/38] Address checkstyle issues. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/analysis/Analyzer.java | 2 +- .../sql/ast/AbstractNodeVisitor.java | 2 +- .../executor/execution/QueryPlanFactory.java | 2 +- .../opensearch/sql/planner/LogicalCursor.java | 3 ++ .../opensearch/sql/analysis/AnalyzerTest.java | 4 +++ .../execution/QueryPlanFactoryTest.java | 1 - .../pagination/PlanSerializerTest.java | 1 - .../planner/physical/ProjectOperatorTest.java | 32 ------------------- .../opensearch/sql/utils/TestOperator.java | 3 +- .../sql/sql/StandalonePaginationIT.java | 2 +- .../request/OpenSearchScrollRequest.java | 2 +- .../request/OpenSearchQueryRequestTest.java | 22 ++++++------- .../request/OpenSearchRequestBuilderTest.java | 1 - .../request/OpenSearchScrollRequestTest.java | 17 ++++------ .../storage/scan/OpenSearchIndexScanTest.java | 19 ++++------- .../org/opensearch/sql/sql/SQLService.java | 4 +-- .../opensearch/sql/sql/SQLServiceTest.java | 6 ++-- 17 files changed, 41 insertions(+), 82 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index 92317f0be4..927f3176b3 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -43,9 +43,9 @@ import org.opensearch.sql.ast.expression.UnresolvedExpression; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Kmeans; diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index db9340f4cc..636355f547 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -41,9 +41,9 @@ import org.opensearch.sql.ast.statement.Statement; import org.opensearch.sql.ast.tree.AD; import org.opensearch.sql.ast.tree.Aggregation; -import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Dedupe; import org.opensearch.sql.ast.tree.Eval; +import org.opensearch.sql.ast.tree.FetchCursor; import org.opensearch.sql.ast.tree.Filter; import org.opensearch.sql.ast.tree.Head; import org.opensearch.sql.ast.tree.Kmeans; diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 2231d5a7c1..5a407ad0c2 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -84,7 +84,7 @@ public AbstractPlan create(String cursor, boolean isExplain, return isExplain ? new ExplainPlan(queryId, plan, explainListener) : plan; } - boolean canConvertToCursor(UnresolvedPlan plan) { + boolean canConvertToCursor(UnresolvedPlan plan) { return plan.accept(new CanPaginateVisitor(), null); } diff --git a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java index 703f4bf0fb..5670e3078d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java +++ b/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java @@ -17,6 +17,9 @@ public class LogicalCursor extends LogicalPlan { @Getter private final StorageEngine engine; + /** + * LogicalCursor constructor. Does not have child plans. + */ public LogicalCursor(String cursor, StorageEngine engine) { super(List.of()); this.cursor = cursor; diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index b9ede4eea7..9f5b37e3c0 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -90,6 +90,7 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.function.OpenSearchFunctions; import org.opensearch.sql.expression.window.WindowDefinition; +import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalMLCommons; @@ -1646,5 +1647,8 @@ public void visit_paginate() { @Test void visit_cursor() { LogicalPlan actual = analyze((new FetchCursor("test"))); + assertTrue(actual instanceof LogicalCursor); + assertEquals(new LogicalCursor("test", + dataSourceService.getDataSource("@opensearch").getStorageEngine()), actual); } } diff --git a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java index e2e1f3684d..c35d506fe7 100644 --- a/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanFactoryTest.java @@ -31,7 +31,6 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.pagination.CanPaginateVisitor; -import org.opensearch.sql.executor.pagination.PlanSerializer; @ExtendWith(MockitoExtension.class) class QueryPlanFactoryTest { diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java index ce3c35c1c0..ace4d29a56 100644 --- a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java @@ -24,7 +24,6 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; -import org.opensearch.sql.ast.dsl.AstDSL; import org.opensearch.sql.exception.NoCursorException; import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index b0422be142..c34b5a56a7 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -235,36 +235,4 @@ public void serializable() { var roundTripPlan = (ProjectOperator) objectInput.readObject(); assertEquals(project, roundTripPlan); } - - @EqualsAndHashCode(callSuper = false) - public static class TestOperator2 extends PhysicalPlan implements ExternalizablePlan { - - @Override - public R accept(PhysicalPlanNodeVisitor visitor, C context) { - return null; - } - - @Override - public boolean hasNext() { - return false; - } - - @Override - public ExprValue next() { - return null; - } - - @Override - public List getChild() { - return null; - } - - @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { - } - - @Override - public void writeExternal(ObjectOutput out) throws IOException { - } - } } diff --git a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java index 1a5ebcc1ed..265800218d 100644 --- a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java +++ b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java @@ -7,7 +7,6 @@ import lombok.Setter; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.executor.pagination.PlanSerializerTest; import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -27,7 +26,7 @@ public TestOperator(int value) { } @Override - public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException { + public void readExternal(ObjectInput in) throws IOException { field = in.readInt(); } diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java index 3488110ac9..aad39c4074 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/StandalonePaginationIT.java @@ -99,7 +99,7 @@ public void onFailure(Exception e) { e.printStackTrace(); fail(e.getMessage()); } - }; + } // arrange { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 8c90542289..dd568a6c9b 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -128,7 +128,7 @@ public void clean(Consumer cleanAction) { * @return true if scroll started */ public boolean isScroll() { - return scrollId != NO_SCROLL_ID; + return !scrollId.equals(NO_SCROLL_ID); } /** diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java index 530ca57779..a92cb44d7a 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequestTest.java @@ -148,15 +148,14 @@ void clean() { void searchRequest() { request.getSourceBuilder().query(QueryBuilders.termQuery("name", "John")); - assertSearchRequest(new SearchRequest() - .indices("test") - .source(new SearchSourceBuilder() - .timeout(DEFAULT_QUERY_TIMEOUT) - .from(0) - .size(200) - .query(QueryBuilders.termQuery("name", "John"))), - request); - + assertSearchRequest(new SearchRequest() + .indices("test") + .source(new SearchSourceBuilder() + .timeout(DEFAULT_QUERY_TIMEOUT) + .from(0) + .size(200) + .query(QueryBuilders.termQuery("name", "John"))), + request); } @Test @@ -176,7 +175,8 @@ void searchCrossClusterRequest() { @Test void writeTo_unsupported() { - assertThrows(UnsupportedOperationException.class, () -> request.writeTo(mock(StreamOutput.class))); + assertThrows(UnsupportedOperationException.class, + () -> request.writeTo(mock(StreamOutput.class))); } private void assertSearchRequest(SearchRequest expected, OpenSearchQueryRequest request) { @@ -184,7 +184,7 @@ private void assertSearchRequest(SearchRequest expected, OpenSearchQueryRequest assertEquals(expected, searchRequest); return when(mock(SearchResponse.class).getHits()) .thenReturn(new SearchHits(new SearchHit[0], - new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)) + new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0.0f)) .getMock(); }; request.search(querySearch, searchScrollRequest -> null); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index c38ba36057..8826295959 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -51,7 +51,6 @@ import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.ScoreSortBuilder; import org.opensearch.search.sort.SortBuilders; -import org.opensearch.sql.common.setting.Settings; import org.opensearch.sql.exception.SemanticCheckException; import org.opensearch.sql.expression.DSL; import org.opensearch.sql.expression.NamedExpression; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index 34aff19cd6..a80f171755 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -6,12 +6,9 @@ package org.opensearch.sql.opensearch.request; -import static org.junit.jupiter.api.Assertions.assertAll; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; @@ -24,8 +21,6 @@ import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.BiConsumer; -import java.util.function.Consumer; import java.util.function.Function; import lombok.SneakyThrows; import org.apache.lucene.search.TotalHits; @@ -54,7 +49,8 @@ @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) class OpenSearchScrollRequestTest { - public static final OpenSearchRequest.IndexName INDEX_NAME = new OpenSearchRequest.IndexName("test"); + public static final OpenSearchRequest.IndexName INDEX_NAME + = new OpenSearchRequest.IndexName("test"); public static final TimeValue SCROLL_TIMEOUT = TimeValue.timeValueMinutes(1); @Mock private Function searchAction; @@ -79,13 +75,14 @@ class OpenSearchScrollRequestTest { private final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); private final OpenSearchScrollRequest request = new OpenSearchScrollRequest( - INDEX_NAME, SCROLL_TIMEOUT, + INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, factory); @Test void constructor() { searchSourceBuilder.fetchSource(new String[] {"test"}, null); - var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, searchSourceBuilder, factory); + var request = new OpenSearchScrollRequest(INDEX_NAME, SCROLL_TIMEOUT, + searchSourceBuilder, factory); assertNotEquals(List.of(), request.getIncludes()); } @@ -280,11 +277,11 @@ void serialize_deserialize_no_needClean() { @Test @SneakyThrows void serialize_deserialize_needClean() { - var stream = new BytesStreamOutput(); lenient().when(searchResponse.getHits()).thenReturn( new SearchHits(new SearchHit[0], new TotalHits(0, TotalHits.Relation.EQUAL_TO), 1F)); lenient().when(searchResponse.getScrollId()).thenReturn(""); + var stream = new BytesStreamOutput(); request.search(searchRequest -> searchResponse, null); request.writeTo(stream); stream.flush(); @@ -326,6 +323,6 @@ void includes() { void assertIncludes(List expected, SearchSourceBuilder sourceBuilder) { assertEquals(expected, new OpenSearchScrollRequest( - INDEX_NAME, SCROLL_TIMEOUT, sourceBuilder, factory).getIncludes()); + INDEX_NAME, SCROLL_TIMEOUT, sourceBuilder, factory).getIncludes()); } } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java index 297c0115c1..e974790629 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanTest.java @@ -20,10 +20,7 @@ import static org.opensearch.search.sort.SortOrder.ASC; import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.util.Arrays; import java.util.HashMap; @@ -62,7 +59,6 @@ import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchIndex; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; -import org.opensearch.sql.storage.StorageEngine; @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) @@ -76,9 +72,6 @@ class OpenSearchIndexScanTest { @Mock private OpenSearchClient client; - @Mock - private Settings settings; - private final OpenSearchExprValueFactory exprValueFactory = new OpenSearchExprValueFactory( Map.of("name", OpenSearchDataType.of(STRING), "department", OpenSearchDataType.of(STRING))); @@ -105,7 +98,7 @@ void throws_no_cursor_exception() { var byteStream = new ByteArrayOutputStream(); var objectStream = new ObjectOutputStream(byteStream)) { assertThrows(NoCursorException.class, () -> objectStream.writeObject(indexScan)); - } + } } @Test @@ -119,14 +112,14 @@ void serialize() { when(engine.getClient()).thenReturn(client); when(engine.getTable(any(), any())).thenReturn(index); var request = new OpenSearchScrollRequest( - INDEX_NAME, CURSOR_KEEP_ALIVE, searchSourceBuilder, factory); + INDEX_NAME, CURSOR_KEEP_ALIVE, searchSourceBuilder, factory); request.setScrollId("valid-id"); try (var indexScan = new OpenSearchIndexScan(client, QUERY_SIZE, request)) { - var planSerializer = new PlanSerializer(engine); - var cursor = planSerializer.convertToCursor(indexScan); - var newPlan = planSerializer.convertToPlan(cursor.toString()); - assertEquals(indexScan, newPlan); + var planSerializer = new PlanSerializer(engine); + var cursor = planSerializer.convertToCursor(indexScan); + var newPlan = planSerializer.convertToPlan(cursor.toString()); + assertEquals(indexScan, newPlan); } } diff --git a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java index 5f9701b187..889f80223f 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/SQLService.java +++ b/sql/src/main/java/org/opensearch/sql/sql/SQLService.java @@ -69,8 +69,8 @@ private AbstractPlan plan( if (request.getCursor().isPresent()) { // Handle v2 cursor here -- legacy cursor was handled earlier. if (isExplainRequest) { - throw new UnsupportedOperationException("Explain of a paged query continuation " + - "is not supported. Use `explain` for the initial query request."); + throw new UnsupportedOperationException("Explain of a paged query continuation " + + "is not supported. Use `explain` for the initial query request."); } return queryExecutionFactory.create(request.getCursor().get(), isExplainRequest, queryListener.orElse(null), explainListener.orElse(null)); diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java index 7eb07b0bd0..f34c95e121 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -8,7 +8,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; 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.doAnswer; @@ -31,7 +30,6 @@ import org.opensearch.sql.executor.ExecutionEngine.ExplainResponseNode; import org.opensearch.sql.executor.QueryService; import org.opensearch.sql.executor.execution.QueryPlanFactory; -import org.opensearch.sql.executor.pagination.PlanSerializer; import org.opensearch.sql.sql.antlr.SQLSyntaxParser; import org.opensearch.sql.sql.domain.SQLQueryRequest; @@ -147,8 +145,8 @@ public void onResponse(ExplainResponse response) { @Override public void onFailure(Exception e) { - assertEquals("Explain of a paged query continuation is not supported." + - " Use `explain` for the initial query request.", e.getMessage()); + assertEquals("Explain of a paged query continuation is not supported." + + " Use `explain` for the initial query request.", e.getMessage()); } }); } From 23bc3ab0e87000262ab6687bd4c0715a0dbdffc7 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 25 May 2023 22:45:25 -0700 Subject: [PATCH 27/38] Better class name. Signed-off-by: MaxKsyunz --- .../src/main/java/org/opensearch/sql/analysis/Analyzer.java | 4 ++-- .../java/org/opensearch/sql/planner/DefaultImplementor.java | 2 +- .../planner/{LogicalCursor.java => LogicalFetchCursor.java} | 4 ++-- .../org/opensearch/sql/planner/logical/LogicalPlanDSL.java | 4 ++-- .../sql/planner/logical/LogicalPlanNodeVisitor.java | 4 ++-- .../test/java/org/opensearch/sql/analysis/AnalyzerTest.java | 6 +++--- .../sql/planner/logical/LogicalPlanNodeVisitorTest.java | 4 ++-- .../opensearch/sql/opensearch/storage/OpenSearchIndex.java | 1 - 8 files changed, 14 insertions(+), 15 deletions(-) rename core/src/main/java/org/opensearch/sql/planner/{LogicalCursor.java => LogicalFetchCursor.java} (85%) diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index 927f3176b3..f9c005fa40 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -81,7 +81,7 @@ import org.opensearch.sql.expression.function.FunctionName; import org.opensearch.sql.expression.function.TableFunctionImplementation; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalDedupe; @@ -215,7 +215,7 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex @Override public LogicalPlan visitCursor(FetchCursor cursor, AnalysisContext context) { - return new LogicalCursor(cursor.getCursor(), + return new LogicalFetchCursor(cursor.getCursor(), dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); } diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index d82ea138b5..534ba19f85 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -150,7 +150,7 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { } @Override - public PhysicalPlan visitCursor(LogicalCursor plan, C context) { + public PhysicalPlan visitCursor(LogicalFetchCursor plan, C context) { return new PlanSerializer(plan.getEngine()).convertToPlan(plan.getCursor()); } diff --git a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java b/core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java similarity index 85% rename from core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java rename to core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java index 5670e3078d..1a5e33192d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/LogicalCursor.java +++ b/core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java @@ -10,7 +10,7 @@ @EqualsAndHashCode(callSuper = false) @ToString -public class LogicalCursor extends LogicalPlan { +public class LogicalFetchCursor extends LogicalPlan { @Getter private final String cursor; @@ -20,7 +20,7 @@ public class LogicalCursor extends LogicalPlan { /** * LogicalCursor constructor. Does not have child plans. */ - public LogicalCursor(String cursor, StorageEngine engine) { + public LogicalFetchCursor(String cursor, StorageEngine engine) { super(List.of()); this.cursor = cursor; this.engine = engine; diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java index 937f346078..35a59882c6 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java @@ -22,7 +22,7 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; @@ -33,7 +33,7 @@ public class LogicalPlanDSL { public static LogicalPlan cursor(String cursor, StorageEngine engine) { - return new LogicalCursor(cursor, engine); + return new LogicalFetchCursor(cursor, engine); } public static LogicalPlan write(LogicalPlan input, Table table, List columns) { diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java index 8aec99988d..5905a55e90 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java @@ -6,7 +6,7 @@ package org.opensearch.sql.planner.logical; -import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; @@ -110,7 +110,7 @@ public R visitPaginate(LogicalPaginate plan, C context) { return visitNode(plan, context); } - public R visitCursor(LogicalCursor plan, C context) { + public R visitCursor(LogicalFetchCursor plan, C context) { return visitNode(plan, context); } } diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index 9f5b37e3c0..d3f6dd33e4 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -90,7 +90,7 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.function.OpenSearchFunctions; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalMLCommons; @@ -1647,8 +1647,8 @@ public void visit_paginate() { @Test void visit_cursor() { LogicalPlan actual = analyze((new FetchCursor("test"))); - assertTrue(actual instanceof LogicalCursor); - assertEquals(new LogicalCursor("test", + assertTrue(actual instanceof LogicalFetchCursor); + assertEquals(new LogicalFetchCursor("test", dataSourceService.getDataSource("@opensearch").getStorageEngine()), actual); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java index dc064abe8d..76d388ad89 100644 --- a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java @@ -36,7 +36,7 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.Aggregator; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalCursor; +import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; @@ -132,7 +132,7 @@ public TableWriteOperator build(PhysicalPlan child) { LogicalNested nested = new LogicalNested(null, nestedArgs, projectList); - LogicalCursor cursor = new LogicalCursor("n:test", mock(StorageEngine.class)); + LogicalFetchCursor cursor = new LogicalFetchCursor("n:test", mock(StorageEngine.class)); return Stream.of( relation, tableScanBuilder, write, tableWriteBuilder, filter, aggregation, rename, project, remove, eval, sort, dedup, window, rareTopN, highlight, mlCommons, ad, ml, paginate, nested, diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java index 4d34fcdaff..532d62333d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java @@ -27,7 +27,6 @@ import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScanBuilder; import org.opensearch.sql.planner.DefaultImplementor; -import org.opensearch.sql.planner.LogicalCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalML; import org.opensearch.sql.planner.logical.LogicalMLCommons; From 264e48338051b1532a407c0ff3ce2dbc0cecacbf Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Thu, 25 May 2023 22:46:44 -0700 Subject: [PATCH 28/38] Update design document to reflect refactor. Signed-off-by: MaxKsyunz --- docs/dev/Pagination-v2.md | 51 +++++++++++++++------------------------ 1 file changed, 19 insertions(+), 32 deletions(-) diff --git a/docs/dev/Pagination-v2.md b/docs/dev/Pagination-v2.md index 60f8e1e708..5379153a4a 100644 --- a/docs/dev/Pagination-v2.md +++ b/docs/dev/Pagination-v2.md @@ -202,7 +202,7 @@ classDiagram ``` When `QueryPlanFactory.create` is passed a subsequent query request, it: -1. Creates an instance of `Cursor` unresolved plan as the sole node in the unresolved query plan. +1. Creates an instance of `FetchCursor` unresolved plan as the sole node in the unresolved query plan. ```mermaid classDiagram @@ -213,11 +213,11 @@ classDiagram -UnresolvedPlan plan -QueryService queryService } - class Cursor { + class FetchCursor { <> -String cursorId } - QueryPlan --* Cursor + QueryPlan --* FetchCursor ``` The examples below show Abstract Query Plan for the same query in different request types: @@ -256,7 +256,7 @@ stateDiagram-v2 } state "Subsequent Query Request" As Sub { - Cursor + FetchCursor } ``` @@ -284,7 +284,7 @@ classDiagram LogicalQueryPlan --* LogicalRelation ``` -For subsequent page requests, `Cursor` unresolved plan is mapped to `LogicalCursor` logical plan. +For subsequent page requests, `FetchCursor` unresolved plan is mapped to `LogicalFetchCursor` logical plan. ```mermaid classDiagram @@ -292,11 +292,11 @@ classDiagram class LogicalQueryPlan { <> } - class LogicalCursor { + class LogicalFetchCursor { <> -String cursorId } - LogicalQueryPlan --* LogicalCursor + LogicalQueryPlan --* LogicalFetchCursor ``` The examples below show logical query plan for the same query in different request types: @@ -331,7 +331,7 @@ stateDiagram-v2 } state "Subsequent Query Request" As Sub { -Cursor +FetchCursor } ``` @@ -500,33 +500,20 @@ Subsequent pages are processed by a new workflow. The key point there: ```mermaid sequenceDiagram - participant SQLService - participant QueryPlanFactory - participant QueryService - participant OpenSearchExecutionEngine - participant LogicalCursor - participant DefaultImplementor - participant PlanSerializer SQLService ->>+ QueryPlanFactory : execute QueryPlanFactory ->>+ QueryService : execute - rect rgb(91, 123, 155) - note over QueryService, PlanSerializer : Deserialization - QueryService ->>+ PlanSerializer: convertToPlan - PlanSerializer -->>- QueryService: Physical Query Plan - end - Note over QueryService : Planner, Optimizer and Implementor
    are skipped - QueryService ->>+ OpenSearchExecutionEngine : execute - rect rgb(91, 123, 155) - note over OpenSearchExecutionEngine, PlanSerializer : Serialization - OpenSearchExecutionEngine ->>+ PlanSerializer : convertToCursor - PlanSerializer -->>- OpenSearchExecutionEngine : cursor - end - rect rgb(91, 123, 155) - Note over OpenSearchExecutionEngine : get total hits - end - OpenSearchExecutionEngine -->>- QueryService: execution completed - QueryService -->>- QueryPlanFactory : execution completed + QueryService ->>+ Analyzer : analyze + Analyzer -->>- QueryService : new LogicalFetchCursor + QueryService ->>+ Planner : plan + Planner ->>+ DefaultImplementor : implement + DefaultImplementor ->>+ PlanSerializer : deserialize + PlanSerializer -->>- DefaultImplementor: physical query plan + DefaultImplementor -->>- Planner : physical query plan + Planner -->>- QueryService : physical query plan + QueryService ->>+ OpenSearchExecutionEngine : execute + OpenSearchExecutionEngine -->>- QueryService: execution completed + QueryService -->>- QueryPlanFactory : execution completed QueryPlanFactory -->>- SQLService : execution completed ``` From 97ef3a5c26da4525371e172b7f4bb39cd76d51e9 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 26 May 2023 00:20:40 -0700 Subject: [PATCH 29/38] Addressed PR feedback. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/analysis/Analyzer.java | 13 ++++----- .../opensearch/sql/ast/tree/FetchCursor.java | 5 ++++ .../sql/planner/DefaultImplementor.java | 1 + .../sql/planner/ExternalizablePlan.java | 15 ---------- .../{ => logical}/LogicalFetchCursor.java | 7 ++++- .../sql/planner/logical/LogicalPlanDSL.java | 1 - .../logical/LogicalPlanNodeVisitor.java | 1 - .../optimizer/LogicalPlanOptimizer.java | 2 +- .../planner/optimizer/PushDownPageSize.java | 4 +-- .../opensearch/sql/analysis/AnalyzerTest.java | 2 +- .../logical/LogicalPlanNodeVisitorTest.java | 1 - docs/dev/query-optimizer-improvement.md | 8 ++---- .../request/ExecutableRequestBuilder.java | 9 ------ .../request/OpenSearchQueryRequest.java | 3 +- .../request/OpenSearchRequestBuilder.java | 3 +- .../request/OpenSearchScrollRequest.java | 3 +- ...OpenSearchIndexScanAggregationBuilder.java | 25 ----------------- .../storage/scan/PushDownQueryBuilder.java | 28 ++++++++++++++----- .../request/OpenSearchRequestBuilderTest.java | 4 +-- 19 files changed, 52 insertions(+), 83 deletions(-) rename core/src/main/java/org/opensearch/sql/planner/{ => logical}/LogicalFetchCursor.java (86%) delete mode 100644 opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index f9c005fa40..54c60d6054 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -81,11 +81,11 @@ import org.opensearch.sql.expression.function.FunctionName; import org.opensearch.sql.expression.function.TableFunctionImplementation; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalAD; import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalDedupe; import org.opensearch.sql.planner.logical.LogicalEval; +import org.opensearch.sql.planner.logical.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalML; @@ -213,12 +213,6 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex tableFunctionImplementation.applyArguments()); } - @Override - public LogicalPlan visitCursor(FetchCursor cursor, AnalysisContext context) { - return new LogicalFetchCursor(cursor.getCursor(), - dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); - } - @Override public LogicalPlan visitLimit(Limit node, AnalysisContext context) { LogicalPlan child = node.getChild().get(0).accept(this, context); @@ -594,4 +588,9 @@ private SortOption analyzeSortOption(List fieldArgs) { return asc ? SortOption.DEFAULT_ASC : SortOption.DEFAULT_DESC; } + @Override + public LogicalPlan visitCursor(FetchCursor cursor, AnalysisContext context) { + return new LogicalFetchCursor(cursor.getCursor(), + dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); + } } diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java index 64655d1100..07070e36df 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.sql.ast.tree; import lombok.EqualsAndHashCode; diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index 534ba19f85..fec136ef7a 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -10,6 +10,7 @@ import org.opensearch.sql.planner.logical.LogicalAggregation; import org.opensearch.sql.planner.logical.LogicalDedupe; import org.opensearch.sql.planner.logical.LogicalEval; +import org.opensearch.sql.planner.logical.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalLimit; import org.opensearch.sql.planner.logical.LogicalNested; diff --git a/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java b/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java index dec64ffdc7..22412679bb 100644 --- a/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java @@ -29,21 +29,6 @@ */ public interface ExternalizablePlan extends Externalizable { - /** - * Argument is an instance of {@link PlanSerializer.CursorDeserializationStream}. - */ - @Override - void readExternal(ObjectInput in) throws IOException, ClassNotFoundException; - - /** - * Each plan which has as a child plan should do. - * 
    {@code
-   * out.writeObject(input.getPlanForSerialization());
-   * }
    - */ - @Override - void writeExternal(ObjectOutput out) throws IOException; - /** * Override to return child or delegated plan, so parent plan should skip this one * for serialization, but it should try to serialize grandchild plan. diff --git a/core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java similarity index 86% rename from core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java rename to core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java index 1a5e33192d..554381335e 100644 --- a/core/src/main/java/org/opensearch/sql/planner/LogicalFetchCursor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java @@ -1,4 +1,9 @@ -package org.opensearch.sql.planner; +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.planner.logical; import java.util.List; import lombok.EqualsAndHashCode; diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java index 35a59882c6..9bd9b58274 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java @@ -22,7 +22,6 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.NamedAggregator; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java index 5905a55e90..04159f4b93 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java @@ -6,7 +6,6 @@ package org.opensearch.sql.planner.logical; -import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.storage.read.TableScanBuilder; import org.opensearch.sql.storage.write.TableWriteBuilder; diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java index 0f6b3e5195..be1227c1da 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java @@ -51,11 +51,11 @@ public static LogicalPlanOptimizer create() { * Phase 2: Transformations that rely on data source push down capability */ new CreateTableScanBuilder(), - new PushDownPageSize(), TableScanPushDown.PUSH_DOWN_FILTER, TableScanPushDown.PUSH_DOWN_AGGREGATION, TableScanPushDown.PUSH_DOWN_SORT, TableScanPushDown.PUSH_DOWN_LIMIT, + new PushDownPageSize(), TableScanPushDown.PUSH_DOWN_HIGHLIGHT, TableScanPushDown.PUSH_DOWN_NESTED, TableScanPushDown.PUSH_DOWN_PROJECT, diff --git a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java index 831428001b..8150de824d 100644 --- a/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java +++ b/core/src/main/java/org/opensearch/sql/planner/optimizer/PushDownPageSize.java @@ -39,8 +39,8 @@ private Optional findTableScanBuilder(LogicalPaginate logicalP Deque plans = new ArrayDeque<>(); plans.add(logicalPaginate); do { - final var plan = plans.removeFirst(); - final var children = plan.getChild(); + var plan = plans.removeFirst(); + var children = plan.getChild(); if (children.stream().anyMatch(TableScanBuilder.class::isInstance)) { if (children.size() > 1) { throw new UnsupportedOperationException( diff --git a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java index d3f6dd33e4..dda359a7df 100644 --- a/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java +++ b/core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java @@ -90,8 +90,8 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.function.OpenSearchFunctions; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalAD; +import org.opensearch.sql.planner.logical.LogicalFetchCursor; import org.opensearch.sql.planner.logical.LogicalFilter; import org.opensearch.sql.planner.logical.LogicalMLCommons; import org.opensearch.sql.planner.logical.LogicalPaginate; diff --git a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java index 76d388ad89..e826a13f6c 100644 --- a/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitorTest.java @@ -36,7 +36,6 @@ import org.opensearch.sql.expression.ReferenceExpression; import org.opensearch.sql.expression.aggregation.Aggregator; import org.opensearch.sql.expression.window.WindowDefinition; -import org.opensearch.sql.planner.LogicalFetchCursor; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.storage.Table; diff --git a/docs/dev/query-optimizer-improvement.md b/docs/dev/query-optimizer-improvement.md index 7503087562..720649b280 100644 --- a/docs/dev/query-optimizer-improvement.md +++ b/docs/dev/query-optimizer-improvement.md @@ -91,6 +91,7 @@ classDiagram +pushDownAggregation(LogicalAggregation) boolean +pushDownSort(LogicalSort) boolean +pushDownLimit(LogicalLimit) boolean + +pushDownPageSize(LogicalPaginate) boolean +pushDownProject(LogicalProject) boolean +pushDownHighlight(LogicalHighlight) boolean +pushDownNested(LogicalNested) boolean @@ -103,16 +104,13 @@ classDiagram +pushDownAggregation(LogicalAggregation) boolean +pushDownSort(LogicalSort) boolean +pushDownLimit(LogicalLimit) boolean + +pushDownPageSize(LogicalPaginate) boolean +pushDownProject(LogicalProject) boolean +pushDownHighlight(LogicalHighlight) boolean +pushDownNested(LogicalNested) boolean +findReferenceExpression(NamedExpression)$ List~ReferenceExpression~ +findReferenceExpressions(List~NamedExpression~)$ Set~ReferenceExpression~ } - class OpenSearchPagedIndexScanBuilder { - +OpenSearchPagedIndexScanBuilder(OpenSearchPagedIndexScan) - +build() TableScanOperator - } class OpenSearchIndexScanBuilder { -TableScanBuilder delegate -boolean isLimitPushedDown @@ -131,7 +129,6 @@ classDiagram LogicalPlan <|-- TableScanBuilder TableScanBuilder <|-- OpenSearchIndexScanQueryBuilder - TableScanBuilder <|-- OpenSearchPagedIndexScanBuilder TableScanBuilder <|-- OpenSearchIndexScanBuilder OpenSearchIndexScanBuilder *-- "1" TableScanBuilder : delegate OpenSearchIndexScanBuilder <.. OpenSearchIndexScanQueryBuilder : creates @@ -159,7 +156,6 @@ classDiagram } class Table { +TableScanBuilder createScanBuilder() - +TableScanBuilder createPagedScanBuilder(int) } class TableScanPushDown~T~ { +Rule~T~ PUSH_DOWN_FILTER$ diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java deleted file mode 100644 index 9b1ec039e2..0000000000 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/ExecutableRequestBuilder.java +++ /dev/null @@ -1,9 +0,0 @@ -package org.opensearch.sql.opensearch.request; - -import org.opensearch.common.unit.TimeValue; - -public interface ExecutableRequestBuilder { - - OpenSearchRequest build(OpenSearchRequest.IndexName indexName, - int maxResultWindow, TimeValue scrollTimeout); -} diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java index 473c30347b..45954a3871 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java @@ -120,6 +120,7 @@ public boolean hasAnotherBatch() { @Override public void writeTo(StreamOutput out) throws IOException { - throw new UnsupportedOperationException("Not necessary"); + throw new UnsupportedOperationException("OpenSearchQueryRequest serialization " + + "is not implemented."); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java index 276b3b0220..97512bec49 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java @@ -49,7 +49,7 @@ @EqualsAndHashCode @Getter @ToString -public class OpenSearchRequestBuilder implements ExecutableRequestBuilder { +public class OpenSearchRequestBuilder { /** * Search request source builder. @@ -92,7 +92,6 @@ public OpenSearchRequestBuilder(int requestedTotalSize, * * @return query request or scroll request */ - @Override public OpenSearchRequest build(OpenSearchRequest.IndexName indexName, int maxResultWindow, TimeValue scrollTimeout) { int size = requestedTotalSize; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index dd568a6c9b..fecc377e7a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -90,7 +90,8 @@ public OpenSearchScrollRequest(IndexName indexName, } - /** Constructor. */ + /** Executes request using either {@param searchAction} or {@param scrollAction} as appropriate. + */ @Override public OpenSearchResponse search(Function searchAction, Function scrollAction) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java index 58946801fb..84883b5209 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScanAggregationBuilder.java @@ -82,31 +82,6 @@ public boolean pushDownSort(LogicalSort sort) { return true; } - @Override - public boolean pushDownLimit(LogicalLimit limit) { - return false; - } - - @Override - public boolean pushDownProject(LogicalProject project) { - return false; - } - - @Override - public boolean pushDownHighlight(LogicalHighlight highlight) { - return false; - } - - @Override - public boolean pushDownPageSize(LogicalPaginate paginate) { - return false; - } - - @Override - public boolean pushDownNested(LogicalNested nested) { - return false; - } - private boolean hasAggregatorInSortBy(LogicalSort sort) { final Set aggregatorNames = aggregatorList.stream().map(NamedAggregator::getName).collect(Collectors.toSet()); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java index 563405049f..274bc4647d 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilder.java @@ -18,19 +18,33 @@ * Translates a logical query plan into OpenSearch DSL and an appropriate request. */ public interface PushDownQueryBuilder { - boolean pushDownFilter(LogicalFilter filter); + default boolean pushDownFilter(LogicalFilter filter) { + return false; + } - boolean pushDownSort(LogicalSort sort); + default boolean pushDownSort(LogicalSort sort) { + return false; + } - boolean pushDownLimit(LogicalLimit limit); + default boolean pushDownLimit(LogicalLimit limit) { + return false; + } - boolean pushDownProject(LogicalProject project); + default boolean pushDownProject(LogicalProject project) { + return false; + } - boolean pushDownHighlight(LogicalHighlight highlight); + default boolean pushDownHighlight(LogicalHighlight highlight) { + return false; + } - boolean pushDownPageSize(LogicalPaginate paginate); + default boolean pushDownPageSize(LogicalPaginate paginate) { + return false; + } - boolean pushDownNested(LogicalNested nested); + default boolean pushDownNested(LogicalNested nested) { + return false; + } OpenSearchRequestBuilder build(); } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java index 8826295959..21618a436d 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilderTest.java @@ -86,8 +86,8 @@ void setup() { @Test void build_query_request() { - int limit = 200; - int offset = 0; + Integer limit = 200; + Integer offset = 0; requestBuilder.pushDownLimit(limit, offset); requestBuilder.pushDownTrackedScore(true); From 2cecaca2f30fa166a19142a6c88fdd1dd82ebe40 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 26 May 2023 00:42:47 -0700 Subject: [PATCH 30/38] Addressed PR feedback. Changed ExternalizablePlan back to SerializablePlan. The old name was better. Signed-off-by: MaxKsyunz --- .../sql/executor/pagination/PlanSerializer.java | 4 ++-- .../{ExternalizablePlan.java => SerializablePlan.java} | 10 +++------- .../sql/planner/physical/ProjectOperator.java | 6 +++--- .../sql/executor/pagination/PlanSerializerTest.java | 4 ++-- ...alizablePlanTest.java => SerializablePlanTest.java} | 4 ++-- .../planner/optimizer/LogicalPlanOptimizerTest.java | 2 +- .../sql/planner/physical/ProjectOperatorTest.java | 5 ----- .../java/org/opensearch/sql/utils/TestOperator.java | 4 ++-- .../executor/protector/ResourceMonitorPlan.java | 8 ++++---- .../opensearch/storage/scan/OpenSearchIndexScan.java | 4 ++-- .../executor/OpenSearchExecutionEngineTest.java | 4 ++-- .../opensearch/executor/ResourceMonitorPlanTest.java | 4 ++-- 12 files changed, 25 insertions(+), 34 deletions(-) rename core/src/main/java/org/opensearch/sql/planner/{ExternalizablePlan.java => SerializablePlan.java} (77%) rename core/src/test/java/org/opensearch/sql/planner/{ExternalizablePlanTest.java => SerializablePlanTest.java} (92%) diff --git a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java index 276d0515dc..07cf174d73 100644 --- a/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java +++ b/core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java @@ -19,7 +19,7 @@ import java.util.zip.GZIPOutputStream; import lombok.RequiredArgsConstructor; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; @@ -40,7 +40,7 @@ public class PlanSerializer { public Cursor convertToCursor(PhysicalPlan plan) { try { return new Cursor(CURSOR_PREFIX - + serialize(((ExternalizablePlan) plan).getPlanForSerialization())); + + serialize(((SerializablePlan) plan).getPlanForSerialization())); // ClassCastException thrown when a plan in the tree doesn't implement SerializablePlan } catch (NotSerializableException | ClassCastException | NoCursorException e) { return Cursor.None; diff --git a/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java similarity index 77% rename from core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java rename to core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java index 22412679bb..83eb0c6fc1 100644 --- a/core/src/main/java/org/opensearch/sql/planner/ExternalizablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java @@ -6,10 +6,6 @@ package org.opensearch.sql.planner; import java.io.Externalizable; -import java.io.IOException; -import java.io.ObjectInput; -import java.io.ObjectOutput; -import org.opensearch.sql.executor.pagination.PlanSerializer; /** * All subtypes of PhysicalPlan which needs to be serialized (in cursor, for pagination feature) @@ -23,11 +19,11 @@ * *
  • * Overwrite {@link #getPlanForSerialization} to return - * another instance of {@link ExternalizablePlan}. + * another instance of {@link SerializablePlan}. *
    • * */ -public interface ExternalizablePlan extends Externalizable { +public interface SerializablePlan extends Externalizable { /** * Override to return child or delegated plan, so parent plan should skip this one @@ -42,7 +38,7 @@ public interface ExternalizablePlan extends Externalizable { * It is needed to skip a `ResourceMonitorPlan` instance only, actually. * @return Next plan for serialization. */ - default ExternalizablePlan getPlanForSerialization() { + default SerializablePlan getPlanForSerialization() { return this; } } diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java index 3a0da4f0dc..1699c97c15 100644 --- a/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java +++ b/core/src/main/java/org/opensearch/sql/planner/physical/ProjectOperator.java @@ -25,7 +25,7 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.NamedExpression; import org.opensearch.sql.expression.parse.ParseExpression; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; /** * Project the fields specified in {@link ProjectOperator#projectList} from input. @@ -33,7 +33,7 @@ @ToString @EqualsAndHashCode(callSuper = false) @AllArgsConstructor -public class ProjectOperator extends PhysicalPlan implements ExternalizablePlan { +public class ProjectOperator extends PhysicalPlan implements SerializablePlan { @Getter private PhysicalPlan input; @Getter @@ -116,6 +116,6 @@ public void readExternal(ObjectInput in) throws IOException, ClassNotFoundExcept @Override public void writeExternal(ObjectOutput out) throws IOException { out.writeObject(projectList); - out.writeObject(((ExternalizablePlan) input).getPlanForSerialization()); + out.writeObject(((SerializablePlan) input).getPlanForSerialization()); } } diff --git a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java index ace4d29a56..8211a3bc12 100644 --- a/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java +++ b/core/src/test/java/org/opensearch/sql/executor/pagination/PlanSerializerTest.java @@ -25,7 +25,7 @@ import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.StorageEngine; import org.opensearch.sql.utils.TestOperator; @@ -144,7 +144,7 @@ void serialize_and_deserialize() { @Test void convertToCursor_and_convertToPlan() { var plan = new TestOperator(100500); - var roundTripPlan = (ExternalizablePlan) + var roundTripPlan = (SerializablePlan) planCache.convertToPlan(planCache.convertToCursor(plan).toString()); assertEquals(plan, roundTripPlan); assertNotSame(plan, roundTripPlan); diff --git a/core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java b/core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java similarity index 92% rename from core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java rename to core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java index 9055fff605..8073445dc0 100644 --- a/core/src/test/java/org/opensearch/sql/planner/ExternalizablePlanTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/SerializablePlanTest.java @@ -18,9 +18,9 @@ @ExtendWith(MockitoExtension.class) @DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class) -public class ExternalizablePlanTest { +public class SerializablePlanTest { @Mock(answer = CALLS_REAL_METHODS) - ExternalizablePlan plan; + SerializablePlan plan; @Test void getPlanForSerialization_defaults_to_self() { diff --git a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java index 5370f0c0be..faedb88111 100644 --- a/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizerTest.java @@ -397,7 +397,7 @@ void table_scan_builder_support_offset_push_down_can_apply_its_rule() { var relation = new LogicalRelation("schema", table); var optimized = LogicalPlanOptimizer.create() .optimize(new LogicalPaginate(42, List.of(project(relation)))); - // `optimized` structure: LogicalPaginate -> LogicalProject -> TableScanBuilder + // `optimized` structure: LogicalProject -> TableScanBuilder // LogicalRelation replaced by a TableScanBuilder instance assertEquals(project(tableScanBuilder), optimized); } diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java index c34b5a56a7..f5ecf76bd0 100644 --- a/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/physical/ProjectOperatorTest.java @@ -23,13 +23,9 @@ import com.google.common.collect.ImmutableMap; import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.IOException; -import java.io.ObjectInput; import java.io.ObjectInputStream; -import java.io.ObjectOutput; import java.io.ObjectOutputStream; import java.util.List; -import lombok.EqualsAndHashCode; import lombok.SneakyThrows; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -40,7 +36,6 @@ import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.expression.DSL; -import org.opensearch.sql.planner.ExternalizablePlan; import org.opensearch.sql.utils.TestOperator; @ExtendWith(MockitoExtension.class) diff --git a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java index 265800218d..6abd31b339 100644 --- a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java +++ b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java @@ -7,11 +7,11 @@ import lombok.Setter; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.exception.NoCursorException; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; -public class TestOperator extends PhysicalPlan implements ExternalizablePlan { +public class TestOperator extends PhysicalPlan implements SerializablePlan { private int field; @Setter private boolean throwNoCursorOnWrite = false; diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java index a65c72ff7e..0ec4d743b3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/protector/ResourceMonitorPlan.java @@ -15,7 +15,7 @@ import lombok.ToString; import org.opensearch.sql.data.model.ExprValue; import org.opensearch.sql.monitor.ResourceMonitor; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -25,7 +25,7 @@ @ToString @RequiredArgsConstructor @EqualsAndHashCode(callSuper = false) -public class ResourceMonitorPlan extends PhysicalPlan implements ExternalizablePlan { +public class ResourceMonitorPlan extends PhysicalPlan implements SerializablePlan { /** * How many method calls to delegate's next() to perform resource check once. @@ -93,8 +93,8 @@ public long getTotalHits() { } @Override - public ExternalizablePlan getPlanForSerialization() { - return (ExternalizablePlan) delegate; + public SerializablePlan getPlanForSerialization() { + return (SerializablePlan) delegate; } /** diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 050f0ae71b..5753495cba 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -24,7 +24,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchScrollRequest; import org.opensearch.sql.opensearch.response.OpenSearchResponse; import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.storage.TableScanOperator; /** @@ -32,7 +32,7 @@ */ @EqualsAndHashCode(onlyExplicitlyIncluded = true, callSuper = false) @ToString(onlyExplicitlyIncluded = true) -public class OpenSearchIndexScan extends TableScanOperator implements ExternalizablePlan { +public class OpenSearchIndexScan extends TableScanOperator implements SerializablePlan { /** OpenSearch client. */ private OpenSearchClient client; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java index 63e0b8e08f..330793a5d6 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngineTest.java @@ -51,7 +51,7 @@ import org.opensearch.sql.opensearch.request.OpenSearchRequest; import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexScan; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.sql.storage.split.Split; @@ -259,7 +259,7 @@ public void onFailure(Exception e) { } @RequiredArgsConstructor - private static class FakePhysicalPlan extends TableScanOperator implements ExternalizablePlan { + private static class FakePhysicalPlan extends TableScanOperator implements SerializablePlan { private final Iterator it; private boolean hasOpen; private boolean hasClosed; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java index 53ddc40cb3..0b9f302ceb 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/executor/ResourceMonitorPlanTest.java @@ -21,7 +21,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.sql.monitor.ResourceMonitor; import org.opensearch.sql.opensearch.executor.protector.ResourceMonitorPlan; -import org.opensearch.sql.planner.ExternalizablePlan; +import org.opensearch.sql.planner.SerializablePlan; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor; @@ -119,7 +119,7 @@ void getTotalHitsSuccess() { @Test void getPlanForSerialization() { - plan = mock(PhysicalPlan.class, withSettings().extraInterfaces(ExternalizablePlan.class)); + plan = mock(PhysicalPlan.class, withSettings().extraInterfaces(SerializablePlan.class)); monitorPlan = new ResourceMonitorPlan(plan, resourceMonitor); assertEquals(plan, monitorPlan.getPlanForSerialization()); } From 840c4a9dbc1464c6b4776dd847f8a1b4ee2be043 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 26 May 2023 09:53:59 -0700 Subject: [PATCH 31/38] Minor cleanup. Signed-off-by: MaxKsyunz --- .../opensearch/sql/utils/TestOperator.java | 5 +++ .../scan/PushDownQueryBuilderTest.java | 42 +++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java diff --git a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java index 6abd31b339..584cf6f3fd 100644 --- a/core/src/test/java/org/opensearch/sql/utils/TestOperator.java +++ b/core/src/test/java/org/opensearch/sql/utils/TestOperator.java @@ -1,3 +1,8 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + package org.opensearch.sql.utils; import java.io.IOException; diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java new file mode 100644 index 0000000000..351814208f --- /dev/null +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java @@ -0,0 +1,42 @@ +package org.opensearch.sql.opensearch.storage.scan; + + +import static org.junit.jupiter.api.Assertions.assertAll; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.mockito.Mockito.mock; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.junit.jupiter.MockitoExtension; +import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder; +import org.opensearch.sql.planner.logical.LogicalFilter; +import org.opensearch.sql.planner.logical.LogicalHighlight; +import org.opensearch.sql.planner.logical.LogicalLimit; +import org.opensearch.sql.planner.logical.LogicalNested; +import org.opensearch.sql.planner.logical.LogicalPaginate; +import org.opensearch.sql.planner.logical.LogicalProject; +import org.opensearch.sql.planner.logical.LogicalSort; + +@ExtendWith(MockitoExtension.class) +class PushDownQueryBuilderTest { + @Test + void default_implementations() { + var sample = new PushDownQueryBuilder() { + @Override + public OpenSearchRequestBuilder build() { + return null; + } + }; + assertAll( + () -> assertFalse(sample.pushDownFilter(mock(LogicalFilter.class))), + () -> assertFalse(sample.pushDownProject(mock(LogicalProject.class))), + () -> assertFalse(sample.pushDownHighlight(mock(LogicalHighlight.class))), + () -> assertFalse(sample.pushDownSort(mock(LogicalSort.class))), + () -> assertFalse(sample.pushDownNested(mock(LogicalNested.class))), + () -> assertFalse(sample.pushDownLimit(mock(LogicalLimit.class))), + () -> assertFalse(sample.pushDownPageSize(mock(LogicalPaginate.class))) + + ); + } + +} From 4c9e958864254eaaba7ee0271dd280eb626d1ad6 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 26 May 2023 11:58:07 -0700 Subject: [PATCH 32/38] Update core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java Co-authored-by: Yury-Fridlyand Signed-off-by: Max Ksyunz --- .../java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java index 9bd9b58274..c0e253ca50 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java @@ -31,7 +31,7 @@ @UtilityClass public class LogicalPlanDSL { - public static LogicalPlan cursor(String cursor, StorageEngine engine) { + public static LogicalPlan fetchCursor(String cursor, StorageEngine engine) { return new LogicalFetchCursor(cursor, engine); } From 102706b1910990a4fbb1f3d5ef20478981bc2b29 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 26 May 2023 11:58:17 -0700 Subject: [PATCH 33/38] Update core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java Co-authored-by: Yury-Fridlyand Signed-off-by: Max Ksyunz --- .../opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java index 04159f4b93..796fb50f26 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java @@ -109,7 +109,7 @@ public R visitPaginate(LogicalPaginate plan, C context) { return visitNode(plan, context); } - public R visitCursor(LogicalFetchCursor plan, C context) { + public R visitFetchCursor(LogicalFetchCursor plan, C context) { return visitNode(plan, context); } } From 2f4b48b8e0e6eeb6a09c21fee8f51e8352b1a811 Mon Sep 17 00:00:00 2001 From: Max Ksyunz Date: Fri, 26 May 2023 11:58:28 -0700 Subject: [PATCH 34/38] Update core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java Co-authored-by: Yury-Fridlyand Signed-off-by: Max Ksyunz --- core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java index 07070e36df..b9761cb442 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java @@ -18,7 +18,7 @@ public class FetchCursor extends UnresolvedPlan { @Override public T accept(AbstractNodeVisitor nodeVisitor, C context) { - return nodeVisitor.visitCursor(this, context); + return nodeVisitor.visitFetchCursor(this, context); } @Override From 33ad6dd9e36954c7e9c6e1fc0ab553743ba9fc2b Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Fri, 26 May 2023 17:59:39 -0700 Subject: [PATCH 35/38] Minor cleanup 2 Signed-off-by: MaxKsyunz --- .../java/org/opensearch/sql/analysis/Analyzer.java | 2 +- .../opensearch/sql/ast/AbstractNodeVisitor.java | 2 +- .../sql/executor/execution/QueryPlanFactory.java | 2 +- .../opensearch/sql/planner/DefaultImplementor.java | 2 +- .../opensearch/sql/planner/SerializablePlan.java | 4 ++++ .../sql/planner/logical/LogicalFetchCursor.java | 2 +- .../sql/planner/DefaultImplementorTest.java | 2 +- .../java/org/opensearch/sql/ppl/StandaloneIT.java | 5 ++--- .../executor/OpenSearchExecutionEngine.java | 1 - .../request/OpenSearchScrollRequest.java | 1 + .../storage/scan/OpenSearchIndexScan.java | 13 +++++-------- .../storage/scan/PushDownQueryBuilderTest.java | 14 +++++++------- .../sql/plugin/config/OpenSearchPluginModule.java | 5 ++--- 13 files changed, 27 insertions(+), 28 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java index 54c60d6054..02b97baa93 100644 --- a/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java +++ b/core/src/main/java/org/opensearch/sql/analysis/Analyzer.java @@ -589,7 +589,7 @@ private SortOption analyzeSortOption(List fieldArgs) { } @Override - public LogicalPlan visitCursor(FetchCursor cursor, AnalysisContext context) { + public LogicalPlan visitFetchCursor(FetchCursor cursor, AnalysisContext context) { return new LogicalFetchCursor(cursor.getCursor(), dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine()); } diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java index 636355f547..beb4833d4d 100644 --- a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java @@ -301,7 +301,7 @@ public T visitPaginate(Paginate paginate, C context) { return visitChildren(paginate, context); } - public T visitCursor(FetchCursor cursor, C context) { + public T visitFetchCursor(FetchCursor cursor, C context) { return visit(cursor, context); } } diff --git a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java index 5a407ad0c2..cc53f5060b 100644 --- a/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java +++ b/core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java @@ -74,7 +74,7 @@ public AbstractPlan create( } /** - * Creates a ContinuePaginatedPlan from a cursor. + * Creates a QueryPlan from a cursor. */ public AbstractPlan create(String cursor, boolean isExplain, ResponseListener queryResponseListener, diff --git a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java index fec136ef7a..a1897245ea 100644 --- a/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java +++ b/core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java @@ -151,7 +151,7 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) { } @Override - public PhysicalPlan visitCursor(LogicalFetchCursor plan, C context) { + public PhysicalPlan visitFetchCursor(LogicalFetchCursor plan, C context) { return new PlanSerializer(plan.getEngine()).convertToPlan(plan.getCursor()); } diff --git a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java index 83eb0c6fc1..ab195da5bf 100644 --- a/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java +++ b/core/src/main/java/org/opensearch/sql/planner/SerializablePlan.java @@ -36,6 +36,10 @@ public interface SerializablePlan extends Externalizable { * * In that case only plans A and C should be attempted to serialize. * It is needed to skip a `ResourceMonitorPlan` instance only, actually. + * + * 
    {@code
+   *    * A.writeObject(B.getPlanForSerialization());
+   *  }
    * @return Next plan for serialization. */ default SerializablePlan getPlanForSerialization() { diff --git a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java index 554381335e..d9a426dfe7 100644 --- a/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java +++ b/core/src/main/java/org/opensearch/sql/planner/logical/LogicalFetchCursor.java @@ -33,6 +33,6 @@ public LogicalFetchCursor(String cursor, StorageEngine engine) { @Override public R accept(LogicalPlanNodeVisitor visitor, C context) { - return visitor.visitCursor(this, context); + return visitor.visitFetchCursor(this, context); } } diff --git a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java index c88fd20d40..d43cb89a3e 100644 --- a/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java +++ b/core/src/test/java/org/opensearch/sql/planner/DefaultImplementorTest.java @@ -229,7 +229,7 @@ void visitLogicalCursor_deserializes_it() { var engine = Mockito.mock(StorageEngine.class); var physicalPlan = new TestOperator(); - var logicalPlan = LogicalPlanDSL.cursor(new PlanSerializer(engine) + var logicalPlan = LogicalPlanDSL.fetchCursor(new PlanSerializer(engine) .convertToCursor(physicalPlan).toString(), engine); assertEquals(physicalPlan, logicalPlan.accept(implementor, null)); } diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java index 5a4e38b6b5..b1fcbf7d1b 100644 --- a/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java @@ -229,13 +229,12 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl } @Provides - public PlanSerializer paginatedPlanCache(StorageEngine storageEngine) { + public PlanSerializer planSerializer(StorageEngine storageEngine) { return new PlanSerializer(storageEngine); } @Provides - public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine, - PlanSerializer planSerializer) { + public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index b2eee3850f..f63eb9e204 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -82,5 +82,4 @@ public ExplainResponseNode visitTableScan(TableScanOperator node, Object context } }); } - } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index fecc377e7a..9c544ba7a3 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -68,6 +68,7 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { /** Default constructor for Externalizable only. */ + @Deprecated(since = "introduction") public OpenSearchScrollRequest() { } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java index 5753495cba..3633e45449 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java @@ -126,18 +126,15 @@ public void readExternal(ObjectInput in) throws IOException { byte[] requestStream = new byte[reqSize]; in.read(requestStream); - var engine = (OpenSearchStorageEngine) ((PlanSerializer.CursorDeserializationStream) in) .resolveObject("engine"); try (BytesStreamInput bsi = new BytesStreamInput(requestStream)) { - request = new OpenSearchScrollRequest(bsi, engine); } maxResponseSize = in.readInt(); client = engine.getClient(); - } @Override @@ -145,16 +142,16 @@ public void writeExternal(ObjectOutput out) throws IOException { if (!request.hasAnotherBatch()) { throw new NoCursorException(); } - + // request is not directly Serializable so.. + // 1. Serialize request to an opensearch byte stream. BytesStreamOutput reqOut = new BytesStreamOutput(); request.writeTo(reqOut); reqOut.flush(); - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - reqOut.bytes().writeTo(baos); - baos.flush(); + // 2. Extract byte[] from the opensearch byte stream + var reqAsBytes = reqOut.bytes().toBytesRef().bytes; - var reqAsBytes = baos.toByteArray(); + // 3. Write out the byte[] to object output stream. out.writeInt(reqAsBytes.length); out.write(reqAsBytes); diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java index 351814208f..0b0568a6b7 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/PushDownQueryBuilderTest.java @@ -28,13 +28,13 @@ public OpenSearchRequestBuilder build() { } }; assertAll( - () -> assertFalse(sample.pushDownFilter(mock(LogicalFilter.class))), - () -> assertFalse(sample.pushDownProject(mock(LogicalProject.class))), - () -> assertFalse(sample.pushDownHighlight(mock(LogicalHighlight.class))), - () -> assertFalse(sample.pushDownSort(mock(LogicalSort.class))), - () -> assertFalse(sample.pushDownNested(mock(LogicalNested.class))), - () -> assertFalse(sample.pushDownLimit(mock(LogicalLimit.class))), - () -> assertFalse(sample.pushDownPageSize(mock(LogicalPaginate.class))) + () -> assertFalse(sample.pushDownFilter(mock(LogicalFilter.class))), + () -> assertFalse(sample.pushDownProject(mock(LogicalProject.class))), + () -> assertFalse(sample.pushDownHighlight(mock(LogicalHighlight.class))), + () -> assertFalse(sample.pushDownSort(mock(LogicalSort.class))), + () -> assertFalse(sample.pushDownNested(mock(LogicalNested.class))), + () -> assertFalse(sample.pushDownLimit(mock(LogicalLimit.class))), + () -> assertFalse(sample.pushDownPageSize(mock(LogicalPaginate.class))) ); } diff --git a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java index 7415c78a8c..f301a242fb 100644 --- a/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java +++ b/plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java @@ -75,7 +75,7 @@ public ExecutionProtector protector(ResourceMonitor resourceMonitor) { } @Provides - public PlanSerializer paginatedPlanCache(StorageEngine storageEngine) { + public PlanSerializer planSerializer(StorageEngine storageEngine) { return new PlanSerializer(storageEngine); } @@ -100,8 +100,7 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl */ @Provides public QueryPlanFactory queryPlanFactory(DataSourceService dataSourceService, - ExecutionEngine executionEngine, - PlanSerializer planSerializer) { + ExecutionEngine executionEngine) { Analyzer analyzer = new Analyzer( new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository); From 0abe2983b6c84264b2f3ce2a164caebb9e81d280 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 29 May 2023 12:39:18 -0700 Subject: [PATCH 36/38] Remove assertions that no longer apply Signed-off-by: MaxKsyunz --- .../java/org/opensearch/sql/sql/PaginationBlackboxIT.java | 5 ----- 1 file changed, 5 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java index 285fc1f989..2a34dabd79 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/PaginationBlackboxIT.java @@ -92,12 +92,7 @@ public void test_pagination_blackbox() { } while(!cursor.isEmpty()); assertTrue("Paged response schema doesn't match to non-paged", schema.similar(response.getJSONArray("schema"))); -// assertEquals(0, response.getInt("total")); -// assertEquals(testReportPrefix + "Last page is not empty", -// 0, response.getInt("size")); -// assertEquals(testReportPrefix + "Last page is not empty", -// 0, response.getJSONArray("datarows").length()); assertEquals(testReportPrefix + "Paged responses return another row count that non-paged", indexSize, rowsReturned); assertTrue(testReportPrefix + "Paged accumulated result has other rows than non-paged", From 1dc3320b32f56cde155633b68a0cf6b1816f1e37 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 29 May 2023 12:40:59 -0700 Subject: [PATCH 37/38] Minor cleanup Add comment and remove unused default constructor. Signed-off-by: MaxKsyunz --- .../org/opensearch/sql/ast/tree/FetchCursor.java | 4 ++++ .../request/OpenSearchScrollRequest.java | 16 +++++----------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java index b9761cb442..aa327c295b 100644 --- a/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java +++ b/core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java @@ -10,6 +10,10 @@ import lombok.RequiredArgsConstructor; import org.opensearch.sql.ast.AbstractNodeVisitor; +/** + * An unresolved plan that represents fetching the next + * batch in paginationed plan. + */ @RequiredArgsConstructor @EqualsAndHashCode(callSuper = false) public class FetchCursor extends UnresolvedPlan { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java index 9c544ba7a3..7173eff171 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequest.java @@ -38,19 +38,19 @@ @Getter @ToString public class OpenSearchScrollRequest implements OpenSearchRequest { - private SearchRequest initialSearchRequest; + private final SearchRequest initialSearchRequest; /** Scroll context timeout. */ - private TimeValue scrollTimeout; + private final TimeValue scrollTimeout; /** * {@link OpenSearchRequest.IndexName}. */ - private IndexName indexName; + private final IndexName indexName; /** Index name. */ @EqualsAndHashCode.Exclude @ToString.Exclude - private OpenSearchExprValueFactory exprValueFactory; + private final OpenSearchExprValueFactory exprValueFactory; /** * Scroll id which is set after first request issued. Because ElasticsearchClient is shared by * multi-thread so this state has to be maintained here. @@ -64,13 +64,7 @@ public class OpenSearchScrollRequest implements OpenSearchRequest { private boolean needClean = false; @Getter - private List includes; - - /** Default constructor for Externalizable only. - */ - @Deprecated(since = "introduction") - public OpenSearchScrollRequest() { - } + private final List includes; /** Constructor. */ public OpenSearchScrollRequest(IndexName indexName, From 701bce77708df19bba543379ce70e6cd88e48002 Mon Sep 17 00:00:00 2001 From: MaxKsyunz Date: Mon, 29 May 2023 14:14:34 -0700 Subject: [PATCH 38/38] Update test to account for prior changes Signed-off-by: MaxKsyunz --- .../sql/opensearch/request/OpenSearchScrollRequestTest.java | 5 ----- 1 file changed, 5 deletions(-) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java index a80f171755..a2585620aa 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/OpenSearchScrollRequestTest.java @@ -297,11 +297,6 @@ void serialize_deserialize_needClean() { assertEquals("", newRequest.getScrollId()); } - @Test - void default_constructor() { - assertEquals(NO_SCROLL_ID, new OpenSearchScrollRequest().getScrollId()); - } - @Test void setScrollId() { request.setScrollId("test");