revert changes in OpenSearchIndexScan

Signed-off-by: Lantao Jin <ltjin@amazon.com>

Author: LantaoJin

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchQueryRequest.java

@@ -81,7 +81,7 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
 
   @ToString.Exclude private Map<String, Object> afterKey;
 
-  @EqualsAndHashCode.Exclude @ToString.Exclude private boolean afterKeyToReset = false;
+  @ToString.Exclude private boolean calciteEnabled;
 
   @TestOnly
   static OpenSearchQueryRequest of(
@@ -108,7 +108,8 @@ public static OpenSearchQueryRequest of(
       SearchSourceBuilder sourceBuilder,
       OpenSearchExprValueFactory factory,
       List<String> includes) {
-    return new OpenSearchQueryRequest(indexName, sourceBuilder, factory, includes, null, null);
+    return new OpenSearchQueryRequest(
+        indexName, sourceBuilder, factory, includes, null, null, false);
   }
 
   /** Build an OpenSearchQueryRequest with PIT support. */
@@ -118,9 +119,10 @@ public static OpenSearchQueryRequest pitOf(
       OpenSearchExprValueFactory factory,
       List<String> includes,
       TimeValue cursorKeepAlive,
-      String pitId) {
+      String pitId,
+      boolean calciteEnabled) {
     return new OpenSearchQueryRequest(
-        indexName, sourceBuilder, factory, includes, cursorKeepAlive, pitId);
+        indexName, sourceBuilder, factory, includes, cursorKeepAlive, pitId, calciteEnabled);
   }
 
   /** Do not new it directly, use of() and pitOf() instead. */
@@ -130,13 +132,15 @@ public static OpenSearchQueryRequest pitOf(
       OpenSearchExprValueFactory factory,
       List<String> includes,
       TimeValue cursorKeepAlive,
-      String pitId) {
+      String pitId,
+      boolean calciteEnabled) {
     this.indexName = indexName;
     this.sourceBuilder = sourceBuilder;
     this.exprValueFactory = factory;
     this.includes = includes;
     this.cursorKeepAlive = cursorKeepAlive;
     this.pitId = pitId;
+    this.calciteEnabled = calciteEnabled;
   }
 
   /**
@@ -181,7 +185,7 @@ public OpenSearchResponse search(
       Function<SearchRequest, SearchResponse> searchAction,
       Function<SearchScrollRequest, SearchResponse> scrollAction) {
     if (this.pitId == null) {
-      return search(searchAction);
+      return calciteEnabled ? search(searchAction) : searchForV2(searchAction);
     } else {
       // Search with PIT instead of scroll API
       return searchWithPIT(searchAction);
@@ -196,29 +200,41 @@ public boolean isCountAggRequest() {
         && sourceBuilder.trackTotalHitsUpTo() == Integer.MAX_VALUE;
   }
 
+  /** Call the old search logic for v2, since we don't support paginating aggregate in v2. */
+  @Deprecated
+  private OpenSearchResponse searchForV2(Function<SearchRequest, SearchResponse> searchAction) {
+    // When SearchRequest doesn't contain PitId, fetch single page request
+    if (needClean) {
+      return new OpenSearchResponse(
+          SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
+    } else {
+      // get the value before set needClean = true
+      boolean isCountAggRequest = isCountAggRequest();
+      needClean = true;
+      return new OpenSearchResponse(
+          searchAction.apply(
+              new SearchRequest().indices(indexName.getIndexNames()).source(sourceBuilder)),
+          exprValueFactory,
+          includes,
+          isCountAggRequest);
+    }
+  }
+
   private OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction) {
     OpenSearchResponse openSearchResponse;
     if (needClean) {
       openSearchResponse =
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
     } else {
-      // On first call: reset builder to clear any afterKey from other requests
+      // Set afterKey to request
       if (this.sourceBuilder.aggregations() != null) {
         this.sourceBuilder.aggregations().getAggregatorFactories().stream()
             .filter(b -> b instanceof CompositeAggregationBuilder)
-            .forEach(
-                c -> {
-                  if (!afterKeyToReset) {
-                    // First call: reset to clear any previous afterKey from shared builder
-                    // Use null instead of empty map to avoid "[after] has 0 value(s)" error
-                    ((CompositeAggregationBuilder) c).aggregateAfter(null);
-                    afterKeyToReset = true;
-                  }
-                  if (afterKey != null && !afterKey.isEmpty()) {
-                    ((CompositeAggregationBuilder) c).aggregateAfter(afterKey);
-                  }
-                });
+            .forEach(c -> ((CompositeAggregationBuilder) c).aggregateAfter(afterKey));
+        if (LOG.isDebugEnabled()) {
+          LOG.debug(sourceBuilder);
+        }
       }
 
       SearchRequest searchRequest =
@@ -230,11 +246,11 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
               this.searchResponse, exprValueFactory, includes, isCountAggRequest());
 
       needClean = openSearchResponse.isEmpty();
-      if (openSearchResponse.isCompositeAggregationResponse()) {
-        InternalComposite compositeAgg =
-            (InternalComposite) this.searchResponse.getAggregations().asList().get(0);
-        // Update afterKey from response
-        afterKey = compositeAgg.afterKey();
+      // Get afterKey from response
+      if (openSearchResponse.isAggregationResponse()) {
+        openSearchResponse.getAggregations().asList().stream()
+            .filter(b -> b instanceof InternalComposite)
+            .forEach(c -> afterKey = ((InternalComposite) c).afterKey());
       }
     }
     return openSearchResponse;

opensearch/src/main/java/org/opensearch/sql/opensearch/request/OpenSearchRequestBuilder.java

@@ -133,13 +133,18 @@ private OpenSearchRequest buildRequestWithPit(
         // Search with PIT request
         String pitId = createPit(indexName, cursorKeepAlive, client);
         return OpenSearchQueryRequest.pitOf(
-            indexName, sourceBuilder, exprValueFactory, includes, cursorKeepAlive, pitId);
+            indexName,
+            sourceBuilder,
+            exprValueFactory,
+            includes,
+            cursorKeepAlive,
+            pitId,
+            isCalciteEnabled());
       } else {
         sourceBuilder.from(startFrom);
         sourceBuilder.size(size);
         // Search with non-Pit request
-        return OpenSearchQueryRequest.pitOf(
-            indexName, sourceBuilder, exprValueFactory, includes, cursorKeepAlive, null);
+        return OpenSearchQueryRequest.of(indexName, sourceBuilder, exprValueFactory, includes);
       }
     } else {
       if (startFrom != 0) {
@@ -149,7 +154,13 @@ private OpenSearchRequest buildRequestWithPit(
       // Search with PIT request
       String pitId = createPit(indexName, cursorKeepAlive, client);
       return OpenSearchQueryRequest.pitOf(
-          indexName, sourceBuilder, exprValueFactory, includes, cursorKeepAlive, pitId);
+          indexName,
+          sourceBuilder,
+          exprValueFactory,
+          includes,
+          cursorKeepAlive,
+          pitId,
+          isCalciteEnabled());
     }
   }
 
