Add Geohex aggregation on geo_shape field (#91956)

This commit adds support for geohex aggregation on geoshape fields using cartesian geometry.

docs/changelog/91956.yaml

@@ -0,0 +1,5 @@
+pr: 91956
+summary: Geohex aggregation on `geo_shape` field
+area: Geo
+type: feature
+issues: []

x-pack/plugin/spatial/src/internalClusterTest/java/org/elasticsearch/xpack/spatial/search/GeoGridAggAndQueryConsistencyIT.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.spatial.LocalStateSpatialPlugin;
+import org.elasticsearch.xpack.spatial.common.H3CartesianUtil;
 import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper;
 import org.elasticsearch.xpack.spatial.index.query.GeoGridQueryBuilder;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder;
@@ -87,6 +88,10 @@ public void testGeoShapeGeoTile() throws IOException {
         );
     }
 
+    public void testGeoShapeGeoHex() throws IOException {
+        doTestGeohexGrid(GeoShapeWithDocValuesFieldMapper.CONTENT_TYPE, () -> GeometryTestUtils.randomGeometryWithoutCircle(0, false));
+    }
+
     private void doTestGeohashGrid(String fieldType, Supplier<Geometry> randomGeometriesSupplier) throws IOException {
         doTestGrid(
             1,
@@ -124,43 +129,13 @@ private void doTestGeohexGrid(String fieldType, Supplier<Geometry> randomGeometr
             }
             return points;
         },
-            this::toGeoHexRectangle,
+            h3 -> H3CartesianUtil.toBoundingBox(H3.stringToH3(h3)),
             GeoHexGridAggregationBuilder::new,
             (s1, s2) -> new GeoGridQueryBuilder(s1).setGridId(GeoGridQueryBuilder.Grid.GEOHEX, s2),
             randomGeometriesSupplier
         );
     }
 
