Fix integ test failures

Signed-off-by: Rupal Mahajan <[email protected]>
opensearch-project · Jun 14, 2024 · 87f611b · 87f611b
1 parent 5bcd134
commit 87f611b
Showing 1 changed file with 22 additions and 21 deletions.
diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java b/legacy/src/main/java/org/opensearch/sql/legacy/executor/multi/MinusExecutor.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.legacy.executor.multi;
 
+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;
 
@@ -36,6 +38,7 @@
 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.query.DefaultQueryAction;
 import org.opensearch.sql.legacy.query.multi.MultiQueryRequestBuilder;
 import org.opensearch.sql.legacy.utils.Util;
@@ -198,7 +201,7 @@ private Set<ComperableHitResult> runWithScrollings() {
             builder.getOriginalSelect(true),
             maxDocsToFetchOnEachScrollShard,
             null,
-            builder.getPit().getPitId());
+            builder.getPit());
     Set<ComperableHitResult> results = new HashSet<>();
 
     SearchHit[] hits = scrollResp.getHits().getHits();
@@ -220,7 +223,7 @@ private Set<ComperableHitResult> runWithScrollings() {
               builder.getOriginalSelect(true),
               maxDocsToFetchOnEachScrollShard,
               scrollResp,
-              builder.getPit().getPitId());
+              builder.getPit());
       hits = scrollResp.getHits().getHits();
     }
     scrollResp =
@@ -230,7 +233,7 @@ private Set<ComperableHitResult> runWithScrollings() {
             builder.getOriginalSelect(false),
             this.maxDocsToFetchOnEachScrollShard,
             null,
-            builder.getPit().getPitId());
+            builder.getPit());
 
     hits = scrollResp.getHits().getHits();
     if (hits == null || hits.length == 0) {
@@ -250,7 +253,7 @@ private Set<ComperableHitResult> runWithScrollings() {
               builder.getOriginalSelect(false),
               maxDocsToFetchOnEachScrollShard,
               scrollResp,
-              builder.getPit().getPitId());
+              builder.getPit());
       hits = scrollResp.getHits().getHits();
     }
 
@@ -321,7 +324,7 @@ private MinusOneFieldAndOptimizationResult runWithScrollingAndAddFilter(
             builder.getOriginalSelect(true),
             maxDocsToFetchOnEachScrollShard,
             null,
-            builder.getPit().getPitId());
+            builder.getPit());
     Set<Object> results = new HashSet<>();
     int currentNumOfResults = 0;
     SearchHit[] hits = scrollResp.getHits().getHits();
@@ -355,7 +358,7 @@ private MinusOneFieldAndOptimizationResult runWithScrollingAndAddFilter(
               secondQuerySelect,
               this.maxDocsToFetchOnEachScrollShard,
               null,
-              builder.getPit().getPitId());
+              builder.getPit());
       SearchHits secondQuerySearchHits = responseForSecondTable.getHits();
 
       SearchHit[] secondQueryHits = secondQuerySearchHits.getHits();
@@ -373,7 +376,7 @@ private MinusOneFieldAndOptimizationResult runWithScrollingAndAddFilter(
                 secondQuerySelect,
                 maxDocsToFetchOnEachScrollShard,
                 responseForSecondTable,
-                builder.getPit().getPitId());
+                builder.getPit());
         secondQueryHits = responseForSecondTable.getHits().getHits();
       }
       results.addAll(currentSetFromResults);
@@ -389,7 +392,7 @@ private MinusOneFieldAndOptimizationResult runWithScrollingAndAddFilter(
               builder.getOriginalSelect(true),
               maxDocsToFetchOnEachScrollShard,
               scrollResp,
-              builder.getPit().getPitId());
+              builder.getPit());
       hits = scrollResp.getHits().getHits();
     }
     return new MinusOneFieldAndOptimizationResult(results, someHit);
@@ -401,27 +404,25 @@ private SearchResponse getResponseWithHits(
       Select select,
       int size,
       SearchResponse previousResponse,
-      String pitId) {
+      PointInTimeHandler pit) {
     // Set Size
     request.setSize(size);
-
-    // Set sort field for search_after
-    boolean ordered = select.isOrderdSelect();
-    if (!ordered) {
-      request.addSort(
-          org.opensearch.search.sort.FieldSortBuilder.DOC_FIELD_NAME,
-          org.opensearch.search.sort.SortOrder.ASC);
-    }
-
     SearchResponse responseWithHits;
-    // Set PIT or scroll
+
     if (LocalClusterState.state().getSettingValue(SQL_PAGINATION_API_SEARCH_AFTER)) {
-      request.setPointInTime(new PointInTimeBuilder(pitId));
-      if (previousResponse != null) {
+      // 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()));
+      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 =