From 824e841fe21c4f318eb767c79e3adef8fe9a2696 Mon Sep 17 00:00:00 2001
From: Sascha Woo <sascha.woo@gmail.com>
Date: Tue, 9 Aug 2022 11:45:18 +0200
Subject: [PATCH] [elc] add convenient methods to improve access to
 aggregations

---
 .../client/elc/ElasticsearchAggregations.java | 42 ++++++++++++-------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/src/main/java/org/springframework/data/elasticsearch/client/elc/ElasticsearchAggregations.java b/src/main/java/org/springframework/data/elasticsearch/client/elc/ElasticsearchAggregations.java
index a67da95bc..2ede2ec99 100644
--- a/src/main/java/org/springframework/data/elasticsearch/client/elc/ElasticsearchAggregations.java
+++ b/src/main/java/org/springframework/data/elasticsearch/client/elc/ElasticsearchAggregations.java
@@ -18,6 +18,7 @@
 import co.elastic.clients.elasticsearch._types.aggregations.Aggregate;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -28,35 +29,46 @@
  * 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 final Map<String, ElasticsearchAggregation> aggregationsAsMap;
 
-	public ElasticsearchAggregations(List<ElasticsearchAggregation> aggregations) {
+	public ElasticsearchAggregations(Map<String, Aggregate> aggregations) {
 
 		Assert.notNull(aggregations, "aggregations must not be null");
 
-		this.aggregations = aggregations;
-	}
-
-	/**
-	 * convenience constructor taking a map as it is returned from the new Elasticsearch client.
-	 *
-	 * @param aggregationsMap aggregate map
-	 */
-	public ElasticsearchAggregations(Map<String, Aggregate> aggregationsMap) {
+		Map<String, ElasticsearchAggregation> elasticsearchAggregationsAsMap = new HashMap<>(aggregations.size());
+		aggregations.forEach((name, aggregate) -> elasticsearchAggregationsAsMap //
+				.put(name, new ElasticsearchAggregation(new Aggregation(name, aggregate))));
 
-		Assert.notNull(aggregationsMap, "aggregationsMap must not be null");
-
-		aggregations = new ArrayList<>(aggregationsMap.size());
-		aggregationsMap
-				.forEach((name, aggregate) -> aggregations.add(new ElasticsearchAggregation(new Aggregation(name, aggregate))));
+		this.aggregationsAsMap = elasticsearchAggregationsAsMap;
+		this.aggregations = new ArrayList<>(aggregationsAsMap.values());
 	}
 
 	@Override
 	public List<ElasticsearchAggregation> aggregations() {
 		return aggregations;
 	}
+
+	/**
+	 * @return the {@link ElasticsearchAggregation}s keyed by aggregation name.
+	 */
+	public Map<String, ElasticsearchAggregation> aggregationsAsMap() {
+		return aggregationsAsMap;
+	}
+
+	/**
+	 * Returns the aggregation that is associated with the specified name.
+	 * 
+	 * @param the name
+	 * @return the aggregation
+	 */
+	public ElasticsearchAggregation get(String name) {
+		return aggregationsAsMap.get(name);
+	}
+
 }