Wire up GeoDistanceAggregation (#55975) (#56042)

This commit is contained in:
Mark Tozzi 2020-04-30 15:43:27 -04:00 committed by GitHub
parent 54dbea6c65
commit d8eb51ed63
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 77 additions and 19 deletions

View File

@ -476,7 +476,9 @@ public class SearchModule {
.addResultReader(InternalAutoDateHistogram::new)
.setAggregatorRegistrar(AutoDateHistogramAggregationBuilder::registerAggregators), builder);
registerAggregation(new AggregationSpec(GeoDistanceAggregationBuilder.NAME, GeoDistanceAggregationBuilder::new,
GeoDistanceAggregationBuilder::parse).addResultReader(InternalGeoDistance::new), builder);
GeoDistanceAggregationBuilder::parse)
.addResultReader(InternalGeoDistance::new)
.setAggregatorRegistrar(GeoDistanceAggregationBuilder::registerAggregators), builder);
registerAggregation(new AggregationSpec(GeoHashGridAggregationBuilder.NAME, GeoHashGridAggregationBuilder::new,
GeoHashGridAggregationBuilder.PARSER)
.addResultReader(InternalGeoHashGrid::new)

View File

@ -39,6 +39,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;
@ -215,6 +216,10 @@ public class GeoDistanceAggregationBuilder extends ValuesSourceAggregationBuilde
}
}
public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
GeoDistanceRangeAggregatorFactory.registerAggregators(builder);
}
private GeoPoint origin;
private List<Range> ranges = new ArrayList<>();
private DistanceUnit unit = DistanceUnit.DEFAULT;
@ -267,9 +272,7 @@ public class GeoDistanceAggregationBuilder extends ValuesSourceAggregationBuilde
@Override
protected ValuesSourceType defaultValueSourceType() {
// TODO: This should probably not be BYTES, but we're not failing tests with BYTES, so needs more tests?
// TODO: this should set defaultValuesSourceType to GEOPOINT
return CoreValuesSourceType.BYTES;
return CoreValuesSourceType.GEOPOINT;
}
@Override

View File

@ -0,0 +1,49 @@
/*
* 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.common.geo.GeoDistance;
import org.elasticsearch.common.unit.DistanceUnit;
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.Map;
public interface GeoDistanceAggregatorSupplier extends AggregatorSupplier {
Aggregator build(
String name,
AggregatorFactories factories,
GeoDistance distanceType,
org.elasticsearch.common.geo.GeoPoint origin,
DistanceUnit units,
ValuesSource valuesSource,
DocValueFormat format,
InternalRange.Factory rangeFactory,
RangeAggregator.Range[] ranges,
boolean keyed,
SearchContext context,
Aggregator parent,
Map<String, Object> metadata) throws IOException;
}

View File

@ -34,18 +34,27 @@ 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.GeoDistanceAggregationBuilder.Range;
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.Map;
import static org.elasticsearch.search.aggregations.support.AggregationUsageService.OTHER_SUBTYPE;
public class GeoDistanceRangeAggregatorFactory extends ValuesSourceAggregatorFactory {
public class GeoDistanceRangeAggregatorFactory
extends ValuesSourceAggregatorFactory {
public static void registerAggregators(ValuesSourceRegistry.Builder builder) {
builder.register(GeoDistanceAggregationBuilder.NAME, CoreValuesSourceType.GEOPOINT,
(GeoDistanceAggregatorSupplier) (name, factories, distanceType, origin, units, valuesSource, format, rangeFactory, ranges, keyed,
context, parent, metadata) -> {
DistanceSource distanceSource = new DistanceSource((ValuesSource.GeoPoint) valuesSource, distanceType, origin, units);
return new RangeAggregator(name, factories, distanceSource, format, rangeFactory, ranges, keyed, context, parent, metadata);
});
}
private final InternalRange.Factory<InternalGeoDistance.Bucket, InternalGeoDistance> rangeFactory = InternalGeoDistance.FACTORY;
private final GeoPoint origin;
@ -81,13 +90,14 @@ public class GeoDistanceRangeAggregatorFactory
Aggregator parent,
boolean collectsFromSingleBucket,
Map<String, Object> metadata) throws IOException {
if (valuesSource instanceof ValuesSource.GeoPoint == false) {
throw new AggregationExecutionException("ValuesSource type " + valuesSource.toString() + "is not supported for aggregation " +
this.name());
AggregatorSupplier aggregatorSupplier = queryShardContext.getValuesSourceRegistry()
.getAggregator(config.valueSourceType(), GeoDistanceAggregationBuilder.NAME);
if (aggregatorSupplier instanceof GeoDistanceAggregatorSupplier == false) {
throw new AggregationExecutionException("Registry miss-match - expected "
+ GeoDistanceAggregatorSupplier.class.getName() + ", found [" + aggregatorSupplier.getClass().toString() + "]");
}
DistanceSource distanceSource = new DistanceSource((ValuesSource.GeoPoint) valuesSource, distanceType, origin, unit);
return new RangeAggregator(name, factories, distanceSource, config.format(), rangeFactory, ranges, keyed, searchContext,
parent, metadata);
return ((GeoDistanceAggregatorSupplier) aggregatorSupplier).build(name, factories, distanceType, origin,
unit, config.toValuesSource(), config.format(), rangeFactory, ranges, keyed, searchContext, parent, metadata);
}
private static class DistanceSource extends ValuesSource.Numeric {
@ -129,10 +139,4 @@ public class GeoDistanceRangeAggregatorFactory
}
}
@Override
public String getStatsSubtype() {
// GeoDistanceRangeAggregatorFactory doesn't register itself with ValuesSourceRegistry
return OTHER_SUBTYPE;
}
}