diff --git a/build.gradle b/build.gradle index dbf0ea65..e404ed48 100644 --- a/build.gradle +++ b/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' diff --git a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java b/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java index 9dc7d6e8..56349622 100644 --- a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java +++ b/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> 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 getAggregations() { + + final var geoHexGridSpec = new AggregationSpec( + GeoHexGridAggregationBuilder.NAME, + GeoHexGridAggregationBuilder::new, + GeoHexGridAggregationBuilder.PARSER + ).addResultReader(GeoHexGrid::new).setAggregatorRegistrar(GeoHexGridAggregationBuilder::registerAggregators); + + return List.of(geoHexGridSpec); + } } diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGrid.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGrid.java new file mode 100644 index 00000000..9b336f86 --- /dev/null +++ b/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 { + + public GeoHexGrid(StreamInput in) throws IOException { + super(in); + } + + @Override + public BaseGeoGrid create(List 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 buckets, Map metadata) { + super(name, requiredSize, buckets, metadata); + } + + @Override + protected Reader getBucketReader() { + return GeoHexGridBucket::new; + } + + @Override + protected BaseGeoGrid create(String name, int requiredSize, List buckets, Map 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; + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregationBuilder.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregationBuilder.java new file mode 100644 index 00000000..c1016261 --- /dev/null +++ b/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 REGISTRY_KEY = new ValuesSourceRegistry.RegistryKey<>( + NAME, + GeoGridAggregatorSupplier.class + ); + public static final ObjectParser 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 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 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 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); + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregator.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregator.java new file mode 100644 index 00000000..813575c9 --- /dev/null +++ b/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 { + + public GeoHexGridAggregator( + String name, + AggregatorFactories factories, + ValuesSource.Numeric valuesSource, + int requiredSize, + int shardSize, + SearchContext aggregationContext, + Aggregator parent, + CardinalityUpperBound cardinality, + Map metadata + ) throws IOException { + super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, cardinality, metadata); + } + + @Override + protected GeoHexGrid buildAggregation(String name, int requiredSize, List buckets, Map metadata) { + return new GeoHexGrid(name, requiredSize, buckets, metadata); + } + + @Override + protected BaseGeoGridBucket newEmptyBucket() { + return new GeoHexGridBucket(0, 0, null); + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorFactory.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorFactory.java new file mode 100644 index 00000000..78fb0edc --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorFactory.java @@ -0,0 +1,129 @@ +/* + * 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.geo.GeoBoundingBox; +import org.opensearch.geo.search.aggregations.bucket.geogrid.CellIdSource; +import org.opensearch.index.query.QueryShardContext; +import org.opensearch.search.aggregations.Aggregator; +import org.opensearch.search.aggregations.AggregatorFactories; +import org.opensearch.search.aggregations.AggregatorFactory; +import org.opensearch.search.aggregations.CardinalityUpperBound; +import org.opensearch.search.aggregations.InternalAggregation; +import org.opensearch.search.aggregations.NonCollectingAggregator; +import org.opensearch.search.aggregations.support.CoreValuesSourceType; +import org.opensearch.search.aggregations.support.ValuesSource; +import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory; +import org.opensearch.search.aggregations.support.ValuesSourceConfig; +import org.opensearch.search.aggregations.support.ValuesSourceRegistry; +import org.opensearch.search.internal.SearchContext; + +/** + * Aggregation Factory for geohex_grid agg + */ +public class GeoHexGridAggregatorFactory extends ValuesSourceAggregatorFactory { + private final int precision; + private final int requiredSize; + private final int shardSize; + private final GeoBoundingBox geoBoundingBox; + + GeoHexGridAggregatorFactory( + String name, + ValuesSourceConfig config, + int precision, + int requiredSize, + int shardSize, + GeoBoundingBox geoBoundingBox, + QueryShardContext queryShardContext, + AggregatorFactory parent, + AggregatorFactories.Builder subFactoriesBuilder, + Map metadata + ) throws IOException { + super(name, config, queryShardContext, parent, subFactoriesBuilder, metadata); + this.precision = precision; + this.requiredSize = requiredSize; + this.shardSize = shardSize; + this.geoBoundingBox = geoBoundingBox; + } + + @Override + protected Aggregator createUnmapped(SearchContext searchContext, Aggregator aggregator, Map map) throws IOException { + final var aggregation = new GeoHexGrid(name, requiredSize, List.of(), metadata); + + return new NonCollectingAggregator(name, searchContext, aggregator, factories, metadata) { + @Override + public InternalAggregation buildEmptyAggregation() { + return aggregation; + } + }; + } + + @Override + protected Aggregator doCreateInternal( + SearchContext searchContext, + Aggregator aggregator, + CardinalityUpperBound cardinalityUpperBound, + Map map + ) throws IOException { + return queryShardContext.getValuesSourceRegistry() + .getAggregator(GeoHexGridAggregationBuilder.REGISTRY_KEY, config) + .build( + name, + factories, + config.getValuesSource(), + precision, + geoBoundingBox, + requiredSize, + shardSize, + searchContext, + aggregator, + cardinalityUpperBound, + metadata + ); + } + + static void registerAggregators(final ValuesSourceRegistry.Builder builder) { + builder.register( + GeoHexGridAggregationBuilder.REGISTRY_KEY, + CoreValuesSourceType.GEOPOINT, + ( + name, + factories, + valuesSource, + precision, + geoBoundingBox, + requiredSize, + shardSize, + aggregationContext, + parent, + cardinality, + metadata) -> { + CellIdSource cellIdSource = new CellIdSource( + (ValuesSource.GeoPoint) valuesSource, + precision, + geoBoundingBox, + GeoHexHelper::longEncode + ); + return new GeoHexGridAggregator( + name, + factories, + cellIdSource, + requiredSize, + shardSize, + aggregationContext, + parent, + cardinality, + metadata + ); + }, + true + ); + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridBucket.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridBucket.java new file mode 100644 index 00000000..29c5b80a --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridBucket.java @@ -0,0 +1,44 @@ +/* + * 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.h3ToGeoPoint; + +import java.io.IOException; + +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket; +import org.opensearch.geospatial.h3.H3; +import org.opensearch.search.aggregations.InternalAggregations; + +/** + * Implementation of geohex grid bucket + */ +public class GeoHexGridBucket extends BaseGeoGridBucket { + + public GeoHexGridBucket(long hashAsLong, long docCount, InternalAggregations aggregations) { + super(hashAsLong, docCount, aggregations); + } + + /** + * Read from a Stream + * @param in {@link StreamInput} contains GridBucket + * @throws IOException + */ + public GeoHexGridBucket(StreamInput in) throws IOException { + super(in); + } + + @Override + public Object getKey() { + return h3ToGeoPoint(hashAsLong); + } + + @Override + public String getKeyAsString() { + return H3.h3ToString(hashAsLong); + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexHelper.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexHelper.java new file mode 100644 index 00000000..f99f892e --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexHelper.java @@ -0,0 +1,77 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import static org.opensearch.geospatial.h3.H3.MIN_H3_RES; +import static org.opensearch.geospatial.h3.H3.h3IsValid; +import static org.opensearch.geospatial.h3.H3.h3ToLatLng; +import static org.opensearch.geospatial.h3.H3.stringToH3; + +import java.util.Locale; + +import lombok.NonNull; + +import org.opensearch.common.geo.GeoPoint; +import org.opensearch.geospatial.h3.H3; + +/** + * Helper class for H3 library + */ +public class GeoHexHelper { + + /** + * Checks whether given precision is within H3 Precision range + * @param precision H3 index precision + */ + public static void checkPrecisionRange(int precision) { + if ((precision < MIN_H3_RES) || (precision > H3.MAX_H3_RES)) { + throw new IllegalArgumentException( + String.format( + Locale.getDefault(), + "Invalid precision of %d . Must be between %d and %d.", + precision, + MIN_H3_RES, + H3.MAX_H3_RES + ) + ); + } + } + + /** + * Converts from long representation of an index to {@link GeoPoint} representation. + * @param h3CellID H3 Cell Id + * @throws IllegalArgumentException if invalid h3CellID is provided + */ + public static GeoPoint h3ToGeoPoint(long h3CellID) { + if (!h3IsValid(h3CellID)) { + throw new IllegalArgumentException(String.format(Locale.getDefault(), "Invalid H3 Cell address: %d", h3CellID)); + } + final var position = h3ToLatLng(h3CellID); + return new GeoPoint(position.getLatDeg(), position.getLonDeg()); + } + + /** + * Converts from {@link String} representation of an index to {@link GeoPoint} representation. + * @param h3CellID H3 Cell Id + * @throws IllegalArgumentException if invalid h3CellID is provided + */ + public static GeoPoint h3ToGeoPoint(@NonNull String h3CellID) { + return h3ToGeoPoint(stringToH3(h3CellID)); + } + + /** + * Encodes longitude/latitude into H3 Cell Address for given precision + * + * @param latitude Latitude in degrees. + * @param longitude Longitude in degrees. + * @param precision Precision, 0 <= res <= 15 + * @return The H3 index. + * @throws IllegalArgumentException latitude, longitude, or precision are out of range. + */ + public static long longEncode(double longitude, double latitude, int precision) { + return H3.geoToH3(latitude, longitude, precision); + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGrid.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGrid.java new file mode 100644 index 00000000..7e0dcda2 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGrid.java @@ -0,0 +1,33 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import java.io.IOException; + +import lombok.NoArgsConstructor; + +import org.opensearch.common.xcontent.ObjectParser; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.geo.search.aggregations.bucket.geogrid.ParsedGeoGrid; + +@NoArgsConstructor +public class ParsedGeoHexGrid extends ParsedGeoGrid { + private static final ObjectParser PARSER = createParser( + ParsedGeoHexGrid::new, + ParsedGeoHexGridBucket::fromXContent, + ParsedGeoHexGridBucket::fromXContent + ); + + public static ParsedGeoGrid fromXContent(XContentParser parser, String name) throws IOException { + final var parsedGeoGrid = PARSER.parse(parser, null); + parsedGeoGrid.setName(name); + return parsedGeoGrid; + } + + public String getType() { + return GeoHexGridAggregationBuilder.NAME; + } +} diff --git a/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGridBucket.java b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGridBucket.java new file mode 100644 index 00000000..9b409f69 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/ParsedGeoHexGridBucket.java @@ -0,0 +1,30 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import java.io.IOException; + +import lombok.NoArgsConstructor; + +import org.opensearch.common.geo.GeoPoint; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.geo.search.aggregations.bucket.geogrid.ParsedGeoGridBucket; + +@NoArgsConstructor +public class ParsedGeoHexGridBucket extends ParsedGeoGridBucket { + + public GeoPoint getKey() { + return GeoHexHelper.h3ToGeoPoint(this.hashAsString); + } + + public String getKeyAsString() { + return this.hashAsString; + } + + static ParsedGeoHexGridBucket fromXContent(XContentParser parser) throws IOException { + return parseXContent(parser, false, ParsedGeoHexGridBucket::new, (p, bucket) -> { bucket.hashAsString = p.textOrNull(); }); + } +} diff --git a/src/test/java/org/opensearch/geospatial/GeospatialTestHelper.java b/src/test/java/org/opensearch/geospatial/GeospatialTestHelper.java index a52a4c82..9486ab0e 100644 --- a/src/test/java/org/opensearch/geospatial/GeospatialTestHelper.java +++ b/src/test/java/org/opensearch/geospatial/GeospatialTestHelper.java @@ -42,6 +42,7 @@ import org.opensearch.common.collect.Tuple; import org.opensearch.geospatial.action.upload.geojson.ContentBuilder; import org.opensearch.geospatial.action.upload.geojson.UploadGeoJSONRequestContent; +import org.opensearch.geospatial.h3.H3; import org.opensearch.geospatial.stats.upload.UploadMetric; import org.opensearch.index.shard.ShardId; import org.opensearch.test.OpenSearchTestCase; @@ -166,4 +167,8 @@ public static double[] toDoubleArray(float[] input) { return IntStream.range(0, input.length).mapToDouble(i -> input[i]).toArray(); } + public static int randomHexGridPrecision() { + return randomIntBetween(H3.MIN_H3_RES, H3.MAX_H3_RES); + } + } diff --git a/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorTests.java b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorTests.java new file mode 100644 index 00000000..c6ff7f41 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridAggregatorTests.java @@ -0,0 +1,342 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import static org.hamcrest.Matchers.equalTo; +import static org.opensearch.geospatial.GeospatialTestHelper.randomLowerCaseString; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.function.Consumer; +import java.util.function.Function; + +import org.apache.lucene.document.LatLonDocValuesField; +import org.apache.lucene.document.SortedSetDocValuesField; +import org.apache.lucene.geo.GeoEncodingUtils; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexableField; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.index.RandomIndexWriter; +import org.apache.lucene.util.BytesRef; +import org.hamcrest.MatcherAssert; +import org.opensearch.common.CheckedConsumer; +import org.opensearch.common.geo.GeoBoundingBox; +import org.opensearch.common.geo.GeoPoint; +import org.opensearch.common.geo.GeoUtils; +import org.opensearch.geo.GeometryTestUtils; +import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket; +import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGrid; +import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGridAggregationBuilder; +import org.opensearch.geometry.Rectangle; +import org.opensearch.geospatial.h3.H3; +import org.opensearch.geospatial.plugin.GeospatialPlugin; +import org.opensearch.index.mapper.GeoPointFieldMapper; +import org.opensearch.index.mapper.MappedFieldType; +import org.opensearch.plugins.SearchPlugin; +import org.opensearch.search.aggregations.Aggregation; +import org.opensearch.search.aggregations.Aggregator; +import org.opensearch.search.aggregations.AggregatorTestCase; +import org.opensearch.search.aggregations.MultiBucketConsumerService; +import org.opensearch.search.aggregations.bucket.terms.StringTerms; +import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; + +// This class is modified from https://github.com/opensearch-project/OpenSearch/blob/main/modules/geo/src/test/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java +// to keep relevant test case required for GeoHex Grid Aggregation. +public class GeoHexGridAggregatorTests extends AggregatorTestCase { + + private static final String GEO_POINT_FIELD_NAME = "location"; + private static final double TOLERANCE = 1E-5D; + + public void testNoDocs() throws IOException { + testCase( + new MatchAllDocsQuery(), + GEO_POINT_FIELD_NAME, + randomPrecision(), + null, + geoGrid -> { assertEquals(0, geoGrid.getBuckets().size()); }, + iw -> { + // Intentionally not writing any docs + } + ); + } + + public void testUnmapped() throws IOException { + testCase( + new MatchAllDocsQuery(), + randomLowerCaseString(), + randomPrecision(), + null, + geoGrid -> { assertEquals(0, geoGrid.getBuckets().size()); }, + iw -> { iw.addDocument(Collections.singleton(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, 10D, 10D))); } + ); + } + + public void testUnmappedMissing() throws IOException { + GeoGridAggregationBuilder builder = createBuilder(randomLowerCaseString()).field(randomLowerCaseString()) + .missing("53.69437,6.475031"); + testCase( + new MatchAllDocsQuery(), + randomPrecision(), + null, + geoGrid -> assertEquals(1, geoGrid.getBuckets().size()), + iw -> iw.addDocument(Collections.singleton(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, 10D, 10D))), + builder + ); + + } + + public void testWithSeveralDocs() throws IOException { + int precision = randomPrecision(); + int numPoints = randomIntBetween(8, 128); + Map expectedCountPerGeoHex = new HashMap<>(); + testCase(new MatchAllDocsQuery(), GEO_POINT_FIELD_NAME, precision, null, geoHexGrid -> { + assertEquals(expectedCountPerGeoHex.size(), geoHexGrid.getBuckets().size()); + for (GeoGrid.Bucket bucket : geoHexGrid.getBuckets()) { + assertEquals((long) expectedCountPerGeoHex.get(bucket.getKeyAsString()), bucket.getDocCount()); + } + assertTrue(hasValue(geoHexGrid)); + }, iw -> { + List points = new ArrayList<>(); + Set distinctAddressPerDoc = new HashSet<>(); + for (int pointId = 0; pointId < numPoints; pointId++) { + double[] latLng = randomLatLng(); + points.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latLng[0], latLng[1])); + String address = h3AddressAsString(latLng[1], latLng[0], precision); + if (distinctAddressPerDoc.contains(address) == false) { + expectedCountPerGeoHex.put(address, expectedCountPerGeoHex.getOrDefault(address, 0) + 1); + } + distinctAddressPerDoc.add(address); + if (usually()) { + iw.addDocument(points); + points.clear(); + distinctAddressPerDoc.clear(); + } + } + if (points.size() != 0) { + iw.addDocument(points); + } + }); + } + + public void testAsSubAgg() throws IOException { + int precision = randomPrecision(); + Map> expectedCountPerTPerGeoHex = new TreeMap<>(); + List> docs = new ArrayList<>(); + for (int i = 0; i < 30; i++) { + String t = randomAlphaOfLength(1); + double[] latLng = randomLatLng(); + + List doc = new ArrayList<>(); + docs.add(doc); + doc.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latLng[0], latLng[1])); + doc.add(new SortedSetDocValuesField("t", new BytesRef(t))); + + String address = h3AddressAsString(latLng[1], latLng[0], precision); + Map expectedCountPerGeoHex = expectedCountPerTPerGeoHex.get(t); + if (expectedCountPerGeoHex == null) { + expectedCountPerGeoHex = new TreeMap<>(); + expectedCountPerTPerGeoHex.put(t, expectedCountPerGeoHex); + } + expectedCountPerGeoHex.put(address, expectedCountPerGeoHex.getOrDefault(address, 0L) + 1); + } + CheckedConsumer buildIndex = iw -> iw.addDocuments(docs); + String aggregation = randomLowerCaseString(); + TermsAggregationBuilder aggregationBuilder = new TermsAggregationBuilder("t").field("t") + .size(expectedCountPerTPerGeoHex.size()) + .subAggregation(createBuilder(aggregation).field(GEO_POINT_FIELD_NAME).precision(precision)); + Consumer verify = (terms) -> { + Map> actual = new TreeMap<>(); + for (StringTerms.Bucket tb : terms.getBuckets()) { + GeoHexGrid gg = tb.getAggregations().get(aggregation); + Map sub = new TreeMap<>(); + for (BaseGeoGridBucket ggb : gg.getBuckets()) { + sub.put(ggb.getKeyAsString(), ggb.getDocCount()); + } + actual.put(tb.getKeyAsString(), sub); + } + MatcherAssert.assertThat(actual, equalTo(expectedCountPerTPerGeoHex)); + }; + testCase(aggregationBuilder, new MatchAllDocsQuery(), buildIndex, verify, keywordField("t"), geoPointField(GEO_POINT_FIELD_NAME)); + } + + public void testBounds() throws IOException { + final int numDocs = randomIntBetween(64, 256); + final GeoHexGridAggregationBuilder builder = createBuilder("_name"); + + expectThrows(IllegalArgumentException.class, () -> builder.precision(-1)); + expectThrows(IllegalArgumentException.class, () -> builder.precision(30)); + + // only consider bounding boxes that are at least TOLERANCE wide and have quantized coordinates + GeoBoundingBox bbox = randomValueOtherThanMany( + (b) -> Math.abs(GeoUtils.normalizeLon(b.right()) - GeoUtils.normalizeLon(b.left())) < TOLERANCE, + GeoHexGridAggregatorTests::randomBBox + ); + Function encodeDecodeLat = (lat) -> GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lat)); + Function encodeDecodeLon = (lon) -> GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lon)); + bbox.topLeft().reset(encodeDecodeLat.apply(bbox.top()), encodeDecodeLon.apply(bbox.left())); + bbox.bottomRight().reset(encodeDecodeLat.apply(bbox.bottom()), encodeDecodeLon.apply(bbox.right())); + + int in = 0, out = 0; + List docs = new ArrayList<>(); + while (in + out < numDocs) { + if (bbox.left() > bbox.right()) { + if (randomBoolean()) { + double lonWithin = randomBoolean() + ? randomDoubleBetween(bbox.left(), 180.0, true) + : randomDoubleBetween(-180.0, bbox.right(), true); + double latWithin = randomDoubleBetween(bbox.bottom(), bbox.top(), true); + in++; + docs.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latWithin, lonWithin)); + } else { + double lonOutside = randomDoubleBetween(bbox.left(), bbox.right(), true); + double latOutside = randomDoubleBetween(bbox.top(), -90, false); + out++; + docs.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latOutside, lonOutside)); + } + } else { + if (randomBoolean()) { + double lonWithin = randomDoubleBetween(bbox.left(), bbox.right(), true); + double latWithin = randomDoubleBetween(bbox.bottom(), bbox.top(), true); + in++; + docs.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latWithin, lonWithin)); + } else { + double lonOutside = GeoUtils.normalizeLon(randomDoubleBetween(bbox.right(), 180.001, false)); + double latOutside = GeoUtils.normalizeLat(randomDoubleBetween(bbox.top(), 90.001, false)); + out++; + docs.add(new LatLonDocValuesField(GEO_POINT_FIELD_NAME, latOutside, lonOutside)); + } + } + + } + + final long numDocsInBucket = in; + final int precision = randomPrecision(); + + testCase(new MatchAllDocsQuery(), GEO_POINT_FIELD_NAME, precision, bbox, geoGrid -> { + assertTrue(hasValue(geoGrid)); + long docCount = 0; + for (int i = 0; i < geoGrid.getBuckets().size(); i++) { + docCount += geoGrid.getBuckets().get(i).getDocCount(); + } + MatcherAssert.assertThat(docCount, equalTo(numDocsInBucket)); + }, iw -> { + for (LatLonDocValuesField docField : docs) { + iw.addDocument(Collections.singletonList(docField)); + } + }); + } + + @Override + public void doAssertReducedMultiBucketConsumer(Aggregation agg, MultiBucketConsumerService.MultiBucketConsumer bucketConsumer) { + /* + * No-op. + */ + } + + /** + * Overriding the Search Plugins list with {@link GeospatialPlugin} so that the testcase will know that this plugin is + * to be loaded during the tests. + * @return List of {@link SearchPlugin} + */ + @Override + protected List getSearchPlugins() { + return Collections.singletonList(new GeospatialPlugin()); + } + + private double[] randomLatLng() { + double lat = (180d * randomDouble()) - 90d; + double lng = (360d * randomDouble()) - 180d; + + // Precision-adjust longitude/latitude to avoid wrong bucket placement + // Internally, lat/lng get converted to 32 bit integers, loosing some precision. + // This does not affect geo hex because it also uses the same algorithm, + // but it does affect other bucketing algos, thus we need to do the same steps here. + lng = GeoEncodingUtils.decodeLongitude(GeoEncodingUtils.encodeLongitude(lng)); + lat = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(lat)); + + return new double[] { lat, lng }; + } + + private void testCase( + Query query, + String field, + int precision, + GeoBoundingBox geoBoundingBox, + Consumer verify, + CheckedConsumer buildIndex + ) throws IOException { + testCase(query, precision, geoBoundingBox, verify, buildIndex, createBuilder("_name").field(field)); + } + + private void testCase( + Query query, + int precision, + GeoBoundingBox geoBoundingBox, + Consumer verify, + CheckedConsumer buildIndex, + GeoGridAggregationBuilder aggregationBuilder + ) throws IOException { + Directory directory = newDirectory(); + RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory); + buildIndex.accept(indexWriter); + indexWriter.close(); + + IndexReader indexReader = DirectoryReader.open(directory); + IndexSearcher indexSearcher = newSearcher(indexReader, true, true); + + aggregationBuilder.precision(precision); + if (geoBoundingBox != null) { + aggregationBuilder.setGeoBoundingBox(geoBoundingBox); + MatcherAssert.assertThat(aggregationBuilder.geoBoundingBox(), equalTo(geoBoundingBox)); + } + + MappedFieldType fieldType = new GeoPointFieldMapper.GeoPointFieldType(GEO_POINT_FIELD_NAME); + + Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType); + aggregator.preCollection(); + indexSearcher.search(query, aggregator); + aggregator.postCollection(); + verify.accept((GeoHexGrid) aggregator.buildTopLevel()); + + indexReader.close(); + directory.close(); + } + + private int randomPrecision() { + return randomIntBetween(H3.MIN_H3_RES, H3.MAX_H3_RES); + } + + private static boolean hasValue(GeoHexGrid agg) { + return agg.getBuckets().stream().anyMatch(bucket -> bucket.getDocCount() > 0); + } + + private static GeoBoundingBox randomBBox() { + Rectangle rectangle = GeometryTestUtils.randomRectangle(); + return new GeoBoundingBox( + new GeoPoint(rectangle.getMaxLat(), rectangle.getMinLon()), + new GeoPoint(rectangle.getMinLat(), rectangle.getMaxLon()) + ); + } + + private String h3AddressAsString(double lng, double lat, int precision) { + return H3.geoToH3Address(lat, lng, precision); + } + + private GeoHexGridAggregationBuilder createBuilder(String name) { + return new GeoHexGridAggregationBuilder(name); + } +} diff --git a/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridParserTests.java b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridParserTests.java new file mode 100644 index 00000000..81b9d551 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridParserTests.java @@ -0,0 +1,122 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.instanceOf; +import static org.opensearch.geospatial.GeospatialTestHelper.randomHexGridPrecision; +import static org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder.NAME; +import static org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGridAggregationBuilder.PARSER; + +import java.util.Locale; + +import org.hamcrest.MatcherAssert; +import org.opensearch.common.xcontent.XContentParseException; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.geo.GeometryTestUtils; +import org.opensearch.geometry.Rectangle; +import org.opensearch.geospatial.h3.H3; +import org.opensearch.test.OpenSearchTestCase; + +public class GeoHexGridParserTests extends OpenSearchTestCase { + + public void testParseValidFromInts() throws Exception { + XContentParser stParser = createParser( + JsonXContent.jsonXContent, + String.format( + Locale.getDefault(), + "{\"field\":\"my_loc\", \"precision\":%d, \"size\": 500, \"shard_size\": 550}", + randomHexGridPrecision() + ) + ); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + // can create a factory + assertNotNull(PARSER.parse(stParser, NAME)); + } + + public void testParseValidFromStrings() throws Exception { + XContentParser stParser = createParser( + JsonXContent.jsonXContent, + String.format( + Locale.getDefault(), + "{\"field\":\"my_loc\", \"precision\":\"%d\", \"size\": \"500\", \"shard_size\": \"550\"}", + randomHexGridPrecision() + ) + ); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + // can create a factory + assertNotNull(PARSER.parse(stParser, NAME)); + } + + public void testParseInvalidUnitPrecision() throws Exception { + XContentParser stParser = createParser( + JsonXContent.jsonXContent, + "{\"field\":\"my_loc\", \"precision\": \"10kg\", \"size\": \"500\", \"shard_size\": \"550\"}" + ); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + XContentParseException ex = expectThrows(XContentParseException.class, () -> PARSER.parse(stParser, NAME)); + MatcherAssert.assertThat(ex.getMessage(), containsString("failed to parse field [precision]")); + MatcherAssert.assertThat(ex.getCause(), instanceOf(NumberFormatException.class)); + assertEquals("For input string: \"10kg\"", ex.getCause().getMessage()); + } + + public void testParseErrorOnBooleanPrecision() throws Exception { + XContentParser stParser = createParser(JsonXContent.jsonXContent, "{\"field\":\"my_loc\", \"precision\":false}"); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + XContentParseException e = expectThrows(XContentParseException.class, () -> PARSER.parse(stParser, NAME)); + MatcherAssert.assertThat(e.getMessage(), containsString("precision doesn't support values of type: VALUE_BOOLEAN")); + } + + public void testParseErrorOnPrecisionOutOfRange() throws Exception { + int invalidPrecision = H3.MAX_H3_RES + 1; + XContentParser stParser = createParser( + JsonXContent.jsonXContent, + String.format(Locale.getDefault(), "{\"field\":\"my_loc\", \"precision\": %d}", invalidPrecision) + ); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + try { + PARSER.parse(stParser, NAME); + fail(); + } catch (XContentParseException ex) { + MatcherAssert.assertThat(ex.getCause(), instanceOf(IllegalArgumentException.class)); + assertEquals( + String.format( + Locale.getDefault(), + "Invalid precision of %d . Must be between %d and %d.", + invalidPrecision, + H3.MIN_H3_RES, + H3.MAX_H3_RES + ), + ex.getCause().getMessage() + ); + } + } + + public void testParseValidBounds() throws Exception { + Rectangle bbox = GeometryTestUtils.randomRectangle(); + XContentParser stParser = createParser( + JsonXContent.jsonXContent, + String.format( + Locale.getDefault(), + "{\"field\":\"my_loc\", \"precision\": 5, \"size\": 500, \"shard_size\": 550,\"bounds\": { \"top\": %s,\"bottom\": %s,\"left\": %s,\"right\": %s}}", + bbox.getMaxY(), + bbox.getMinY(), + bbox.getMinX(), + bbox.getMaxX() + ) + ); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + // can create a factory + assertNotNull(PARSER.parse(stParser, NAME)); + } +} diff --git a/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridTests.java b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridTests.java new file mode 100644 index 00000000..d12bfbbb --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/search/aggregations/bucket/geogrid/GeoHexGridTests.java @@ -0,0 +1,184 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.search.aggregations.bucket.geogrid; + +import static org.hamcrest.Matchers.equalTo; +import static org.opensearch.geospatial.GeospatialTestHelper.randomHexGridPrecision; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.apache.lucene.index.IndexWriter; +import org.hamcrest.MatcherAssert; +import org.opensearch.common.ParseField; +import org.opensearch.common.xcontent.ContextParser; +import org.opensearch.common.xcontent.NamedXContentRegistry; +import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGrid; +import org.opensearch.geo.search.aggregations.bucket.geogrid.BaseGeoGridBucket; +import org.opensearch.geo.search.aggregations.bucket.geogrid.GeoGrid; +import org.opensearch.geospatial.h3.H3; +import org.opensearch.geospatial.plugin.GeospatialPlugin; +import org.opensearch.plugins.SearchPlugin; +import org.opensearch.search.aggregations.Aggregation; +import org.opensearch.search.aggregations.InternalAggregations; +import org.opensearch.search.aggregations.ParsedMultiBucketAggregation; +import org.opensearch.test.InternalMultiBucketAggregationTestCase; + +// This class is modified from https://github.com/opensearch-project/OpenSearch/blob/main/modules/geo/src/test/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoGridTestCase.java +// to keep relevant test case required for GeoHex Grid. +public class GeoHexGridTests extends InternalMultiBucketAggregationTestCase { + + private static final double LATITUDE_MIN = -90.0; + private static final double LATITUDE_MAX = 90.0; + private static final double LONGITUDE_MIN = -180.0; + private static final double LONGITUDE_MAX = 180.0; + private static final int MIN_BUCKET_SIZE = 1; + private static final int MAX_BUCKET_SIZE = 3; + + public void testCreateFromBuckets() { + BaseGeoGrid original = createTestInstance(); + MatcherAssert.assertThat(original, equalTo(original.create(original.getBuckets()))); + } + + @Override + protected int minNumberOfBuckets() { + return MIN_BUCKET_SIZE; + } + + @Override + protected int maxNumberOfBuckets() { + return MAX_BUCKET_SIZE; + } + + /** + * Overriding the method so that tests can get the aggregation specs for namedWriteable. + * + * @return GeoSpatialPlugin + */ + @Override + protected SearchPlugin registerPlugin() { + return new GeospatialPlugin(); + } + + /** + * Overriding with the {@link ParsedGeoHexGrid} so that it can be parsed. We need to do this as {@link GeospatialPlugin} + * is registering this Aggregation. + * + * @return a List of {@link NamedXContentRegistry.Entry} + */ + @Override + protected List getNamedXContents() { + final List namedXContents = new ArrayList<>(getDefaultNamedXContents()); + final ContextParser hexGridParser = (p, c) -> ParsedGeoHexGrid.fromXContent(p, (String) c); + namedXContents.add( + new NamedXContentRegistry.Entry(Aggregation.class, new ParseField(GeoHexGridAggregationBuilder.NAME), hexGridParser) + ); + return namedXContents; + } + + @Override + protected GeoHexGrid createTestInstance(String name, Map metadata, InternalAggregations aggregations) { + final int precision = randomHexGridPrecision(); + int size = randomNumberOfBuckets(); + List buckets = new ArrayList<>(size); + for (int i = 0; i < size; i++) { + double latitude = randomDoubleBetween(LATITUDE_MIN, LATITUDE_MAX, false); + double longitude = randomDoubleBetween(LONGITUDE_MIN, LONGITUDE_MAX, false); + + long addressAsLong = longEncode(longitude, latitude, precision); + buckets.add(createInternalGeoGridBucket(addressAsLong, randomInt(IndexWriter.MAX_DOCS), aggregations)); + } + return createInternalGeoGrid(name, size, buckets, metadata); + } + + @Override + protected void assertReduced(GeoHexGrid reduced, List inputs) { + Map> map = new HashMap<>(); + for (GeoHexGrid input : inputs) { + for (GeoGrid.Bucket bucketBase : input.getBuckets()) { + GeoHexGridBucket bucket = (GeoHexGridBucket) bucketBase; + List buckets = map.computeIfAbsent(bucket.hashAsLong(), k -> new ArrayList<>()); + buckets.add(bucket); + } + } + List expectedBuckets = new ArrayList<>(); + for (Map.Entry> entry : map.entrySet()) { + long docCount = 0; + for (GeoHexGridBucket bucket : entry.getValue()) { + docCount += bucket.getDocCount(); + } + expectedBuckets.add(createInternalGeoGridBucket(entry.getKey(), docCount, InternalAggregations.EMPTY)); + } + expectedBuckets.sort((first, second) -> { + int cmp = Long.compare(second.getDocCount(), first.getDocCount()); + if (cmp == 0) { + return second.compareTo(first); + } + return cmp; + }); + int requestedSize = inputs.get(0).getRequiredSize(); + expectedBuckets = expectedBuckets.subList(0, Math.min(requestedSize, expectedBuckets.size())); + assertEquals(expectedBuckets.size(), reduced.getBuckets().size()); + for (int i = 0; i < reduced.getBuckets().size(); i++) { + GeoGrid.Bucket expected = expectedBuckets.get(i); + GeoGrid.Bucket actual = reduced.getBuckets().get(i); + assertEquals(expected.getDocCount(), actual.getDocCount()); + assertEquals(expected.getKey(), actual.getKey()); + } + } + + @Override + protected Class implementationClass() { + return ParsedGeoHexGrid.class; + } + + @Override + protected GeoHexGrid mutateInstance(GeoHexGrid instance) { + String name = instance.getName(); + int size = instance.getRequiredSize(); + List buckets = instance.getBuckets(); + Map metadata = instance.getMetadata(); + switch (between(0, 3)) { + case 0: + name += randomAlphaOfLength(5); + break; + case 1: + buckets = new ArrayList<>(buckets); + buckets.add( + createInternalGeoGridBucket(randomNonNegativeLong(), randomInt(IndexWriter.MAX_DOCS), InternalAggregations.EMPTY) + ); + break; + case 2: + size = size + between(1, 10); + break; + case 3: + if (metadata == null) { + metadata = new HashMap<>(1); + } else { + metadata = new HashMap<>(instance.getMetadata()); + } + metadata.put(randomAlphaOfLength(15), randomInt()); + break; + default: + throw new AssertionError("Illegal randomisation branch"); + } + return createInternalGeoGrid(name, size, buckets, metadata); + } + + private GeoHexGrid createInternalGeoGrid(String name, int size, List buckets, Map metadata) { + return new GeoHexGrid(name, size, buckets, metadata); + } + + private GeoHexGridBucket createInternalGeoGridBucket(Long key, long docCount, InternalAggregations aggregations) { + return new GeoHexGridBucket(key, docCount, aggregations); + } + + private long longEncode(double lng, double lat, int precision) { + return H3.geoToH3(lng, lat, precision); + } +} diff --git a/src/yamlRestTest/resources/rest-api-spec/test/20_geohex_grid.yml b/src/yamlRestTest/resources/rest-api-spec/test/20_geohex_grid.yml new file mode 100644 index 00000000..da689826 --- /dev/null +++ b/src/yamlRestTest/resources/rest-api-spec/test/20_geohex_grid.yml @@ -0,0 +1,61 @@ +setup: + - do: + indices.create: + index: cities + body: + settings: + number_of_replicas: 0 + mappings: + properties: + location: + type: geo_point + +--- +"Basic test": + - do: + bulk: + refresh: true + body: + - index: + _index: cities + _id: 1 + - location: "52.374081,4.912350" + - index: + _index: cities + _id: 2 + - location: "52.369219,4.901618" + - index: + _index: cities + _id: 3 + - location: "52.371667,4.914722" + - index: + _index: cities + _id: 4 + - location: "51.222900,4.405200" + - index: + _index: cities + _id: 5 + - location: "48.861111,2.336389" + - index: + _index: cities + _id: 6 + - location: "48.860000,2.327000" + + - do: + search: + rest_total_hits_as_int: true + body: + aggregations: + grid: + geohex_grid: + field: location + precision: 4 + + + - match: { hits.total: 6 } + - match: { aggregations.grid.buckets.0.key: 841969dffffffff } + - match: { aggregations.grid.buckets.0.doc_count: 3 } + - match: { aggregations.grid.buckets.1.key: 841fb47ffffffff } + - match: { aggregations.grid.buckets.1.doc_count: 2 } + - match: { aggregations.grid.buckets.2.key: 841fa4dffffffff } + - match: { aggregations.grid.buckets.2.doc_count: 1 }