@@ -451,4 +462,10 @@ private Supplier<NestedQueryBuilder> createEmptyNestedQuery(String path) {
   private BoolQueryBuilder query() {
     return (BoolQueryBuilder) sourceBuilder.query();
   }
+
+  private boolean isCalciteEnabled() {
+    return settings == null
+        || settings.getSettingValue(Settings.Key.CALCITE_ENGINE_ENABLED) == null
+        || Boolean.TRUE.equals(settings.getSettingValue(Settings.Key.CALCITE_ENGINE_ENABLED));
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java

@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import lombok.EqualsAndHashCode;
+import lombok.Getter;
 import lombok.ToString;
 import org.jetbrains.annotations.TestOnly;
 import org.opensearch.action.search.SearchResponse;
@@ -47,7 +48,7 @@ public class OpenSearchResponse implements Iterable<ExprValue> {
   private final SearchHits hits;
 
   /** Search aggregation result. */
-  private final Aggregations aggregations;
+  @Getter private final Aggregations aggregations;
 
   /** List of requested include fields. */
   private final List<String> includes;

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java

@@ -228,8 +228,6 @@ public TableScanBuilder createScanBuilder() {
             new OpenSearchIndexScan(
                 client,
                 requestBuilder.getMaxResponseSize(),
-                requestBuilder.getMaxResultWindow(),
-                this.getQueryBucketSize(),
                 requestBuilder.build(
                     indexName, cursorKeepAlive, client, cachedFieldOpenSearchTypes.isEmpty()));
     return new OpenSearchIndexScanBuilder(builder, createScanOperator);

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexScan.java

@@ -22,6 +22,7 @@
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.request.OpenSearchQueryRequest;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
+import org.opensearch.sql.opensearch.response.OpenSearchResponse;
 import org.opensearch.sql.opensearch.storage.OpenSearchStorageEngine;
 import org.opensearch.sql.planner.SerializablePlan;
 import org.opensearch.sql.storage.TableScanOperator;
@@ -43,33 +44,28 @@ public class OpenSearchIndexScan extends TableScanOperator implements Serializab
   /** Number of rows returned. */
   private Integer queryCount;
 
+  /** Search response for current batch. */
   private Iterator<ExprValue> iterator;
 
-  private BackgroundSearchScanner bgScanner;
-
   /** Creates index scan based on a provided OpenSearchRequestBuilder. */
   public OpenSearchIndexScan(
-      OpenSearchClient client,
-      int maxResponseSize,
-      int maxResultWindow,
-      int queryBucketSize,
-      OpenSearchRequest request) {
+      OpenSearchClient client, int maxResponseSize, OpenSearchRequest request) {
     this.maxResponseSize = maxResponseSize;
     this.client = client;
     this.request = request;
-    this.bgScanner = new BackgroundSearchScanner(client, maxResultWindow, queryBucketSize);
   }
 
   @TestOnly
   public OpenSearchIndexScan(OpenSearchClient client, OpenSearchRequest request) {
-    this(client, Integer.MAX_VALUE, 10000, 10000, request);
+    this(client, Integer.MAX_VALUE, request);
   }
 
   @Override
   public void open() {
     super.open();
+    iterator = Collections.emptyIterator();
     queryCount = 0;
-    this.bgScanner.startScanning(request);
+    fetchNextBatch();
   }
 
   @Override
@@ -84,8 +80,8 @@ public boolean hasNext() {
       return false;
     }
 
-    if (iterator == null || (!iterator.hasNext() && !this.bgScanner.isScanDone())) {
-      iterator = fetchNextBatch();
+    if (!iterator.hasNext()) {
+      fetchNextBatch();
     }
     return iterator.hasNext();
   }
@@ -101,16 +97,17 @@ public ExprValue next() {
     return iterator.next();
   }
 
-  private Iterator<ExprValue> fetchNextBatch() {
-    BackgroundSearchScanner.SearchBatchResult result = bgScanner.fetchNextBatch(request);
-    return result.iterator();
+  private void fetchNextBatch() {
+    OpenSearchResponse response = client.search(request);
+    if (!response.isEmpty()) {
+      iterator = response.iterator();
+    }
   }
 
   @Override
   public void close() {
     super.close();
-    iterator = Collections.emptyIterator();
-    bgScanner.close();
+
     client.cleanup(request);
   }