Create bucket aggregation parsers that supports parsing nested sub aggregations

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

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

@@ -5,9 +5,9 @@
 
 package org.opensearch.sql.calcite.remote;
 
-import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WEBLOGS;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WEBLOGS;
 import static org.opensearch.sql.util.MatcherUtils.rows;
 import static org.opensearch.sql.util.MatcherUtils.schema;
 import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
@@ -264,10 +264,7 @@ public void testCaseRangeAggregationPushdown() throws IOException {
                     + ") | stats count() as total by range_bucket | sort range_bucket",
                 TEST_INDEX_WEBLOGS));
 
-    verifySchema(
-        actual,
-        schema("range_bucket", "string"),
-        schema("total", "long"));
+    verifySchema(actual, schema("range_bucket", "string"), schema("total", "long"));
 
     // This should work but won't be optimized due to implicit NULL bucket
     assertTrue(actual.getJSONArray("datarows").length() > 0);
@@ -279,12 +276,10 @@ public void testCaseRangeAggregationWithMetrics() throws IOException {
     JSONObject actual =
         executeQuery(
             String.format(
-                "source=%s | eval size_category = case("
-                    + "  cast(bytes as int) < 2000, 'small',"
-                    + "  cast(bytes as int) >= 2000 AND cast(bytes as int) < 5000, 'medium',"
-                    + "  cast(bytes as int) >= 5000, 'large'"
-                    + ") | stats count() as total, avg(cast(bytes as int)) as avg_bytes by size_category"
-                    + " | sort size_category",
+                "source=%s | eval size_category = case(  cast(bytes as int) < 2000, 'small', "
+                    + " cast(bytes as int) >= 2000 AND cast(bytes as int) < 5000, 'medium', "
+                    + " cast(bytes as int) >= 5000, 'large') | stats count() as total,"
+                    + " avg(cast(bytes as int)) as avg_bytes by size_category | sort size_category",
                 TEST_INDEX_WEBLOGS));
 
     verifySchema(
@@ -304,19 +299,14 @@ public void testCaseRangeAggregationWithElse() throws IOException {
     JSONObject actual =
         executeQuery(
             String.format(
-                "source=%s | eval status_category = case("
-                    + "  cast(response as int) < 300, 'success',"
-                    + "  cast(response as int) >= 300 AND cast(response as int) < 400, 'redirect',"
-                    + "  cast(response as int) >= 400 AND cast(response as int) < 500, 'client_error',"
-                    + "  cast(response as int) >= 500, 'server_error'"
-                    + "  else 'unknown'"
-                    + ") | stats count() by status_category | sort status_category",
+                "source=%s | eval status_category = case(  cast(response as int) < 300, 'success', "
+                    + " cast(response as int) >= 300 AND cast(response as int) < 400, 'redirect', "
+                    + " cast(response as int) >= 400 AND cast(response as int) < 500,"
+                    + " 'client_error',  cast(response as int) >= 500, 'server_error'  else"
+                    + " 'unknown') | stats count() by status_category | sort status_category",
                 TEST_INDEX_WEBLOGS));
 
-    verifySchema(
-        actual,
-        schema("status_category", "string"),
-        schema("count()", "long"));
+    verifySchema(actual, schema("status_category", "string"), schema("count()", "long"));
 
     // Should handle the ELSE case for null/non-numeric responses
     assertTrue(actual.getJSONArray("datarows").length() > 0);
@@ -335,10 +325,7 @@ public void testNonOptimizableCaseExpression() throws IOException {
                     + ") | stats count() by mixed_condition",
                 TEST_INDEX_WEBLOGS));
 
-    verifySchema(
-        actual,
-        schema("mixed_condition", "string"),
-        schema("count()", "long"));
+    verifySchema(actual, schema("mixed_condition", "string"), schema("count()", "long"));
 
     // This should work but won't be optimized
     assertTrue(actual.getJSONArray("datarows").length() > 0);
@@ -356,10 +343,7 @@ public void testCaseWithNonLiteralResult() throws IOException {
                     + ") | stats count() by computed_result | head 3",
                 TEST_INDEX_WEBLOGS));
 
-    verifySchema(
-        actual,
-        schema("computed_result", "string"),
-        schema("count()", "long"));
+    verifySchema(actual, schema("computed_result", "string"), schema("count()", "long"));
 
     // This should work but won't be optimized to range aggregation
     assertTrue(actual.getJSONArray("datarows").length() > 0);
@@ -379,10 +363,7 @@ public void testOptimizableCaseRangeAggregation() throws IOException {
                     + ") | stats count() by size_bucket | sort size_bucket",
                 TEST_INDEX_WEBLOGS));
 
