Create a parent abstract classes for BucketAggregationParsers

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

Author: yuancu

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.List;
+import java.util.Map;
+import org.opensearch.search.SearchHits;
+import org.opensearch.search.aggregations.bucket.composite.CompositeAggregation;
+import org.opensearch.search.aggregations.bucket.range.Range;
+
+/**
+ * Abstract base class for parsing bucket aggregations from OpenSearch responses. Provides common
+ * functionality for extracting key-value pairs from different types of buckets.
+ */
+public abstract class AbstractBucketAggregationParser
+    implements OpenSearchAggregationResponseParser {
+  /**
+   * Extracts key-value pairs from a composite aggregation bucket without processing its
+   * sub-aggregations.
+   *
+   * <p>For example, for the following CompositeAggregation bucket in response:
+   *
+   * <pre>{@code
+   * {
+   *   "key": {
+   *     "firstname": "William",
+   *     "lastname": "Shakespeare"
+   *   },
+   *   "sub_agg_name": {
+   *     "buckets": []
+   *   }
+   * }
+   * }</pre>
+   *
+   * It returns {@code {"firstname": "William", "lastname": "Shakespeare"}} as the response.
+   *
+   * @param bucket the composite aggregation bucket to extract data from
+   * @return a map containing the bucket's key-value pairs
+   */
+  protected Map<String, Object> extract(CompositeAggregation.Bucket bucket) {
+    return bucket.getKey();
+  }
+
+  /**
+   * Extracts key-value pairs from a range aggregation bucket without processing its
+   * sub-aggregations.
+   *
+   * @param bucket the range aggregation bucket to extract data from
+   * @param name the name to use as the key in the returned map
+   * @return a map containing the bucket's key mapped to the provided name
+   */
+  protected Map<String, Object> extract(Range.Bucket bucket, String name) {
+    return Map.of(name, bucket.getKey());
+  }
+
+  @Override
+  public List<Map<String, Object>> parse(SearchHits hits) {
+    throw new UnsupportedOperationException(this.getClass() + " doesn't support parse(SearchHits)");
+  }
+}

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

@@ -5,27 +5,48 @@
 
 package org.opensearch.sql.opensearch.response.agg;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import org.opensearch.search.SearchHits;
 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;
 import org.opensearch.search.aggregations.bucket.range.Range;
 
+/**
+ * Parser for bucket aggregations that contain sub-aggregations. This parser handles multiple levels
+ * of multi-bucket aggregations by delegates sublevels to sub-parsers.
+ *
+ * <p>Please note that it does not handle metric or value count responses -- they should be parsed
+ * only in {@link LeafBucketAggregationParser}.
+ */
 @Getter
-@EqualsAndHashCode
-public class BucketAggregationParser implements OpenSearchAggregationResponseParser {
-  private final OpenSearchAggregationResponseParser subAggParser;
+@EqualsAndHashCode(callSuper = false)
+public class BucketAggregationParser extends AbstractBucketAggregationParser {
+  /** The sub-aggregation parser used to process nested aggregations within each bucket. */
+  private final AbstractBucketAggregationParser subAggParser;
 
-  public BucketAggregationParser(OpenSearchAggregationResponseParser subAggParser) {
+  /**
+   * Constructs a new BucketAggregationParser with the specified sub-aggregation parser.
+   *
+   * @param subAggParser the parser to handle sublevel multi-bucket aggregations within each bucket
+   */
+  public BucketAggregationParser(AbstractBucketAggregationParser subAggParser) {
     this.subAggParser = subAggParser;
   }
 
+  /**
+   * Parses the provided aggregations into a list of maps containing the aggregated data. This
+   * method handles multi-bucket aggregations by processing each bucket and merging the results with
+   * bucket-specific key information.
+   *
+   * @param aggregations the aggregations to parse
+   * @return a list of maps containing the parsed aggregation data
+   * @throws IllegalStateException if the aggregation type is not supported or if the sub-parser
+   *     type is invalid
+   */
   @Override
   public List<Map<String, Object>> parse(Aggregations aggregations) {
     if (subAggParser instanceof BucketAggregationParser) {
@@ -51,21 +72,16 @@ public List<Map<String, Object>> parse(Aggregations aggregations) {
   private List<Map<String, Object>> parse(MultiBucketsAggregation.Bucket bucket, String name) {
     List<Map<String, Object>> results = subAggParser.parse(bucket.getAggregations());
     if (bucket instanceof CompositeAggregation.Bucket compositeBucket) {
-      Map<String, Object> common = new HashMap<>(compositeBucket.getKey());
+      Map<String, Object> common = extract(compositeBucket);
       for (Map<String, Object> r : results) {
         r.putAll(common);
       }
-    } else if (bucket instanceof Range.Bucket) {
+    } else if (bucket instanceof Range.Bucket rangeBucket) {
+      Map<String, Object> common = extract(rangeBucket, name);
       for (Map<String, Object> r : results) {
-        r.put(name, bucket.getKey());
+        r.putAll(common);
       }
     }
     return results;
   }
-
-  @Override
-  public List<Map<String, Object>> parse(SearchHits hits) {
-    throw new UnsupportedOperationException(
-        "BucketAggregationParser doesn't support parse(SearchHits)");
-  }
 }

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

@@ -11,19 +11,20 @@
 import java.util.Objects;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import org.opensearch.search.SearchHits;
 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;
 import org.opensearch.search.aggregations.bucket.range.Range;
 
 /**
- * Use LeafBucketAggregationParser only when there is a single group-by key, it returns multiple
- * buckets. {@link BucketAggregationParser} is used for multiple group by keys
+ * Parser for leaf-level bucket aggregations that may contain metric information but no nested
+ * multi-bucket aggregations.
+ *
+ * <p>For aggregations with nested bucket structures, use {@link BucketAggregationParser} instead.
  */
-@EqualsAndHashCode
-public class LeafBucketAggregationParser implements OpenSearchAggregationResponseParser {
+@EqualsAndHashCode(callSuper = false)
+public class LeafBucketAggregationParser extends AbstractBucketAggregationParser {
   @Getter private final MetricParserHelper metricsParser;
   // countAggNameList dedicated the list of count aggregations which are filled by doc_count
   private List<String> countAggNameList = List.of();
@@ -50,16 +51,10 @@ public List<Map<String, Object>> parse(Aggregations aggregations) {
         .getBuckets().stream().map(b -> parse(b, agg.getName())).filter(Objects::nonNull).toList();
   }
 
-  @Override
-  public List<Map<String, Object>> parse(SearchHits hits) {
-    throw new UnsupportedOperationException(
-        "LeafBucketAggregationParser doesn't support parse(SearchHits)");
-  }
-
   private Map<String, Object> parse(MultiBucketsAggregation.Bucket bucket, String name) {
     Map<String, Object> result = metricsParser.parse(bucket.getAggregations());
     if (bucket instanceof CompositeAggregation.Bucket compositeBucket) {
-      result.putAll(compositeBucket.getKey());
+      result.putAll(extract(compositeBucket));
     } else if (bucket instanceof Range.Bucket) {
       if (bucket.getDocCount() == 0) {
         return null;