Skip to content

Commit ce75d61

Browse files
VijayanBVijayanB
committed
Add geohex aggregation
This aggregation will use uber's h3 to group coordinates into H3 cell. Created new aggregation type geohex_grid. The precision will be between 0 and 15. This aggreation has default precision as 5, similar to geohash and geotile. Signed-off-by: Vijayan Balasubramanian <balasvij@amazon.com>
1 parent a2eaf96 commit ce75d61

15 files changed

Lines changed: 1319 additions & 0 deletions

File tree

‎build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,8 @@ publishing {
135135
// Dependencies
136136
//****************************************************************************/
137137
dependencies {
138+
implementation "org.opensearch.plugin:geo:${opensearch_version}"
139+
api project(":libs:h3")
138140
yamlRestTestRuntimeOnly "org.apache.logging.log4j:log4j-core:${versions.log4j}"
139141
testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}"
140142
testImplementation 'org.json:json:20211205'

‎src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import org.opensearch.geospatial.index.query.xyshape.XYShapeQueryBuilder;
3535
import org.opensearch.geospatial.processor.FeatureProcessor;
3636
import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction;
37+
import org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGrid;
38+
import org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder;
3739
import org.opensearch.geospatial.stats.upload.RestUploadStatsAction;
3840
import org.opensearch.geospatial.stats.upload.UploadStats;
3941
import org.opensearch.geospatial.stats.upload.UploadStatsAction;
@@ -120,4 +122,19 @@ public List<QuerySpec<?>> getQueries() {
120122
// Register XYShapeQuery Builder to be delegated for query type: xy_shape
121123
return List.of(new QuerySpec<>(XYShapeQueryBuilder.NAME, XYShapeQueryBuilder::new, XYShapeQueryBuilder::fromXContent));
122124
}
125+
126+
/**
127+
* Registering {@link GeoHexGrid} aggregation on GeoPoint field.
128+
*/
129+
@Override
130+
public List<AggregationSpec> getAggregations() {
131+
132+
final var geoHexGridSpec = new AggregationSpec(
133+
GeoHexGridAggregationBuilder.NAME,
134+
GeoHexGridAggregationBuilder::new,
135+
GeoHexGridAggregationBuilder.PARSER
136+
).addResultReader(GeoHexGrid::new).setAggregatorRegistrar(GeoHexGridAggregationBuilder::registerAggregators);
137+
138+
return List.of(geoHexGridSpec);
139+
}
123140
}

‎src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGrid.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
7+
8+
import java.io.IOException;
9+
import java.util.List;
10+
import java.util.Map;
11+
12+
import org.opensearch.common.io.stream.StreamInput;
13+
import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGrid;
14+
import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket;
15+
import org.opensearch.search.aggregations.InternalAggregations;
16+
17+
/**
18+
* Represents a grid of cells where each cell's location is determined by a h3 cell address.
19+
* All h3CellAddress in a grid are of the same precision
20+
*/
21+
public final class GeoHexGrid extends BaseGeoGrid<GeoHexGridBucket> {
22+
23+
public GeoHexGrid(StreamInput in) throws IOException {
24+
super(in);
25+
}
26+
27+
@Override
28+
public BaseGeoGrid create(List<BaseGeoGridBucket> list) {
29+
return new GeoHexGrid(name, requiredSize, buckets, metadata);
30+
}
31+
32+
@Override
33+
public BaseGeoGridBucket createBucket(InternalAggregations internalAggregations, BaseGeoGridBucket baseGeoGridBucket) {
34+
return new GeoHexGridBucket(baseGeoGridBucket.hashAsLong(), baseGeoGridBucket.getDocCount(), internalAggregations);
35+
}
36+
37+
@Override
38+
public String getWriteableName() {
39+
return GeoHexGridAggregationBuilder.NAME;
40+
}
41+
42+
protected GeoHexGrid(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
43+
super(name, requiredSize, buckets, metadata);
44+
}
45+
46+
@Override
47+
protected Reader<GeoHexGridBucket> getBucketReader() {
48+
return GeoHexGridBucket::new;
49+
}
50+
51+
@Override
52+
protected BaseGeoGrid create(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
53+
return new GeoHexGrid(name, requiredSize, buckets, metadata);
54+
}
55+
56+
@Override
57+
protected GeoHexGridBucket createBucket(long address, long docCount, InternalAggregations internalAggregations) {
58+
return new GeoHexGridBucket(address, docCount, internalAggregations);
59+
}
60+
61+
int getRequiredSize() {
62+
return requiredSize;
63+
}
64+
}