-    verifySchema(
-        actual,
-        schema("size_bucket", "string"),
-        schema("count()", "long"));
+    verifySchema(actual, schema("size_bucket", "string"), schema("count()", "long"));
 
     // This should work - the explicit ELSE makes it potentially optimizable
     assertTrue(actual.getJSONArray("datarows").length() > 0);

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

@@ -41,7 +41,6 @@
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-
 import lombok.RequiredArgsConstructor;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.rel.core.Aggregate;
@@ -60,10 +59,10 @@
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.aggregations.AggregatorFactories.Builder;
 import org.opensearch.search.aggregations.BucketOrder;
-import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder;
 import org.opensearch.search.aggregations.bucket.composite.CompositeValuesSourceBuilder;
 import org.opensearch.search.aggregations.bucket.composite.TermsValuesSourceBuilder;
 import org.opensearch.search.aggregations.bucket.missing.MissingOrder;
+import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder;
 import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
 import org.opensearch.search.aggregations.metrics.ExtendedStats;
 import org.opensearch.search.aggregations.metrics.PercentilesAggregationBuilder;
@@ -81,11 +80,10 @@
 import org.opensearch.sql.opensearch.request.PredicateAnalyzer.NamedFieldExpression;
 import org.opensearch.sql.opensearch.response.agg.ArgMaxMinParser;
 import org.opensearch.sql.opensearch.response.agg.BucketAggregationParser;
-import org.opensearch.sql.opensearch.response.agg.CompositeAggregationParser;
+import org.opensearch.sql.opensearch.response.agg.LeafBucketAggregationParser;
 import org.opensearch.sql.opensearch.response.agg.MetricParser;
 import org.opensearch.sql.opensearch.response.agg.NoBucketAggregationParser;
 import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser;
-import org.opensearch.sql.opensearch.response.agg.RangeParser;
 import org.opensearch.sql.opensearch.response.agg.SinglePercentileParser;
 import org.opensearch.sql.opensearch.response.agg.SingleValueParser;
 import org.opensearch.sql.opensearch.response.agg.StatsParser;
@@ -204,11 +202,19 @@ public static Pair<List<AggregationBuilder>, OpenSearchAggregationResponseParser
       Builder metricBuilder = builderAndParser.getLeft();
       List<MetricParser> metricParsers = builderAndParser.getRight();
       // Find group by fields derived from CASE functions and convert them to range queries
-      CaseRangeAnalyzer rangeAnalyzer = CaseRangeAnalyzer.create(rowType);
-      List<Pair<Integer, RangeAggregationBuilder>> groupsByCase = groupList.stream()
-              .map(i -> Pair.of(i, project.getProjects().get(i)))
-              .filter(p -> p.getRight() instanceof RexCall rexCall && rexCall.getKind() == SqlKind.CASE)
-              .map(p -> Pair.of(p.getLeft(), rangeAnalyzer.analyze((RexCall) p.getRight())))
+      List<Pair<Integer, RangeAggregationBuilder>> groupsByCase =
+          groupList.stream()
+              .map(i -> Pair.of(i, project.getNamedProjects().get(i)))
+              .filter(
+                  p ->
+                      p.getRight().getKey() instanceof RexCall rexCall
+                          && rexCall.getKind() == SqlKind.CASE)
+              .map(
+                  p ->
+                      Pair.of(
+                          p.getLeft(),
+                          CaseRangeAnalyzer.create(p.getRight().getValue(), rowType)
+                              .analyze((RexCall) p.getRight().getKey())))
               .filter(p -> p.getRight().isPresent())
               .map(p -> Pair.of(p.getLeft(), p.getRight().get()))
               .toList();
@@ -220,28 +226,32 @@ public static Pair<List<AggregationBuilder>, OpenSearchAggregationResponseParser
       // Note that but a composite aggregation can not be a sub aggregation of range aggregation,
       // but range aggregation can be a sub aggregation of a composite aggregation.
       AggregationBuilder rangeAggregationBuilder = null;
