From bebbc375ae83145ffb4f0f080d7d5233d81d0b73 Mon Sep 17 00:00:00 2001 From: Mark Tozzi Date: Tue, 28 Apr 2020 12:10:21 -0400 Subject: [PATCH] Wire up IpRangeAggregation to ValuesSourceRegistry (#55831) (#55859) --- .../elasticsearch/search/SearchModule.java | 4 +- .../bucket/range/BinaryRangeAggregator.java | 4 +- .../range/BinaryRangeAggregatorFactory.java | 21 ++++++--- .../range/IpRangeAggregationBuilder.java | 5 +++ .../range/IpRangeAggregatorSupplier.java | 44 +++++++++++++++++++ 5 files changed, 69 insertions(+), 9 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregatorSupplier.java diff --git a/server/src/main/java/org/elasticsearch/search/SearchModule.java b/server/src/main/java/org/elasticsearch/search/SearchModule.java index e34f3623558..8f677f77f77 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchModule.java +++ b/server/src/main/java/org/elasticsearch/search/SearchModule.java @@ -460,7 +460,9 @@ public class SearchModule { .addResultReader(InternalDateRange::new) .setAggregatorRegistrar(DateRangeAggregationBuilder::registerAggregators), builder); registerAggregation(new AggregationSpec(IpRangeAggregationBuilder.NAME, IpRangeAggregationBuilder::new, - IpRangeAggregationBuilder.PARSER).addResultReader(InternalBinaryRange::new), builder); + IpRangeAggregationBuilder.PARSER) + .addResultReader(InternalBinaryRange::new) + .setAggregatorRegistrar(IpRangeAggregationBuilder::registerAggregators), builder); registerAggregation(new AggregationSpec(HistogramAggregationBuilder.NAME, HistogramAggregationBuilder::new, HistogramAggregationBuilder.PARSER) .addResultReader(InternalHistogram::new) diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java index 66a06b67d69..fab0dad2f6d 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java @@ -77,11 +77,11 @@ public final class BinaryRangeAggregator extends BucketsAggregator { final Range[] ranges; public BinaryRangeAggregator(String name, AggregatorFactories factories, - ValuesSource.Bytes valuesSource, DocValueFormat format, + ValuesSource valuesSource, DocValueFormat format, List ranges, boolean keyed, SearchContext context, Aggregator parent, Map metadata) throws IOException { super(name, factories, context, parent, metadata); - this.valuesSource = valuesSource; + this.valuesSource = (ValuesSource.Bytes) valuesSource; this.format = format; this.keyed = keyed; this.ranges = ranges.toArray(new Range[0]); diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregatorFactory.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregatorFactory.java index 06585389446..84f2b8def05 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregatorFactory.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregatorFactory.java @@ -23,17 +23,23 @@ import org.elasticsearch.search.aggregations.AggregationExecutionException; import org.elasticsearch.search.aggregations.Aggregator; import org.elasticsearch.search.aggregations.AggregatorFactories.Builder; 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; +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 BinaryRangeAggregatorFactory - extends ValuesSourceAggregatorFactory { +public class BinaryRangeAggregatorFactory extends ValuesSourceAggregatorFactory { + + public static void registerAggregators(ValuesSourceRegistry.Builder builder) { + builder.register(IpRangeAggregationBuilder.NAME, CoreValuesSourceType.IP, (IpRangeAggregatorSupplier) BinaryRangeAggregator::new); + } private final List ranges; private final boolean keyed; @@ -60,11 +66,14 @@ public class BinaryRangeAggregatorFactory SearchContext searchContext, Aggregator parent, boolean collectsFromSingleBucket, Map 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(), + IpRangeAggregationBuilder.NAME); + + if (aggregatorSupplier instanceof IpRangeAggregatorSupplier == false) { + throw new AggregationExecutionException("Registry miss-match - expected IpRangeAggregatorSupplier, found [" + + aggregatorSupplier.getClass().toString() + "]"); } - return new BinaryRangeAggregator(name, factories, (ValuesSource.Bytes) valuesSource, config.format(), + return ((IpRangeAggregatorSupplier) aggregatorSupplier).build(name, factories, valuesSource, config.format(), ranges, keyed, searchContext, parent, metadata); } diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregationBuilder.java index 50141441352..0dae1d23dad 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregationBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregationBuilder.java @@ -40,6 +40,7 @@ import org.elasticsearch.search.aggregations.support.CoreValuesSourceType; import org.elasticsearch.search.aggregations.support.ValuesSourceAggregationBuilder; 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.aggregations.support.ValuesSourceType; import java.io.IOException; @@ -210,6 +211,10 @@ public final class IpRangeAggregationBuilder } } + public static void registerAggregators(ValuesSourceRegistry.Builder builder) { + BinaryRangeAggregatorFactory.registerAggregators(builder); + } + private boolean keyed = false; private List ranges = new ArrayList<>(); diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregatorSupplier.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregatorSupplier.java new file mode 100644 index 00000000000..e9f528432b3 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/IpRangeAggregatorSupplier.java @@ -0,0 +1,44 @@ +/* + * 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.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 IpRangeAggregatorSupplier extends AggregatorSupplier { + + Aggregator build(String name, + AggregatorFactories factories, + ValuesSource valuesSource, + DocValueFormat format, + List ranges, + boolean keyed, + SearchContext context, + Aggregator parent, + Map metadata) throws IOException; +}