Merge OpenSearchPagedIndexScan and OpenSearchIndexScan (#1600)

Signed-off-by: MaxKsyunz <[email protected]>
Signed-off-by: Max Ksyunz <[email protected]>
Co-authored-by: Yury-Fridlyand <[email protected]>

Author: Max Ksyunz

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;
@@ -44,6 +45,7 @@
 import org.opensearch.sql.ast.tree.Aggregation;
 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;
@@ -64,7 +66,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;
@@ -84,6 +85,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.LogicalML;
@@ -211,7 +213,6 @@ public LogicalPlan visitTableFunction(TableFunction node, AnalysisContext contex
         tableFunctionImplementation.applyArguments());
   }
 
-
   @Override
   public LogicalPlan visitLimit(Limit node, AnalysisContext context) {
     LogicalPlan child = node.getChild().get(0).accept(this, context);
@@ -587,4 +588,9 @@ private SortOption analyzeSortOption(List<Argument> fieldArgs) {
     return asc ? SortOption.DEFAULT_ASC : SortOption.DEFAULT_DESC;
   }
 
+  @Override
+  public LogicalPlan visitFetchCursor(FetchCursor cursor, AnalysisContext context) {
+    return new LogicalFetchCursor(cursor.getCursor(),
+      dataSourceService.getDataSource(DEFAULT_DATASOURCE_NAME).getStorageEngine());
+  }
 }

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -43,6 +43,7 @@
 import org.opensearch.sql.ast.tree.Aggregation;
 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;
@@ -299,4 +300,8 @@ public T visitExplain(Explain node, C context) {
   public T visitPaginate(Paginate paginate, C context) {
     return visitChildren(paginate, context);
   }
+
+  public T visitFetchCursor(FetchCursor cursor, C context) {
+    return visit(cursor, context);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/tree/FetchCursor.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+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 {
+  @Getter
+  final String cursor;
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitFetchCursor(this, context);
+  }
+
+  @Override
+  public UnresolvedPlan attach(UnresolvedPlan child) {
+    throw new UnsupportedOperationException("Cursor unresolved plan does not support children");
+  }
+}

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<ExecutionEngine.QueryResponse> listener) {
-    executionEngine.execute(plan, ExecutionContext.emptyExecutionContext(), listener);
-  }
-
   /**
    * Execute the {@link UnresolvedPlan}, with {@link PlanContext} and using {@link ResponseListener}
    * to get response.

core/src/main/java/org/opensearch/sql/executor/execution/ContinuePaginatedPlan.java

@@ -17,12 +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.FetchCursor;
+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.PlanSerializer;
+import org.opensearch.sql.executor.pagination.CanPaginateVisitor;
 
 /**
  * QueryExecution Factory.
@@ -39,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
@@ -65,25 +66,28 @@ public void onFailure(Exception e) {
   /**
    * Create QueryExecution from Statement.
    */
-  public AbstractPlan createContinuePaginatedPlan(
+  public AbstractPlan create(
       Statement statement,
       Optional<ResponseListener<ExecutionEngine.QueryResponse>> queryListener,
       Optional<ResponseListener<ExecutionEngine.ExplainResponse>> explainListener) {
     return statement.accept(this, Pair.of(queryListener, explainListener));
   }
 
   /**
-   * Creates a ContinuePaginatedPlan from a cursor.
+   * Creates a QueryPlan from a cursor.
    */
-  public AbstractPlan createContinuePaginatedPlan(String cursor, boolean isExplain,
-      ResponseListener<ExecutionEngine.QueryResponse> queryResponseListener,
-      ResponseListener<ExecutionEngine.ExplainResponse> explainListener) {
+  public AbstractPlan create(String cursor, boolean isExplain,
+                             ResponseListener<ExecutionEngine.QueryResponse> queryResponseListener,
+                             ResponseListener<ExecutionEngine.ExplainResponse> explainListener) {
     QueryId queryId = QueryId.queryId();
-    var plan = new ContinuePaginatedPlan(queryId, cursor, queryService,
-            planSerializer, queryResponseListener);
+    var plan = new QueryPlan(queryId, new FetchCursor(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 +98,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());
@@ -119,7 +123,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());
   }

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.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.SerializablePlan;
 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.

core/src/main/java/org/opensearch/sql/executor/pagination/PlanSerializer.java

@@ -6,9 +6,11 @@
 
 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;
+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;
@@ -148,6 +150,11 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) {
         + "implementing and optimizing logical plan with relation involved");
   }
 
+  @Override
+  public PhysicalPlan visitFetchCursor(LogicalFetchCursor 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);

core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.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)
@@ -29,21 +25,6 @@
  */
 public interface SerializablePlan 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.
-   * <pre>{@code
-   * out.writeObject(input.getPlanForSerialization());
-   * }</pre>
-   */
-  @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.
@@ -55,6 +36,10 @@ public interface SerializablePlan extends Externalizable {
    * </pre>
    * In that case only plans A and C should be attempted to serialize.
    * It is needed to skip a `ResourceMonitorPlan` instance only, actually.
+   *
+   *  <pre>{@code
+   *    * A.writeObject(B.getPlanForSerialization());
+   *  }</pre>
    * @return Next plan for serialization.
    */
   default SerializablePlan getPlanForSerialization() {

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

@@ -0,0 +1,38 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.logical;
+
+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 LogicalFetchCursor extends LogicalPlan {
+  @Getter
+  private final String cursor;
+
+  @Getter
+  private final StorageEngine engine;
+
+  /**
+   * LogicalCursor constructor. Does not have child plans.
+   */
+  public LogicalFetchCursor(String cursor, StorageEngine engine) {
+    super(List.of());
+    this.cursor = cursor;
+    this.engine = engine;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitFetchCursor(this, context);
+  }
+}