Add pit for join queries

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -23,6 +23,7 @@ public enum Key {
     SQL_SLOWLOG("plugins.sql.slowlog"),
     SQL_CURSOR_KEEP_ALIVE("plugins.sql.cursor.keep_alive"),
     SQL_DELETE_ENABLED("plugins.sql.delete.enabled"),
+    SQL_PAGINATION_API_SEARCH_AFTER("plugins.sql.pagination.api"),
 
     /** PPL Settings. */
     PPL_ENABLED("plugins.ppl.enabled"),

integ-test/src/test/java/org/opensearch/sql/legacy/JoinIT.java

@@ -288,6 +288,8 @@ public void hintMultiSearchCanRunFewTimesNL() throws IOException {
     Assert.assertThat(hits.length(), equalTo(42));
   }
 
+  // TODO: Fix joinWithGeoIntersectNL test when SQL_PAGINATION_API_SEARCH_AFTER is true
+  @Ignore
   @Test
   public void joinWithGeoIntersectNL() throws IOException {
 
@@ -455,7 +457,7 @@ public void joinParseCheckSelectedFieldsSplitNLConditionOrderGT() throws IOExcep
             "SELECT /*! USE_NL*/ a.firstname, a.lastname, a.gender, d.firstname, d.age  FROM %s a"
                 + " JOIN %s d on a.age < d.age WHERE (d.firstname = 'Lynn' OR d.firstname ="
                 + " 'Obrien') AND a.firstname = 'Mcgee'",
-            TEST_INDEX_PEOPLE,
+            TEST_INDEX_PEOPLE2,
             TEST_INDEX_ACCOUNT);
 
     JSONObject result = executeQuery(query);
@@ -501,7 +503,7 @@ public void joinParseCheckSelectedFieldsSplitNLConditionOrderLT() throws IOExcep
             "SELECT /*! USE_NL*/ a.firstname, a.lastname, a.gender, d.firstname, d.age  FROM %s a"
                 + " JOIN %s d on a.age > d.age WHERE (d.firstname = 'Sandoval' OR d.firstname ="
                 + " 'Hewitt') AND a.firstname = 'Fulton'",
-            TEST_INDEX_PEOPLE,
+            TEST_INDEX_PEOPLE2,
             TEST_INDEX_ACCOUNT);
 
     JSONObject result = executeQuery(query);

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

@@ -45,6 +45,7 @@ public static SearchHits executeJoinSearchAction(
     ElasticJoinExecutor executor =
         ElasticJoinExecutor.createJoinExecutor(client, joinRequestBuilder);
     executor.run();
+    joinQueryAction.getPointInTimeHandler().deletePointInTime();
     return executor.getHits();
   }
 

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

@@ -5,6 +5,9 @@
 
 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;
 import java.util.Collection;
 import java.util.HashMap;
@@ -28,9 +31,11 @@
 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.query.SqlElasticRequestBuilder;
@@ -256,23 +261,44 @@ protected void updateMetaSearchResults(SearchResponse searchResponse) {
     this.metaResults.updateTimeOut(searchResponse.isTimedOut());
   }
 
