[elc] add convenient methods to improve access to aggregations

Author: xhaggi

src/main/java/org/springframework/data/elasticsearch/client/elc/ElasticsearchAggregations.java

@@ -18,6 +18,8 @@
 import co.elastic.clients.elasticsearch._types.aggregations.Aggregate;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -28,35 +30,72 @@
  * AggregationsContainer implementation for the Elasticsearch aggregations.
  *
  * @author Peter-Josef Meisch
+ * @author Sascha Woo
  * @since 4.4
  */
 public class ElasticsearchAggregations implements AggregationsContainer<List<ElasticsearchAggregation>> {
 
 	private final List<ElasticsearchAggregation> aggregations;
+	private Map<String, ElasticsearchAggregation> aggregationsAsMap;
 
 	public ElasticsearchAggregations(List<ElasticsearchAggregation> aggregations) {
 
 		Assert.notNull(aggregations, "aggregations must not be null");
 
 		this.aggregations = aggregations;
+
+		if (aggregations.isEmpty()) {
+			aggregationsAsMap = Collections.emptyMap();
+		}
 	}
 
 	/**
-	 * convenience constructor taking a map as it is returned from the new Elasticsearch client.
+	 * Convenience constructor taking a map as it is returned from the new Elasticsearch client.
 	 *
-	 * @param aggregationsMap aggregate map
+	 * @param aggregationsAsMap aggregate map
 	 */
-	public ElasticsearchAggregations(Map<String, Aggregate> aggregationsMap) {
+	public ElasticsearchAggregations(Map<String, Aggregate> aggregationsAsMap) {
 
-		Assert.notNull(aggregationsMap, "aggregationsMap must not be null");
+		Assert.notNull(aggregationsAsMap, "aggregationsAsMap must not be null");
 
-		aggregations = new ArrayList<>(aggregationsMap.size());
-		aggregationsMap
+		aggregations = new ArrayList<>(aggregationsAsMap.size());
+		aggregationsAsMap
 				.forEach((name, aggregate) -> aggregations.add(new ElasticsearchAggregation(new Aggregation(name, aggregate))));
+
+		if (aggregations.isEmpty()) {
+			this.aggregationsAsMap = Collections.emptyMap();
+		}
 	}
 
 	@Override
 	public List<ElasticsearchAggregation> aggregations() {
 		return aggregations;
 	}
+
+	/**
+	 * @return the {@link ElasticsearchAggregation}s keyed by aggregation name.
+	 */
+	public Map<String, ElasticsearchAggregation> asMap() {
+
+		if (aggregationsAsMap == null) {
+			Map<String, ElasticsearchAggregation> newAggregationsAsMap = new HashMap<>(aggregations.size());
+			for (ElasticsearchAggregation aggregation : aggregations) {
+				newAggregationsAsMap.put(aggregation.aggregation().getName(), aggregation);
+			}
+			aggregationsAsMap = Collections.unmodifiableMap(newAggregationsAsMap);
+		}
+
+		return aggregationsAsMap;
+	}
+
+	/**
+	 * Returns the aggregation that is associated with the specified name.
+	 * 
+	 * @param the name
+	 * @return the aggregation
+	 */
+	public ElasticsearchAggregation get(String name) {
+		return asMap().get(name);
+	}
+
 }