Pushdown system limit automatically for data-intensive operations

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

Author: LantaoJin

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

@@ -40,6 +40,10 @@ public enum Key {
     /** Common Settings for SQL and PPL. */
     QUERY_MEMORY_LIMIT("plugins.query.memory_limit"),
     QUERY_SIZE_LIMIT("plugins.query.size_limit"),
+    QUERY_SYSTEM_LIMIT("plugins.query.system_limit"),
+    // currently, only join could bloat data, so SYSTEM_LIMIT_JOIN equals to SYSTEM_LIMIT
+    // TODO add more specific configs in future when we have other data-bloat operations.
+    QUERY_SYSTEM_LIMIT_JOIN("plugins.query.system_limit.join"),
     ENCYRPTION_MASTER_KEY("plugins.query.datasources.encryption.masterkey"),
     DATASOURCES_URI_HOSTS_DENY_LIST("plugins.query.datasources.uri.hosts.denylist"),
     DATASOURCES_LIMIT("plugins.query.datasources.limit"),

core/src/main/java/org/opensearch/sql/calcite/utils/PlanUtils.java

@@ -21,6 +21,7 @@
 import java.util.List;
 import javax.annotation.Nullable;
 import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.rex.RexWindowBound;
@@ -353,4 +354,14 @@ public Void visitChildren(Node node, Object context) {
         };
     node.accept(leafVisitor, null);
   }
+
+  /**
+   * Extract the integer value from a RexNode, which is used to extract value from offset and limit
+   *
+   * @param r RexNode to extract value from
+   * @param defaultValue default value could be Integer or null
+   */
+  static Integer intValue(RexNode r, Integer defaultValue) {
+    return r instanceof RexLiteral ? ((RexLiteral) r).getValueAs(Integer.class) : defaultValue;
+  }
 }

docs/user/admin/settings.rst

@@ -166,7 +166,7 @@ plugins.query.size_limit
 Description
 -----------
 
-The new engine fetches a default size of index from OpenSearch set by this setting, the default value equals to max result window in index level (10000 by default). You can change the value to any value not greater than the max result window value in index level (`index.max_result_window`), here is an example::
+The size configures the maximum amount of rows to be fetched from query execution results. The default value is: 10000. You can change the value::
 
 	>> curl -H 'Content-Type: application/json' -X PUT localhost:9200/_plugins/_query/settings -d '{
 	  "transient" : {
@@ -188,6 +188,51 @@ Result set::
       }
     }
 
+plugins.query.system_limit
+==========================
+
+Description
+-----------
+
+The size configures the maximum amount of documents to be pull from OpenSearch for data-intensive operations (e.g. join, lookup). The default value is: 50000.
+
+From v3.0.0, PPL introduces commands that may increase data volume. To prevent out-of-memory problem, the system automatically enforces this ``system limit`` in the pushdown context of the physical index scan operator for such commands.
+
+Version
+-------
+3.1.0
+
+Example
+-------
+
+Change the system_limit to 10000::
+
+    sh$ curl -sS -H 'Content-Type: application/json' \
+    ... -X PUT localhost:9200/_plugins/_query/settings \
+    ... -d '{"persistent" : {"plugins.query.system_limit" : "10000"}}'
+    {
+      "acknowledged": true,
+      "persistent": {
+        "plugins": {
+          "query": {
+            "system_limit": "10000"
+          }
+        }
+      },
+      "transient": {}
+    }
+
+Rollback to default value::
+
+    sh$ curl -sS -H 'Content-Type: application/json' \
+    ... -X PUT localhost:9200/_plugins/_query/settings \
+    ... -d '{"persistent" : {"plugins.query.system_limit" : null}}'
+    {
+      "acknowledged": true,
+      "persistent": {},
+      "transient": {}
+    }
+
 plugins.query.memory_limit
 ==========================
 

docs/user/ppl/admin/settings.rst

