Src files for GeoTile and GeoHash Aggregations on GeoShape

Signed-off-by: Navneet Verma <navneev@amazon.com>

Author: navneet1v

modules/geo/src/main/java/org/opensearch/geo/GeoModulePlugin.java

@@ -39,21 +39,17 @@
 import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoTileGrid;
 import org.opensearch.geo.search.aggregations.metrics.GeoBounds;
 import org.opensearch.geo.search.aggregations.metrics.GeoBoundsAggregationBuilder;
-import org.opensearch.geo.search.aggregations.metrics.GeoBoundsGeoShapeAggregator;
 import org.opensearch.geo.search.aggregations.metrics.InternalGeoBounds;
 import org.opensearch.index.mapper.GeoShapeFieldMapper;
 import org.opensearch.index.mapper.Mapper;
 import org.opensearch.plugins.MapperPlugin;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.plugins.SearchPlugin;
 import org.opensearch.search.aggregations.bucket.composite.CompositeAggregation;
-import org.opensearch.search.aggregations.support.CoreValuesSourceType;
-import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
 
 public class GeoModulePlugin extends Plugin implements MapperPlugin, SearchPlugin {
 
@@ -63,7 +59,8 @@ public Map<String, Mapper.TypeParser> getMappers() {
     }
 
     /**
-     * Registering {@link GeoBounds} aggregation on GeoPoint field.
+     * Registering {@link GeoBounds}, {@link GeoHashGrid}, {@link GeoTileGrid} aggregation on GeoPoint and GeoShape
+     * fields.
      */
     @Override
     public List<AggregationSpec> getAggregations() {
@@ -105,23 +102,4 @@ public List<CompositeAggregationSpec> getCompositeAggregations() {
             )
         );
     }
-
-    /**
-     * Registering the GeoBounds Aggregation on the GeoShape Field. This function allows plugins to register new
-     * aggregations using aggregation names that are already defined in Core, as long as the new aggregations target
-     * different ValuesSourceTypes.
-     *
-     * @return A list of the new registrar functions
-     */
-    @Override
-    public List<Consumer<ValuesSourceRegistry.Builder>> getAggregationExtentions() {
-        final Consumer<ValuesSourceRegistry.Builder> geoShapeConsumer = builder -> builder.register(
-            GeoBoundsAggregationBuilder.REGISTRY_KEY,
-            CoreValuesSourceType.GEO_SHAPE,
-            GeoBoundsGeoShapeAggregator::new,
-            true
-        );
-        return Collections.singletonList(geoShapeConsumer);
-    }
-
 }

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/composite/GeoTileGridValuesSourceBuilder.java

@@ -42,7 +42,7 @@
 import org.opensearch.common.xcontent.ObjectParser;
 import org.opensearch.common.xcontent.XContentBuilder;
 import org.opensearch.common.xcontent.XContentParser;
-import org.opensearch.geo.search.aggregations.bucket.geogrid.CellIdSource;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.cells.CellIdSource;
 import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoTileGridAggregationBuilder;
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.query.QueryShardContext;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/metrics/GeoGridAggregatorSupplier.java renamed to modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoGridAggregatorSupplier.java

@@ -30,10 +30,9 @@
  * GitHub history for details.
  */
 
-package org.opensearch.geo.search.aggregations.metrics;
+package org.opensearch.geo.search.aggregations.bucket.geogrid;
 
 import org.opensearch.common.geo.GeoBoundingBox;
-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;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoHashGridAggregationBuilder.java

@@ -40,7 +40,6 @@
 import org.opensearch.search.aggregations.AggregationBuilder;
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.aggregations.AggregatorFactory;
-import org.opensearch.geo.search.aggregations.metrics.GeoGridAggregatorSupplier;
 import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;
 import org.opensearch.search.aggregations.support.ValuesSourceRegistry;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoHashGridAggregatorFactory.java

@@ -33,6 +33,9 @@
 package org.opensearch.geo.search.aggregations.bucket.geogrid;
 
 import org.opensearch.common.geo.GeoBoundingBox;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.cells.CellIdSource;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.cells.GeoShapeCellIdSource;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.util.GeoShapeHashUtil;
 import org.opensearch.geometry.utils.Geohash;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.aggregations.Aggregator;