-  protected SearchResponse scrollOneTimeWithMax(
-      Client client, TableInJoinRequestBuilder tableRequest) {
-    SearchRequestBuilder scrollRequest =
-        tableRequest
-            .getRequestBuilder()
-            .setScroll(new TimeValue(60000))
-            .setSize(MAX_RESULTS_ON_ONE_FETCH);
+  public SearchResponse getResponseWithHits(
+      Client client,
+      TableInJoinRequestBuilder tableRequest,
+      int size,
+      SearchResponse previousResponse) {
+    // Set Size
+    SearchRequestBuilder request = tableRequest.getRequestBuilder().setSize(size);
+
+    // Set sort field for search_after
     boolean ordered = tableRequest.getOriginalSelect().isOrderdSelect();
     if (!ordered) {
-      scrollRequest.addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC);
+      request.addSort(FieldSortBuilder.DOC_FIELD_NAME, SortOrder.ASC);
     }
-    SearchResponse responseWithHits = scrollRequest.get();
-    // on ordered select - not using SCAN , elastic returns hits on first scroll
-    // es5.0 elastic always return docs on scan
-    //  if(!ordered)
-    //  responseWithHits = client.prepareSearchScroll(responseWithHits.getScrollId())
-    //  .setScroll(new TimeValue(600000)).get();
+
+    SearchResponse responseWithHits;
+    // Set PIT or scroll
+    if (LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER)) {
+      request.setPointInTime(new PointInTimeBuilder(tableRequest.getPitId()));
+      if (previousResponse != null) {
+        request.searchAfter(previousResponse.getHits().getSortFields());
+      }
+      responseWithHits = request.get();
+    } else {
+      if (previousResponse == null) {
+        responseWithHits =
+            client
+                .prepareSearchScroll(previousResponse.getScrollId())
+                .setScroll(
+                    new TimeValue(LocalClusterState.state().getSettingValue(SQL_CURSOR_KEEP_ALIVE)))
+                .execute()
+                .actionGet();
+      } else {
+        request.setScroll(
+            new TimeValue(LocalClusterState.state().getSettingValue(SQL_CURSOR_KEEP_ALIVE)));
+        responseWithHits = request.get();
+      }
+    }
+
     return responseWithHits;
   }
 }

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

