Add pit for multi query

legacy/src/main/java/org/opensearch/sql/legacy/executor/ElasticHitsExecutor.java

@@ -5,13 +5,77 @@
 
 package org.opensearch.sql.legacy.executor;
 
+import static org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME;
+import static org.opensearch.search.sort.SortOrder.ASC;
+import static org.opensearch.sql.common.setting.Settings.Key.SQL_CURSOR_KEEP_ALIVE;
+import static org.opensearch.sql.common.setting.Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER;
+
 import java.io.IOException;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.action.search.SearchRequestBuilder;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.client.Client;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.search.SearchHits;
+import org.opensearch.search.builder.PointInTimeBuilder;
+import org.opensearch.sql.legacy.domain.Select;
+import org.opensearch.sql.legacy.esdomain.LocalClusterState;
 import org.opensearch.sql.legacy.exception.SqlParseException;
+import org.opensearch.sql.legacy.pit.PointInTimeHandler;
 
 /** Created by Eliran on 21/8/2016. */
-public interface ElasticHitsExecutor {
-  void run() throws IOException, SqlParseException;
+public abstract class ElasticHitsExecutor {
+
+  protected static final Logger LOG = LogManager.getLogger();
+  protected PointInTimeHandler pit;
+  protected Client client;
+
+  protected abstract void run() throws IOException, SqlParseException;
+
+  protected abstract SearchHits getHits();
+
+  public SearchResponse getResponseWithHits(
+      Client client,
+      SearchRequestBuilder request,
+      Select select,
+      int size,
+      SearchResponse previousResponse,
+      PointInTimeHandler pit) {
+    // Set Size
+    request.setSize(size);
+    SearchResponse responseWithHits;
+
+    if (LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER)) {
+      // Set sort field for search_after
+      boolean ordered = select.isOrderdSelect();
+      if (!ordered) {
+        request.addSort(DOC_FIELD_NAME, ASC);
+      }
+      // Set PIT
+      request.setPointInTime(new PointInTimeBuilder(pit.getPitId()));
+      // from and size is alternate method to paginate result.
+      // If select has from clause, search after is not required.
+      if (previousResponse != null && select.getFrom().isEmpty()) {
+        request.searchAfter(previousResponse.getHits().getSortFields());
+      }
+      responseWithHits = request.get();
+    } else {
+      // Set scroll
+      TimeValue keepAlive = LocalClusterState.state().getSettingValue(SQL_CURSOR_KEEP_ALIVE);
+      if (previousResponse != null) {
+        responseWithHits =
+            client
+                .prepareSearchScroll(previousResponse.getScrollId())
+                .setScroll(keepAlive)
+                .execute()
+                .actionGet();
+      } else {
+        request.setScroll(keepAlive);
+        responseWithHits = request.get();
+      }
+    }
 
-  SearchHits getHits();
+    return responseWithHits;
+  }
 }

legacy/src/main/java/org/opensearch/sql/legacy/executor/join/ElasticJoinExecutor.java

@@ -5,7 +5,6 @@
 
 package org.opensearch.sql.legacy.executor.join;
 
-import static org.opensearch.sql.common.setting.Settings.Key.SQL_CURSOR_KEEP_ALIVE;
 import static org.opensearch.sql.common.setting.Settings.Key.SQL_PAGINATION_API_SEARCH_AFTER;
 
 import java.io.IOException;
@@ -16,30 +15,22 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Stream;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.search.TotalHits.Relation;
-import org.opensearch.action.search.SearchRequestBuilder;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.Client;
 import org.opensearch.common.document.DocumentField;
-import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.rest.BytesRestResponse;
 import org.opensearch.rest.RestChannel;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
-import org.opensearch.search.builder.PointInTimeBuilder;
-import org.opensearch.search.sort.FieldSortBuilder;
-import org.opensearch.search.sort.SortOrder;
 import org.opensearch.sql.legacy.domain.Field;
 import org.opensearch.sql.legacy.esdomain.LocalClusterState;
 import org.opensearch.sql.legacy.exception.SqlParseException;
 import org.opensearch.sql.legacy.executor.ElasticHitsExecutor;
-import org.opensearch.sql.legacy.pit.PointInTimeHandler;
 import org.opensearch.sql.legacy.pit.PointInTimeHandlerImpl;
 import org.opensearch.sql.legacy.query.SqlElasticRequestBuilder;
 import org.opensearch.sql.legacy.query.join.HashJoinElasticRequestBuilder;
@@ -49,17 +40,14 @@
 import org.opensearch.sql.legacy.query.planner.HashJoinQueryPlanRequestBuilder;
 
 /** Created by Eliran on 15/9/2015. */
-public abstract class ElasticJoinExecutor implements ElasticHitsExecutor {
-  private static final Logger LOG = LogManager.getLogger();
+public abstract class ElasticJoinExecutor extends ElasticHitsExecutor {
 
   protected List<SearchHit> results; // Keep list to avoid copy to new array in SearchHits
   protected MetaSearchResult metaResults;
   protected final int MAX_RESULTS_ON_ONE_FETCH = 10000;
   private Set<String> aliasesOnReturn;
   private boolean allFieldsReturn;
-  protected Client client;
   protected String[] indices;
-  protected PointInTimeHandler pit;
 
   protected ElasticJoinExecutor(Client client, JoinRequestBuilder requestBuilder) {
     metaResults = new MetaSearchResult();
@@ -283,38 +271,14 @@ protected void updateMetaSearchResults(SearchResponse searchResponse) {
 
   public SearchResponse getResponseWithHits(
       TableInJoinRequestBuilder tableRequest, int size, SearchResponse previousResponse) {
-    // Set Size
-    SearchRequestBuilder request = tableRequest.getRequestBuilder().setSize(size);
-    SearchResponse responseWithHits;
-    if (LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER)) {
-      // Set sort field for search_after
-      boolean ordered = tableRequest.getOriginalSelect().isOrderdSelect();
-      if (!ordered) {
-        request.addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC);
-      }
-      // Set PIT
-      request.setPointInTime(new PointInTimeBuilder(pit.getPitId()));
-      if (previousResponse != null) {
-        request.searchAfter(previousResponse.getHits().getSortFields());
-      }
-      responseWithHits = request.get();
-    } else {
-      // Set scroll
-      TimeValue keepAlive = LocalClusterState.state().getSettingValue(SQL_CURSOR_KEEP_ALIVE);
-      if (previousResponse != null) {
-        responseWithHits =
-            client
-                .prepareSearchScroll(previousResponse.getScrollId())
-                .setScroll(keepAlive)
-                .execute()
-                .actionGet();
-      } else {
-        request.setScroll(keepAlive);
-        responseWithHits = request.get();
-      }
-    }
 
-    return responseWithHits;
+    return getResponseWithHits(
+        client,
+        tableRequest.getRequestBuilder(),
+        tableRequest.getOriginalSelect(),
+        size,
+        previousResponse,
+        pit);
   }
 
   public String[] getIndices(JoinRequestBuilder joinRequestBuilder) {