Skip to content

Commit 3b74015

Browse files
talevytalevy
authored
Add geo_shape support for the geo_centroid aggregation (#55602)
this commit leverages the new geo_shape doc values to register a new geo_centroid aggregator that works on geo_shape field.
1 parent 029a925 commit 3b74015

File tree

17 files changed

+541
-36
lines changed

17 files changed

+541
-36
lines changed

‎server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidAggregatorSupplier.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,6 @@
3030
@FunctionalInterface
3131
public interface GeoCentroidAggregatorSupplier extends AggregatorSupplier {
3232

33-
GeoCentroidAggregator build(String name, SearchContext context, Aggregator parent,
33+
MetricsAggregator build(String name, SearchContext context, Aggregator parent,
3434
ValuesSource valuesSource, Map<String, Object> metadata) throws IOException;
3535
}

‎server/src/main/java/org/elasticsearch/search/aggregations/metrics/InternalGeoCentroid.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ public static double decodeLongitude(long encodedLatLon) {
5353
return GeoEncodingUtils.decodeLongitude((int) (encodedLatLon & 0xFFFFFFFFL));
5454
}
5555

56-
InternalGeoCentroid(String name, GeoPoint centroid, long count, Map<String, Object> metadata) {
56+
public InternalGeoCentroid(String name, GeoPoint centroid, long count, Map<String, Object> metadata) {
5757
super(name, metadata);
5858
assert (centroid == null) == (count == 0);
5959
this.centroid = centroid;

‎x-pack/plugin/core/src/main/java/org/elasticsearch/license/XPackLicenseState.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ public enum Feature {
4545
SECURITY_TOKEN_SERVICE(OperationMode.GOLD, false),
4646
SECURITY_API_KEY_SERVICE(OperationMode.MISSING, false),
4747
SECURITY_AUTHORIZATION_REALM(OperationMode.PLATINUM, true),
48-
SECURITY_AUTHORIZATION_ENGINE(OperationMode.PLATINUM, true);
48+
SECURITY_AUTHORIZATION_ENGINE(OperationMode.PLATINUM, true),
49+
SPATIAL_GEO_CENTROID(OperationMode.GOLD, true);
4950

5051
final OperationMode minimumOperationMode;
5152
final boolean needsActive;

‎x-pack/plugin/spatial/build.gradle‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,15 @@ dependencies {
1818

1919
restResources {
2020
restApi {
21-
includeCore '_common', 'indices', 'index', 'search'
21+
includeCore '_common', 'bulk', 'indices', 'index', 'search'
2222
}
2323
restTests {
2424
includeCore 'geo_shape'
2525
}
2626
}
2727

2828
testClusters.integTest {
29+
setting 'xpack.license.self_generated.type', 'trial'
2930
testDistribution = 'DEFAULT'
3031
}
3132

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,21 +10,27 @@
1010
import org.elasticsearch.geo.GeoPlugin;
1111
import org.elasticsearch.index.mapper.Mapper;
1212
import org.elasticsearch.ingest.Processor;
13+
import org.elasticsearch.license.LicenseUtils;
14+
import org.elasticsearch.license.XPackLicenseState;
1315
import org.elasticsearch.plugins.ActionPlugin;
1416
import org.elasticsearch.plugins.IngestPlugin;
1517
import org.elasticsearch.plugins.MapperPlugin;
1618
import org.elasticsearch.plugins.SearchPlugin;
1719
import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregationBuilder;
1820
import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregatorSupplier;
21+
import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
22+
import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregatorSupplier;
1923
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
24+
import org.elasticsearch.xpack.core.XPackPlugin;
2025
import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction;
2126
import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction;
22-
import org.elasticsearch.xpack.spatial.search.aggregations.metrics.GeoShapeBoundsAggregator;
27+
import org.elasticsearch.xpack.spatial.aggregations.metrics.GeoShapeCentroidAggregator;
2328
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper;
2429
import org.elasticsearch.xpack.spatial.index.mapper.PointFieldMapper;
2530
import org.elasticsearch.xpack.spatial.index.mapper.ShapeFieldMapper;
2631
import org.elasticsearch.xpack.spatial.index.query.ShapeQueryBuilder;
2732
import org.elasticsearch.xpack.spatial.ingest.CircleProcessor;
33+
import org.elasticsearch.xpack.spatial.search.aggregations.metrics.GeoShapeBoundsAggregator;
2834
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSource;
2935
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
3036

@@ -39,6 +45,11 @@
3945

4046
public class SpatialPlugin extends GeoPlugin implements ActionPlugin, MapperPlugin, SearchPlugin, IngestPlugin {
4147

48+
// to be overriden by tests
49+
protected XPackLicenseState getLicenseState() {
50+
return XPackPlugin.getSharedLicenseState();
51+
}
52+
4253
@Override
4354
public List<ActionPlugin.ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
4455
return Arrays.asList(
@@ -62,18 +73,29 @@ public List<QuerySpec<?>> getQueries() {
6273

6374
@Override
6475
public List<Consumer<ValuesSourceRegistry.Builder>> getAggregationExtentions() {
65-
return List.of(SpatialPlugin::registerGeoShapeBoundsAggregator);
76+
return List.of(this::registerGeoShapeBoundsAggregator, this::registerGeoShapeCentroidAggregator);
6677
}
6778

6879
@Override
6980
public Map<String, Processor.Factory> getProcessors(Processor.Parameters parameters) {
7081
return Map.of(CircleProcessor.TYPE, new CircleProcessor.Factory());
7182
}
7283

73-
public static void registerGeoShapeBoundsAggregator(ValuesSourceRegistry.Builder builder) {
84+
public void registerGeoShapeBoundsAggregator(ValuesSourceRegistry.Builder builder) {
7485
builder.register(GeoBoundsAggregationBuilder.NAME, GeoShapeValuesSourceType.instance(),
7586
(GeoBoundsAggregatorSupplier) (name, aggregationContext, parent, valuesSource, wrapLongitude, metadata)
7687
-> new GeoShapeBoundsAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource,
7788
wrapLongitude, metadata));
7889
}
90+
91+
public void registerGeoShapeCentroidAggregator(ValuesSourceRegistry.Builder builder) {
92+
builder.register(GeoCentroidAggregationBuilder.NAME, GeoShapeValuesSourceType.instance(),
93+
(GeoCentroidAggregatorSupplier) (name, aggregationContext, parent, valuesSource, metadata)
94+
-> {
95+
if (getLicenseState().isAllowed(XPackLicenseState.Feature.SPATIAL_GEO_CENTROID)) {
96+
return new GeoShapeCentroidAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource, metadata);
97+
}
98+
throw LicenseUtils.newComplianceException("geo_centroid aggregation on geo_shape fields");
99+
});
100+
}
79101
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/aggregations/metrics/GeoShapeCentroidAggregator.java‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
7+
8+
package org.elasticsearch.xpack.spatial.aggregations.metrics;
9+
10+
import org.apache.lucene.index.LeafReaderContext;
11+
import org.elasticsearch.common.geo.GeoPoint;
12+
import org.elasticsearch.common.lease.Releasables;
13+
import org.elasticsearch.common.util.BigArrays;
14+
import org.elasticsearch.common.util.ByteArray;
15+
import org.elasticsearch.common.util.DoubleArray;
16+
import org.elasticsearch.common.util.LongArray;
17+
import org.elasticsearch.search.aggregations.Aggregator;
18+
import org.elasticsearch.search.aggregations.InternalAggregation;
19+
import org.elasticsearch.search.aggregations.LeafBucketCollector;
20+
import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
21+
import org.elasticsearch.search.aggregations.metrics.CompensatedSum;
22+
import org.elasticsearch.search.aggregations.metrics.InternalGeoCentroid;
23+
import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
24+
import org.elasticsearch.search.internal.SearchContext;
25+
import org.elasticsearch.xpack.spatial.index.fielddata.DimensionalShapeType;
26+
import org.elasticsearch.xpack.spatial.index.fielddata.MultiGeoShapeValues;
27+
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSource;
28+
29+
import java.io.IOException;
30+
import java.util.Map;
31+
32+
/**
33+
* A geo metric aggregator that computes a geo-centroid from a {@code geo_shape} type field
34+
*/
35+
public final class GeoShapeCentroidAggregator extends MetricsAggregator {
36+
private final GeoShapeValuesSource valuesSource;
37+
private DoubleArray lonSum, lonCompensations, latSum, latCompensations, weightSum, weightCompensations;
38+
private LongArray counts;
39+
private ByteArray dimensionalShapeTypes;
40+
41+
public GeoShapeCentroidAggregator(String name, SearchContext context, Aggregator parent,
42+
GeoShapeValuesSource valuesSource, Map<String, Object> metadata) throws IOException {
43+
super(name, context, parent, metadata);
44+
this.valuesSource = valuesSource;
45+
if (valuesSource != null) {
46+
final BigArrays bigArrays = context.bigArrays();
47+
lonSum = bigArrays.newDoubleArray(1, true);
48+
lonCompensations = bigArrays.newDoubleArray(1, true);
49+
latSum = bigArrays.newDoubleArray(1, true);
50+
latCompensations = bigArrays.newDoubleArray(1, true);
51+
weightSum = bigArrays.newDoubleArray(1, true);
52+
weightCompensations = bigArrays.newDoubleArray(1, true);
53+
counts = bigArrays.newLongArray(1, true);
54+
dimensionalShapeTypes = bigArrays.newByteArray(1, true);
55+
}
56+
}
57+
58+
@Override
59+
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
60+
if (valuesSource == null) {
61+
return LeafBucketCollector.NO_OP_COLLECTOR;
62+
}
63+
final BigArrays bigArrays = context.bigArrays();
64+
final MultiGeoShapeValues values = valuesSource.geoShapeValues(ctx);
65+
final CompensatedSum compensatedSumLat = new CompensatedSum(0, 0);
66+
final CompensatedSum compensatedSumLon = new CompensatedSum(0, 0);
67+
final CompensatedSum compensatedSumWeight = new CompensatedSum(0, 0);
68+
69+
return new LeafBucketCollectorBase(sub, values) {
70+
@Override
71+
public void collect(int doc, long bucket) throws IOException {
72+
latSum = bigArrays.grow(latSum, bucket + 1);
73+
lonSum = bigArrays.grow(lonSum, bucket + 1);
74+
weightSum = bigArrays.grow(weightSum, bucket + 1);
75+
lonCompensations = bigArrays.grow(lonCompensations, bucket + 1);
76+
latCompensations = bigArrays.grow(latCompensations, bucket + 1);
77+
weightCompensations = bigArrays.grow(weightCompensations, bucket + 1);
78+
counts = bigArrays.grow(counts, bucket + 1);
79+
dimensionalShapeTypes = bigArrays.grow(dimensionalShapeTypes, bucket + 1);
80+
81+
if (values.advanceExact(doc)) {
82+
final int valueCount = values.docValueCount();
83+
// increment by the number of points for this document
84+
counts.increment(bucket, valueCount);
85+
// Compute the sum of double values with Kahan summation algorithm which is more
86+
// accurate than naive summation.
87+
DimensionalShapeType shapeType = DimensionalShapeType.fromOrdinalByte(dimensionalShapeTypes.get(bucket));
88+
double sumLat = latSum.get(bucket);
89+
double compensationLat = latCompensations.get(bucket);
90+
double sumLon = lonSum.get(bucket);
91+
double compensationLon = lonCompensations.get(bucket);
92+
double sumWeight = weightSum.get(bucket);
93+
double compensatedWeight = weightCompensations.get(bucket);
94+
95+
compensatedSumLat.reset(sumLat, compensationLat);
96+
compensatedSumLon.reset(sumLon, compensationLon);
97+
compensatedSumWeight.reset(sumWeight, compensatedWeight);
98+
99+
// update the sum
100+
for (int i = 0; i < valueCount; ++i) {
101+
MultiGeoShapeValues.GeoShapeValue value = values.nextValue();
102+
int compares = shapeType.compareTo(value.dimensionalShapeType());
103+
if (compares < 0) {
104+
double coordinateWeight = value.weight();
105+
compensatedSumLat.reset(coordinateWeight * value.lat(), 0.0);
106+
compensatedSumLon.reset(coordinateWeight * value.lon(), 0.0);
107+
compensatedSumWeight.reset(coordinateWeight, 0.0);
108+
dimensionalShapeTypes.set(bucket, (byte) value.dimensionalShapeType().ordinal());
109+
} else if (compares == 0) {
110+
double coordinateWeight = value.weight();
111+
// weighted latitude
112+
compensatedSumLat.add(coordinateWeight * value.lat());
113+
// weighted longitude
114+
compensatedSumLon.add(coordinateWeight * value.lon());
115+
// weight
116+
compensatedSumWeight.add(coordinateWeight);
117+
}
118+
// else (compares > 0)
119+
// do not modify centroid calculation since shape is of lower dimension than the running dimension
120+
121+
}
122+
lonSum.set(bucket, compensatedSumLon.value());
123+
lonCompensations.set(bucket, compensatedSumLon.delta());
124+
latSum.set(bucket, compensatedSumLat.value());
125+
latCompensations.set(bucket, compensatedSumLat.delta());
126+
weightSum.set(bucket, compensatedSumWeight.value());
127+
weightCompensations.set(bucket, compensatedSumWeight.delta());
128+
}
129+
}
130+
};
131+
}
132+
133+
@Override
134+
public InternalAggregation buildAggregation(long bucket) {
135+
if (valuesSource == null || bucket >= counts.size()) {
136+
return buildEmptyAggregation();
137+
}
138+
final long bucketCount = counts.get(bucket);
139+
final double bucketWeight = weightSum.get(bucket);
140+
final GeoPoint bucketCentroid = (bucketWeight > 0)
141+
? new GeoPoint(latSum.get(bucket) / bucketWeight, lonSum.get(bucket) / bucketWeight)
142+
: null;
143+
return new InternalGeoCentroid(name, bucketCentroid , bucketCount, metadata());
144+
}
145+
146+
@Override
147+
public InternalAggregation buildEmptyAggregation() {
148+
return new InternalGeoCentroid(name, null, 0L, metadata());
149+
}
150+
151+
@Override
152+
public void doClose() {
153+
Releasables.close(latSum, latCompensations, lonSum, lonCompensations, counts, weightSum, weightCompensations,
154+
dimensionalShapeTypes);
155+
}
156+
}

‎x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/LocalStateSpatialPlugin.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
7+
package org.elasticsearch.xpack.spatial;
8+
9+
import org.apache.lucene.util.LuceneTestCase;
10+
import org.elasticsearch.license.License;
11+
import org.elasticsearch.license.TestUtils;
12+
import org.elasticsearch.license.XPackLicenseState;
13+
import org.elasticsearch.test.VersionUtils;
14+
15+
/**
16+
* This class overrides the {@link SpatialPlugin} in order
17+
* to provide the integration test clusters a hook into a real
18+
* {@link XPackLicenseState}. In the cases that this is used, the
19+
* actual license's operation mode is not important
20+
*/
21+
public class LocalStateSpatialPlugin extends SpatialPlugin {
22+
protected XPackLicenseState getLicenseState() {
23+
TestUtils.UpdatableLicenseState licenseState = new TestUtils.UpdatableLicenseState();
24+
License.OperationMode operationMode = License.OperationMode.TRIAL;
25+
licenseState.update(operationMode, true, VersionUtils.randomVersion(LuceneTestCase.random()));
26+
return licenseState;
27+
}
28+
}

‎x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/SpatialPluginTests.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.spatial;
7+
8+
import org.elasticsearch.ElasticsearchSecurityException;
9+
import org.elasticsearch.license.License;
10+
import org.elasticsearch.license.TestUtils;
11+
import org.elasticsearch.license.XPackLicenseState;
12+
import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
13+
import org.elasticsearch.search.aggregations.metrics.GeoCentroidAggregatorSupplier;
14+
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
15+
import org.elasticsearch.test.ESTestCase;
16+
import org.elasticsearch.test.VersionUtils;
17+
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
18+
19+
import java.util.List;
20+
import java.util.function.Consumer;
21+
22+
import static org.hamcrest.Matchers.equalTo;
23+
24+
public class SpatialPluginTests extends ESTestCase {
25+
26+
public void testGeoCentroidLicenseCheck() {
27+
for (License.OperationMode operationMode : License.OperationMode.values()) {
28+
SpatialPlugin plugin = getPluginWithOperationMode(operationMode);
29+
ValuesSourceRegistry.Builder registryBuilder = new ValuesSourceRegistry.Builder();
30+
List<Consumer<ValuesSourceRegistry.Builder>> registrar = plugin.getAggregationExtentions();
31+
registrar.forEach(c -> c.accept(registryBuilder));
32+
ValuesSourceRegistry registry = registryBuilder.build();
33+
GeoCentroidAggregatorSupplier centroidSupplier = (GeoCentroidAggregatorSupplier) registry.getAggregator(
34+
GeoShapeValuesSourceType.instance(), GeoCentroidAggregationBuilder.NAME);
35+
if (License.OperationMode.TRIAL != operationMode &&
36+
License.OperationMode.compare(operationMode, License.OperationMode.GOLD) < 0) {
37+
ElasticsearchSecurityException exception = expectThrows(ElasticsearchSecurityException.class,
38+
() -> centroidSupplier.build(null, null, null, null, null));
39+
assertThat(exception.getMessage(),
40+
equalTo("current license is non-compliant for [geo_centroid aggregation on geo_shape fields]"));
41+
}
42+
}
43+
}
44+
45+
private SpatialPlugin getPluginWithOperationMode(License.OperationMode operationMode) {
46+
return new SpatialPlugin() {
47+
protected XPackLicenseState getLicenseState() {
48+
TestUtils.UpdatableLicenseState licenseState = new TestUtils.UpdatableLicenseState();
49+
licenseState.update(operationMode, true, VersionUtils.randomVersion(random()));
50+
return licenseState;
51+
}
52+
};
53+
}
54+
}

0 commit comments

Comments
 (0)