Skip to content

Commit

Permalink
Src files for GeoTile and GeoHash Aggregations on GeoShape
Browse files Browse the repository at this point in the history
Signed-off-by: Navneet Verma <navneev@amazon.com>
  • Loading branch information
navneet1v committed Dec 15, 2022
1 parent cb26035 commit 3a9fdf1
Show file tree
Hide file tree
Showing 18 changed files with 468 additions and 33 deletions.
26 changes: 2 additions & 24 deletions modules/geo/src/main/java/org/opensearch/geo/GeoModulePlugin.java
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand All @@ -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() {
Expand Down Expand Up @@ -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);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -120,6 +123,7 @@ protected Aggregator doCreateInternal(
}

static void registerAggregators(ValuesSourceRegistry.Builder builder) {
// register GeoPoint Aggregation
builder.register(
GeoHashGridAggregationBuilder.REGISTRY_KEY,
CoreValuesSourceType.GEOPOINT,
Expand Down Expand Up @@ -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
);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}
Loading

0 comments on commit 3a9fdf1

Please sign in to comment.