Fix unit tests

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

Author: yuancu

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

@@ -204,6 +204,7 @@ public static Pair<List<AggregationBuilder>, OpenSearchAggregationResponseParser
       // Find group by fields derived from CASE functions and convert them to range queries
       List<Pair<Integer, RangeAggregationBuilder>> groupsByCase =
           groupList.stream()
+              .filter(i -> project != null && i < project.getProjects().size())
               .map(i -> Pair.of(i, project.getNamedProjects().get(i)))
               .filter(
                   p ->

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

@@ -10,13 +10,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import lombok.Getter;
 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;
 
 public class BucketAggregationParser implements OpenSearchAggregationResponseParser {
-  private final OpenSearchAggregationResponseParser subAggParser;
+  @Getter private final OpenSearchAggregationResponseParser subAggParser;
 
   public BucketAggregationParser(OpenSearchAggregationResponseParser subAggParser) {
     this.subAggParser = subAggParser;

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

@@ -7,11 +7,13 @@
 
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
 import lombok.EqualsAndHashCode;
+import lombok.Getter;
 import org.opensearch.search.aggregations.Aggregation;
 import org.opensearch.search.aggregations.Aggregations;
 import org.opensearch.search.aggregations.bucket.MultiBucketsAggregation;
@@ -24,7 +26,7 @@
  */
 @EqualsAndHashCode
 public class LeafBucketAggregationParser implements OpenSearchAggregationResponseParser {
-  private final MetricParserHelper metricsParser;
+  @Getter private final MetricParserHelper metricsParser;
 
   public LeafBucketAggregationParser(MetricParser... metricParserList) {
     metricsParser = new MetricParserHelper(Arrays.asList(metricParserList));
@@ -47,25 +49,19 @@ public List<Map<String, Object>> parse(Aggregations aggregations) {
   private Map<String, Object> parse(MultiBucketsAggregation.Bucket bucket, String name) {
     if (bucket instanceof CompositeAggregation.Bucket compositeBucket) {
       return parse(compositeBucket);
-    } else if (bucket instanceof Range.Bucket rangeBucket) {
-      return parse(rangeBucket, name);
     }
-    return metricsParser.parse(bucket.getAggregations());
-  }
-
-  private Map<String, Object> parse(CompositeAggregation.Bucket bucket) {
-    Map<String, Object> resultMap = new HashMap<>();
-    resultMap.putAll(bucket.getKey());
+    if (bucket instanceof Range.Bucket && bucket.getDocCount() == 0) {
+      return null;
+    }
+    Map<String, Object> resultMap = new LinkedHashMap<>();
+    resultMap.put(name, bucket.getKey());
     resultMap.putAll(metricsParser.parse(bucket.getAggregations()));
     return resultMap;
   }
 
-  private Map<String, Object> parse(Range.Bucket bucket, String name) {
-    if (bucket.getDocCount() == 0) {
-      return null;
-    }
+  private Map<String, Object> parse(CompositeAggregation.Bucket bucket) {
     Map<String, Object> resultMap = new HashMap<>();
-    resultMap.put(name, bucket.getKey());
+    resultMap.putAll(bucket.getKey());
     resultMap.putAll(metricsParser.parse(bucket.getAggregations()));
     return resultMap;
   }

opensearch/src/test/java/org/opensearch/sql/opensearch/request/AggregateAnalyzerTest.java

@@ -46,8 +46,9 @@
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType;
 import org.opensearch.sql.opensearch.request.AggregateAnalyzer.ExpressionNotAnalyzableException;
-import org.opensearch.sql.opensearch.response.agg.CompositeAggregationParser;
+import org.opensearch.sql.opensearch.response.agg.BucketAggregationParser;
 import org.opensearch.sql.opensearch.response.agg.FilterParser;
+import org.opensearch.sql.opensearch.response.agg.LeafBucketAggregationParser;
 import org.opensearch.sql.opensearch.response.agg.MetricParserHelper;
 import org.opensearch.sql.opensearch.response.agg.NoBucketAggregationParser;
 import org.opensearch.sql.opensearch.response.agg.OpenSearchAggregationResponseParser;
@@ -281,9 +282,11 @@ void analyze_groupBy() throws ExpressionNotAnalyzableException {
             + "{\"b\":{\"terms\":{\"field\":\"b.keyword\",\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},"
             + "\"aggregations\":{\"cnt\":{\"value_count\":{\"field\":\"_index\"}}}}}]",
         result.getLeft().toString());
-    assertInstanceOf(CompositeAggregationParser.class, result.getRight());
+    assertInstanceOf(BucketAggregationParser.class, result.getRight());
     MetricParserHelper metricsParser =
-        ((CompositeAggregationParser) result.getRight()).getMetricsParser();
+        ((LeafBucketAggregationParser)
+                ((BucketAggregationParser) result.getRight()).getSubAggParser())
+            .getMetricsParser();
     assertEquals(1, metricsParser.getMetricParserMap().size());
     metricsParser
         .getMetricParserMap()
@@ -592,8 +595,11 @@ private Project createMockProject(List<Integer> refIndex) {
       when(ref.getType()).thenReturn(typeFactory.createSqlType(SqlTypeName.INTEGER));
       rexNodes.add(ref);
     }
+    List<org.apache.calcite.util.Pair<RexNode, String>> namedProjects =
+        rexNodes.stream().map(n -> org.apache.calcite.util.Pair.of(n, n.toString())).toList();
     when(project.getProjects()).thenReturn(rexNodes);
     when(project.getRowType()).thenReturn(rowType);
+    when(project.getNamedProjects()).thenReturn(namedProjects);
     return project;
   }