Support construct AggregationResponseParser during Aggregator build stage (#108)

* Support construct AggregationResponseParser during Aggregator build stage

* modify the doc

Signed-off-by: penghuo <[email protected]>

Author: penghuo

opensearch/src/main/java/org/opensearch/sql/opensearch/data/value/OpenSearchExprValueFactory.java

@@ -63,7 +63,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
-import lombok.AllArgsConstructor;
+import lombok.Getter;
 import lombok.Setter;
 import org.opensearch.common.time.DateFormatters;
 import org.opensearch.sql.data.model.ExprBooleanValue;
@@ -86,18 +86,22 @@
 import org.opensearch.sql.opensearch.data.utils.Content;
 import org.opensearch.sql.opensearch.data.utils.ObjectContent;
 import org.opensearch.sql.opensearch.data.utils.OpenSearchJsonContent;
+import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser;
 
 /**
  * Construct ExprValue from OpenSearch response.
  */
-@AllArgsConstructor
 public class OpenSearchExprValueFactory {
   /**
    * The Mapping of Field and ExprType.
    */
   @Setter
   private Map<String, ExprType> typeMapping;
 
+  @Getter
+  @Setter
+  private OpenSearchAggregationResponseParser parser;
+
   private static final DateTimeFormatter DATE_TIME_FORMATTER =
       new DateTimeFormatterBuilder()
           .appendOptional(SQL_LITERAL_DATE_TIME_FORMAT)
@@ -131,6 +135,14 @@ public class OpenSearchExprValueFactory {
           .put(OPENSEARCH_BINARY, c -> new OpenSearchExprBinaryValue(c.stringValue()))
           .build();
 
+  /**
+   * Constructor of OpenSearchExprValueFactory.
+   */
+  public OpenSearchExprValueFactory(
+      Map<String, ExprType> typeMapping) {
+    this.typeMapping = typeMapping;
+  }
+
   /**
    * The struct construction has the following assumption. 1. The field has OpenSearch Object
    * data type. https://www.elastic.co/guide/en/elasticsearch/reference/current/object.html 2. The

opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchAggregationResponseParser.java

@@ -103,7 +103,7 @@ public boolean isAggregationResponse() {
    */
   public Iterator<ExprValue> iterator() {
     if (isAggregationResponse()) {
-      return OpenSearchAggregationResponseParser.parse(aggregations).stream().map(entry -> {
+      return exprValueFactory.getParser().parse(aggregations).stream().map(entry -> {
         ImmutableMap.Builder<String, ExprValue> builder = new ImmutableMap.Builder<>();
         for (Map.Entry<String, Object> value : entry.entrySet()) {
           builder.put(value.getKey(), exprValueFactory.construct(value.getKey(), value.getValue()));

opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java

@@ -0,0 +1,51 @@
+/*
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.opensearch.search.aggregations.Aggregations;
+import org.opensearch.search.aggregations.bucket.composite.CompositeAggregation;
+
+/**
+ * Composite Aggregation Parser which include composite aggregation and metric parsers.
+ */
+public class CompositeAggregationParser implements OpenSearchAggregationResponseParser {
+
+  private final MetricParserHelper metricsParser;
+
+  public CompositeAggregationParser(MetricParser... metricParserList) {
+    metricsParser = new MetricParserHelper(Arrays.asList(metricParserList));
+  }
+
+  public CompositeAggregationParser(List<MetricParser> metricParserList) {
+    metricsParser = new MetricParserHelper(metricParserList);
+  }
+
+  @Override
+  public List<Map<String, Object>> parse(Aggregations aggregations) {
+    return ((CompositeAggregation) aggregations.asList().get(0))
+        .getBuckets().stream().map(this::parse).collect(Collectors.toList());
+  }
+
+  private Map<String, Object> parse(CompositeAggregation.Bucket bucket) {
+    Map<String, Object> resultMap = new HashMap<>();
+    resultMap.putAll(bucket.getKey());
+    resultMap.putAll(metricsParser.parse(bucket.getAggregations()));
+    return resultMap;
+  }
+}

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

@@ -0,0 +1,38 @@
+/*
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.Map;
+import lombok.Builder;
+import lombok.Getter;
+import org.opensearch.search.aggregations.Aggregation;
+import org.opensearch.search.aggregations.bucket.filter.Filter;
+
+/**
+ * {@link Filter} Parser.
+ * The current use case is filter aggregation, e.g. avg(age) filter(balance>0). The filter parser
+ * do nothing and return the result from metricsParser.
+ */
+@Builder
+public class FilterParser implements MetricParser {
+
+  private final MetricParser metricsParser;
+
+  @Getter private final String name;
+
+  @Override
+  public Map<String, Object> parse(Aggregation aggregations) {
+    return metricsParser.parse(((Filter) aggregations).getAggregations().asList().get(0));
+  }
+}

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

@@ -0,0 +1,36 @@
+/*
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.Map;
+import org.opensearch.search.aggregations.Aggregation;
+
+/**
+ * Metric Aggregation Parser.
+ */
+public interface MetricParser {
+
+  /**
+   * Get the name of metric parser.
+   */
+  String getName();
+
+  /**
+   * Parse the {@link Aggregation}.
+   *
+   * @param aggregation {@link Aggregation}
+   * @return the map between metric name and metric value.
+   */
+  Map<String, Object> parse(Aggregation aggregation);
+}

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

@@ -0,0 +1,56 @@
+/*
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.search.aggregations.Aggregation;
+import org.opensearch.search.aggregations.Aggregations;
+import org.opensearch.sql.common.utils.StringUtils;
+
+/**
+ * Parse multiple metrics in one bucket.
+ */
+@RequiredArgsConstructor
+public class MetricParserHelper {
+
+  private final Map<String, MetricParser> metricParserMap;
+
+  public MetricParserHelper(List<MetricParser> metricParserList) {
+    metricParserMap =
+        metricParserList.stream().collect(Collectors.toMap(MetricParser::getName, m -> m));
+  }
+
+  /**
+   * Parse {@link Aggregations}.
+   *
+   * @param aggregations {@link Aggregations}
+   * @return the map between metric name and metric value.
+   */
+  public Map<String, Object> parse(Aggregations aggregations) {
+    Map<String, Object> resultMap = new HashMap<>();
+    for (Aggregation aggregation : aggregations) {
+      if (metricParserMap.containsKey(aggregation.getName())) {
+        resultMap.putAll(metricParserMap.get(aggregation.getName()).parse(aggregation));
+      } else {
+        throw new RuntimeException(StringUtils.format("couldn't parse field %s in aggregation "
+            + "response", aggregation.getName()));
+      }
+    }
+    return resultMap;
+  }
+}

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

@@ -0,0 +1,41 @@
+/*
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package org.opensearch.sql.opensearch.response.agg;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import org.opensearch.search.aggregations.Aggregations;
+
+/**
+ * No Bucket Aggregation Parser which include only metric parsers.
+ */
+public class NoBucketAggregationParser implements OpenSearchAggregationResponseParser {
+
+  private final MetricParserHelper metricsParser;
+
+  public NoBucketAggregationParser(MetricParser... metricParserList) {
+    metricsParser = new MetricParserHelper(Arrays.asList(metricParserList));
+  }
+
+  public NoBucketAggregationParser(List<MetricParser> metricParserList) {
+    metricsParser = new MetricParserHelper(metricParserList);
+  }
+
+  @Override
+  public List<Map<String, Object>> parse(Aggregations aggregations) {
+    return Collections.singletonList(metricsParser.parse(aggregations));
+  }
+}