refactor: throw exception if pushdown cannot be applied

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

Author: LantaoJin

core/src/main/java/org/opensearch/sql/ast/expression/Argument.java

@@ -21,6 +21,7 @@
 @EqualsAndHashCode(callSuper = false)
 public class Argument extends UnresolvedExpression {
   public static final String BUCKET_NULLABLE = "bucket_nullable";
+  public static final String NESTED = "nested";
 
   private final String argName;
   private final Literal value;

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -1047,7 +1047,9 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     //           \- Scan t
     List<RexInputRef> trimmedRefs = new ArrayList<>();
     trimmedRefs.addAll(PlanUtils.getInputRefs(resolvedGroupByList)); // group-by keys first
-    trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolvedAggCallList));
+    List<RexInputRef> aggCallRefs = PlanUtils.getInputRefsFromAggCall(resolvedAggCallList);
+    trimmedRefs.addAll(aggCallRefs);
+    List<Boolean> nestedList = isNestedAggregation(context.relBuilder, aggCallRefs);
     context.relBuilder.project(trimmedRefs);
 
     // Re-resolve all attributes based on adding trimmed Project.
@@ -1060,6 +1062,9 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(reResolved.getLeft()), reResolved.getRight());
     if (hintBucketNonNull) PlanUtils.addIgnoreNullBucketHintToAggregate(context.relBuilder);
+    if (nestedList.stream().anyMatch(b -> b)) {
+      PlanUtils.addNestedHintToAggregate(context.relBuilder, nestedList);
+    }
     // During aggregation, Calcite projects both input dependencies and output group-by fields.
     // When names conflict, Calcite adds numeric suffixes (e.g., "value0").
     // Apply explicit renaming to restore the intended aliases.
@@ -1068,6 +1073,18 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     return Pair.of(reResolved.getLeft(), reResolved.getRight());
   }
 