@@ -125,7 +125,7 @@ plugins.query.size_limit
 Description
 -----------
 
-The size configures the maximum amount of rows to be fetched from PPL execution results. The default value is: 10000
+The size configures the maximum amount of rows to be fetched from query execution results. The default value is: 10000
 
 Example
 -------
@@ -159,3 +159,48 @@ Rollback to default value::
     }
 
 Note: the legacy settings of ``opendistro.query.size_limit`` is deprecated, it will fallback to the new settings if you request an update with the legacy name.
+
+plugins.query.system_limit
+==========================
+
+Description
+-----------
+
+The size configures the maximum amount of documents to be pull from OpenSearch for data-intensive operations (e.g. join, lookup). The default value is: 50000.
+
+From v3.0.0, PPL introduces commands that may increase data volume. To prevent out-of-memory problem, the system automatically enforces this ``system limit`` in the pushdown context of the physical index scan operator for such commands.
+
+Version
+-------
+3.1.0
+
+Example
+-------
+
+Change the system_limit to 10000::
+
+    sh$ curl -sS -H 'Content-Type: application/json' \
+    ... -X PUT localhost:9200/_plugins/_query/settings \
+    ... -d '{"persistent" : {"plugins.query.system_limit" : "10000"}}'
+    {
+      "acknowledged": true,
+      "persistent": {
+        "plugins": {
+          "query": {
+            "system_limit": "10000"
+          }
+        }
+      },
+      "transient": {}
+    }
+
+Rollback to default value::
+
+    sh$ curl -sS -H 'Content-Type: application/json' \
+    ... -X PUT localhost:9200/_plugins/_query/settings \
+    ... -d '{"persistent" : {"plugins.query.system_limit" : null}}'
+    {
+      "acknowledged": true,
+      "persistent": {},
+      "transient": {}
+    }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -5,8 +5,14 @@
 
 package org.opensearch.sql.calcite.remote;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_HOBBIES;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_OCCUPATION;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_STATE_COUNTRY;
+import static org.opensearch.sql.util.MatcherUtils.assertJsonEqualsIgnoreId;
+
 import java.io.IOException;
 import org.junit.Ignore;
+import org.junit.jupiter.api.Test;
 import org.opensearch.sql.ppl.ExplainIT;
 
 public class CalciteExplainIT extends ExplainIT {
@@ -15,6 +21,9 @@ public void init() throws Exception {
     super.init();
     enableCalcite();
     disallowCalciteFallback();
+    loadIndex(Index.STATE_COUNTRY);
+    loadIndex(Index.OCCUPATION);
+    loadIndex(Index.HOBBIES);
   }
 
   @Override
@@ -43,6 +52,65 @@ public void testTrendlineWithSortPushDownExplain() throws Exception {
         "https://github.com/opensearch-project/sql/issues/3466");
   }
 
