Fix doctest and IT

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

Author: LantaoJin

core/src/main/java/org/opensearch/sql/exception/NonFallbackCalciteException.java

@@ -11,4 +11,8 @@ public class NonFallbackCalciteException extends QueryEngineException {
   public NonFallbackCalciteException(String message) {
     super(message);
   }
+
+  public NonFallbackCalciteException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLIntegTestCase.java

@@ -113,6 +113,7 @@ private Settings defaultSettings() {
           new ImmutableMap.Builder<Key, Object>()
               .put(Key.QUERY_SIZE_LIMIT, 200)
               .put(Key.QUERY_BUCKET_SIZE, 1000)
+              .put(Key.SEARCH_MAX_BUCKETS, 65535)
               .put(Key.SQL_CURSOR_KEEP_ALIVE, TimeValue.timeValueMinutes(1))
               .put(Key.FIELD_TYPE_TOLERANCE, true)
               .put(Key.CALCITE_ENGINE_ENABLED, true)

integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java

@@ -167,6 +167,7 @@ private Settings defaultSettings() {
           new ImmutableMap.Builder<Key, Object>()
               .put(Key.QUERY_SIZE_LIMIT, 200)
               .put(Key.QUERY_BUCKET_SIZE, 1000)
+              .put(Key.SEARCH_MAX_BUCKETS, 65535)
               .put(Key.SQL_CURSOR_KEEP_ALIVE, TimeValue.timeValueMinutes(1))
               .put(Key.FIELD_TYPE_TOLERANCE, true)
               .build();

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

@@ -79,7 +79,9 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
 
   private SearchResponse searchResponse = null;
 
-  @ToString.Exclude private Map<String, Object> afterKey = Collections.emptyMap();
+  @ToString.Exclude private Map<String, Object> afterKey;
+
+  @EqualsAndHashCode.Exclude @ToString.Exclude private boolean afterKeyToReset = false;
 
   @TestOnly
   static OpenSearchQueryRequest of(
@@ -201,13 +203,22 @@ private OpenSearchResponse search(Function<SearchRequest, SearchResponse> search
           new OpenSearchResponse(
               SearchHits.empty(), exprValueFactory, includes, isCountAggRequest());
     } else {
+      // On first call: reset builder to clear any afterKey from other requests
       if (this.sourceBuilder.aggregations() != null) {
         this.sourceBuilder.aggregations().getAggregatorFactories().stream()
             .filter(b -> b instanceof CompositeAggregationBuilder)
-            .forEach(c -> ((CompositeAggregationBuilder) c).aggregateAfter(afterKey));
-        if (LOG.isDebugEnabled()) {
-          LOG.debug(sourceBuilder);
-        }
+            .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);
+                  }
+                });
       }
 
       SearchRequest searchRequest =

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

@@ -122,11 +122,16 @@ public boolean isCompositeAggregationResponse() {
         && (aggregations.asList().get(0) instanceof InternalComposite);
   }
 
+  /**
+   * Get the size of composite aggregation bucket. Must be called after
+   * isCompositeAggregationResponse() is true.
+   */
   public int getCompositeBucketSize() {
     if (isCompositeAggregationResponse()) {
       return ((InternalComposite) aggregations.asList().get(0)).getBuckets().size();
     }
-    throw new RuntimeException("Should not go to here");
+    assert false : "Should never call here";
+    return -1;
   }
 
   public boolean isCountResponse() {

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

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.opensearch.storage;
 
+import static org.opensearch.search.aggregations.MultiBucketConsumerService.DEFAULT_MAX_BUCKETS;
+
 import com.google.common.annotations.VisibleForTesting;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -199,9 +201,15 @@ public Integer getMaxResultWindow() {
   }
 
   public Integer getQueryBucketSize() {
-    return Math.min(
-        settings.getSettingValue(Settings.Key.QUERY_BUCKET_SIZE),
-        settings.getSettingValue(Settings.Key.SEARCH_MAX_BUCKETS));
+    return Math.min(settings.getSettingValue(Settings.Key.QUERY_BUCKET_SIZE), getMaxBuckets());
+  }
+
+  public Integer getMaxBuckets() {
+    try {
+      return settings.getSettingValue(Settings.Key.SEARCH_MAX_BUCKETS);
+    } catch (Exception e) {
+      return DEFAULT_MAX_BUCKETS;
+    }
   }
 
   /** TODO: Push down operations to index scan operator as much as possible in future. */

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

@@ -123,7 +123,8 @@ private OpenSearchResponse getCurrentResponse(OpenSearchRequest request) {
         throw new NonFallbackCalciteException(
             "Failed to fetch data from the index: the background task failed or interrupted.\n"
                 + "  Inner error: "
-                + e.getMessage());
+                + e.getMessage(),
+            e);
       }
     } else {
       return client.search(request);

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/AggPushDownAction.java

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.opensearch.storage.scan.context;
 
+import static org.opensearch.search.aggregations.MultiBucketConsumerService.DEFAULT_MAX_BUCKETS;
+
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -43,8 +45,6 @@
 @Getter
 @EqualsAndHashCode
 public class AggPushDownAction implements OSRequestBuilderAction {
-  private static final int MAX_BUCKET_SIZE = 65535;
-
   private Pair<List<AggregationBuilder>, OpenSearchAggregationResponseParser> builderAndParser;
   private final Map<String, OpenSearchDataType> extendedTypeMapping;
   private final long scriptCount;
@@ -186,13 +186,13 @@ public void rePushDownRareTop(RareTopDigest digest) {
           for (int i = 0; i < composite.sources().size(); i++) {
             TermsValuesSourceBuilder terms = (TermsValuesSourceBuilder) composite.sources().get(i);
             if (i == 0) { // first
-              aggregationBuilder = buildTermsAggregationBuilder(terms, null, MAX_BUCKET_SIZE);
+              aggregationBuilder = buildTermsAggregationBuilder(terms, null, DEFAULT_MAX_BUCKETS);
             } else if (i == composite.sources().size() - 1) { // last
               aggregationBuilder.subAggregation(
                   buildTermsAggregationBuilder(terms, bucketOrder, digest.number()));
             } else {
               aggregationBuilder.subAggregation(
-                  buildTermsAggregationBuilder(terms, null, MAX_BUCKET_SIZE));
+                  buildTermsAggregationBuilder(terms, null, DEFAULT_MAX_BUCKETS));
             }
           }
         } else {