Vs refactor wire up ranges and date ranges (#52918)

not-napoleon · web-flow · commit 29463c43a23b · 2020-03-04T17:05:34.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
@@ -417,9 +417,13 @@ private void registerAggregations(List<SearchPlugin> plugins) {
         registerAggregation(new AggregationSpec(SignificantTextAggregationBuilder.NAME, SignificantTextAggregationBuilder::new,
                 SignificantTextAggregationBuilder::parse));
         registerAggregation(new AggregationSpec(RangeAggregationBuilder.NAME, RangeAggregationBuilder::new,
-                RangeAggregationBuilder::parse).addResultReader(InternalRange::new));
+                RangeAggregationBuilder::parse)
+                    .addResultReader(InternalRange::new)
+                    .setAggregatorRegistrar(RangeAggregationBuilder::registerAggregators));
         registerAggregation(new AggregationSpec(DateRangeAggregationBuilder.NAME, DateRangeAggregationBuilder::new,
-                DateRangeAggregationBuilder::parse).addResultReader(InternalDateRange::new));
+                DateRangeAggregationBuilder::parse)
+                    .addResultReader(InternalDateRange::new)
+                    .setAggregatorRegistrar(DateRangeAggregationBuilder::registerAggregators));
         registerAggregation(new AggregationSpec(IpRangeAggregationBuilder.NAME, IpRangeAggregationBuilder::new,
                 IpRangeAggregationBuilder::parse).addResultReader(InternalBinaryRange::new));
         registerAggregation(new AggregationSpec(HistogramAggregationBuilder.NAME, HistogramAggregationBuilder::new,
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/AbstractRangeAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/AbstractRangeAggregatorFactory.java
@@ -20,17 +20,21 @@
 package org.elasticsearch.search.aggregations.bucket.range;
 
 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.bucket.range.RangeAggregator.Range;
 import org.elasticsearch.search.aggregations.bucket.range.RangeAggregator.Unmapped;
 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.ValuesSource.Numeric;
 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;
@@ -42,20 +46,44 @@ public class AbstractRangeAggregatorFactory<R extends Range> extends ValuesSourc
     private final InternalRange.Factory<?, ?> rangeFactory;
     private final R[] ranges;
     private final boolean keyed;
+    private final String aggregationTypeName;
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry, String aggregationName) {
+        valuesSourceRegistry.register(aggregationName,
+            List.of(CoreValuesSourceType.NUMERIC, CoreValuesSourceType.DATE, CoreValuesSourceType.BOOLEAN),
+            new RangeAggregatorSupplier() {
+                @Override
+                public Aggregator build(String name,
+                                        AggregatorFactories factories,
+                                        Numeric valuesSource,
+                                        DocValueFormat format,
+                                        InternalRange.Factory rangeFactory,
+                                        Range[] ranges,
+                                        boolean keyed,
+                                        SearchContext context,
+                                        Aggregator parent,
+                                        List<PipelineAggregator> pipelineAggregators,
+                                        Map<String, Object> metaData) throws IOException {
+                    return new RangeAggregator(name, factories, valuesSource, format, rangeFactory, ranges, keyed, context, parent,
+                        pipelineAggregators, metaData);
+                }
+            });
+    }
     public AbstractRangeAggregatorFactory(String name,
-                                            ValuesSourceConfig config,
-                                            R[] ranges,
-                                            boolean keyed,
-                                            InternalRange.Factory<?, ?> rangeFactory,
-                                            QueryShardContext queryShardContext,
-                                            AggregatorFactory parent,
-                                            AggregatorFactories.Builder subFactoriesBuilder,
-                                            Map<String, Object> metaData) throws IOException {
+                                          String aggregationTypeName,
+                                          ValuesSourceConfig config,
+                                          R[] ranges,
+                                          boolean keyed,
+                                          InternalRange.Factory<?, ?> rangeFactory,
+                                          QueryShardContext queryShardContext,
+                                          AggregatorFactory parent,
+                                          AggregatorFactories.Builder subFactoriesBuilder,
+                                          Map<String, Object> metaData) throws IOException {
         super(name, config, queryShardContext, parent, subFactoriesBuilder, metaData);
         this.ranges = ranges;
         this.keyed = keyed;
         this.rangeFactory = rangeFactory;
+        this.aggregationTypeName = aggregationTypeName;
     }
 
     @Override
@@ -73,13 +101,14 @@ protected Aggregator doCreateInternal(ValuesSource valuesSource,
                                             boolean collectsFromSingleBucket,
                                             List<PipelineAggregator> pipelineAggregators,
                                             Map<String, Object> metaData) throws IOException {
-        if (valuesSource instanceof Numeric == false) {
-            throw new AggregationExecutionException("ValuesSource type " + valuesSource.toString() + "is not supported for aggregation " +
-            this.name());
+
+        AggregatorSupplier aggregatorSupplier = queryShardContext.getValuesSourceRegistry().getAggregator(config.valueSourceType(),
+            aggregationTypeName);
+        if (aggregatorSupplier instanceof RangeAggregatorSupplier == false) {
+            throw new AggregationExecutionException("Registry miss-match - expected RangeAggregatorSupplier, found [" +
+                aggregatorSupplier.getClass().toString() + "]");
         }
-        return new RangeAggregator(name, factories, (Numeric) valuesSource, config.format(), rangeFactory, ranges, keyed, searchContext,
-            parent, pipelineAggregators, metaData);
+        return ((RangeAggregatorSupplier)aggregatorSupplier).build(name, factories, (Numeric) valuesSource, config.format(), rangeFactory,
+            ranges, keyed, searchContext, parent, pipelineAggregators, metaData);
     }
-
-
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregationBuilder.java
@@ -30,6 +30,7 @@
 import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 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;
@@ -60,6 +61,10 @@ private static RangeAggregator.Range parseRange(XContentParser parser) throws IO
         return RangeAggregator.Range.fromXContent(parser);
     }
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        AbstractRangeAggregatorFactory.registerAggregators(valuesSourceRegistry, NAME);
+    }
+
     public DateRangeAggregationBuilder(String name) {
         super(name, InternalDateRange.FACTORY);
     }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregatorFactory.java