+  /**
+   * Return a list of the neste aggregation flag. For example: aggCalls: [count(), count(a.b),
+   * avg(a.c)] -> aggCallRefs [1, 2] -> nestedList [true, true]
+   */
+  private List<Boolean> isNestedAggregation(RelBuilder relBuilder, List<RexInputRef> aggCallRefs) {
+    return aggCallRefs.stream()
+        .map(r -> relBuilder.peek().getRowType().getFieldNames().get(r.getIndex()))
+        .map(name -> org.apache.commons.lang3.StringUtils.substringBefore(name, "."))
+        .map(root -> relBuilder.field(root).getType().getSqlTypeName() == SqlTypeName.ARRAY)
+        .toList();
+  }
+
   /**
    * Imitates {@code Registrar.registerExpression} of {@link RelBuilder} to derive the output order
    * of group-by keys after aggregation.

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

@@ -52,7 +52,6 @@
 import org.apache.calcite.jdbc.CalcitePrepare;
 import org.apache.calcite.jdbc.CalciteSchema;
 import org.apache.calcite.jdbc.Driver;
-import org.apache.calcite.linq4j.function.Function0;
 import org.apache.calcite.plan.Context;
 import org.apache.calcite.plan.Contexts;
 import org.apache.calcite.plan.Convention;
@@ -175,8 +174,11 @@ public Connection connect(
     }
 
     @Override
-    protected Function0<CalcitePrepare> createPrepareFactory() {
-      return OpenSearchPrepareImpl::new;
+    public CalcitePrepare createPrepare() {
+      if (prepareFactory != null) {
+        return prepareFactory.get();
+      }
+      return new OpenSearchPrepareImpl();
     }
   }
 
@@ -298,10 +300,10 @@ public OpenSearchCalcitePreparingStmt(
 
     @Override
     protected PreparedResult implement(RelRoot root) {
-      Hook.PLAN_BEFORE_IMPLEMENTATION.run(root);
-      RelDataType resultType = root.rel.getRowType();
-      boolean isDml = root.kind.belongsTo(SqlKind.DML);
       if (root.rel instanceof Scannable scannable) {
+        Hook.PLAN_BEFORE_IMPLEMENTATION.run(root);
+        RelDataType resultType = root.rel.getRowType();
+        boolean isDml = root.kind.belongsTo(SqlKind.DML);
         final Bindable bindable = dataContext -> scannable.scan();
 
         return new PreparedResultImpl(

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

@@ -19,7 +19,12 @@ public class PPLHintStrategyTable {
           () ->
               HintStrategyTable.builder()
                   .hintStrategy(
-                      "stats_args",
+                      "agg_args",
+                      (hint, rel) -> {
+                        return rel instanceof LogicalAggregate;
+                      })
+                  .hintStrategy(
+                      "nested_agg",
                       (hint, rel) -> {
                         return rel instanceof LogicalAggregate;
                       })

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

@@ -64,6 +64,7 @@
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
+import org.opensearch.sql.utils.Utils;
 
 public interface PlanUtils {
 
@@ -603,11 +604,30 @@ static void replaceTop(RelBuilder relBuilder, RelNode relNode) {
   }
 
   static void addIgnoreNullBucketHintToAggregate(RelBuilder relBuilder) {
-    final RelHint statHits =
-        RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
     assert relBuilder.peek() instanceof LogicalAggregate
         : "Stats hits should be added to LogicalAggregate";
-    relBuilder.hints(statHits);
+    final RelHint statHint =
+        RelHint.builder("agg_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
+    relBuilder.hints(statHint);
+    relBuilder.getCluster().setHintStrategies(PPLHintStrategyTable.getHintStrategyTable());
+  }
+
+  static void addNestedHintToAggregate(RelBuilder relBuilder, List<Boolean> nestedList) {
+    assert relBuilder.peek() instanceof LogicalAggregate
+        : "Stats hits should be added to LogicalAggregate";
+    LogicalAggregate aggregate = (LogicalAggregate) relBuilder.peek();
+    List<Integer> indicesWithArgList =
+        Utils.zipWithIndex(aggregate.getAggCallList()).stream()
+            .filter(p -> !p.getKey().getArgList().isEmpty())
+            .map(org.apache.commons.lang3.tuple.Pair::getValue)
+            .toList();
+    assert indicesWithArgList.size() == nestedList.size();
+    RelHint.Builder builder = RelHint.builder("nested_agg");
+    for (int i = 0; i < indicesWithArgList.size(); i++) {
+      builder.hintOption(indicesWithArgList.get(i).toString(), nestedList.get(i).toString());
+    }
+    final RelHint statHint = builder.build();
+    relBuilder.hints(statHint);
     relBuilder.getCluster().setHintStrategies(PPLHintStrategyTable.getHintStrategyTable());
   }
 

core/src/main/java/org/opensearch/sql/utils/Utils.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.utils;
+
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import org.apache.commons.lang3.tuple.Pair;
+
+public interface Utils {
+  static <I> List<Pair<I, Integer>> zipWithIndex(List<I> input) {
+    LinkedList<Pair<I, Integer>> result = new LinkedList<>();
+    Iterator<I> iter = input.iterator();
+    int index = 0;
+    while (iter.hasNext()) {
+      result.add(Pair.of(iter.next(), index++));
+    }
+    return result;
+  }
+}

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

@@ -19,11 +19,13 @@
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WORK_INFORMATION;
 import static org.opensearch.sql.util.MatcherUtils.assertJsonEqualsIgnoreId;
 import static org.opensearch.sql.util.MatcherUtils.assertYamlEqualsIgnoreId;
+import static org.opensearch.sql.util.MatcherUtils.verifyErrorMessageContains;
 
 import java.io.IOException;
 import java.util.Locale;
 import org.junit.Ignore;
 import org.junit.Test;
+import org.opensearch.client.ResponseException;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.ppl.ExplainIT;
 
@@ -2215,4 +2217,33 @@ public void testRexStandardizationForScript() throws IOException {
                 TEST_INDEX_BANK),
             true));
   }
+
+  @Test
+  public void testNestedAggPushDownExplain() throws Exception {
+    enabledOnlyWhenPushdownIsEnabled();
+    String expected = loadExpectedPlan("explain_nested_agg_push.yaml");
+
+    assertYamlEqualsIgnoreId(
+        expected,
+        explainQueryYaml(
+            "source=opensearch-sql_test_index_nested_simple | stats count(address.area) as"
+                + " count_area, min(address.area) as min_area, max(address.area) as max_area,"
+                + " avg(address.area) as avg_area, avg(age) as avg_age by name"));
+  }
+
+  @Test
+  public void testNestedAggExplainWhenPushdownNotApplied() throws Exception {
+    enabledOnlyWhenPushdownIsEnabled();
+    Exception e =
+        assertThrows(
+            ResponseException.class,
+            () ->
+                explainQueryYaml(
+                    "source=opensearch-sql_test_index_nested_simple | head 10000 | stats"
+                        + " count(address.area) as count_area, min(address.area) as min_area,"
+                        + " max(address.area) as max_area, avg(address.area) as avg_area, avg(age)"
+                        + " as avg_age by name"));
+    verifyErrorMessageContains(
+        e, "Nested aggregate is unsupported when pushdown cannot be applied");
+  }
 }

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

@@ -7,11 +7,15 @@
 
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_NESTED_SIMPLE;
 import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.schema;
 import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+import static org.opensearch.sql.util.MatcherUtils.verifySchemaInOrder;
 
 import java.io.IOException;
 import org.json.JSONObject;
+import org.junit.Ignore;
 import org.junit.jupiter.api.Test;
 import org.opensearch.sql.ppl.StatsCommandIT;
 
@@ -21,6 +25,8 @@ public void init() throws Exception {
     super.init();
     enableCalcite();
     setQuerySizeLimit(2000);
+
+    loadIndex(Index.NESTED_SIMPLE);
   }
 
   @Test
@@ -106,4 +112,60 @@ public void testPaginatingStatsForHeadFrom() throws IOException {
       resetQueryBucketSize();
     }
   }
+
+  @Test
+  public void testNestedAggregation() throws IOException {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | stats count(address.area) as count_area, min(address.area) as"
+                    + " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
+                    + " avg(age) as avg_age",
+                TEST_INDEX_NESTED_SIMPLE));
+    verifySchemaInOrder(
+        actual,
+        isCalciteEnabled() ? schema("count_area", "bigint") : schema("count_area", "int"),
+        schema("min_area", "double"),
+        schema("max_area", "double"),
+        schema("avg_area", "double"),
+        schema("avg_age", "double"));
+    verifyDataRows(actual, rows(9, 9.99, 1000.99, 300.11555555555555, 25.2));
+  }
+
+  @Test
+  public void testNestedAggregationBy() throws IOException {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | stats count(address.area) as count_area, min(address.area) as"
+                    + " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
+                    + " avg(age) as avg_age by name",
+                TEST_INDEX_NESTED_SIMPLE));
+    verifySchemaInOrder(
+        actual,
+        isCalciteEnabled() ? schema("count_area", "bigint") : schema("count_area", "int"),
+        schema("min_area", "double"),
+        schema("max_area", "double"),
+        schema("avg_area", "double"),
+        schema("avg_age", "double"),
+        schema("name", "string"));
+    verifyDataRows(
+        actual,
+        rows(4, 10.24, 400.99, 209.69, 24, "abbas"),
+        rows(0, null, null, null, 19, "andy"),
+        rows(2, 9.99, 1000.99, 505.49, 32, "chen"),
+        rows(1, 190.5, 190.5, 190.5, 25, "david"),
+        rows(2, 231.01, 429.79, 330.4, 26, "peng"));
+  }
+
+  @Ignore("https://github.com/opensearch-project/sql/issues/3384")
+  public void testNestedAggregationBySpan() throws IOException {
+    JSONObject actual =
+        executeQuery(
+            String.format(
+                "source=%s | stats count(address.area) as count_area, min(address.area) as"
+                    + " min_area, max(address.area) as max_area, avg(address.area) as avg_area,"
+                    + " avg(age) as avg_age by span(age, 10)",
+                TEST_INDEX_NESTED_SIMPLE));
+  }
 }