Skip to content

Commit 4431ed7

Browse files
talevytalevy
authored
Add geo_bounds aggregation support for geo_shape (#55328)
This commit adds a new GeoShapeBoundsAggregator to the spatial plugin and registers it with the GeoShapeValuesSourceType. This enables geo_bounds aggregations on geo_shape fields
1 parent 7fafec0 commit 4431ed7

File tree

9 files changed

+438
-14
lines changed

9 files changed

+438
-14
lines changed

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

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

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

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -32,16 +32,16 @@
3232
import java.util.Objects;
3333

3434
public class InternalGeoBounds extends InternalAggregation implements GeoBounds {
35-
final double top;
36-
final double bottom;
37-
final double posLeft;
38-
final double posRight;
39-
final double negLeft;
40-
final double negRight;
41-
final boolean wrapLongitude;
42-
43-
InternalGeoBounds(String name, double top, double bottom, double posLeft, double posRight,
44-
double negLeft, double negRight, boolean wrapLongitude, Map<String, Object> metadata) {
35+
public final double top;
36+
public final double bottom;
37+
public final double posLeft;
38+
public final double posRight;
39+
public final double negLeft;
40+
public final double negRight;
41+
public final boolean wrapLongitude;
42+
43+
public InternalGeoBounds(String name, double top, double bottom, double posLeft, double posRight,
44+
double negLeft, double negRight, boolean wrapLongitude, Map<String, Object> metadata) {
4545
super(name, metadata);
4646
this.top = top;
4747
this.bottom = bottom;

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,14 @@
1414
import org.elasticsearch.plugins.IngestPlugin;
1515
import org.elasticsearch.plugins.MapperPlugin;
1616
import org.elasticsearch.plugins.SearchPlugin;
17+
import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregationBuilder;
18+
import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregatorSupplier;
19+
import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
1720
import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction;
1821
import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction;
22+
import org.elasticsearch.xpack.spatial.aggregations.metrics.GeoShapeBoundsAggregator;
23+
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSource;
24+
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSourceType;
1925
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper;
2026
import org.elasticsearch.xpack.spatial.index.mapper.PointFieldMapper;
2127
import org.elasticsearch.xpack.spatial.index.mapper.ShapeFieldMapper;
@@ -27,6 +33,7 @@
2733
import java.util.HashMap;
2834
import java.util.List;
2935
import java.util.Map;
36+
import java.util.function.Consumer;
3037

3138
import static java.util.Collections.singletonList;
3239

@@ -53,8 +60,20 @@ public List<QuerySpec<?>> getQueries() {
5360
return singletonList(new QuerySpec<>(ShapeQueryBuilder.NAME, ShapeQueryBuilder::new, ShapeQueryBuilder::fromXContent));
5461
}
5562

63+
@Override
64+
public List<Consumer<ValuesSourceRegistry>> getBareAggregatorRegistrar() {
65+
return List.of(SpatialPlugin::registerGeoShapeBoundsAggregator);
66+
}
67+
5668
@Override
5769
public Map<String, Processor.Factory> getProcessors(Processor.Parameters parameters) {
5870
return Map.of(CircleProcessor.TYPE, new CircleProcessor.Factory());
5971
}
72+
73+
public static void registerGeoShapeBoundsAggregator(ValuesSourceRegistry valuesSourceRegistry) {
74+
valuesSourceRegistry.register(GeoBoundsAggregationBuilder.NAME, GeoShapeValuesSourceType.INSTANCE,
75+
(GeoBoundsAggregatorSupplier) (name, aggregationContext, parent, valuesSource, wrapLongitude, metadata)
76+
-> new GeoShapeBoundsAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource,
77+
wrapLongitude, metadata));
78+
}
6079
}

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

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
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.aggregations.metrics;
8+
9+
import org.apache.lucene.index.LeafReaderContext;
10+
import org.elasticsearch.common.lease.Releasables;
11+
import org.elasticsearch.common.util.BigArrays;
12+
import org.elasticsearch.common.util.DoubleArray;
13+
import org.elasticsearch.search.aggregations.Aggregator;
14+
import org.elasticsearch.search.aggregations.InternalAggregation;
15+
import org.elasticsearch.search.aggregations.LeafBucketCollector;
16+
import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
17+
import org.elasticsearch.search.aggregations.metrics.InternalGeoBounds;
18+
import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
19+
import org.elasticsearch.search.internal.SearchContext;
20+
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSource;
21+
import org.elasticsearch.xpack.spatial.index.mapper.MultiGeoShapeValues;
22+
23+
import java.io.IOException;
24+
import java.util.Map;
25+
26+
public final class GeoShapeBoundsAggregator extends MetricsAggregator {
27+
private final GeoShapeValuesSource valuesSource;
28+
private final boolean wrapLongitude;
29+
private DoubleArray tops;
30+
private DoubleArray bottoms;
31+
private DoubleArray posLefts;
32+
private DoubleArray posRights;
33+
private DoubleArray negLefts;
34+
private DoubleArray negRights;
35+
36+
public GeoShapeBoundsAggregator(String name, SearchContext aggregationContext, Aggregator parent,
37+
GeoShapeValuesSource valuesSource, boolean wrapLongitude, Map<String, Object> metadata) throws IOException {
38+
super(name, aggregationContext, parent, metadata);
39+
this.valuesSource = valuesSource;
40+
this.wrapLongitude = wrapLongitude;
41+
if (valuesSource != null) {
42+
final BigArrays bigArrays = context.bigArrays();
43+
tops = bigArrays.newDoubleArray(1, false);
44+
tops.fill(0, tops.size(), Double.NEGATIVE_INFINITY);
45+
bottoms = bigArrays.newDoubleArray(1, false);
46+
bottoms.fill(0, bottoms.size(), Double.POSITIVE_INFINITY);
47+
posLefts = bigArrays.newDoubleArray(1, false);
48+
posLefts.fill(0, posLefts.size(), Double.POSITIVE_INFINITY);
49+
posRights = bigArrays.newDoubleArray(1, false);
50+
posRights.fill(0, posRights.size(), Double.NEGATIVE_INFINITY);
51+
negLefts = bigArrays.newDoubleArray(1, false);
52+
negLefts.fill(0, negLefts.size(), Double.POSITIVE_INFINITY);
53+
negRights = bigArrays.newDoubleArray(1, false);
54+
negRights.fill(0, negRights.size(), Double.NEGATIVE_INFINITY);
55+
}
56+
}
57+
58+
@Override
59+
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
60+
LeafBucketCollector sub) {
61+
if (valuesSource == null) {
62+
return LeafBucketCollector.NO_OP_COLLECTOR;
63+
}
64+
final BigArrays bigArrays = context.bigArrays();
65+
final MultiGeoShapeValues values = valuesSource.geoShapeValues(ctx);
66+
return new LeafBucketCollectorBase(sub, values) {
67+
@Override
68+
public void collect(int doc, long bucket) throws IOException {
69+
if (bucket >= tops.size()) {
70+
long from = tops.size();
71+
tops = bigArrays.grow(tops, bucket + 1);
72+
tops.fill(from, tops.size(), Double.NEGATIVE_INFINITY);
73+
bottoms = bigArrays.resize(bottoms, tops.size());
74+
bottoms.fill(from, bottoms.size(), Double.POSITIVE_INFINITY);
75+
posLefts = bigArrays.resize(posLefts, tops.size());
76+
posLefts.fill(from, posLefts.size(), Double.POSITIVE_INFINITY);
77+
posRights = bigArrays.resize(posRights, tops.size());
78+
posRights.fill(from, posRights.size(), Double.NEGATIVE_INFINITY);
79+
negLefts = bigArrays.resize(negLefts, tops.size());
80+
negLefts.fill(from, negLefts.size(), Double.POSITIVE_INFINITY);
81+
negRights = bigArrays.resize(negRights, tops.size());
82+
negRights.fill(from, negRights.size(), Double.NEGATIVE_INFINITY);
83+
}
84+
85+
if (values.advanceExact(doc)) {
86+
final int valuesCount = values.docValueCount();
87+
88+
for (int i = 0; i < valuesCount; ++i) {
89+
MultiGeoShapeValues.GeoShapeValue value = values.nextValue();
90+
MultiGeoShapeValues.BoundingBox bounds = value.boundingBox();
91+
double top = Math.max(tops.get(bucket), bounds.top);
92+
double bottom = Math.min(bottoms.get(bucket), bounds.bottom);
93+
double posLeft = Math.min(posLefts.get(bucket), bounds.posLeft);
94+
double posRight = Math.max(posRights.get(bucket), bounds.posRight);
95+
double negLeft = Math.min(negLefts.get(bucket), bounds.negLeft);
96+
double negRight = Math.max(negRights.get(bucket), bounds.negRight);
97+
tops.set(bucket, top);
98+
bottoms.set(bucket, bottom);
99+
posLefts.set(bucket, posLeft);
100+
posRights.set(bucket, posRight);
101+
negLefts.set(bucket, negLeft);
102+
negRights.set(bucket, negRight);
103+
}
104+
}
105+
}
106+
};
107+
}
108+
109+
110+
@Override
111+
public InternalAggregation buildAggregation(long owningBucketOrdinal) {
112+
if (valuesSource == null) {
113+
return buildEmptyAggregation();
114+
}
115+
double top = tops.get(owningBucketOrdinal);
116+
double bottom = bottoms.get(owningBucketOrdinal);
117+
double posLeft = posLefts.get(owningBucketOrdinal);
118+
double posRight = posRights.get(owningBucketOrdinal);
119+
double negLeft = negLefts.get(owningBucketOrdinal);
120+
double negRight = negRights.get(owningBucketOrdinal);
121+
return new InternalGeoBounds(name, top, bottom, posLeft, posRight, negLeft, negRight, wrapLongitude, metadata());
122+
}
123+
124+
@Override
125+
public InternalAggregation buildEmptyAggregation() {
126+
return new InternalGeoBounds(name, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.POSITIVE_INFINITY,
127+
Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY, wrapLongitude, metadata());
128+
}
129+
130+
@Override
131+
public void doClose() {
132+
Releasables.close(tops, bottoms, posLefts, posRights, negLefts, negRights);
133+
}
134+
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeValuesSource.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOExc
2929

3030
};
3131

32-
abstract MultiGeoShapeValues geoShapeValues(LeafReaderContext context);
32+
public abstract MultiGeoShapeValues geoShapeValues(LeafReaderContext context);
3333

3434
@Override
3535
public DocValueBits docsWithValue(LeafReaderContext context) throws IOException {

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeValuesSourceType.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525

2626
public class GeoShapeValuesSourceType implements Writeable, ValuesSourceType {
2727

28-
static GeoShapeValuesSourceType INSTANCE = new GeoShapeValuesSourceType();
28+
public static GeoShapeValuesSourceType INSTANCE = new GeoShapeValuesSourceType();
2929

3030
@Override
3131
public ValuesSource getEmpty() {
@@ -58,7 +58,7 @@ public ValuesSource replaceMissing(ValuesSource valuesSource, Object rawMissing,
5858
final MultiGeoShapeValues.GeoShapeValue missing = MultiGeoShapeValues.GeoShapeValue.missing(rawMissing.toString());
5959
return new GeoShapeValuesSource() {
6060
@Override
61-
MultiGeoShapeValues geoShapeValues(LeafReaderContext context) {
61+
public MultiGeoShapeValues geoShapeValues(LeafReaderContext context) {
6262
MultiGeoShapeValues values = geoShapeValuesSource.geoShapeValues(context);
6363
return new MultiGeoShapeValues() {
6464

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
import org.elasticsearch.index.mapper.TypeParsers;
3636
import org.elasticsearch.index.query.QueryShardContext;
3737
import org.elasticsearch.index.query.VectorGeoShapeQueryProcessor;
38+
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
3839

3940
import java.io.IOException;
4041
import java.util.HashMap;
@@ -170,6 +171,11 @@ public Query existsQuery(QueryShardContext context) {
170171
}
171172
}
172173

174+
@Override
175+
public ValuesSourceType getValuesSourceType() {
176+
return GeoShapeValuesSourceType.INSTANCE;
177+
}
178+
173179
@Override
174180
public GeoShapeWithDocValuesFieldType clone() {
175181
return new GeoShapeWithDocValuesFieldType(this);

0 commit comments

Comments
 (0)