+      BucketAggregationParser bucketAggregationParser = null;
       if (!groupsByCase.isEmpty()) {
         for (int i = 0; i < groupsByCase.size(); i++) {
           Pair<Integer, RangeAggregationBuilder> pair = groupsByCase.get(i);
           if (i == 0) {
             rangeAggregationBuilder = pair.getRight();
+            bucketAggregationParser =
+                new BucketAggregationParser(new LeafBucketAggregationParser(metricParsers));
           } else {
             groupsByCase.get(i - 1).getRight().subAggregation(pair.getRight());
+            bucketAggregationParser = new BucketAggregationParser(bucketAggregationParser);
           }
         }
         groupsByCase.getLast().getRight().subAggregations(metricBuilder);
-        metricParsers.add(new RangeParser("case_range"));
       }
 
       // Remove groups that are converted to ranges from groupList
       Set<Integer> toRemove = groupsByCase.stream().map(Pair::getLeft).collect(Collectors.toSet());
-      List<Integer> filteredGroupList = groupList.stream().filter(i -> !toRemove.contains(i)).toList();
+      List<Integer> filteredGroupList =
+          groupList.stream().filter(i -> !toRemove.contains(i)).toList();
 
       // The top-level query is a range query: stats count() by range_field
       // RangeAgg
       //   Metric
       if (!groupsByCase.isEmpty() && filteredGroupList.isEmpty()) {
-        return Pair.of(List.of(rangeAggregationBuilder), new BucketAggregationParser(metricParsers));
+        return Pair.of(List.of(rangeAggregationBuilder), bucketAggregationParser);
       }
       // No parent composite aggregation or range aggregation is attached: stats count()
       // Metric
@@ -250,21 +260,23 @@ else if (aggregate.getGroupSet().isEmpty() && filteredGroupList.isEmpty()) {
             ImmutableList.copyOf(metricBuilder.getAggregatorFactories()),
             new NoBucketAggregationParser(metricParsers));
       }
-      // It has both composite aggregation and range aggregation: stats count() by range_field, non_range_field
+      // It has both composite aggregation and range aggregation: stats count() by range_field,
+      // non_range_field
       // CompositeAgg
       //   RangeAgg
       //     Metric
       else if (!groupsByCase.isEmpty()) {
         List<CompositeValuesSourceBuilder<?>> buckets =
-                createCompositeBuckets(filteredGroupList, project, helper);
+            createCompositeBuckets(filteredGroupList, project, helper);
         return Pair.of(
-                Collections.singletonList(
-                        AggregationBuilders.composite("composite_buckets", buckets)
-                                .subAggregation(rangeAggregationBuilder)
-                                .size(AGGREGATION_BUCKET_SIZE)),
-                new CompositeAggregationParser(metricParsers));
+            Collections.singletonList(
+                AggregationBuilders.composite("composite_buckets", buckets)
+                    .subAggregation(rangeAggregationBuilder)
+                    .size(AGGREGATION_BUCKET_SIZE)),
+            new BucketAggregationParser(bucketAggregationParser));
       }