‎src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregationBuilder.java‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
7+
8+
import static org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexHelper.checkPrecisionRange;
9+
10+
import java.io.IOException;
11+
import java.util.Map;
12+
13+
import org.opensearch.OpenSearchParseException;
14+
import org.opensearch.common.geo.GeoBoundingBox;
15+
import org.opensearch.common.io.stream.StreamInput;
16+
import org.opensearch.common.xcontent.ObjectParser;
17+
import org.opensearch.common.xcontent.XContentParser;
18+
import org.opensearch.common.xcontent.support.XContentMapValues;
19+
import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder;
20+
import org.opensearch.geo.search.aggregations.metrics.GeoGridAggregatorSupplier;
21+
import org.opensearch.index.query.QueryShardContext;
22+
import org.opensearch.search.aggregations.AggregationBuilder;
23+
import org.opensearch.search.aggregations.AggregatorFactories;
24+
import org.opensearch.search.aggregations.AggregatorFactory;
25+
import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
26+
import org.opensearch.search.aggregations.support.ValuesSourceConfig;
27+
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
28+
29+
/**
30+
* Aggregation Builder for geo hex grid
31+
*/
32+
public class GeoHexGridAggregationBuilder extends GeoGridAggregationBuilder {
33+
34+
/**
35+
* Aggregation context name
36+
*/
37+
public static final String NAME = "geohex_grid";
38+
public static final ValuesSourceRegistry.RegistryKey<GeoGridAggregatorSupplier> REGISTRY_KEY = new ValuesSourceRegistry.RegistryKey<>(
39+
NAME,
40+
GeoGridAggregatorSupplier.class
41+
);
42+
public static final ObjectParser<GeoHexGridAggregationBuilder, String> PARSER = createParser(
43+
NAME,
44+
GeoHexGridAggregationBuilder::parsePrecision,
45+
GeoHexGridAggregationBuilder::new
46+
);
47+
private static final int DEFAULT_MAX_NUM_CELLS = 10000;
48+
private static final int DEFAULT_PRECISION = 5;
49+
private static final int DEFAULT_SHARD_SIZE = -1;
50+
51+
public GeoHexGridAggregationBuilder(String name) {
52+
super(name);
53+
precision(DEFAULT_PRECISION);
54+
size(DEFAULT_MAX_NUM_CELLS);
55+
shardSize = DEFAULT_SHARD_SIZE;
56+
}
57+
58+
public GeoHexGridAggregationBuilder(StreamInput in) throws IOException {
59+
super(in);
60+
}
61+
62+
@Override
63+
public String getType() {
64+
return NAME;
65+
}
66+
67+
/**
68+
* Register's Geo Hex Aggregation
69+
* @param builder Builder to register new Aggregation
70+
*/
71+
public static void registerAggregators(final ValuesSourceRegistry.Builder builder) {
72+
GeoHexGridAggregatorFactory.registerAggregators(builder);
73+
}
74+
75+
@Override
76+
public GeoGridAggregationBuilder precision(int precision) {
77+
checkPrecisionRange(precision);
78+
this.precision = precision;
79+
return this;
80+
}
81+
82+
protected GeoHexGridAggregationBuilder(
83+
GeoGridAggregationBuilder clone,
84+
AggregatorFactories.Builder factoriesBuilder,
85+
Map<String, Object> metadata
86+
) {
87+
super(clone, factoriesBuilder, metadata);
88+
}
89+
90+
@Override
91+
protected ValuesSourceAggregatorFactory createFactory(
92+
String name,
93+
ValuesSourceConfig config,
94+
int precision,
95+
int requiredSize,
96+
int shardSize,
97+
GeoBoundingBox geoBoundingBox,
98+
QueryShardContext queryShardContext,
99+
AggregatorFactory aggregatorFactory,
100+
AggregatorFactories.Builder builder,
101+
Map<String, Object> metadata
102+
) throws IOException {
103+
return new GeoHexGridAggregatorFactory(
104+
name,
105+
config,
106+
precision,
107+
requiredSize,
108+
shardSize,
109+
geoBoundingBox,
110+
queryShardContext,
111+
aggregatorFactory,
112+
builder,
113+
metadata
114+
);
115+
}
116+
117+
@Override
118+
protected ValuesSourceRegistry.RegistryKey<?> getRegistryKey() {
119+
return REGISTRY_KEY;
120+
}
121+
122+
@Override
123+
protected AggregationBuilder shallowCopy(AggregatorFactories.Builder builder, Map<String, Object> metadata) {
124+
return new GeoHexGridAggregationBuilder(this, builder, metadata);
125+
}
126+
127+
private static int parsePrecision(final XContentParser parser) throws IOException, OpenSearchParseException {
128+
final var token = parser.currentToken();
129+
if (token.equals(XContentParser.Token.VALUE_NUMBER)) {
130+
return XContentMapValues.nodeIntegerValue(parser.intValue());
131+
}
132+
final var precision = parser.text();
133+
return XContentMapValues.nodeIntegerValue(precision);
134+
}
135+
}

‎src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregator.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
7+
8+
import java.io.IOException;
9+
import java.util.List;
10+
import java.util.Map;
11+
12+
import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket;
13+
import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGridAggregator;
14+
import org.opensearch.search.aggregations.Aggregator;
15+
import org.opensearch.search.aggregations.AggregatorFactories;
16+
import org.opensearch.search.aggregations.CardinalityUpperBound;
17+
import org.opensearch.search.aggregations.support.ValuesSource;
18+
import org.opensearch.search.internal.SearchContext;
19+
20+
/**
21+
* Aggregates data expressed as H3 Cell ID.
22+
*/
23+
public class GeoHexGridAggregator extends GeoGridAggregator<GeoHexGrid> {
24+
25+
public GeoHexGridAggregator(
26+
String name,
27+
AggregatorFactories factories,
28+
ValuesSource.Numeric valuesSource,
29+
int requiredSize,
30+
int shardSize,
31+
SearchContext aggregationContext,
32+
Aggregator parent,
33+
CardinalityUpperBound cardinality,
34+
Map<String, Object> metadata
35+
) throws IOException {
36+
super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, cardinality, metadata);
37+
}
38+
39+
@Override
40+
protected GeoHexGrid buildAggregation(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
41+
return new GeoHexGrid(name, requiredSize, buckets, metadata);
42+
}
43+
44+
@Override
45+
protected BaseGeoGridBucket newEmptyBucket() {
46+
return new GeoHexGridBucket(0, 0, null);
47+
}
48+
}

0 commit comments

Comments
 (0)