-    private Rectangle toGeoHexRectangle(String bucketKey) {
-        final long h3 = H3.stringToH3(bucketKey);
-        final CellBoundary boundary = H3.h3ToGeoBoundary(h3);
-        double minLat = Double.POSITIVE_INFINITY;
-        double minLon = Double.POSITIVE_INFINITY;
-        double maxLat = Double.NEGATIVE_INFINITY;
-        double maxLon = Double.NEGATIVE_INFINITY;
-        for (int i = 0; i < boundary.numPoints(); i++) {
-            final double boundaryLat = boundary.getLatLon(i).getLatDeg();
-            final double boundaryLon = boundary.getLatLon(i).getLonDeg();
-            minLon = Math.min(minLon, boundaryLon);
-            maxLon = Math.max(maxLon, boundaryLon);
-            minLat = Math.min(minLat, boundaryLat);
-            maxLat = Math.max(maxLat, boundaryLat);
-        }
-        final int resolution = H3.getResolution(h3);
-        if (H3.geoToH3(90, 0, resolution) == h3) {
-            // north pole
-            return new Rectangle(-180d, 180d, 90, minLat);
-        } else if (H3.geoToH3(-90, 0, resolution) == h3) {
-            // south pole
-            return new Rectangle(-180d, 180d, maxLat, -90);
-        } else if (maxLon - minLon > 180d) {
-            // crosses dateline
-            return new Rectangle(maxLon, minLon, maxLat, minLat);
-        } else {
-            return new Rectangle(minLon, maxLon, maxLat, minLat);
-        }
-    }
-
     private void doTestGrid(
         int minPrecision,
         int maxPrecision,

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

@@ -51,16 +51,19 @@
 import org.elasticsearch.xpack.spatial.search.aggregations.GeoLineAggregationBuilder;
 import org.elasticsearch.xpack.spatial.search.aggregations.InternalGeoLine;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.BoundedGeoHashGridTiler;
+import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.BoundedGeoHexGridTiler;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.BoundedGeoTileGridTiler;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoGridTiler;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoHexCellIdSource;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoHexGridAggregator;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoShapeCellIdSource;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoShapeHashGridAggregator;
+import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoShapeHexGridAggregator;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.GeoShapeTileGridAggregator;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.InternalGeoHexGrid;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.UnboundedGeoHashGridTiler;
+import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.UnboundedGeoHexGridTiler;
 import org.elasticsearch.xpack.spatial.search.aggregations.bucket.geogrid.UnboundedGeoTileGridTiler;
 import org.elasticsearch.xpack.spatial.search.aggregations.metrics.CartesianBoundsAggregationBuilder;
 import org.elasticsearch.xpack.spatial.search.aggregations.metrics.CartesianBoundsAggregator;
@@ -310,12 +313,9 @@ private void registerGeoShapeGridAggregators(ValuesSourceRegistry.Builder builde
                 collectsFromSingleBucket,
                 metadata) -> {
                 if (GEO_GRID_AGG_FEATURE.check(getLicenseState())) {
-                    final GeoGridTiler tiler;
-                    if (geoBoundingBox.isUnbounded()) {
-                        tiler = new UnboundedGeoHashGridTiler(precision);
-                    } else {
-                        tiler = new BoundedGeoHashGridTiler(precision, geoBoundingBox);
-                    }
+                    final GeoGridTiler tiler = geoBoundingBox.isUnbounded()
+                        ? new UnboundedGeoHashGridTiler(precision)
+                        : new BoundedGeoHashGridTiler(precision, geoBoundingBox);
                     GeoShapeCellIdSource cellIdSource = new GeoShapeCellIdSource((GeoShapeValuesSource) valuesSource, tiler);
                     GeoShapeHashGridAggregator agg = new GeoShapeHashGridAggregator(
                         name,
@@ -353,12 +353,9 @@ private void registerGeoShapeGridAggregators(ValuesSourceRegistry.Builder builde
                 collectsFromSingleBucket,
                 metadata) -> {
                 if (GEO_GRID_AGG_FEATURE.check(getLicenseState())) {
-                    final GeoGridTiler tiler;
-                    if (geoBoundingBox.isUnbounded()) {
-                        tiler = new UnboundedGeoTileGridTiler(precision);
-                    } else {
-                        tiler = new BoundedGeoTileGridTiler(precision, geoBoundingBox);
-                    }
+                    final GeoGridTiler tiler = geoBoundingBox.isUnbounded()
+                        ? new UnboundedGeoTileGridTiler(precision)
+                        : new BoundedGeoTileGridTiler(precision, geoBoundingBox);
                     GeoShapeCellIdSource cellIdSource = new GeoShapeCellIdSource((GeoShapeValuesSource) valuesSource, tiler);
                     GeoShapeTileGridAggregator agg = new GeoShapeTileGridAggregator(
                         name,
@@ -379,6 +376,46 @@ private void registerGeoShapeGridAggregators(ValuesSourceRegistry.Builder builde
             },
             true
         );
+
+        builder.register(
+            GeoHexGridAggregationBuilder.REGISTRY_KEY,
+            GeoShapeValuesSourceType.instance(),
+            (
+                name,
+                factories,
+                valuesSource,
+                precision,
+                geoBoundingBox,
+                requiredSize,
+                shardSize,
+                context,
+                parent,
+                collectsFromSingleBucket,
+                metadata) -> {
+                if (GEO_GRID_AGG_FEATURE.check(getLicenseState())) {
+                    final GeoGridTiler tiler = geoBoundingBox.isUnbounded()
+                        ? new UnboundedGeoHexGridTiler(precision)
+                        : new BoundedGeoHexGridTiler(precision, geoBoundingBox);
+                    GeoShapeCellIdSource cellIdSource = new GeoShapeCellIdSource((GeoShapeValuesSource) valuesSource, tiler);
+                    GeoShapeHexGridAggregator agg = new GeoShapeHexGridAggregator(
+                        name,
+                        factories,
+                        cellIdSource,
+                        requiredSize,
+                        shardSize,
+                        context,
+                        parent,
+                        collectsFromSingleBucket,
+                        metadata
+                    );
+                    // this would ideally be something set in an immutable way on the ValuesSource
+                    cellIdSource.setCircuitBreakerConsumer(agg::addRequestBytes);
+                    return agg;
+                }
+                throw LicenseUtils.newComplianceException("geohex_grid aggregation on geo_shape fields");
+            },
+            true
+        );
     }
 
     private static void registerValueCountAggregator(ValuesSourceRegistry.Builder builder) {