+  @Test
+  public void testPushDownSystemLimitForJoinExplain() throws Exception {
+    // the SYSTEM LIMIT should apply to each table of join
+    String expected = loadFromFile("expectedOutput/calcite/explain_join_push_system_limit.json");
+
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | inner join left=a, right=b ON a.name = b.name %s | fields "
+                    + "a.name, a.age, a.state, a.country, b.occupation, b.country, b.salary",
+                TEST_INDEX_STATE_COUNTRY, TEST_INDEX_OCCUPATION)));
+  }
+
+  @Test
+  public void testPushDownSystemLimitForMultipleJoinExplain() throws Exception {
+    // the SYSTEM LIMIT should apply to each table of multi-join
+    String expected =
+        loadFromFile("expectedOutput/calcite/explain_multi_join_push_system_limit.json");
+
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source=%s | inner join left=a, right=b ON a.name = b.name %s"
+                    + " | left join left=a, right=b ON a.name = b.name %s",
+                TEST_INDEX_STATE_COUNTRY, TEST_INDEX_OCCUPATION, TEST_INDEX_HOBBIES)));
+  }
+
+  @Test
+  public void testPushDownSystemLimitForExistsSubqueryExplain() throws Exception {
+    // the SYSTEM LIMIT should apply to each table of exists subquery
+    // since exists subquery translates to hash join
+    String expected =
+        loadFromFile("expectedOutput/calcite/explain_exists_subquery_push_system_limit.json");
+
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source = %s exists [ source = %s | where name = %s.name ]",
+                TEST_INDEX_STATE_COUNTRY, TEST_INDEX_OCCUPATION, TEST_INDEX_STATE_COUNTRY)));
+  }
+
+  @Test
+  public void testPushDownSystemLimitForInSubqueryExplain() throws Exception {
+    // the SYSTEM LIMIT should apply to each table of in subquery
+    // since in subquery translates to hash join
+    String expected =
+        loadFromFile("expectedOutput/calcite/explain_in_subquery_push_system_limit.json");
+
+    assertJsonEqualsIgnoreId(
+        expected,
+        explainQueryToString(
+            String.format(
+                "source = %s | where name in [ source = %s | fields name ]",
+                TEST_INDEX_STATE_COUNTRY, TEST_INDEX_OCCUPATION)));
+  }
+
   @Override
   @Ignore("test only in v2")
   public void testExplainModeUnsupportedInV2() throws IOException {}

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteSettingsIT.java

@@ -5,7 +5,14 @@
 
 package org.opensearch.sql.calcite.remote;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+import static org.opensearch.sql.util.MatcherUtils.verifyNumOfRows;
+
 import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.jupiter.api.Test;
 import org.opensearch.sql.common.setting.Settings.Key;
 import org.opensearch.sql.ppl.SettingsIT;
 
@@ -30,4 +37,52 @@ public void testQuerySizeLimit_NoPushdown() throws IOException {
           }
         });
   }
+
+  @Test
+  public void testQuerySystemLimit() throws IOException {
+    // system limit won't impact data-intensive operations
+    setQuerySystemLimit(2);
+    JSONObject result =
+        executeQuery(String.format("search source=%s age>35 | fields firstname", TEST_INDEX_BANK));
+    verifyDataRows(result, rows("Hattie"), rows("Elinor"), rows("Virginia"));
+
+    // for non data-intensive operations, the rows still control by query.size_limit
+    setQuerySizeLimit(1);
+    result =
+        executeQuery(String.format("search source=%s age>35 | fields firstname", TEST_INDEX_BANK));
+    verifyDataRows(result, rows("Hattie"));
+    resetQuerySizeLimit();
+    resetQuerySystemLimit();
+  }
+
+  @Test
+  public void testQuerySystemLimitWithJoin() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "search source=%s age>35 | fields firstname"
+                    + " | full join left=l right=r on l.firstname=r.firstname"
+                    + " [ search source=%s | fields firstname ]",
+                TEST_INDEX_BANK, TEST_INDEX_BANK));
+    verifyNumOfRows(result, 7);
+
+    // system limit will impact data-intensive operations
+    setQuerySystemLimit(2);
+    result =
+        executeQuery(
+            String.format(
+                "search source=%s age>35 | fields firstname"
+                    + " | full join left=l right=r on l.firstname=r.firstname"
+                    + " [ search source=%s | fields firstname ]",
+                TEST_INDEX_BANK, TEST_INDEX_BANK));
+    verifyNumOfRows(result, 3);
+
+    // amount of final result should equals to query.size_limit
+    setQuerySizeLimit(1);
+    result =
+        executeQuery(String.format("search source=%s age>35 | fields firstname", TEST_INDEX_BANK));
+    verifyNumOfRows(result, 1);
+    resetQuerySizeLimit();
+    resetQuerySystemLimit();
+  }
 }

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

