Removed ValuesSourceRegistry.registerAny() (#55747)

* Removed VSRegistry.registerAny()
* All ValuesSourceTypes must be registered
explicitly
* Removed lambdas in ValuesSourceRegistry

server/src/main/java/org/elasticsearch/search/aggregations/bucket/missing/MissingAggregator.java

@@ -37,8 +37,13 @@ public class MissingAggregator extends BucketsAggregator implements SingleBucket
 
     private final ValuesSource valuesSource;
 
-    public MissingAggregator(String name, AggregatorFactories factories, ValuesSource valuesSource,
-            SearchContext aggregationContext, Aggregator parent, Map<String, Object> metadata) throws IOException {
+    public MissingAggregator(
+            String name,
+            AggregatorFactories factories,
+            ValuesSource valuesSource,
+            SearchContext aggregationContext,
+            Aggregator parent,
+            Map<String, Object> metadata) throws IOException {
         super(name, factories, aggregationContext, parent, metadata);
         this.valuesSource = valuesSource;
     }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/missing/MissingAggregatorFactory.java

@@ -33,25 +33,13 @@ import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
-import java.util.List;
 import java.util.Map;
 
 public class MissingAggregatorFactory extends ValuesSourceAggregatorFactory {
 
     public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
-        builder.register(
-            MissingAggregationBuilder.NAME,
-            List.of(
-                CoreValuesSourceType.NUMERIC,
-                CoreValuesSourceType.BYTES,
-                CoreValuesSourceType.GEOPOINT,
-                CoreValuesSourceType.RANGE,
-                CoreValuesSourceType.IP,
-                CoreValuesSourceType.BOOLEAN,
-                CoreValuesSourceType.DATE
-            ),
-            (MissingAggregatorSupplier) MissingAggregator::new
-        );
+        builder.register(MissingAggregationBuilder.NAME, CoreValuesSourceType.ALL_CORE,
+            (MissingAggregatorSupplier) MissingAggregator::new);
     }
 
     public MissingAggregatorFactory(String name, ValuesSourceConfig config, QueryShardContext queryShardContext,

server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregator.java

@@ -49,7 +49,7 @@ import java.util.Map;
 /**
  * An aggregator that computes approximate counts of unique values.
  */
-class CardinalityAggregator extends NumericMetricsAggregator.SingleValue {
+public class CardinalityAggregator extends NumericMetricsAggregator.SingleValue {
 
     private final int precision;
     private final ValuesSource valuesSource;
@@ -60,12 +60,13 @@ class CardinalityAggregator extends NumericMetricsAggregator.SingleValue {
 
     private Collector collector;
 
-    CardinalityAggregator(String name,
-                            ValuesSource valuesSource,
-                            int precision,
-                            SearchContext context,
-                            Aggregator parent,
-                            Map<String, Object> metadata) throws IOException {
+    public CardinalityAggregator(
+            String name,
+            ValuesSource valuesSource,
+            int precision,
+            SearchContext context,
+            Aggregator parent,
+            Map<String, Object> metadata) throws IOException {
         super(name, context, parent, metadata);
         this.valuesSource = valuesSource;
         this.precision = precision;

server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorFactory.java

@@ -25,6 +25,7 @@ import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.support.AggregatorSupplier;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
@@ -48,22 +49,9 @@ class CardinalityAggregatorFactory extends ValuesSourceAggregatorFactory {
         this.precisionThreshold = precisionThreshold;
     }
 
-    static void registerAggregators(ValuesSourceRegistry.Builder builder) {
-        builder.registerAny(CardinalityAggregationBuilder.NAME, cardinalityAggregatorSupplier());
-    }
-
-    private static CardinalityAggregatorSupplier cardinalityAggregatorSupplier(){
-        return new CardinalityAggregatorSupplier() {
-            @Override
-            public Aggregator build(String name,
-                                    ValuesSource valuesSource,
-                                    int precision,
-                                    SearchContext context,
-                                    Aggregator parent,
-                                    Map<String, Object> metadata) throws IOException {
-                return new CardinalityAggregator(name, valuesSource, precision, context, parent, metadata);
-            }
-        };
+    public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
+        builder.register(CardinalityAggregationBuilder.NAME, CoreValuesSourceType.ALL_CORE,
+            (CardinalityAggregatorSupplier) CardinalityAggregator::new);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregator.java

@@ -41,15 +41,19 @@ import java.util.Map;
  * This aggregator works in a multi-bucket mode, that is, when serves as a sub-aggregator, a single aggregator instance aggregates the
  * counts for all buckets owned by the parent aggregator)
  */
-class ValueCountAggregator extends NumericMetricsAggregator.SingleValue {
+public class ValueCountAggregator extends NumericMetricsAggregator.SingleValue {
 
     final ValuesSource valuesSource;
 
     // a count per bucket
     LongArray counts;
 
-    ValueCountAggregator(String name, ValuesSource valuesSource,
-            SearchContext aggregationContext, Aggregator parent, Map<String, Object> metadata) throws IOException {
+    public ValueCountAggregator(
+            String name,
+            ValuesSource valuesSource,
+            SearchContext aggregationContext,
+            Aggregator parent,
+            Map<String, Object> metadata) throws IOException {
         super(name, aggregationContext, parent, metadata);
         this.valuesSource = valuesSource;
         if (valuesSource != null) {

server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorFactory.java

@@ -25,6 +25,7 @@ import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.support.AggregatorSupplier;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
@@ -37,17 +38,8 @@ import java.util.Map;
 class ValueCountAggregatorFactory extends ValuesSourceAggregatorFactory {
 
     public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
-        builder.registerAny(ValueCountAggregationBuilder.NAME,
-            new ValueCountAggregatorSupplier() {
-                @Override
-                public Aggregator build(String name,
-                                        ValuesSource valuesSource,
-                                        SearchContext aggregationContext,
-                                        Aggregator parent,
-                                        Map<String, Object> metadata) throws IOException {
-                    return new ValueCountAggregator(name, valuesSource, aggregationContext, parent, metadata);
-                }
-            });
+        builder.register(ValueCountAggregationBuilder.NAME, CoreValuesSourceType.ALL_CORE,
+            (ValueCountAggregatorSupplier) ValueCountAggregator::new);
     }
 
     ValueCountAggregatorFactory(String name, ValuesSourceConfig config, QueryShardContext queryShardContext,

server/src/main/java/org/elasticsearch/search/aggregations/support/CoreValuesSourceType.java

@@ -35,6 +35,8 @@ import org.elasticsearch.search.aggregations.AggregationExecutionException;
 
 import java.time.ZoneId;
 import java.time.ZoneOffset;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Locale;
 import java.util.function.LongSupplier;
 
@@ -42,6 +44,7 @@ import java.util.function.LongSupplier;
  * {@link CoreValuesSourceType} holds the {@link ValuesSourceType} implementations for the core aggregations package.
  */
 public enum CoreValuesSourceType implements ValuesSourceType {
+
     NUMERIC() {
         @Override
         public ValuesSource getEmpty() {
@@ -288,4 +291,6 @@ public enum CoreValuesSourceType implements ValuesSourceType {
         return name().toLowerCase(Locale.ROOT);
     }
 
+    /** List containing all members of the enumeration. */
+    public static List<ValuesSourceType> ALL_CORE = Arrays.asList(CoreValuesSourceType.values());
 }

server/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSourceRegistry.java

@@ -32,7 +32,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Predicate;
 
 /**
  * {@link ValuesSourceRegistry} holds the mapping from {@link ValuesSourceType}s to {@link AggregatorSupplier}s.  DO NOT directly
@@ -42,40 +41,28 @@ import java.util.function.Predicate;
  */
 public class ValuesSourceRegistry {
     public static class Builder {
-        private Map<String, List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>>> aggregatorRegistry = new HashMap<>();
+        private Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry = new HashMap<>();
+
         /**
-         * Register a ValuesSource to Aggregator mapping.
-         *
+         * Register a ValuesSource to Aggregator mapping. This method registers mappings that only apply to a
+         * single {@link ValuesSourceType}
          * @param aggregationName The name of the family of aggregations, typically found via
          *                        {@link ValuesSourceAggregationBuilder#getType()}
-         * @param appliesTo A predicate which accepts the resolved {@link ValuesSourceType} and decides if the given aggregator can be
-         *                  applied to that type.
+         * @param valuesSourceType The ValuesSourceType this mapping applies to.
          * @param aggregatorSupplier An Aggregation-specific specialization of AggregatorSupplier which will construct the mapped aggregator
+         *                           from the aggregation standard set of parameters
          */
-        private void register(String aggregationName, Predicate<ValuesSourceType> appliesTo,
+        public void register(String aggregationName, ValuesSourceType valuesSourceType,
                                           AggregatorSupplier aggregatorSupplier) {
             if (aggregatorRegistry.containsKey(aggregationName) == false) {
                 aggregatorRegistry.put(aggregationName, new ArrayList<>());
             }
-            aggregatorRegistry.get(aggregationName).add( new AbstractMap.SimpleEntry<>(appliesTo, aggregatorSupplier));
-        }
-
-        /**
-         * Register a ValuesSource to Aggregator mapping.  This version provides a convenience method for mappings that only apply to a
-         * single {@link ValuesSourceType}, to allow passing in the type and auto-wrapping it in a predicate
-         * @param aggregationName The name of the family of aggregations, typically found via
-         *                        {@link ValuesSourceAggregationBuilder#getType()}
-         * @param valuesSourceType The ValuesSourceType this mapping applies to.
-         * @param aggregatorSupplier An Aggregation-specific specialization of AggregatorSupplier which will construct the mapped aggregator
-         *                           from the aggregation standard set of parameters
-         */
-        public void register(String aggregationName, ValuesSourceType valuesSourceType, AggregatorSupplier aggregatorSupplier) {
-            register(aggregationName, (candidate) -> valuesSourceType.equals(candidate), aggregatorSupplier);
+            aggregatorRegistry.get(aggregationName).add(new AbstractMap.SimpleEntry<>(valuesSourceType, aggregatorSupplier));
         }
 
         /**
-         * Register a ValuesSource to Aggregator mapping.  This version provides a convenience method for mappings that only apply to a
-         * known list of {@link ValuesSourceType}, to allow passing in the type and auto-wrapping it in a predicate
+         * Register a ValuesSource to Aggregator mapping. This version provides a convenience method for mappings that apply to a
+         * known list of {@link ValuesSourceType}
          *  @param aggregationName The name of the family of aggregations, typically found via
          *                         {@link ValuesSourceAggregationBuilder#getType()}
          * @param valuesSourceTypes The ValuesSourceTypes this mapping applies to.
@@ -83,48 +70,28 @@ public class ValuesSourceRegistry {
          *                           from the aggregation standard set of parameters
          */
         public void register(String aggregationName, List<ValuesSourceType> valuesSourceTypes, AggregatorSupplier aggregatorSupplier) {
-            register(aggregationName, (candidate) -> {
-                for (ValuesSourceType valuesSourceType : valuesSourceTypes) {
-                    if (valuesSourceType.equals(candidate)) {
-                        return true;
-                    }
-                }
-                return false;
-            }, aggregatorSupplier);
-        }
-
-        /**
-         * Register an aggregator that applies to any values source type.  This is a convenience method for aggregations that do not care at
-         * all about the types of their inputs.  Aggregations using this version of registration should not make any other registrations, as
-         * the aggregator registered using this function will be applied in all cases.
-         *
-         * @param aggregationName The name of the family of aggregations, typically found via
-         *                        {@link ValuesSourceAggregationBuilder#getType()}
-         * @param aggregatorSupplier An Aggregation-specific specialization of AggregatorSupplier which will construct the mapped aggregator
-         *                           from the aggregation standard set of parameters.
-         */
-        public void registerAny(String aggregationName, AggregatorSupplier aggregatorSupplier) {
-            register(aggregationName, (ignored) -> true, aggregatorSupplier);
+            for (ValuesSourceType valuesSourceType : valuesSourceTypes) {
+                register(aggregationName, valuesSourceType, aggregatorSupplier);
+            }
         }
 
-
         public ValuesSourceRegistry build() {
             return new ValuesSourceRegistry(aggregatorRegistry);
         }
     }
 
     /** Maps Aggregation names to (ValuesSourceType, Supplier) pairs, keyed by ValuesSourceType */
-    private Map<String, List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>>> aggregatorRegistry;
-    public ValuesSourceRegistry(Map<String, List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>>> aggregatorRegistry) {
+    private Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry;
+    public ValuesSourceRegistry(Map<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> aggregatorRegistry) {
         /*
-         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
+         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()];
         int i = 0;
-        for (Map.Entry<String, List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>>> entry : aggregatorRegistry.entrySet()) {
+        for (Map.Entry<String, List<Map.Entry<ValuesSourceType, AggregatorSupplier>>> entry : aggregatorRegistry.entrySet()) {
             String aggName = entry.getKey();
-            List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>> values = entry.getValue();
+            List<Map.Entry<ValuesSourceType, AggregatorSupplier>> values = entry.getValue();
             Map.Entry newEntry = Map.entry(aggName, List.of(values.toArray()));
             copiedEntries[i++] = newEntry;
         }
@@ -132,9 +99,9 @@ public class ValuesSourceRegistry {
     }
 
     private AggregatorSupplier findMatchingSuppier(ValuesSourceType valuesSourceType,
-                                                   List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>> supportedTypes) {
-        for (Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier> candidate : supportedTypes) {
-            if (candidate.getKey().test(valuesSourceType)) {
+                                                   List<Map.Entry<ValuesSourceType, AggregatorSupplier>> supportedTypes) {
+        for (Map.Entry<ValuesSourceType, AggregatorSupplier> candidate : supportedTypes) {
+            if (candidate.getKey().equals(valuesSourceType)) {
                 return candidate.getValue();
             }
         }

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java

@@ -16,10 +16,16 @@ import org.elasticsearch.plugins.ActionPlugin;
 import org.elasticsearch.plugins.IngestPlugin;
 import org.elasticsearch.plugins.MapperPlugin;
 import org.elasticsearch.plugins.SearchPlugin;
+import org.elasticsearch.search.aggregations.metrics.CardinalityAggregationBuilder;
+import org.elasticsearch.search.aggregations.metrics.CardinalityAggregator;
+import org.elasticsearch.search.aggregations.metrics.CardinalityAggregatorSupplier;
 import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregatorSupplier;
 import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregatorSupplier;
+import org.elasticsearch.search.aggregations.metrics.ValueCountAggregationBuilder;
+import org.elasticsearch.search.aggregations.metrics.ValueCountAggregator;
+import org.elasticsearch.search.aggregations.metrics.ValueCountAggregatorSupplier;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.xpack.core.XPackPlugin;
 import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction;
@@ -73,7 +79,12 @@ public class SpatialPlugin extends GeoPlugin implements ActionPlugin, MapperPlug
 
     @Override
     public List<Consumer<ValuesSourceRegistry.Builder>> getAggregationExtentions() {
-        return List.of(this::registerGeoShapeBoundsAggregator, this::registerGeoShapeCentroidAggregator);
+        return List.of(
+            this::registerGeoShapeBoundsAggregator,
+            this::registerGeoShapeCentroidAggregator,
+            SpatialPlugin::registerValueCountAggregator,
+            SpatialPlugin::registerCardinalityAggregator
+        );
     }
 
     @Override
@@ -92,10 +103,21 @@ public class SpatialPlugin extends GeoPlugin implements ActionPlugin, MapperPlug
         builder.register(GeoCentroidAggregationBuilder.NAME, GeoShapeValuesSourceType.instance(),
             (GeoCentroidAggregatorSupplier) (name, aggregationContext, parent, valuesSource, metadata)
                 -> {
-            if (getLicenseState().isAllowed(XPackLicenseState.Feature.SPATIAL_GEO_CENTROID)) {
-                return new GeoShapeCentroidAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource, metadata);
-            }
+                if (getLicenseState().isAllowed(XPackLicenseState.Feature.SPATIAL_GEO_CENTROID)) {
+                    return new GeoShapeCentroidAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource, metadata);
+                }
                 throw LicenseUtils.newComplianceException("geo_centroid aggregation on geo_shape fields");
             });
     }
+
+    public static void registerValueCountAggregator(ValuesSourceRegistry.Builder builder) {
+        builder.register(ValueCountAggregationBuilder.NAME, GeoShapeValuesSourceType.instance(),
+            (ValueCountAggregatorSupplier) ValueCountAggregator::new
+        );
+    }
+
+    public static void registerCardinalityAggregator(ValuesSourceRegistry.Builder builder) {
+        builder.register(CardinalityAggregationBuilder.NAME, GeoShapeValuesSourceType.instance(),
+            (CardinalityAggregatorSupplier) CardinalityAggregator::new);
+    }
 }