ValuesSource refactoring: Wire up string_stats aggregation (#52875)

csoulios · web-flow · commit d434a9d2658d · 2020-02-27T18:13:15.000+02:00
diff --git a/test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java b/test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java
@@ -135,7 +135,7 @@ public abstract class AggregatorTestCase extends ESTestCase {
     private static final String NESTEDFIELD_PREFIX = "nested_";
     private List<Releasable> releasables = new ArrayList<>();
     private static final String TYPE_NAME = "type";
-    private static ValuesSourceRegistry valuesSourceRegistry;
+    protected static ValuesSourceRegistry valuesSourceRegistry;
 
     // A list of field types that should not be tested, or are not currently supported
     private static List<String> TYPE_TEST_BLACKLIST = List.of(
@@ -683,7 +683,7 @@ private void writeTestDoc(MappedFieldType fieldType, String fieldName, RandomInd
 
         String typeName = fieldType.typeName();
         ValuesSourceType vst = fieldType.getValuesSourceType();
-        
+
         if (vst.equals(CoreValuesSourceType.NUMERIC)) {
             if (typeName.equals(NumberFieldMapper.NumberType.DOUBLE.typeName())) {
                 long encoded = NumericUtils.doubleToSortableLong(Math.abs(randomDouble()));
diff --git a/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/AnalyticsPlugin.java b/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/AnalyticsPlugin.java
@@ -66,7 +66,9 @@ public List<AggregationSpec> getAggregations() {
             new AggregationSpec(
                 StringStatsAggregationBuilder.NAME,
                 StringStatsAggregationBuilder::new,
-                StringStatsAggregationBuilder.PARSER).addResultReader(InternalStringStats::new),
+                StringStatsAggregationBuilder.PARSER)
+                    .addResultReader(InternalStringStats::new)
+                    .setAggregatorRegistrar(StringStatsAggregationBuilder::registerAggregators),
             new AggregationSpec(
                 BoxplotAggregationBuilder.NAME,
                 BoxplotAggregationBuilder::new,
diff --git a/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregationBuilder.java b/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregationBuilder.java
@@ -18,6 +18,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;
@@ -108,6 +109,10 @@ public StringStatsAggregationBuilder showDistribution(boolean showDistribution)
         return this;
     }
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        StringStatsAggregatorFactory.registerAggregators(valuesSourceRegistry);
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(super.hashCode(), showDistribution);
diff --git a/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorFactory.java b/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorFactory.java
@@ -7,14 +7,18 @@
 package org.elasticsearch.xpack.analytics.stringstats;
 
 import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.search.DocValueFormat;
 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.CoreValuesSourceType;
 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;
@@ -34,6 +38,24 @@ class StringStatsAggregatorFactory extends ValuesSourceAggregatorFactory {
         this.showDistribution = showDistribution;
     }
 
+    static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        valuesSourceRegistry.register(StringStatsAggregationBuilder.NAME,
+            CoreValuesSourceType.BYTES, new StringStatsAggregatorSupplier() {
+                @Override
+                public Aggregator build(String name,
+                                        ValuesSource valuesSource,
+                                        boolean showDistribution,
+                                        DocValueFormat format,
+                                        SearchContext context,
+                                        Aggregator parent,
+                                        List<PipelineAggregator> pipelineAggregators,
+                                        Map<String, Object> metaData) throws IOException {
+                    return new StringStatsAggregator(name, showDistribution, (ValuesSource.Bytes) valuesSource,
+                        format, context, parent, pipelineAggregators, metaData);
+                }
+            });
+    }
+
     @Override
     protected Aggregator createUnmapped(SearchContext searchContext,
                                             Aggregator parent,
@@ -50,12 +72,15 @@ protected Aggregator doCreateInternal(ValuesSource valuesSource,
                                           boolean collectsFromSingleBucket,
                                           List<PipelineAggregator> pipelineAggregators,
                                           Map<String, Object> metaData) throws IOException {
-        if (valuesSource instanceof ValuesSource.Bytes == false) {
-            throw new AggregationExecutionException("ValuesSource type " + valuesSource.toString() + "is not supported for aggregation " +
-                this.name());
+        AggregatorSupplier aggregatorSupplier = queryShardContext.getValuesSourceRegistry().getAggregator(config.valueSourceType(),
+            StringStatsAggregationBuilder.NAME);
+
+        if (aggregatorSupplier instanceof StringStatsAggregatorSupplier == false) {
+            throw new AggregationExecutionException("Registry miss-match - expected StringStatsAggregatorSupplier, found [" +
+                aggregatorSupplier.getClass().toString() + "]");
         }
-        return new StringStatsAggregator(name, showDistribution, (ValuesSource.Bytes) valuesSource, config.format(), searchContext, parent,
-                                         pipelineAggregators, metaData);
+        return ((StringStatsAggregatorSupplier) aggregatorSupplier).build(name, valuesSource, showDistribution, config.format(),
+            searchContext, parent, pipelineAggregators, metaData);
     }
 
 }
diff --git a/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorSupplier.java b/x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorSupplier.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.analytics.stringstats;
+
+import org.elasticsearch.search.DocValueFormat;
+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 StringStatsAggregatorSupplier extends AggregatorSupplier {
+
+    Aggregator build(String name,
+                     ValuesSource valuesSource,
+                     boolean showDistribution,
+                     DocValueFormat format,
+                     SearchContext context,
+                     Aggregator parent,
+                     List<PipelineAggregator> pipelineAggregators,
+                     Map<String, Object> metaData) throws IOException;
+}
diff --git a/x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorTests.java b/x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorTests.java
@@ -27,6 +27,7 @@
 import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregator;
 import org.elasticsearch.search.aggregations.support.ValueType;
+import org.junit.BeforeClass;
 
 import java.io.IOException;
 import java.util.List;
@@ -36,6 +37,11 @@
 
 public class StringStatsAggregatorTests extends AggregatorTestCase {
 
+    @BeforeClass()
+    public static void registerBuilder() {
+        StringStatsAggregationBuilder.registerAggregators(valuesSourceRegistry);
+    }
+
     private void testCase(Query query,
                           CheckedConsumer<RandomIndexWriter, IOException> buildIndex,
                           Consumer<InternalStringStats> verify) throws IOException {
@@ -258,5 +264,4 @@ public void testNestedAggregation() throws IOException {
         indexReader.close();
         directory.close();
     }
-
 }
