Extend query size limit using scroll

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java

@@ -18,7 +18,6 @@
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.NamedExpression;
-import org.opensearch.sql.expression.ParseExpression;
 import org.opensearch.sql.expression.ReferenceExpression;
 import org.opensearch.sql.expression.aggregation.NamedAggregator;
 import org.opensearch.sql.expression.window.WindowDefinition;

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

@@ -31,6 +31,7 @@
 import java.nio.file.Paths;
 import java.util.Locale;
 
+import static com.google.common.base.Strings.isNullOrEmpty;
 import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient;
 import static org.opensearch.sql.legacy.TestUtils.getAccountIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getBankIndexMapping;
@@ -71,6 +72,8 @@ public abstract class SQLIntegTestCase extends OpenSearchSQLRestTestCase {
   public static final String TRANSIENT = "transient";
   public static final Integer DEFAULT_QUERY_SIZE_LIMIT =
       Integer.parseInt(System.getProperty("defaultQuerySizeLimit", "200"));
+  public static final Integer DEFAULT_MAX_RESULT_WINDOW =
+      Integer.parseInt(System.getProperty("defaultMaxResultWindow", "10000"));
 
   public boolean shouldResetQuerySizeLimit() {
     return true;
@@ -161,6 +164,15 @@ protected static void wipeAllClusterSettings() throws IOException {
     updateClusterSettings(new ClusterSetting("transient", "*", null));
   }
 
+  protected void setMaxResultWindow(String indexName, Integer window) throws IOException {
+    updateIndexSettings(indexName, "{ \"index\": { \"max_result_window\":" + window + " } }");
+  }
+
+  protected void resetMaxResultWindow(String indexName) throws IOException {
+    updateIndexSettings(indexName,
+        "{ \"index\": { \"max_result_window\": " + DEFAULT_MAX_RESULT_WINDOW + " } }");
+  }
+
   /**
    * Provide for each test to load test index, data and other setup work
    */
@@ -378,6 +390,18 @@ public String toString() {
     }
   }
 
+  protected static JSONObject updateIndexSettings(String indexName, String setting)
+      throws IOException {
+    Request request = new Request("PUT", "/" + indexName + "/_settings");
+    if (!isNullOrEmpty(setting)) {
+      request.setJsonEntity(setting);
+    }
+    RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
+    restOptionsBuilder.addHeader("Content-Type", "application/json");
+    request.setOptions(restOptionsBuilder);
+    return new JSONObject(executeRequest(request));
+  }
+
   protected String makeRequest(String query) {
     return makeRequest(query, 0);
   }

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

@@ -14,6 +14,7 @@
 import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.jupiter.api.Test;
 
 public class HeadCommandIT extends PPLIntegTestCase {
@@ -26,6 +27,7 @@ public void beforeTest() throws IOException {
   @After
   public void afterTest() throws IOException {
     resetQuerySizeLimit();
+    resetMaxResultWindow(TEST_INDEX_ACCOUNT);
   }
 
   @Override
@@ -60,6 +62,76 @@ public void testHeadWithNumber() throws IOException {
         rows("Nanette", 28));
   }
 
+  @Ignore("Fix https://github.com/opensearch-project/sql/issues/703#issuecomment-1211422130")
+  @Test
+  public void testHeadWithNumberLargerThanQuerySizeLimit() throws IOException {
+    setQuerySizeLimit(5);
+    JSONObject result =
+        executeQuery(String.format(
+            "source=%s | fields firstname, age | head 10", TEST_INDEX_ACCOUNT));
+    verifyDataRows(result,
+        rows("Amber", 32),
+        rows("Hattie", 36),
+        rows("Nanette", 28),
+        rows("Dale", 33),
+        rows("Elinor", 36),
+        rows("Virginia", 39),
+        rows("Dillard", 34),
+        rows("Mcgee", 39),
+        rows("Aurelia", 37),
+        rows("Fulton", 23));
+  }
+
+  @Test
+  public void testHeadWithNumberLargerThanMaxResultWindow() throws IOException {
+    setMaxResultWindow(TEST_INDEX_ACCOUNT, 10);
+    JSONObject result =
+        executeQuery(String.format(
+            "source=%s | fields firstname, age | head 15", TEST_INDEX_ACCOUNT));
+    verifyDataRows(result,
+        rows("Amber", 32),
+        rows("Hattie", 36),
+        rows("Nanette", 28),
+        rows("Dale", 33),
+        rows("Elinor", 36),
+        rows("Virginia", 39),
+        rows("Dillard", 34),
+        rows("Mcgee", 39),
+        rows("Aurelia", 37),
+        rows("Fulton", 23),
+        rows("Burton", 31),
+        rows("Josie", 32),
+        rows("Hughes", 30),
+        rows("Hall", 25),
+        rows("Deidre", 33));
+  }
+
+  @Ignore("Fix https://github.com/opensearch-project/sql/issues/703#issuecomment-1211422130")
+  @Test
+  public void testHeadWithLargeNumber() throws IOException {
+    setQuerySizeLimit(5);
+    setMaxResultWindow(TEST_INDEX_ACCOUNT, 10);
+    JSONObject result =
+        executeQuery(String.format(
+            "source=%s | fields firstname, age | head 15", TEST_INDEX_ACCOUNT));
+    verifyDataRows(result,
+        rows("Amber", 32),
+        rows("Hattie", 36),
+        rows("Nanette", 28),
+        rows("Dale", 33),
+        rows("Elinor", 36),
+        rows("Virginia", 39),
+        rows("Dillard", 34),
+        rows("Mcgee", 39),
+        rows("Aurelia", 37),
+        rows("Fulton", 23),
+        rows("Burton", 31),
+        rows("Josie", 32),
+        rows("Hughes", 30),
+        rows("Hall", 25),
+        rows("Deidre", 33));
+  }
+
   @Test
   public void testHeadWithNumberAndFrom() throws IOException {
     JSONObject result =

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchClient.java

@@ -30,6 +30,14 @@ public interface OpenSearchClient {
    */
   Map<String, IndexMapping> getIndexMappings(String... indexExpression);
 
+  /**
+   * Fetch index.max_result_window settings according to index expression given.
+   *
+   * @param indexExpression index expression
+   * @return map from index name to its max result window
+   */
+  Map<String, Integer> getIndexMaxResultWindows(String... indexExpression);
+
   /**
    * Perform search query in the search request.
    *

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java

@@ -24,11 +24,14 @@
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.AliasMetadata;
+import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.metadata.MappingMetadata;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.collect.ImmutableOpenMap;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.index.IndexSettings;
 import org.opensearch.sql.opensearch.mapping.IndexMapping;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
@@ -86,6 +89,29 @@ public Map<String, IndexMapping> getIndexMappings(String... indexExpression) {
     }
   }
 
+  /**
+   * Fetch index.max_result_window settings according to index expression given.
+   *
+   * @param indexExpression index expression
+   * @return map from index name to its max result window
+   */
+  @Override
+  public Map<String, Integer> getIndexMaxResultWindows(String... indexExpression) {
+    ClusterState state = clusterService.state();
+    ImmutableOpenMap<String, IndexMetadata> indicesMetadata = state.metadata().getIndices();
+    String[] concreteIndices = resolveIndexExpression(state, indexExpression);
+
+    ImmutableMap.Builder<String, Integer> result = ImmutableMap.builder();
+    for (String index : concreteIndices) {
+      Settings settings = indicesMetadata.get(index).getSettings();
+      Integer maxResultWindow = settings.getAsInt("index.max_result_window",
+          IndexSettings.MAX_RESULT_WINDOW_SETTING.getDefault(settings));
+      result.put(index, maxResultWindow);
+    }
+
+    return result.build();
+  }
+
   /**
    * TODO: Scroll doesn't work for aggregation. Support aggregation later.
    */

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java

@@ -11,12 +11,15 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import lombok.RequiredArgsConstructor;
 import org.opensearch.action.admin.cluster.settings.ClusterGetSettingsRequest;
+import org.opensearch.action.admin.indices.settings.get.GetSettingsRequest;
+import org.opensearch.action.admin.indices.settings.get.GetSettingsResponse;
 import org.opensearch.action.search.ClearScrollRequest;
 import org.opensearch.client.RequestOptions;
 import org.opensearch.client.RestHighLevelClient;
@@ -26,6 +29,7 @@
 import org.opensearch.client.indices.GetMappingsResponse;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.cluster.metadata.AliasMetadata;
+import org.opensearch.common.collect.ImmutableOpenMap;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.sql.opensearch.mapping.IndexMapping;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
@@ -54,6 +58,36 @@ public Map<String, IndexMapping> getIndexMappings(String... indexExpression) {
     }
   }
 
+  @Override
+  public Map<String, Integer> getIndexMaxResultWindows(String... indexExpression) {
+    GetSettingsRequest request = new GetSettingsRequest()
+        .indices(indexExpression).includeDefaults(true);
+    try {
+      GetSettingsResponse response = client.indices().getSettings(request, RequestOptions.DEFAULT);
+      ImmutableOpenMap<String, Settings> settings = response.getIndexToSettings();
+      ImmutableOpenMap<String, Settings> defaultSettings = response.getIndexToDefaultSettings();
+      Map<String, Integer> result = new HashMap<>();
+
+      defaultSettings.forEach(entry -> {
+        Integer maxResultWindow = entry.value.getAsInt("index.max_result_window", null);
+        if (maxResultWindow != null) {
+          result.put(entry.key, maxResultWindow);
+        }
+      });
+
+      settings.forEach(entry -> {
+        Integer maxResultWindow = entry.value.getAsInt("index.max_result_window", null);
+        if (maxResultWindow != null) {
+          result.put(entry.key, maxResultWindow);
+        }
+      });
+
+      return result;
+    } catch (IOException e) {
+      throw new IllegalStateException("Failed to get max result window for " + indexExpression, e);
+    }
+  }
+
   @Override
   public OpenSearchResponse search(OpenSearchRequest request) {
     return request.search(

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

@@ -49,7 +49,7 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
 
 
   /**
-   * ElasticsearchExprValueFactory.
+   * OpenSearchExprValueFactory.
    */
   @EqualsAndHashCode.Exclude
   @ToString.Exclude
@@ -61,15 +61,15 @@ public class OpenSearchQueryRequest implements OpenSearchRequest {
   private boolean searchDone = false;
 
   /**
-   * Constructor of ElasticsearchQueryRequest.
+   * Constructor of OpenSearchQueryRequest.
    */
   public OpenSearchQueryRequest(String indexName, int size,
                                 OpenSearchExprValueFactory factory) {
     this(new IndexName(indexName), size, factory);
   }
 
   /**
-   * Constructor of ElasticsearchQueryRequest.
+   * Constructor of OpenSearchQueryRequest.
    */
   public OpenSearchQueryRequest(IndexName indexName, int size,
       OpenSearchExprValueFactory factory) {
@@ -81,6 +81,16 @@ public OpenSearchQueryRequest(IndexName indexName, int size,
     this.exprValueFactory = factory;
   }
 
+  /**
+   * Constructor of OpenSearchQueryRequest.
+   */
+  public OpenSearchQueryRequest(IndexName indexName, SearchSourceBuilder sourceBuilder,
+                                OpenSearchExprValueFactory factory) {
+    this.indexName = indexName;
+    this.sourceBuilder = sourceBuilder;
+    this.exprValueFactory = factory;
+  }
+
   @Override
   public OpenSearchResponse search(Function<SearchRequest, SearchResponse> searchAction,
                                    Function<SearchScrollRequest, SearchResponse> scrollAction) {