Skip to content

Commit 289106a

Browse files
talevytalevy
authored
Refactor GeoHashGrid to be abstract and re-usable (#37742)
This change split out all the specific GeoHash classes for the geohash_grid aggregation into abstract GeoGrid classes that can be re-used for specific hashing types, like `geohash`
1 parent 33cac52 commit 289106a

29 files changed

+778
-410
lines changed

‎client/rest-high-level/src/main/java/org/elasticsearch/client/RestHighLevelClient.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@
9393
import org.elasticsearch.search.aggregations.bucket.filter.FiltersAggregationBuilder;
9494
import org.elasticsearch.search.aggregations.bucket.filter.ParsedFilter;
9595
import org.elasticsearch.search.aggregations.bucket.filter.ParsedFilters;
96-
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder;
96+
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGridAggregationBuilder;
9797
import org.elasticsearch.search.aggregations.bucket.geogrid.ParsedGeoHashGrid;
9898
import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
9999
import org.elasticsearch.search.aggregations.bucket.global.ParsedGlobal;
@@ -1758,7 +1758,7 @@ static List<NamedXContentRegistry.Entry> getDefaultNamedXContents() {
17581758
map.put(GlobalAggregationBuilder.NAME, (p, c) -> ParsedGlobal.fromXContent(p, (String) c));
17591759
map.put(FilterAggregationBuilder.NAME, (p, c) -> ParsedFilter.fromXContent(p, (String) c));
17601760
map.put(InternalSampler.PARSER_NAME, (p, c) -> ParsedSampler.fromXContent(p, (String) c));
1761-
map.put(GeoGridAggregationBuilder.NAME, (p, c) -> ParsedGeoHashGrid.fromXContent(p, (String) c));
1761+
map.put(GeoHashGridAggregationBuilder.NAME, (p, c) -> ParsedGeoHashGrid.fromXContent(p, (String) c));
17621762
map.put(RangeAggregationBuilder.NAME, (p, c) -> ParsedRange.fromXContent(p, (String) c));
17631763
map.put(DateRangeAggregationBuilder.NAME, (p, c) -> ParsedDateRange.fromXContent(p, (String) c));
17641764
map.put(GeoDistanceAggregationBuilder.NAME, (p, c) -> ParsedGeoDistance.fromXContent(p, (String) c));

‎server/src/main/java/org/elasticsearch/search/SearchModule.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@
108108
import org.elasticsearch.search.aggregations.bucket.filter.FiltersAggregationBuilder;
109109
import org.elasticsearch.search.aggregations.bucket.filter.InternalFilter;
110110
import org.elasticsearch.search.aggregations.bucket.filter.InternalFilters;
111-
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder;
111+
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGridAggregationBuilder;
112112
import org.elasticsearch.search.aggregations.bucket.geogrid.InternalGeoHashGrid;
113113
import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
114114
import org.elasticsearch.search.aggregations.bucket.global.InternalGlobal;
@@ -420,8 +420,8 @@ private void registerAggregations(List<SearchPlugin> plugins) {
420420
AutoDateHistogramAggregationBuilder::parse).addResultReader(InternalAutoDateHistogram::new));
421421
registerAggregation(new AggregationSpec(GeoDistanceAggregationBuilder.NAME, GeoDistanceAggregationBuilder::new,
422422
GeoDistanceAggregationBuilder::parse).addResultReader(InternalGeoDistance::new));
423-
registerAggregation(new AggregationSpec(GeoGridAggregationBuilder.NAME, GeoGridAggregationBuilder::new,
424-
GeoGridAggregationBuilder::parse).addResultReader(InternalGeoHashGrid::new));
423+
registerAggregation(new AggregationSpec(GeoHashGridAggregationBuilder.NAME, GeoHashGridAggregationBuilder::new,
424+
GeoHashGridAggregationBuilder::parse).addResultReader(InternalGeoHashGrid::new));
425425
registerAggregation(new AggregationSpec(NestedAggregationBuilder.NAME, NestedAggregationBuilder::new,
426426
NestedAggregationBuilder::parse).addResultReader(InternalNested::new));
427427
registerAggregation(new AggregationSpec(ReverseNestedAggregationBuilder.NAME, ReverseNestedAggregationBuilder::new,

‎server/src/main/java/org/elasticsearch/search/aggregations/AggregationBuilders.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,8 @@
2828
import org.elasticsearch.search.aggregations.bucket.filter.Filters;
2929
import org.elasticsearch.search.aggregations.bucket.filter.FiltersAggregationBuilder;
3030
import org.elasticsearch.search.aggregations.bucket.filter.FiltersAggregator.KeyedFilter;
31-
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder;
32-
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGrid;
31+
import org.elasticsearch.search.aggregations.bucket.geogrid.InternalGeoHashGrid;
32+
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGridAggregationBuilder;
3333
import org.elasticsearch.search.aggregations.bucket.global.Global;
3434
import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
3535
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder;
@@ -244,10 +244,10 @@ public static HistogramAggregationBuilder histogram(String name) {
244244
}
245245

246246
/**
247-
* Create a new {@link GeoHashGrid} aggregation with the given name.
247+
* Create a new {@link InternalGeoHashGrid} aggregation with the given name.
248248
*/
249-
public static GeoGridAggregationBuilder geohashGrid(String name) {
250-
return new GeoGridAggregationBuilder(name);
249+
public static GeoHashGridAggregationBuilder geohashGrid(String name) {
250+
return new GeoHashGridAggregationBuilder(name);
251251
}
252252

253253
/**

‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/BucketPriorityQueue.java‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.search.aggregations.bucket.geogrid;
20+
21+
import org.apache.lucene.util.PriorityQueue;
22+
23+
class BucketPriorityQueue<B extends InternalGeoGridBucket> extends PriorityQueue<B> {
24+
25+
BucketPriorityQueue(int size) {
26+
super(size);
27+
}
28+
29+
@Override
30+
protected boolean lessThan(InternalGeoGridBucket o1, InternalGeoGridBucket o2) {
31+
int cmp = Long.compare(o2.getDocCount(), o1.getDocCount());
32+
if (cmp == 0) {
33+
cmp = o2.compareTo(o1);
34+
if (cmp == 0) {
35+
cmp = System.identityHashCode(o2) - System.identityHashCode(o1);
36+
}
37+
}
38+
return cmp > 0;
39+
}
40+
}

‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/CellIdSource.java‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020

2121
import org.apache.lucene.index.LeafReaderContext;
2222
import org.apache.lucene.index.SortedNumericDocValues;
23-
import org.elasticsearch.common.geo.GeoHashUtils;
2423
import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
2524
import org.elasticsearch.index.fielddata.MultiGeoPointValues;
2625
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
@@ -36,11 +35,13 @@
3635
class CellIdSource extends ValuesSource.Numeric {
3736
private final ValuesSource.GeoPoint valuesSource;
3837
private final int precision;
38+
private final GeoPointLongEncoder encoder;
3939

40-
CellIdSource(GeoPoint valuesSource, int precision) {
40+
CellIdSource(GeoPoint valuesSource, int precision, GeoPointLongEncoder encoder) {
4141
this.valuesSource = valuesSource;
4242
//different GeoPoints could map to the same or different geohash cells.
4343
this.precision = precision;
44+
this.encoder = encoder;
4445
}
4546

4647
public int precision() {
@@ -54,7 +55,7 @@ public boolean isFloatingPoint() {
5455

5556
@Override
5657
public SortedNumericDocValues longValues(LeafReaderContext ctx) {
57-
return new CellValues(valuesSource.geoPointValues(ctx), precision);
58+
return new CellValues(valuesSource.geoPointValues(ctx), precision, encoder);
5859
}
5960

6061
@Override
@@ -67,13 +68,24 @@ public SortedBinaryDocValues bytesValues(LeafReaderContext ctx) {
6768
throw new UnsupportedOperationException();
6869
}
6970

71+
/**
72+
* The encoder to use to convert a geopoint's (lon, lat, precision) into
73+
* a long-encoded bucket key for aggregating.
74+
*/
75+
@FunctionalInterface
76+
public interface GeoPointLongEncoder {
77+
long encode(double lon, double lat, int precision);
78+
}
79+
7080
private static class CellValues extends AbstractSortingNumericDocValues {
7181
private MultiGeoPointValues geoValues;
7282
private int precision;
83+
private GeoPointLongEncoder encoder;
7384

74-
protected CellValues(MultiGeoPointValues geoValues, int precision) {
85+
protected CellValues(MultiGeoPointValues geoValues, int precision, GeoPointLongEncoder encoder) {
7586
this.geoValues = geoValues;
7687
this.precision = precision;
88+
this.encoder = encoder;
7789
}
7890

7991
@Override
@@ -82,8 +94,7 @@ public boolean advanceExact(int docId) throws IOException {
8294
resize(geoValues.docValueCount());
8395
for (int i = 0; i < docValueCount(); ++i) {
8496
org.elasticsearch.common.geo.GeoPoint target = geoValues.nextValue();
85-
values[i] = GeoHashUtils.longEncode(target.getLon(), target.getLat(),
86-
precision);
97+
values[i] = encoder.encode(target.getLon(), target.getLat(), precision);
8798
}
8899
sort();
89100
return true;

‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoHashGrid.java‎ renamed to ‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGrid.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,19 +23,20 @@
2323
import java.util.List;
2424

2525
/**
26-
* A {@code geohash_grid} aggregation. Defines multiple buckets, each representing a cell in a geo-grid of a specific
26+
* A geo-grid aggregation. Defines multiple buckets, each representing a cell in a geo-grid of a specific
2727
* precision.
2828
*/
29-
public interface GeoHashGrid extends MultiBucketsAggregation {
29+
public interface GeoGrid extends MultiBucketsAggregation {
3030

3131
/**
32-
* A bucket that is associated with a {@code geohash_grid} cell. The key of the bucket is the {@code geohash} of the cell
32+
* A bucket that is associated with a geo-grid cell. The key of the bucket is
33+
* the {@link InternalGeoGridBucket#getKeyAsString()} of the cell
3334
*/
3435
interface Bucket extends MultiBucketsAggregation.Bucket {
3536
}
3637

3738
/**
38-
* @return The buckets of this aggregation (each bucket representing a geohash grid cell)
39+
* @return The buckets of this aggregation (each bucket representing a geo-grid cell)
3940
*/
4041
@Override
4142
List<? extends Bucket> getBuckets();

‎server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregationBuilder.java‎

Lines changed: 32 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,12 @@
2020
package org.elasticsearch.search.aggregations.bucket.geogrid;
2121

2222
import org.elasticsearch.ElasticsearchException;
23-
import org.elasticsearch.common.geo.GeoUtils;
23+
import org.elasticsearch.common.ParseField;
2424
import org.elasticsearch.common.io.stream.StreamInput;
2525
import org.elasticsearch.common.io.stream.StreamOutput;
2626
import org.elasticsearch.common.xcontent.ObjectParser;
2727
import org.elasticsearch.common.xcontent.XContentBuilder;
2828
import org.elasticsearch.common.xcontent.XContentParser;
29-
import org.elasticsearch.search.aggregations.AggregationBuilder;
3029
import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
3130
import org.elasticsearch.search.aggregations.AggregatorFactory;
3231
import org.elasticsearch.search.aggregations.bucket.BucketUtils;
@@ -44,31 +43,31 @@
4443
import java.util.Map;
4544
import java.util.Objects;
4645

47-
import static org.elasticsearch.common.geo.GeoUtils.parsePrecision;
48-
49-
public class GeoGridAggregationBuilder extends ValuesSourceAggregationBuilder<ValuesSource.GeoPoint, GeoGridAggregationBuilder>
46+
public abstract class GeoGridAggregationBuilder extends ValuesSourceAggregationBuilder<ValuesSource.GeoPoint, GeoGridAggregationBuilder>
5047
implements MultiBucketAggregationBuilder {
51-
public static final String NAME = "geohash_grid";
52-
public static final int DEFAULT_PRECISION = 5;
53-
public static final int DEFAULT_MAX_NUM_CELLS = 10000;
54-
55-
private static final ObjectParser<GeoGridAggregationBuilder, Void> PARSER;
56-
static {
57-
PARSER = new ObjectParser<>(GeoGridAggregationBuilder.NAME);
58-
ValuesSourceParserHelper.declareGeoFields(PARSER, false, false);
59-
PARSER.declareField((parser, builder, context) -> builder.precision(parsePrecision(parser)), GeoHashGridParams.FIELD_PRECISION,
60-
org.elasticsearch.common.xcontent.ObjectParser.ValueType.INT);
61-
PARSER.declareInt(GeoGridAggregationBuilder::size, GeoHashGridParams.FIELD_SIZE);
62-
PARSER.declareInt(GeoGridAggregationBuilder::shardSize, GeoHashGridParams.FIELD_SHARD_SIZE);
63-
}
48+
/* recognized field names in JSON */
49+
static final ParseField FIELD_PRECISION = new ParseField("precision");
50+
static final ParseField FIELD_SIZE = new ParseField("size");
51+
static final ParseField FIELD_SHARD_SIZE = new ParseField("shard_size");
6452

65-
public static GeoGridAggregationBuilder parse(String aggregationName, XContentParser parser) throws IOException {
66-
return PARSER.parse(parser, new GeoGridAggregationBuilder(aggregationName), null);
53+
protected int precision;
54+
protected int requiredSize;
55+
protected int shardSize;
56+
57+
@FunctionalInterface
58+
protected interface PrecisionParser {
59+
int parse(XContentParser parser) throws IOException;
6760
}
6861

69-
private int precision = DEFAULT_PRECISION;
70-
private int requiredSize = DEFAULT_MAX_NUM_CELLS;
71-
private int shardSize = -1;
62+
public static ObjectParser<GeoGridAggregationBuilder, Void> createParser(String name, PrecisionParser precisionParser) {
63+
ObjectParser<GeoGridAggregationBuilder, Void> parser = new ObjectParser<>(name);
64+
ValuesSourceParserHelper.declareGeoFields(parser, false, false);
65+
parser.declareField((p, builder, context) -> builder.precision(precisionParser.parse(p)), FIELD_PRECISION,
66+
org.elasticsearch.common.xcontent.ObjectParser.ValueType.INT);
67+
parser.declareInt(GeoGridAggregationBuilder::size, FIELD_SIZE);
68+
parser.declareInt(GeoGridAggregationBuilder::shardSize, FIELD_SHARD_SIZE);
69+
return parser;
70+
}
7271

7372
public GeoGridAggregationBuilder(String name) {
7473
super(name, ValuesSourceType.GEOPOINT, ValueType.GEOPOINT);
@@ -79,11 +78,7 @@ protected GeoGridAggregationBuilder(GeoGridAggregationBuilder clone, Builder fac
7978
this.precision = clone.precision;
8079
this.requiredSize = clone.requiredSize;
8180
this.shardSize = clone.shardSize;
82-
}
8381

84-
@Override
85-
protected AggregationBuilder shallowCopy(Builder factoriesBuilder, Map<String, Object> metaData) {
86-
return new GeoGridAggregationBuilder(this, factoriesBuilder, metaData);
8782
}
8883

8984
/**
@@ -103,10 +98,12 @@ protected void innerWriteTo(StreamOutput out) throws IOException {
10398
out.writeVInt(shardSize);
10499
}
105100

106-
public GeoGridAggregationBuilder precision(int precision) {
107-
this.precision = GeoUtils.checkPrecisionRange(precision);
108-
return this;
109-
}
101+
/**
102+
* method to validate and set the precision value
103+
* @param precision the precision to set for the aggregation
104+
* @return the {@link GeoGridAggregationBuilder} builder
105+
*/
106+
public abstract GeoGridAggregationBuilder precision(int precision);
110107

111108
public int precision() {
112109
return precision;
@@ -154,7 +151,7 @@ public int shardSize() {
154151

155152
if (requiredSize <= 0 || shardSize <= 0) {
156153
throw new ElasticsearchException(
157-
"parameters [required_size] and [shard_size] must be >0 in geohash_grid aggregation [" + name + "].");
154+
"parameters [required_size] and [shard_size] must be > 0 in " + getType() + " aggregation [" + name + "].");
158155
}
159156

160157
if (shardSize < requiredSize) {
@@ -166,10 +163,10 @@ public int shardSize() {
166163

167164
@Override
168165
protected XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
169-
builder.field(GeoHashGridParams.FIELD_PRECISION.getPreferredName(), precision);
170-
builder.field(GeoHashGridParams.FIELD_SIZE.getPreferredName(), requiredSize);
166+
builder.field(FIELD_PRECISION.getPreferredName(), precision);
167+
builder.field(FIELD_SIZE.getPreferredName(), requiredSize);
171168
if (shardSize > -1) {
172-
builder.field(GeoHashGridParams.FIELD_SHARD_SIZE.getPreferredName(), shardSize);
169+
builder.field(FIELD_SHARD_SIZE.getPreferredName(), shardSize);
173170
}
174171
return builder;
175172
}
@@ -193,11 +190,4 @@ protected boolean innerEquals(Object obj) {
193190
protected int innerHashCode() {
194191
return Objects.hash(precision, requiredSize, shardSize);
195192
}
196-
197-
@Override
198-
public String getType() {
199-
return NAME;
200-
}
201-
202-
203193
}

0 commit comments

Comments
 (0)