Add geohex aggregation (#160)

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 <[email protected]>

build.gradle

@@ -135,6 +135,8 @@ publishing {
 // Dependencies
 //****************************************************************************/
 dependencies {
+    implementation "org.opensearch.plugin:geo:${opensearch_version}"
+    api project(":libs:h3")
     yamlRestTestRuntimeOnly "org.apache.logging.log4j:log4j-core:${versions.log4j}"
     testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}"
     testImplementation 'org.json:json:20211205'

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

@@ -34,6 +34,8 @@
 import org.opensearch.geospatial.index.query.xyshape.XYShapeQueryBuilder;
 import org.opensearch.geospatial.processor.FeatureProcessor;
 import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction;
+import org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGrid;
+import org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder;
 import org.opensearch.geospatial.stats.upload.RestUploadStatsAction;
 import org.opensearch.geospatial.stats.upload.UploadStats;
 import org.opensearch.geospatial.stats.upload.UploadStatsAction;
@@ -120,4 +122,19 @@ public List<QuerySpec<?>> getQueries() {
         // Register XYShapeQuery Builder to be delegated for query type: xy_shape
         return List.of(new QuerySpec<>(XYShapeQueryBuilder.NAME, XYShapeQueryBuilder::new, XYShapeQueryBuilder::fromXContent));
     }
+
+    /**
+     * Registering {@link GeoHexGrid} aggregation on GeoPoint field.
+     */
+    @Override
+    public List<AggregationSpec> getAggregations() {
+
+        final var geoHexGridSpec = new AggregationSpec(
+            GeoHexGridAggregationBuilder.NAME,
+            GeoHexGridAggregationBuilder::new,
+            GeoHexGridAggregationBuilder.PARSER
+        ).addResultReader(GeoHexGrid::new).setAggregatorRegistrar(GeoHexGridAggregationBuilder::registerAggregators);
+
+        return List.of(geoHexGridSpec);
+    }
 }

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGrid;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket;
+import org.opensearch.search.aggregations.InternalAggregations;
+
+/**
+ * Represents a grid of cells where each cell's location is determined by a h3 cell address.
+ * All h3CellAddress in a grid are of the same precision
+ */
+public final class GeoHexGrid extends BaseGeoGrid<GeoHexGridBucket> {
+
+    public GeoHexGrid(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public BaseGeoGrid create(List<BaseGeoGridBucket> list) {
+        return new GeoHexGrid(name, requiredSize, buckets, metadata);
+    }
+
+    @Override
+    public BaseGeoGridBucket createBucket(InternalAggregations internalAggregations, BaseGeoGridBucket baseGeoGridBucket) {
+        return new GeoHexGridBucket(baseGeoGridBucket.hashAsLong(), baseGeoGridBucket.getDocCount(), internalAggregations);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return GeoHexGridAggregationBuilder.NAME;
+    }
+
+    protected GeoHexGrid(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
+        super(name, requiredSize, buckets, metadata);
+    }
+
+    @Override
+    protected Reader<GeoHexGridBucket> getBucketReader() {
+        return GeoHexGridBucket::new;
+    }
+
+    @Override
+    protected BaseGeoGrid create(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
+        return new GeoHexGrid(name, requiredSize, buckets, metadata);
+    }
+
+    @Override
+    protected GeoHexGridBucket createBucket(long address, long docCount, InternalAggregations internalAggregations) {
+        return new GeoHexGridBucket(address, docCount, internalAggregations);
+    }
+
+    int getRequiredSize() {
+        return requiredSize;
+    }
+}

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

@@ -0,0 +1,135 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
+
+import static org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexHelper.checkPrecisionRange;
+
+import java.io.IOException;
+import java.util.Map;
+
+import org.opensearch.OpenSearchParseException;
+import org.opensearch.common.geo.GeoBoundingBox;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.xcontent.ObjectParser;
+import org.opensearch.common.xcontent.XContentParser;
+import org.opensearch.common.xcontent.support.XContentMapValues;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder;
+import org.opensearch.geo.search.aggregations.metrics.GeoGridAggregatorSupplier;
+import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.search.aggregations.AggregationBuilder;
+import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.AggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.opensearch.search.aggregations.support.ValuesSourceConfig;
+import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
+
+/**
+ * Aggregation Builder for geo hex grid
+ */
+public class GeoHexGridAggregationBuilder extends GeoGridAggregationBuilder {
+
+    /**
+     * Aggregation context name
+     */
+    public static final String NAME = "geohex_grid";
+    public static final ValuesSourceRegistry.RegistryKey<GeoGridAggregatorSupplier> REGISTRY_KEY = new ValuesSourceRegistry.RegistryKey<>(
+        NAME,
+        GeoGridAggregatorSupplier.class
+    );
+    public static final ObjectParser<GeoHexGridAggregationBuilder, String> PARSER = createParser(
+        NAME,
+        GeoHexGridAggregationBuilder::parsePrecision,
+        GeoHexGridAggregationBuilder::new
+    );
+    private static final int DEFAULT_MAX_NUM_CELLS = 10000;
+    private static final int DEFAULT_PRECISION = 5;
+    private static final int DEFAULT_SHARD_SIZE = -1;
+
+    public GeoHexGridAggregationBuilder(String name) {
+        super(name);
+        precision(DEFAULT_PRECISION);
+        size(DEFAULT_MAX_NUM_CELLS);
+        shardSize = DEFAULT_SHARD_SIZE;
+    }
+
+    public GeoHexGridAggregationBuilder(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public String getType() {
+        return NAME;
+    }
+
+    /**
+     * Register's Geo Hex Aggregation
+     * @param builder Builder to register new Aggregation
+     */
+    public static void registerAggregators(final ValuesSourceRegistry.Builder builder) {
+        GeoHexGridAggregatorFactory.registerAggregators(builder);
+    }
+
+    @Override
+    public GeoGridAggregationBuilder precision(int precision) {
+        checkPrecisionRange(precision);
+        this.precision = precision;
+        return this;
+    }
+
+    protected GeoHexGridAggregationBuilder(
+        GeoGridAggregationBuilder clone,
+        AggregatorFactories.Builder factoriesBuilder,
+        Map<String, Object> metadata
+    ) {
+        super(clone, factoriesBuilder, metadata);
+    }
+
+    @Override
+    protected ValuesSourceAggregatorFactory createFactory(
+        String name,
+        ValuesSourceConfig config,
+        int precision,
+        int requiredSize,
+        int shardSize,
+        GeoBoundingBox geoBoundingBox,
+        QueryShardContext queryShardContext,
+        AggregatorFactory aggregatorFactory,
+        AggregatorFactories.Builder builder,
+        Map<String, Object> metadata
+    ) throws IOException {
+        return new GeoHexGridAggregatorFactory(
+            name,
+            config,
+            precision,
+            requiredSize,
+            shardSize,
+            geoBoundingBox,
+            queryShardContext,
+            aggregatorFactory,
+            builder,
+            metadata
+        );
+    }
+
+    @Override
+    protected ValuesSourceRegistry.RegistryKey<?> getRegistryKey() {
+        return REGISTRY_KEY;
+    }
+
+    @Override
+    protected AggregationBuilder shallowCopy(AggregatorFactories.Builder builder, Map<String, Object> metadata) {
+        return new GeoHexGridAggregationBuilder(this, builder, metadata);
+    }
+
+    private static int parsePrecision(final XContentParser parser) throws IOException, OpenSearchParseException {
+        final var token = parser.currentToken();
+        if (token.equals(XContentParser.Token.VALUE_NUMBER)) {
+            return XContentMapValues.nodeIntegerValue(parser.intValue());
+        }
+        final var precision = parser.text();
+        return XContentMapValues.nodeIntegerValue(precision);
+    }
+}

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.geospatial.search.aggregations.bucket.geogrid;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGridAggregator;
+import org.opensearch.search.aggregations.Aggregator;
+import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.aggregations.CardinalityUpperBound;
+import org.opensearch.search.aggregations.support.ValuesSource;
+import org.opensearch.search.internal.SearchContext;
+
+/**
+ * Aggregates data expressed as H3 Cell ID.
+ */
+public class GeoHexGridAggregator extends GeoGridAggregator<GeoHexGrid> {
+
+    public GeoHexGridAggregator(
+        String name,
+        AggregatorFactories factories,
+        ValuesSource.Numeric valuesSource,
+        int requiredSize,
+        int shardSize,
+        SearchContext aggregationContext,
+        Aggregator parent,
+        CardinalityUpperBound cardinality,
+        Map<String, Object> metadata
+    ) throws IOException {
+        super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, cardinality, metadata);
+    }
+
+    @Override
+    protected GeoHexGrid buildAggregation(String name, int requiredSize, List<BaseGeoGridBucket> buckets, Map<String, Object> metadata) {
+        return new GeoHexGrid(name, requiredSize, buckets, metadata);
+    }
+
+    @Override
+    protected BaseGeoGridBucket newEmptyBucket() {
+        return new GeoHexGridBucket(0, 0, null);
+    }
+}