Wire up Value Count (#52225)

not-napoleon · web-flow · commit e567ecb48311 · 2020-02-24T15:50:39.000-05:00
diff --git a/server/src/main/java/org/elasticsearch/search/SearchModule.java b/server/src/main/java/org/elasticsearch/search/SearchModule.java
@@ -353,7 +353,9 @@ private void registerAggregations(List<SearchPlugin> plugins) {
         registerAggregation(new AggregationSpec(ExtendedStatsAggregationBuilder.NAME, ExtendedStatsAggregationBuilder::new,
                 ExtendedStatsAggregationBuilder::parse).addResultReader(InternalExtendedStats::new));
         registerAggregation(new AggregationSpec(ValueCountAggregationBuilder.NAME, ValueCountAggregationBuilder::new,
-                ValueCountAggregationBuilder::parse).addResultReader(InternalValueCount::new));
+                ValueCountAggregationBuilder::parse)
+                    .addResultReader(InternalValueCount::new)
+                    .setAggregatorRegistrar(ValueCountAggregationBuilder::registerAggregators));
         registerAggregation(new AggregationSpec(PercentilesAggregationBuilder.NAME, PercentilesAggregationBuilder::new,
                 PercentilesAggregationBuilder::parse)
                     .addResultReader(InternalTDigestPercentiles.NAME, InternalTDigestPercentiles::new)
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorFactory.java
@@ -51,7 +51,7 @@ class CardinalityAggregatorFactory extends ValuesSourceAggregatorFactory {
     }
 
     static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
-        valuesSourceRegistry.register(CardinalityAggregationBuilder.NAME, (ignored) -> true, cardinalityAggregatorSupplier());
+        valuesSourceRegistry.registerAny(CardinalityAggregationBuilder.NAME, cardinalityAggregatorSupplier());
     }
 
     private static CardinalityAggregatorSupplier cardinalityAggregatorSupplier(){
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregationBuilder.java
@@ -34,6 +34,7 @@
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 import org.elasticsearch.search.aggregations.support.ValuesSourceParserHelper;
+import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 
 import java.io.IOException;
@@ -52,6 +53,10 @@ public static AggregationBuilder parse(String aggregationName, XContentParser pa
         return PARSER.parse(parser, new ValueCountAggregationBuilder(aggregationName), null);
     }
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        ValueCountAggregatorFactory.registerAggregators(valuesSourceRegistry);
+    }
+
     public ValueCountAggregationBuilder(String name) {
         super(name);
     }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorFactory.java
@@ -20,13 +20,16 @@
 package org.elasticsearch.search.aggregations.metrics;
 
 import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.AggregatorFactory;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.AggregatorSupplier;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
+import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -35,6 +38,21 @@
 
 class ValueCountAggregatorFactory extends ValuesSourceAggregatorFactory {
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        valuesSourceRegistry.registerAny(ValueCountAggregationBuilder.NAME,
+            new ValueCountAggregatorSupplier() {
+                @Override
+                public Aggregator build(String name,
+                                        ValuesSource valuesSource,
+                                        SearchContext aggregationContext,
+                                        Aggregator parent,
+                                        List<PipelineAggregator> pipelineAggregators,
+                                        Map<String, Object> metaData) throws IOException {
+                    return new ValueCountAggregator(name, valuesSource, aggregationContext, parent, pipelineAggregators, metaData);
+                }
+            });
+    }
+
     ValueCountAggregatorFactory(String name, ValuesSourceConfig config, QueryShardContext queryShardContext,
                                 AggregatorFactory parent, AggregatorFactories.Builder subFactoriesBuilder,
                                 Map<String, Object> metaData) throws IOException {
@@ -56,6 +74,13 @@ protected Aggregator doCreateInternal(ValuesSource valuesSource,
                                             boolean collectsFromSingleBucket,
                                             List<PipelineAggregator> pipelineAggregators,
                                             Map<String, Object> metaData) throws IOException {
-        return new ValueCountAggregator(name, valuesSource, searchContext, parent, pipelineAggregators, metaData);
+        AggregatorSupplier aggregatorSupplier = queryShardContext.getValuesSourceRegistry().getAggregator(config.valueSourceType(),
+            ValueCountAggregationBuilder.NAME);
+        if (aggregatorSupplier instanceof ValueCountAggregatorSupplier == false) {
+            throw new AggregationExecutionException("Registry miss-match - expected ValueCountAggregatorSupplier, found [" +
+                aggregatorSupplier.getClass().toString() + "]");
+        }
+        return ((ValueCountAggregatorSupplier) aggregatorSupplier)
+            .build(name, valuesSource, searchContext, parent, pipelineAggregators,metaData);
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorSupplier.java b/server/src/main/java/org/elasticsearch/search/aggregations/metrics/ValueCountAggregatorSupplier.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.search.aggregations.metrics;
+
+import org.elasticsearch.search.aggregations.Aggregator;
+import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.AggregatorSupplier;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
+import org.elasticsearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+public interface ValueCountAggregatorSupplier extends AggregatorSupplier {
+    Aggregator build(String name,
+                     ValuesSource valuesSource,
+                     SearchContext aggregationContext,
+                     Aggregator parent,
+                     List<PipelineAggregator> pipelineAggregators,
+                     Map<String, Object> metaData) throws IOException;
+}
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
@@ -54,7 +54,7 @@ public class ValuesSourceRegistry {
      * different worker threads.  Thus we want to optimize the read case to be thread safe and fast, which the immutable
      * collections do well.  Using immutable collections requires a copy on write mechanic, thus the somewhat non-intuitive
      * implementation of this method.
-     * @param aggregationName The name of the family of aggregations, typically found via ValuesSourceAggregationBuilder.getType()
+     * @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 aggregatorSupplier An Aggregation-specific specialization of AggregatorSupplier which will construct the mapped aggregator
@@ -75,7 +75,7 @@ public synchronized void register(String aggregationName, Predicate<ValuesSource
     /**
      * 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 ValuesSourceAggregationBuilder.getType()
+     *  @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
@@ -87,7 +87,7 @@ public void register(String aggregationName, ValuesSourceType valuesSourceType,
     /**
      * 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
-     *  @param aggregationName The name of the family of aggregations, typically found via ValuesSourceAggregationBuilder.getType()
+     *  @param aggregationName The name of the family of aggregations, typically found via {@link ValuesSourceAggregationBuilder#getType()}
      * @param valuesSourceTypes The ValuesSourceTypes 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
@@ -103,6 +103,19 @@ public void register(String aggregationName, List<ValuesSourceType> valuesSource
         }, 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);
+    }
+
     private AggregatorSupplier findMatchingSuppier(ValuesSourceType valuesSourceType,
                                                    List<Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier>> supportedTypes) {
         for (Map.Entry<Predicate<ValuesSourceType>, AggregatorSupplier> candidate : supportedTypes) {

Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ class CardinalityAggregatorFactory extends ValuesSourceAggregatorFactory {`
`51`	`51`	`}`
`52`	`52`
`53`	`53`	`static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {`
`54`		`- valuesSourceRegistry.register(CardinalityAggregationBuilder.NAME, (ignored) -> true, cardinalityAggregatorSupplier());`
	`54`	`+ valuesSourceRegistry.registerAny(CardinalityAggregationBuilder.NAME, cardinalityAggregatorSupplier());`
`55`	`55`	`}`
`56`	`56`
`57`	`57`	`private static CardinalityAggregatorSupplier cardinalityAggregatorSupplier(){`