Support ResourceMonitor with Calcite (opensearch-project#3738)

* Support ResourceMonitor with Calcite

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

* fix IT

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

* fix UT related to decimal literal

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

---------

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

Author: LantaoJin

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

@@ -0,0 +1,14 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.exception;
+
+/** Non-fallback to v2 exception for Calcite. */
+public class NonFallbackCalciteException extends QueryEngineException {
+
+  public NonFallbackCalciteException(String message) {
+    super(message);
+  }
+}

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -41,6 +41,7 @@
 import org.opensearch.sql.common.setting.Settings.Key;
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
+import org.opensearch.sql.exception.NonFallbackCalciteException;
 import org.opensearch.sql.planner.PlanContext;
 import org.opensearch.sql.planner.Planner;
 import org.opensearch.sql.planner.logical.LogicalPaginate;
@@ -107,7 +108,7 @@ public void executeWithCalcite(
                 return null;
               });
     } catch (Throwable t) {
-      if (isCalciteFallbackAllowed()) {
+      if (isCalciteFallbackAllowed() && !(t instanceof NonFallbackCalciteException)) {
         log.warn("Fallback to V2 query engine since got exception", t);
         executeWithLegacy(plan, queryType, listener, Optional.of(t));
       } else {

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

@@ -5,10 +5,8 @@
 
 package org.opensearch.sql.calcite.remote;
 
-import org.junit.Ignore;
 import org.opensearch.sql.ppl.ResourceMonitorIT;
 
-@Ignore("https://github.com/opensearch-project/sql/issues/3454")
 public class CalciteResourceMonitorIT extends ResourceMonitorIT {
   @Override
   public void init() throws Exception {

opensearch/src/main/java/org/opensearch/sql/opensearch/monitor/OpenSearchResourceMonitor.java

@@ -48,6 +48,11 @@ public OpenSearchResourceMonitor(Settings settings, OpenSearchMemoryHealthy memo
   public boolean isHealthy() {
     try {
       ByteSizeValue limit = settings.getSettingValue(Settings.Key.QUERY_MEMORY_LIMIT);
+      if (limit == null) {
+        // undefined, be always healthy, this is useful in Calcite standalone ITs
+        // since AlwaysHealthyMonitor is not work within Calcite tests.
+        return true;
+      }
       Supplier<Boolean> booleanSupplier =
           Retry.decorateSupplier(retry, () -> memoryMonitor.isMemoryHealthy(limit.getBytes()));
       return booleanSupplier.get();

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

@@ -25,6 +25,8 @@
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
+import org.opensearch.sql.opensearch.monitor.OpenSearchMemoryHealthy;
+import org.opensearch.sql.opensearch.monitor.OpenSearchResourceMonitor;
 import org.opensearch.sql.opensearch.planner.physical.ADOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLOperator;
@@ -261,6 +263,10 @@ public OpenSearchRequestBuilder createRequestBuilder() {
     return new OpenSearchRequestBuilder(createExprValueFactory(), settings);
   }
 
+  public OpenSearchResourceMonitor createOpenSearchResourceMonitor() {
+    return new OpenSearchResourceMonitor(getSettings(), new OpenSearchMemoryHealthy());
+  }
+
   public OpenSearchRequest buildRequest(OpenSearchRequestBuilder requestBuilder) {
     final TimeValue cursorKeepAlive = settings.getSettingValue(Settings.Key.SQL_CURSOR_KEEP_ALIVE);
     return requestBuilder.build(

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

@@ -101,7 +101,8 @@ public Enumerator<Object> enumerator() {
             osIndex.getClient(),
             getFieldPath(),
             requestBuilder.getMaxResponseSize(),
-            osIndex.buildRequest(requestBuilder));
+            osIndex.buildRequest(requestBuilder),
+            osIndex.createOpenSearchResourceMonitor());
       }
     };
   }

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

@@ -13,6 +13,8 @@
 import org.apache.calcite.linq4j.Enumerator;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
+import org.opensearch.sql.exception.NonFallbackCalciteException;
+import org.opensearch.sql.monitor.ResourceMonitor;
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.response.OpenSearchResponse;
@@ -37,6 +39,12 @@ public class OpenSearchIndexEnumerator implements Enumerator<Object> {
   /** Largest number of rows allowed in the response. */
   @EqualsAndHashCode.Include @ToString.Include private final int maxResponseSize;
 
+  /** How many moveNext() calls to perform resource check once. */
+  private static final long NUMBER_OF_NEXT_CALL_TO_CHECK = 1000;
+
+  /** ResourceMonitor. */
+  private final ResourceMonitor monitor;
+
   /** Number of rows returned. */
   private Integer queryCount;
 
@@ -49,13 +57,18 @@ public OpenSearchIndexEnumerator(
       OpenSearchClient client,
       List<String> fields,
       int maxResponseSize,
-      OpenSearchRequest request) {
+      OpenSearchRequest request,
+      ResourceMonitor monitor) {
     this.client = client;
     this.fields = fields;
     this.request = request;
     this.maxResponseSize = maxResponseSize;
+    this.monitor = monitor;
     this.queryCount = 0;
     this.current = null;
+    if (!this.monitor.isHealthy()) {
+      throw new NonFallbackCalciteException("insufficient resources to run the query, quit.");
+    }
   }
 
   private void fetchNextBatch() {
@@ -88,6 +101,11 @@ public boolean moveNext() {
       return false;
     }
 
+    boolean shouldCheck = (queryCount % NUMBER_OF_NEXT_CALL_TO_CHECK == 0);
+    if (shouldCheck && !this.monitor.isHealthy()) {
+      throw new NonFallbackCalciteException("insufficient resources to load next row, quit.");
+    }
+
     if (iterator == null || !iterator.hasNext()) {
       fetchNextBatch();
     }

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLAppendcolTest.java

@@ -145,9 +145,9 @@ public void testAppendcolStats() {
     verifyLogical(root, expectedLogical);
     String expectedResult =
         ""
-            + "count()=5; DEPTNO=20; avg(SAL)=2175.00\n"
-            + "count()=3; DEPTNO=10; avg(SAL)=2916.66\n"
-            + "count()=6; DEPTNO=30; avg(SAL)=1566.66\n";
+            + "count()=5; DEPTNO=20; avg(SAL)=2175.\n"
+            + "count()=3; DEPTNO=10; avg(SAL)=2916.666666\n"
+            + "count()=6; DEPTNO=30; avg(SAL)=1566.666666\n";
     verifyResult(root, expectedResult);
 
     String expectedSparkSql =
@@ -185,9 +185,9 @@ public void testAppendcolStatsOverride() {
     verifyLogical(root, expectedLogical);
     String expectedResult =
         ""
-            + "count()=5; DEPTNO=20; avg(SAL)=2175.00\n"
-            + "count()=3; DEPTNO=10; avg(SAL)=2916.66\n"
-            + "count()=6; DEPTNO=30; avg(SAL)=1566.66\n";
+            + "count()=5; DEPTNO=20; avg(SAL)=2175.\n"
+            + "count()=3; DEPTNO=10; avg(SAL)=2916.666666\n"
+            + "count()=6; DEPTNO=30; avg(SAL)=1566.666666\n";
     verifyResult(root, expectedResult);
 
     String expectedSparkSql =