@@ -30,15 +30,16 @@
 public class DateRangeAggregatorFactory extends AbstractRangeAggregatorFactory<RangeAggregator.Range> {
 
     public DateRangeAggregatorFactory(String name,
-                                        ValuesSourceConfig config,
-                                        RangeAggregator.Range[] ranges,
-                                        boolean keyed,
-                                        InternalRange.Factory<?, ?> rangeFactory,
-                                        QueryShardContext queryShardContext,
-                                        AggregatorFactory parent,
-                                        AggregatorFactories.Builder subFactoriesBuilder,
-                                        Map<String, Object> metaData) throws IOException {
-        super(name, config, ranges, keyed, rangeFactory, queryShardContext, parent, subFactoriesBuilder, metaData);
+                                      ValuesSourceConfig config,
+                                      RangeAggregator.Range[] ranges,
+                                      boolean keyed,
+                                      InternalRange.Factory<?, ?> rangeFactory,
+                                      QueryShardContext queryShardContext,
+                                      AggregatorFactory parent,
+                                      AggregatorFactories.Builder subFactoriesBuilder,
+                                      Map<String, Object> metaData) throws IOException {
+        super(name, DateRangeAggregationBuilder.NAME, config, ranges, keyed, rangeFactory, queryShardContext, parent, subFactoriesBuilder,
+            metaData);
     }
 
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregationBuilder.java
@@ -30,6 +30,7 @@
 import org.elasticsearch.search.aggregations.bucket.range.RangeAggregator.Range;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 import org.elasticsearch.search.aggregations.support.ValuesSourceParserHelper;
+import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 
 import java.io.IOException;
 import java.util.Map;
@@ -58,6 +59,10 @@ private static Range parseRange(XContentParser parser) throws IOException {
         return Range.fromXContent(parser);
     }
 
+    public static void registerAggregators(ValuesSourceRegistry valuesSourceRegistry) {
+        AbstractRangeAggregatorFactory.registerAggregators(valuesSourceRegistry, NAME);
+    }
+
     public RangeAggregationBuilder(String name) {
         super(name, InternalRange.FACTORY);
     }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorFactory.java
@@ -31,10 +31,17 @@
 
 public class RangeAggregatorFactory extends AbstractRangeAggregatorFactory<RangeAggregator.Range> {
 
-    public RangeAggregatorFactory(String name, ValuesSourceConfig config, Range[] ranges, boolean keyed,
-                                  Factory<?, ?> rangeFactory, QueryShardContext queryShardContext, AggregatorFactory parent,
-                                  AggregatorFactories.Builder subFactoriesBuilder, Map<String, Object> metaData) throws IOException {
-        super(name, config, ranges, keyed, rangeFactory, queryShardContext, parent, subFactoriesBuilder, metaData);
+    public RangeAggregatorFactory(String name,
+                                  ValuesSourceConfig config,
+                                  Range[] ranges,
+                                  boolean keyed,
+                                  Factory<?, ?> rangeFactory,
+                                  QueryShardContext queryShardContext,
+                                  AggregatorFactory parent,
+                                  AggregatorFactories.Builder subFactoriesBuilder,
+                                  Map<String, Object> metaData) throws IOException {
+        super(name, RangeAggregationBuilder.NAME, config, ranges, keyed, rangeFactory, queryShardContext, parent, subFactoriesBuilder,
+            metaData);
     }
 
 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorSupplier.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorSupplier.java
@@ -0,0 +1,45 @@
+/*
+ * 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.bucket.range;
+
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.Aggregator;
+import org.elasticsearch.search.aggregations.AggregatorFactories;
+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 RangeAggregatorSupplier extends AggregatorSupplier {
+    Aggregator build(String name,
+                     AggregatorFactories factories,
+                     ValuesSource.Numeric valuesSource,
+                     DocValueFormat format,
+                     InternalRange.Factory rangeFactory,
+                     RangeAggregator.Range[] ranges,
+                     boolean keyed,
+                     SearchContext context,
+                     Aggregator parent,
+                     List<PipelineAggregator> pipelineAggregators,
+                     Map<String, Object> metaData) throws IOException;
+}
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregatorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/range/DateRangeAggregatorTests.java
@@ -224,8 +224,8 @@ public void testKeywordField() {
                 () -> testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
                     iw.addDocument(singleton(new SortedSetDocValuesField("string", new BytesRef("foo"))));
                 }, range -> fail("Should have thrown exception"), fieldType));
-            // I believe this error is coming from improperly parsing the range, not the field.
-            assertEquals("For input string: \"2015-11-13\"", e.getMessage());
+            assertEquals("Field [not_a_number] of type [keyword(indexed,tokenized)] is not supported for aggregation [date_range]",
+                e.getMessage());
         }
 
         public void testBadMissingField() {
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregatorTests.java
@@ -238,7 +238,7 @@ public void testUnsupportedType() {
             () -> testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
                 iw.addDocument(singleton(new SortedSetDocValuesField("string", new BytesRef("foo"))));
             }, range -> fail("Should have thrown exception"), fieldType));
-        assertEquals(e.getMessage(), "Expected numeric type on field [not_a_number], but got [keyword]");
+        assertEquals("Field [not_a_number] of type [keyword(indexed,tokenized)] is not supported for aggregation [range]", e.getMessage());
     }
 
     public void testBadMissingField() {

Original file line number	Diff line number	Diff line change
`@@ -238,7 +238,7 @@ public void testUnsupportedType() {`
`238`	`238`	`() -> testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {`
`239`	`239`	`iw.addDocument(singleton(new SortedSetDocValuesField("string", new BytesRef("foo"))));`
`240`	`240`	`}, range -> fail("Should have thrown exception"), fieldType));`
`241`		`- assertEquals(e.getMessage(), "Expected numeric type on field [not_a_number], but got [keyword]");`
	`241`	`+ assertEquals("Field [not_a_number] of type [keyword(indexed,tokenized)] is not supported for aggregation [range]", e.getMessage());`
`242`	`242`	`}`
`243`	`243`
`244`	`244`	`public void testBadMissingField() {`