@@ -120,6 +123,7 @@ protected Aggregator doCreateInternal(
     }
 
     static void registerAggregators(ValuesSourceRegistry.Builder builder) {
+        // register GeoPoint Aggregation
         builder.register(
             GeoHashGridAggregationBuilder.REGISTRY_KEY,
             CoreValuesSourceType.GEOPOINT,
@@ -155,5 +159,41 @@ static void registerAggregators(ValuesSourceRegistry.Builder builder) {
             },
             true
         );
+        // register GeoShape Aggregation
+        builder.register(
+            GeoHashGridAggregationBuilder.REGISTRY_KEY,
+            CoreValuesSourceType.GEO_SHAPE,
+            (
+                name,
+                factories,
+                valuesSource,
+                precision,
+                geoBoundingBox,
+                requiredSize,
+                shardSize,
+                aggregationContext,
+                parent,
+                cardinality,
+                metadata) -> {
+                final GeoShapeCellIdSource cellIdSource = new GeoShapeCellIdSource(
+                    (ValuesSource.GeoShape) valuesSource,
+                    precision,
+                    geoBoundingBox,
+                    GeoShapeHashUtil::encodeShape
+                );
+                return new GeoHashGridAggregator(
+                    name,
+                    factories,
+                    cellIdSource,
+                    requiredSize,
+                    shardSize,
+                    aggregationContext,
+                    parent,
+                    cardinality,
+                    metadata
+                );
+            },
+            true
+        );
     }
 }

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoTileGridAggregationBuilder.java

@@ -39,7 +39,6 @@
 import org.opensearch.search.aggregations.AggregationBuilder;
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.aggregations.AggregatorFactory;
-import org.opensearch.geo.search.aggregations.metrics.GeoGridAggregatorSupplier;
 import org.opensearch.search.aggregations.bucket.GeoTileUtils;
 import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 import org.opensearch.search.aggregations.support.ValuesSourceConfig;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoTileGridAggregatorFactory.java

@@ -33,6 +33,8 @@
 package org.opensearch.geo.search.aggregations.bucket.geogrid;
 
 import org.opensearch.common.geo.GeoBoundingBox;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.cells.CellIdSource;
+import org.opensearch.geo.search.aggregations.bucket.geogrid.cells.GeoShapeCellIdSource;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.search.aggregations.Aggregator;
 import org.opensearch.search.aggregations.AggregatorFactories;
@@ -154,5 +156,42 @@ static void registerAggregators(ValuesSourceRegistry.Builder builder) {
             },
             true
         );
+
+        // registers Aggregation on GeoShape
+        builder.register(
+            GeoTileGridAggregationBuilder.REGISTRY_KEY,
+            CoreValuesSourceType.GEO_SHAPE,
+            (
+                name,
+                factories,
+                valuesSource,
+                precision,
+                geoBoundingBox,
+                requiredSize,
+                shardSize,
+                aggregationContext,
+                parent,
+                cardinality,
+                metadata) -> {
+                GeoShapeCellIdSource cellIdSource = new GeoShapeCellIdSource(
+                    (ValuesSource.GeoShape) valuesSource,
+                    precision,
+                    geoBoundingBox,
+                    GeoTileUtils::encodeShape
+                );
+                return new GeoTileGridAggregator(
+                    name,
+                    factories,
+                    cellIdSource,
+                    requiredSize,
+                    shardSize,
+                    aggregationContext,
+                    parent,
+                    cardinality,
+                    metadata
+                );
+            },
+            true
+        );
     }
 }

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/BoundedCellValues.java renamed to modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/cells/BoundedCellValues.java

@@ -29,7 +29,7 @@
  * GitHub history for details.
  */
 
-package org.opensearch.geo.search.aggregations.bucket.geogrid;
+package org.opensearch.geo.search.aggregations.bucket.geogrid.cells;
 
 import org.opensearch.common.geo.GeoBoundingBox;
 import org.opensearch.index.fielddata.MultiGeoPointValues;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/CellIdSource.java renamed to modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/cells/CellIdSource.java

@@ -29,7 +29,7 @@
  * GitHub history for details.
  */
 
-package org.opensearch.geo.search.aggregations.bucket.geogrid;
+package org.opensearch.geo.search.aggregations.bucket.geogrid.cells;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/CellValues.java renamed to modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/cells/CellValues.java

@@ -29,7 +29,7 @@
  * GitHub history for details.
  */
 