@@ -117,6 +117,7 @@ private Settings defaultSettings() {
               .put(Key.CALCITE_FALLBACK_ALLOWED, false)
               .put(Key.CALCITE_PUSHDOWN_ENABLED, false)
               .put(Key.CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR, 0.9)
+              .put(Key.QUERY_SYSTEM_LIMIT, 5000)
               .put(Key.DEFAULT_PATTERN_METHOD, "SIMPLE_PATTERN")
               .build();
 
@@ -144,6 +145,7 @@ protected Settings enablePushdown() {
               .put(Key.CALCITE_FALLBACK_ALLOWED, false)
               .put(Key.CALCITE_PUSHDOWN_ENABLED, true)
               .put(Key.CALCITE_PUSHDOWN_ROWCOUNT_ESTIMATION_FACTOR, 0.9)
+              .put(Key.QUERY_SYSTEM_LIMIT, 5000)
               .put(Key.DEFAULT_PATTERN_METHOD, "SIMPLE_PATTERN")
               .build();
 

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

@@ -92,6 +92,8 @@ public abstract class SQLIntegTestCase extends OpenSearchSQLRestTestCase {
       Integer.parseInt(System.getProperty("defaultQuerySizeLimit", "200"));
   public static final Integer DEFAULT_MAX_RESULT_WINDOW =
       Integer.parseInt(System.getProperty("defaultMaxResultWindow", "10000"));
+  public static final Integer DEFAULT_QUERY_SYSTEM_LIMIT =
+      Integer.parseInt(System.getProperty("defaultMaxResultWindow", "50000"));
 
   public boolean shouldResetQuerySizeLimit() {
     return true;
@@ -187,6 +189,20 @@ protected void resetQuerySizeLimit() throws IOException {
             DEFAULT_QUERY_SIZE_LIMIT.toString()));
   }
 
+  protected void setQuerySystemLimit(Integer limit) throws IOException {
+    updateClusterSettings(
+        new ClusterSetting(
+            "transient", Settings.Key.QUERY_SYSTEM_LIMIT.getKeyValue(), limit.toString()));
+  }
+
+  protected void resetQuerySystemLimit() throws IOException {
+    updateClusterSettings(
+        new ClusterSetting(
+            "transient",
+            Settings.Key.QUERY_SYSTEM_LIMIT.getKeyValue(),
+            DEFAULT_QUERY_SYSTEM_LIMIT.toString()));
+  }
+
   @SneakyThrows
   protected void setDataSourcesEnabled(String clusterSettingType, boolean value) {
     updateClusterSettings(

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

@@ -25,6 +25,11 @@ public void init() throws Exception {
     loadIndex(Index.ACCOUNT);
   }
 
+  protected String loadFromFile(String filename) throws Exception {
+    URI uri = Resources.getResource(filename).toURI();
+    return new String(Files.readAllBytes(Paths.get(uri)));
+  }
+
   @Test
   public void testExplain() throws Exception {
     String expected =
@@ -247,11 +252,6 @@ public void testTrendlineWithSortPushDownExplain() throws Exception {
                 + "| fields ageTrend"));
   }
 
-  String loadFromFile(String filename) throws Exception {
-    URI uri = Resources.getResource(filename).toURI();
-    return new String(Files.readAllBytes(Paths.get(uri)));
-  }
-
   @Test
   public void testExplainModeUnsupportedInV2() throws IOException {
     try {

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

@@ -33,6 +33,7 @@ public void testQuerySizeLimit() throws IOException {
     result =
         executeQuery(String.format("search source=%s age>35 | fields firstname", TEST_INDEX_BANK));
     verifyDataRows(result, rows("Hattie"));
+    resetQuerySizeLimit();
   }
 
   @Test
@@ -62,5 +63,6 @@ public void testQuerySizeLimit_NoPushdown() throws IOException {
                 "search source=%s | eval a = 1 | where age>35 | fields firstname",
                 TEST_INDEX_BANK));
     verifyDataRows(result, rows("Hattie"));
+    resetQuerySizeLimit();
   }
 }