Simplify the ValuesSourceRegistry structure

Follow up to elastic#55747.
imotov · May 4, 2020 · ad56f7b · ad56f7b
1 parent e9a873a
commit ad56f7b
Showing 1 changed file with 7 additions and 10 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java b/server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java
@@ -98,33 +98,30 @@ public ValuesSourceRegistry build() {
 
     /** Maps Aggregation names to (ValuesSourceType, Supplier) pairs, keyed by ValuesSourceType */
     private final AggregationUsageService usageService;
-    private Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry;
+    private Map<String, Map<ValuesSourceType, AggregatorSupplier>> aggregatorRegistry;
     public ValuesSourceRegistry(Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry,
                                 AggregationUsageService usageService) {
         /*
          Make an immutatble copy of our input map. Since this is write once, read many, we'll spend a bit of extra time to shape this
          into a Map.of(), which is more read optimized than just using a hash map.
          */
-        Map.Entry[] copiedEntries = new Map.Entry[aggregatorRegistry.size()];
+        @SuppressWarnings("unchecked")
+        Map.Entry<String, Map<ValuesSourceType, AggregatorSupplier>>[] copiedEntries = new Map.Entry[aggregatorRegistry.size()];
         int i = 0;
         for (Map.Entry<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> entry : aggregatorRegistry.entrySet()) {
             String aggName = entry.getKey();
             List<Map.Entry<ValuesSourceType, AggregatorSupplier>> values = entry.getValue();
-            Map.Entry newEntry = Map.entry(aggName, List.of(values.toArray()));
+            @SuppressWarnings("unchecked") Map.Entry<String, Map<ValuesSourceType, AggregatorSupplier>> newEntry =
+                Map.entry(aggName, Map.ofEntries(values.toArray(new Map.Entry[0])));
             copiedEntries[i++] = newEntry;
         }
         this.aggregatorRegistry = Map.ofEntries(copiedEntries);
         this.usageService = usageService;
     }
 
     private AggregatorSupplier findMatchingSuppier(ValuesSourceType valuesSourceType,
-                                                   List<Map.Entry<ValuesSourceType, AggregatorSupplier>> supportedTypes) {
-        for (Map.Entry<ValuesSourceType, AggregatorSupplier> candidate : supportedTypes) {
-            if (candidate.getKey().equals(valuesSourceType)) {
-                return candidate.getValue();
-            }
-        }
-        return null;
+                                                   Map<ValuesSourceType, AggregatorSupplier> supportedTypes) {
+        return supportedTypes.get(valuesSourceType);
     }
 
     public AggregatorSupplier getAggregator(ValuesSourceType valuesSourceType, String aggregationName) {