-package org.opensearch.geo.search.aggregations.bucket.geogrid;
+package org.opensearch.geo.search.aggregations.bucket.geogrid.cells;
 
 import org.opensearch.index.fielddata.AbstractSortingNumericDocValues;
 import org.opensearch.index.fielddata.MultiGeoPointValues;
@@ -54,12 +54,16 @@ protected CellValues(MultiGeoPointValues geoValues, int precision, CellIdSource.
         this.encoder = encoder;
     }
 
+    // we will need a way to clear the values array, as it is storing the states
     @Override
     public boolean advanceExact(int docId) throws IOException {
         if (geoValues.advanceExact(docId)) {
             int docValueCount = geoValues.docValueCount();
             resize(docValueCount);
             int j = 0;
+            // we will need a way to convert the shape into a set of GeoTiles. Now, rather than iterating over the
+            // points as in GeoPoints, we need to add the list of tiles which shape is intersecting with like we are
+            // doing for the points.
             for (int i = 0; i < docValueCount; i++) {
                 j = advanceValue(geoValues.nextValue(), j);
             }

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/cells/GeoShapeCellIdSource.java

@@ -0,0 +1,99 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.geo.search.aggregations.bucket.geogrid.cells;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.opensearch.common.geo.GeoBoundingBox;
+import org.opensearch.common.geo.GeoShapeDocValue;
+import org.opensearch.index.fielddata.GeoShapeValue;
+import org.opensearch.index.fielddata.SortedBinaryDocValues;
+import org.opensearch.index.fielddata.SortedNumericDoubleValues;
+import org.opensearch.search.aggregations.support.ValuesSource;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * ValueSource class which converts the {@link GeoShapeValue} to numeric long values for bucketing. This class uses the
+ * {@link GeoShapeCellIdSource.GeoShapeLongEncoder} to encode the geo_shape to {@link Long} values which can be iterated
+ * to do the bucket aggregation.
+ *
+ * @opensearch.internal
+ */
+public class GeoShapeCellIdSource extends ValuesSource.Numeric {
+
+    private final ValuesSource.GeoShape geoShape;
+    private final int precision;
+    private final GeoBoundingBox geoBoundingBox;
+    private final GeoShapeCellIdSource.GeoShapeLongEncoder encoder;
+
+    public GeoShapeCellIdSource(
+        final ValuesSource.GeoShape geoShape,
+        final int precision,
+        final GeoBoundingBox geoBoundingBox,
+        final GeoShapeCellIdSource.GeoShapeLongEncoder encoder
+    ) {
+        this.geoShape = geoShape;
+        this.geoBoundingBox = geoBoundingBox;
+        this.precision = precision;
+        this.encoder = encoder;
+    }
+
+    /**
+     * Get the current {@link SortedBinaryDocValues}.
+     *
+     * @param context {@link LeafReaderContext}
+     */
+    @Override
+    public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOException {
+        throw new UnsupportedOperationException("The bytesValues operation is not supported on GeoShapeCellIdSource");
+    }
+
+    /**
+     * Whether the underlying data is floating-point or not.
+     */
+    @Override
+    public boolean isFloatingPoint() {
+        return false;
+    }
+
+    /**
+     * Get the current {@link SortedNumericDocValues}.
+     *
+     * @param context {@link LeafReaderContext}
+     */
+    @Override
+    public SortedNumericDocValues longValues(final LeafReaderContext context) {
+        if (geoBoundingBox.isUnbounded()) {
+            return new GeoShapeCellValues.UnboundedCellValues(geoShape.getGeoShapeValues(context), precision, encoder);
+        }
+        return new GeoShapeCellValues.BoundedCellValues(geoShape.getGeoShapeValues(context), precision, encoder, geoBoundingBox);
+    }
+
+    /**
+     * Get the current {@link SortedNumericDoubleValues}.
+     *
+     * @param context {@link LeafReaderContext}
+     */
+    @Override
+    public SortedNumericDoubleValues doubleValues(LeafReaderContext context) {
+        throw new UnsupportedOperationException("The doubleValues operation is not supported on GeoShapeCellIdSource");
+    }
+
+    /**
+     * Encoder to encode the GeoShapes to the specific long values for the aggregation.
+     *
+     * @opensearch.internal
+     */
+    @FunctionalInterface
+    public interface GeoShapeLongEncoder {
+        List<Long> encode(final GeoShapeDocValue geoShapeDocValue, final int precision);
+    }
+}