-      // It does not have range aggregation, but has composite aggregation: stats count() by non_range_field
+      // It does not have range aggregation, but has composite aggregation: stats count() by
+      // non_range_field
       // CompositeAgg
       //   Metric
       else {
@@ -275,7 +287,7 @@ else if (!groupsByCase.isEmpty()) {
                 AggregationBuilders.composite("composite_buckets", buckets)
                     .subAggregations(metricBuilder)
                     .size(AGGREGATION_BUCKET_SIZE)),
-            new CompositeAggregationParser(metricParsers));
+            new BucketAggregationParser(new LeafBucketAggregationParser(metricParsers)));
       }
     } catch (Throwable e) {
       Throwables.throwIfInstanceOf(e, UnsupportedOperationException.class);

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

@@ -40,17 +40,14 @@ public class CaseRangeAnalyzer {
   /** The default key to use if there isn't a key specified for the else case */
   public static final String DEFAULT_ELSE_KEY = "null";
 
-  /** The name for the range aggregation */
-  public static final String NAME = "case_range";
-
   private final RelDataType rowType;
   private final RangeSet<Double> takenRange;
   private final RangeAggregationBuilder builder;
 
-  public CaseRangeAnalyzer(RelDataType rowType) {
+  public CaseRangeAnalyzer(String name, RelDataType rowType) {
     this.rowType = rowType;
     this.takenRange = TreeRangeSet.create();
-    this.builder = AggregationBuilders.range(NAME);
+    this.builder = AggregationBuilders.range(name).keyed(true);
   }
 
   /**
@@ -59,8 +56,8 @@ public CaseRangeAnalyzer(RelDataType rowType) {
    * @param rowType the row type information for field resolution
    * @return a new CaseRangeAnalyzer instance
    */
-  public static CaseRangeAnalyzer create(RelDataType rowType) {
-    return new CaseRangeAnalyzer(rowType);
+  public static CaseRangeAnalyzer create(String name, RelDataType rowType) {
+    return new CaseRangeAnalyzer(name, rowType);
   }
 
   /**

opensearch/src/main/java/org/opensearch/sql/opensearch/response/agg/BucketAggregationParser.java

@@ -5,43 +5,67 @@
 
 package org.opensearch.sql.opensearch.response.agg;
 
-import java.util.Arrays;
-import java.util.LinkedHashMap;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
-import lombok.EqualsAndHashCode;
 import org.opensearch.search.aggregations.Aggregation;
 import org.opensearch.search.aggregations.Aggregations;
 import org.opensearch.search.aggregations.bucket.MultiBucketsAggregation;
+import org.opensearch.search.aggregations.bucket.composite.CompositeAggregation;
 
-/**
- * Use BucketAggregationParser only when there is a single group-by key, it returns multiple
- * buckets. {@link CompositeAggregationParser} is used for multiple group by keys
- */
-@EqualsAndHashCode
 public class BucketAggregationParser implements OpenSearchAggregationResponseParser {
-  private final MetricParserHelper metricsParser;
-
-  public BucketAggregationParser(MetricParser... metricParserList) {
-    metricsParser = new MetricParserHelper(Arrays.asList(metricParserList));
-  }
+  private final OpenSearchAggregationResponseParser subAggParser;
 
-  public BucketAggregationParser(List<MetricParser> metricParserList) {
-    metricsParser = new MetricParserHelper(metricParserList);
+  public BucketAggregationParser(OpenSearchAggregationResponseParser subAggParser) {
+    this.subAggParser = subAggParser;
   }
 
   @Override
   public List<Map<String, Object>> parse(Aggregations aggregations) {
-    Aggregation agg = aggregations.asList().getFirst();
-    return ((MultiBucketsAggregation) agg)
-        .getBuckets().stream().map(b -> parse(b, agg.getName())).collect(Collectors.toList());
+    if (subAggParser instanceof BucketAggregationParser) {
+      return aggregations.asList().stream()
+          .map(
+              aggregation -> {
+                if (aggregation instanceof CompositeAggregation) {
+                  return (CompositeAggregation) aggregation;
+                } else {
+                  return (MultiBucketsAggregation) aggregation;
+                }
+              })
+          .map(MultiBucketsAggregation::getBuckets)
+          .flatMap(List::stream)
+          .map(this::parse)
+          .flatMap(List::stream)
+          .collect(Collectors.toList());
+    } else if (subAggParser instanceof LeafBucketAggregationParser) {
+      return subAggParser.parse(aggregations);
+    } else {
+      throw new IllegalStateException(
+          "Sub parsers of a BucketAggregationParser can only be either BucketAggregationParser or"
+              + " LeafBucketAggregationParser");
+    }
+  }
+
+  private List<Map<String, Object>> parse(MultiBucketsAggregation.Bucket bucket) {
+    if (bucket instanceof CompositeAggregation.Bucket compositeBucket) {
+      return parse(compositeBucket);
+    }
+    List<Map<String, Object>> results = new ArrayList<>();
+    for (Aggregation subAgg : bucket.getAggregations()) {
+      var sub = (Aggregations) subAgg;
+      results.addAll(subAggParser.parse(sub));
+    }
+    return results;
   }
 
-  private Map<String, Object> parse(MultiBucketsAggregation.Bucket bucket, String keyName) {
-    Map<String, Object> resultMap = new LinkedHashMap<>();
-    resultMap.put(keyName, bucket.getKey());
-    resultMap.putAll(metricsParser.parse(bucket.getAggregations()));
-    return resultMap;
+  private List<Map<String, Object>> parse(CompositeAggregation.Bucket bucket) {
+    Map<String, Object> common = new HashMap<>(bucket.getKey());
+    List<Map<String, Object>> results = subAggParser.parse(bucket.getAggregations());
+    for (Map<String, Object> r : results) {
+      r.putAll(common);
+    }
+    return results;
   }
 }