@@ -20,7 +20,6 @@
 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.index.mapper.MapperService;
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
@@ -124,16 +123,13 @@ private List<SearchHit> createCombinedResults(TableInJoinRequestBuilder secondTa
     Integer hintLimit = secondTableRequest.getHintLimit();
     SearchResponse searchResponse;
     boolean finishedScrolling;
+
     if (hintLimit != null && hintLimit < MAX_RESULTS_ON_ONE_FETCH) {
-      searchResponse = secondTableRequest.getRequestBuilder().setSize(hintLimit).get();
+      searchResponse = getResponseWithHits(client, secondTableRequest, hintLimit, null);
       finishedScrolling = true;
     } else {
       searchResponse =
-          secondTableRequest
-              .getRequestBuilder()
-              .setScroll(new TimeValue(60000))
-              .setSize(MAX_RESULTS_ON_ONE_FETCH)
-              .get();
+          getResponseWithHits(client, secondTableRequest, MAX_RESULTS_ON_ONE_FETCH, null);
       // es5.0 no need to scroll again!
       //            searchResponse = client.prepareSearchScroll(searchResponse.getScrollId())
       //            .setScroll(new TimeValue(600000)).get();
@@ -214,11 +210,8 @@ private List<SearchHit> createCombinedResults(TableInJoinRequestBuilder secondTa
         if (secondTableHits.length > 0
             && (hintLimit == null || fetchedSoFarFromSecondTable >= hintLimit)) {
           searchResponse =
-              client
-                  .prepareSearchScroll(searchResponse.getScrollId())
-                  .setScroll(new TimeValue(600000))
-                  .execute()
-                  .actionGet();
+              getResponseWithHits(
+                  client, secondTableRequest, MAX_RESULTS_ON_ONE_FETCH, searchResponse);
         } else {
           break;
         }
@@ -292,12 +285,13 @@ private List<SearchHit> fetchAllHits(TableInJoinRequestBuilder tableInJoinReques
 
   private List<SearchHit> scrollTillLimit(
       TableInJoinRequestBuilder tableInJoinRequest, Integer hintLimit) {
-    SearchResponse scrollResp = scrollOneTimeWithMax(client, tableInJoinRequest);
+    SearchResponse response =
+        getResponseWithHits(client, tableInJoinRequest, MAX_RESULTS_ON_ONE_FETCH, null);
 
-    updateMetaSearchResults(scrollResp);
+    updateMetaSearchResults(response);
     List<SearchHit> hitsWithScan = new ArrayList<>();
     int curentNumOfResults = 0;
-    SearchHit[] hits = scrollResp.getHits().getHits();
+    SearchHit[] hits = response.getHits().getHits();
 
     if (hintLimit == null) {
       hintLimit = MAX_RESULTS_FOR_FIRST_TABLE;
@@ -311,13 +305,9 @@ private List<SearchHit> scrollTillLimit(
         System.out.println("too many results for first table, stoping at:" + curentNumOfResults);
         break;
       }
-      scrollResp =
-          client
-              .prepareSearchScroll(scrollResp.getScrollId())
-              .setScroll(new TimeValue(600000))
-              .execute()
-              .actionGet();
-      hits = scrollResp.getHits().getHits();
+      response =
+          getResponseWithHits(client, tableInJoinRequest, MAX_RESULTS_FOR_FIRST_TABLE, response);
+      hits = response.getHits().getHits();
     }
     return hitsWithScan;
   }

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

@@ -18,7 +18,6 @@
 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.index.mapper.MapperService;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
@@ -111,11 +110,19 @@ protected List<SearchHit> innerRun() throws SqlParseException {
           if (!BackOffRetryStrategy.isHealthy()) {
             throw new IllegalStateException("Memory circuit is broken");
           }
-          firstTableResponse =
-              client
-                  .prepareSearchScroll(firstTableResponse.getScrollId())
-                  .setScroll(new TimeValue(600000))
-                  .get();
+          Integer hintLimit = nestedLoopsRequest.getFirstTable().getHintLimit();
+          if (hintLimit != null && hintLimit < MAX_RESULTS_ON_ONE_FETCH) {
+            firstTableResponse =
+                getResponseWithHits(
+                    client, nestedLoopsRequest.getFirstTable(), hintLimit, firstTableResponse);
+          } else {
+            firstTableResponse =
+                getResponseWithHits(
+                    client,
+                    nestedLoopsRequest.getFirstTable(),
+                    MAX_RESULTS_ON_ONE_FETCH,
+                    firstTableResponse);
+          }
         } else {
           finishedWithFirstTable = true;
         }
@@ -287,12 +294,11 @@ private FetchWithScrollResponse firstFetch(TableInJoinRequestBuilder tableReques
     boolean needScrollForFirstTable = false;
     SearchResponse responseWithHits;
     if (hintLimit != null && hintLimit < MAX_RESULTS_ON_ONE_FETCH) {
-
       responseWithHits = tableRequest.getRequestBuilder().setSize(hintLimit).get();
       needScrollForFirstTable = false;
     } else {
       // scroll request with max.
-      responseWithHits = scrollOneTimeWithMax(client, tableRequest);
+      responseWithHits = getResponseWithHits(client, tableRequest, MAX_RESULTS_ON_ONE_FETCH, null);
       if (responseWithHits.getHits().getTotalHits() != null
           && responseWithHits.getHits().getTotalHits().value < MAX_RESULTS_ON_ONE_FETCH) {
         needScrollForFirstTable = true;

legacy/src/main/java/org/opensearch/sql/legacy/pit/PointInTimeHandler.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.legacy.pit;
+
+import lombok.Getter;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.action.search.CreatePitRequest;
+import org.opensearch.action.search.CreatePitResponse;
+import org.opensearch.action.search.DeletePitRequest;
+import org.opensearch.client.Client;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.core.action.ActionListener;
+
+public class PointInTimeHandler {
+  private Client client;
+  @Getter private String pitId;
+  private static final Logger LOG = LogManager.getLogger();
+
+  public PointInTimeHandler(Client client, String[] indices) {
+    this.client = client;
+
+    CreatePitRequest createPitRequest = new CreatePitRequest(new TimeValue(600000), false, indices);
+    client.createPit(
+        createPitRequest,
+        new ActionListener<>() {
+          @Override
+          public void onResponse(CreatePitResponse createPitResponse) {
+            pitId = createPitResponse.getId();
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            LOG.error("Error occurred while creating PIT", e);
+          }
+        });
+  }
+
+  public void deletePointInTime() {
+    DeletePitRequest deletePitRequest = new DeletePitRequest(pitId);
+    client.deletePits(
+        deletePitRequest,
+        new ActionListener<>() {
+          @Override
+          public void onResponse(org.opensearch.action.search.DeletePitResponse deletePitResponse) {
+            LOG.debug(deletePitResponse);
+          }
+
+          @Override
+          public void onFailure(Exception e) {
+            LOG.error("Error occurred while deleting PIT", e);
+          }
+        });
+  }
+}

legacy/src/main/java/org/opensearch/sql/legacy/query/join/JoinRequestBuilder.java

@@ -7,6 +7,7 @@
 
 import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
 import java.io.IOException;
+import lombok.Setter;
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionRequestBuilder;
 import org.opensearch.action.search.MultiSearchRequest;
@@ -25,6 +26,7 @@ public class JoinRequestBuilder implements SqlElasticRequestBuilder {
   private TableInJoinRequestBuilder secondTable;
   private SQLJoinTableSource.JoinType joinType;
   private int totalLimit;
+  @Setter private String pitId;
 
   public JoinRequestBuilder() {
     firstTable = new TableInJoinRequestBuilder();

legacy/src/main/java/org/opensearch/sql/legacy/query/join/OpenSearchJoinQueryAction.java

@@ -7,12 +7,14 @@
 
 import java.util.List;
 import org.opensearch.client.Client;
+import org.opensearch.common.util.ArrayUtils;
 import org.opensearch.sql.legacy.domain.Field;
 import org.opensearch.sql.legacy.domain.JoinSelect;
 import org.opensearch.sql.legacy.domain.Select;
 import org.opensearch.sql.legacy.domain.TableOnJoinSelect;
 import org.opensearch.sql.legacy.domain.hints.Hint;
 import org.opensearch.sql.legacy.exception.SqlParseException;
+import org.opensearch.sql.legacy.pit.PointInTimeHandler;
 import org.opensearch.sql.legacy.query.DefaultQueryAction;
 import org.opensearch.sql.legacy.query.QueryAction;
 import org.opensearch.sql.legacy.query.SqlElasticRequestBuilder;
@@ -23,6 +25,7 @@
 public abstract class OpenSearchJoinQueryAction extends QueryAction {
 
   protected JoinSelect joinSelect;
+  private PointInTimeHandler pointInTimeHandler;
 
   public OpenSearchJoinQueryAction(Client client, JoinSelect joinSelect) {
     super(client, joinSelect);
@@ -34,9 +37,14 @@ public SqlElasticRequestBuilder explain() throws SqlParseException {
     JoinRequestBuilder requestBuilder = createSpecificBuilder();
     fillBasicJoinRequestBuilder(requestBuilder);
     fillSpecificRequestBuilder(requestBuilder);
+    updateRequestWithPitId(requestBuilder);
     return requestBuilder;
   }
 
+  public PointInTimeHandler getPointInTimeHandler() {
+    return pointInTimeHandler;
+  }
+
   protected abstract void fillSpecificRequestBuilder(JoinRequestBuilder requestBuilder)
       throws SqlParseException;
 
@@ -92,6 +100,18 @@ protected void updateRequestWithHints(JoinRequestBuilder requestBuilder) {
     }
   }
 
+  private void updateRequestWithPitId(JoinRequestBuilder requestBuilder) {
+    String[] indices =
+        ArrayUtils.concat(
+            requestBuilder.getFirstTable().getOriginalSelect().getIndexArr(),
+            requestBuilder.getSecondTable().getOriginalSelect().getIndexArr());
+    this.pointInTimeHandler = new PointInTimeHandler(client, indices);
+    String pitId = pointInTimeHandler.getPitId();
+    requestBuilder.setPitId(pitId);
+    requestBuilder.getFirstTable().setPitId(pitId);
+    requestBuilder.getSecondTable().setPitId(pitId);
+  }
+
   private Config queryPlannerConfig(JoinRequestBuilder requestBuilder) {
     return ((HashJoinQueryPlanRequestBuilder) requestBuilder).getConfig();
   }

legacy/src/main/java/org/opensearch/sql/legacy/query/join/TableInJoinRequestBuilder.java

@@ -17,6 +17,7 @@ public class TableInJoinRequestBuilder {
   private List<Field> returnedFields;
   private Select originalSelect;
   private Integer hintLimit;
+  private String pitId;
 
   public TableInJoinRequestBuilder() {}
 
@@ -59,4 +60,12 @@ public Integer getHintLimit() {
   public void setHintLimit(Integer hintLimit) {
     this.hintLimit = hintLimit;
   }
+
+  public String getPitId() {
+    return pitId;
+  }
+
+  public void setPitId(String pitId) {
+    this.pitId = pitId;
+  }
 }