diff --git a/core/src/main/java/org/elasticsearch/search/SearchModule.java b/core/src/main/java/org/elasticsearch/search/SearchModule.java index 99b3d4c8894e..ca3be20d5335 100644 --- a/core/src/main/java/org/elasticsearch/search/SearchModule.java +++ b/core/src/main/java/org/elasticsearch/search/SearchModule.java @@ -171,6 +171,8 @@ import org.elasticsearch.search.aggregations.metrics.geocentroid.GeoCentroidAggregationBuilder; import org.elasticsearch.search.aggregations.metrics.geocentroid.GeoCentroidParser; import org.elasticsearch.search.aggregations.metrics.geocentroid.InternalGeoCentroid; +import org.elasticsearch.search.aggregations.metrics.geoheatmap.GeoHeatmapAggregationBuilder; +import org.elasticsearch.search.aggregations.metrics.geoheatmap.InternalGeoHeatmap; import org.elasticsearch.search.aggregations.metrics.max.InternalMax; import org.elasticsearch.search.aggregations.metrics.max.MaxAggregationBuilder; import org.elasticsearch.search.aggregations.metrics.max.MaxParser; @@ -438,6 +440,8 @@ private void registerAggregations(List plugins) { new GeoDistanceParser()).addResultReader(InternalGeoDistance::new)); registerAggregation(new AggregationSpec(GeoGridAggregationBuilder.NAME, GeoGridAggregationBuilder::new, new GeoHashGridParser()) .addResultReader(InternalGeoHashGrid::new)); + registerAggregation(new AggregationSpec(GeoHeatmapAggregationBuilder.NAME, GeoHeatmapAggregationBuilder::new, + GeoHeatmapAggregationBuilder::parse).addResultReader(InternalGeoHeatmap::new)); registerAggregation(new AggregationSpec(NestedAggregationBuilder.NAME, NestedAggregationBuilder::new, NestedAggregationBuilder::parse).addResultReader(InternalNested::new)); registerAggregation(new AggregationSpec(ReverseNestedAggregationBuilder.NAME, ReverseNestedAggregationBuilder::new, diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmap.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmap.java new file mode 100644 index 000000000000..1bffd7310452 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmap.java @@ -0,0 +1,44 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +/** + * The internal representation of a heatmap + */ +public interface GeoHeatmap { + String getName(); + + int getGridLevel(); + + int getColumns(); + + int getRows(); + + int[] getCounts(); + + double getMinX(); + + double getMinY(); + + double getMaxX(); + + double getMaxY(); + +} diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationBuilder.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationBuilder.java new file mode 100644 index 000000000000..50f36ee27b50 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationBuilder.java @@ -0,0 +1,305 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import org.elasticsearch.common.ParseField; +import org.elasticsearch.common.ParsingException; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.search.aggregations.AbstractAggregationBuilder; +import org.elasticsearch.search.aggregations.AggregatorFactories.Builder; +import org.elasticsearch.search.aggregations.AggregatorFactory; +import org.elasticsearch.search.aggregations.InternalAggregation.Type; +import org.elasticsearch.search.aggregations.support.AggregationContext; +import org.locationtech.spatial4j.shape.Shape; + +import java.io.IOException; +import java.util.Objects; +import java.util.Optional; + +/** + * Collects the various parameters for a heatmap aggregation and builds a + * factory + */ +public class GeoHeatmapAggregationBuilder extends AbstractAggregationBuilder { + public static final String NAME = "heatmap"; + public static final Type TYPE = new Type(NAME); + public static final ParseField AGGREGATION_NAME_FIELD = new ParseField(NAME); + public static final ParseField GEOM_FIELD = new ParseField("geom"); + public static final ParseField MAX_CELLS_FIELD = new ParseField("max_cells"); + public static final ParseField DIST_ERR_FIELD = new ParseField("dist_err"); + public static final ParseField DIST_ERR_PCT_FIELD = new ParseField("dist_err_pct"); + public static final ParseField GRID_LEVEL_FIELD = new ParseField("grid_level"); + + private QueryBuilder geom; + private Double distErr; + private Double distErrPct; + private Integer gridLevel; + private Integer maxCells; + private String field; + + /** + * Creates a blank builder + * + * @param name + * the name that was given this aggregation instance + */ + public GeoHeatmapAggregationBuilder(String name) { + super(name, TYPE); + } + + /** + * A utility method to construct a blank builder through the Java API + * + * @param name + * a name for the aggregator instance + * @return a new blank builder + */ + public static GeoHeatmapAggregationBuilder heatmap(String name) { + return new GeoHeatmapAggregationBuilder(name); + } + + /** + * Read from a stream + */ + public GeoHeatmapAggregationBuilder(StreamInput in) throws IOException { + super(in, TYPE); + field = in.readString(); + geom = in.readOptionalNamedWriteable(QueryBuilder.class); + distErr = in.readOptionalDouble(); + distErrPct = in.readOptionalDouble(); + gridLevel = in.readOptionalVInt(); + maxCells = in.readOptionalVInt(); + } + + @Override + protected void doWriteTo(StreamOutput out) throws IOException { + out.writeString(field); + out.writeOptionalNamedWriteable(geom); + out.writeOptionalDouble(distErr); + out.writeOptionalDouble(distErrPct); + out.writeOptionalVInt(gridLevel); + out.writeOptionalVInt(maxCells); + } + + @Override + public String getWriteableName() { + return NAME; + } + + /** + * Construct a builder from XContent, which usually comes from the JSON + * query API + */ + public static GeoHeatmapAggregationBuilder parse(String aggregationName, QueryParseContext context) throws IOException { + XContentParser parser = context.parser(); + + XContentParser.Token token = null; + String currentFieldName = null; + String field = null; + GeoShapeQueryBuilder geom = null; + Integer maxCells = null; + Double distErr = null; + Double distErrPct = null; + Integer gridLevel = null; + + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + currentFieldName = parser.currentName(); + } else if (token == XContentParser.Token.VALUE_STRING) { + if ("field".equals(currentFieldName)) { + field = parser.text(); + } else if (context.getParseFieldMatcher().match(currentFieldName, DIST_ERR_FIELD)) { + distErr = DistanceUnit.parse(parser.text(), DistanceUnit.DEFAULT, DistanceUnit.DEFAULT); + } + } else if (token.isValue()) { + if (context.getParseFieldMatcher().match(currentFieldName, MAX_CELLS_FIELD)) { + maxCells = parser.intValue(); + } else if (context.getParseFieldMatcher().match(currentFieldName, DIST_ERR_PCT_FIELD)) { + distErrPct = parser.doubleValue(); + } else if (context.getParseFieldMatcher().match(currentFieldName, GRID_LEVEL_FIELD)) { + gridLevel = parser.intValue(); + } else { + throw new ParsingException(parser.getTokenLocation(), + "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "]."); + } + } else if (token == XContentParser.Token.START_OBJECT) { + if (context.getParseFieldMatcher().match(currentFieldName, GEOM_FIELD)) { + geom = (GeoShapeQueryBuilder) context.parseInnerQueryBuilder() + .filter(qb -> qb.getWriteableName().equals(GeoShapeQueryBuilder.NAME)).orElse(null); + } else { + throw new ParsingException(parser.getTokenLocation(), "Unknown key for a " + token + " in [" + currentFieldName + "].", + parser.getTokenLocation()); + } + } else if (token == XContentParser.Token.VALUE_NULL) { + if (context.getParseFieldMatcher().match(currentFieldName, MAX_CELLS_FIELD) + || context.getParseFieldMatcher().match(currentFieldName, DIST_ERR_PCT_FIELD) + || context.getParseFieldMatcher().match(currentFieldName, GRID_LEVEL_FIELD) + || context.getParseFieldMatcher().match(currentFieldName, GEOM_FIELD)) { + continue; + } else { + throw new ParsingException(parser.getTokenLocation(), + "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "]."); + } + } else { + throw new ParsingException(parser.getTokenLocation(), + "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "]."); + } + } + + if (field == null) { + throw new ParsingException(null, "Missing required field [field] for geo_heatmap aggregation [" + aggregationName + "]"); + } + return new GeoHeatmapAggregationBuilder(aggregationName).geom(geom).field(field).maxCells(maxCells) + .distErr(distErr).distErrPct(distErrPct).gridLevel(gridLevel); + } + + @Override + protected AggregatorFactory doBuild(AggregationContext context, AggregatorFactory parent, Builder subFactoriesBuilder) + throws IOException { + Shape inputShape = null; + if (geom != null) { + GeoShapeQueryBuilder shapeBuilder = (GeoShapeQueryBuilder) geom; + inputShape = shapeBuilder.shape().build(); + } + GeoHeatmapAggregatorFactory factory = new GeoHeatmapAggregatorFactory(name, type, field, Optional.ofNullable(inputShape), + Optional.ofNullable(maxCells), Optional.ofNullable(distErr), Optional.ofNullable(distErrPct), + Optional.ofNullable(gridLevel), context, parent, subFactoriesBuilder, metaData); + return factory; + } + + @Override + protected XContentBuilder internalXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + if (geom != null) { + builder.field(GEOM_FIELD.getPreferredName(), geom); + } + builder.field("field", field); + builder.field(MAX_CELLS_FIELD.getPreferredName(), maxCells); + if (distErr != null) { + builder.field(DIST_ERR_FIELD.getPreferredName(), distErr); + } + builder.field(DIST_ERR_PCT_FIELD.getPreferredName(), distErrPct); + if (gridLevel != null) { + builder.field(GRID_LEVEL_FIELD.getPreferredName(), gridLevel); + } + builder.endObject(); + return builder; + } + + /** + * @param field + * the field on which to build the heatmap; must be a geo_shape + * field. + * @return this builder + */ + public GeoHeatmapAggregationBuilder field(String field) { + if (field == null) { + throw new IllegalArgumentException("[field] must not be null: [" + name + "]"); + } + this.field = field; + return this; + } + + /** + * Sets a bounding geometry for the heatmap. The heatmap itself will be + * rectangular, but hits outside of this geometry will not be counted + * + * @param geom + * the bounding geometry; defaults to the world rectangle + * @return this builder + */ + public GeoHeatmapAggregationBuilder geom(GeoShapeQueryBuilder geom) { + this.geom = geom; + return this; + } + + /** + * Sets the maximum allowable error for determining where an indexed shape is + * relative to the heatmap cells + * + * @param distErr + * The distance in meters + * @return this builder + */ + public GeoHeatmapAggregationBuilder distErr(Double distErr) { + if (distErr != null) this.distErr = distErr; + return this; + } + + /** + * Sets the maximum allowable error for determining where an indexed shape is + * relative to the heatmap cells, specified as a fraction of the shape size + * + * @param distErrPct + * A fraction from 0.0 to 0.5 + * @return this builder + */ + public GeoHeatmapAggregationBuilder distErrPct(Double distErrPct) { + if (distErrPct != null) this.distErrPct = distErrPct; + return this; + } + + /** + * Sets the grid level (granularity) of the heatmap + * + * @param gridLevel + * higher numbers mean higher granularity; defaults to 7 + * @return this builder + */ + public GeoHeatmapAggregationBuilder gridLevel(Integer gridLevel) { + if (gridLevel != null) this.gridLevel = gridLevel; + return this; + } + + /** + * Set the maximum number of cells that can be returned in the heatmap + * + * @param maxCells + * defaults to 100,000 + * @return this builder + */ + public GeoHeatmapAggregationBuilder maxCells(Integer maxCells) { + if (maxCells != null) this.maxCells = maxCells; + return this; + } + + @Override + protected int doHashCode() { + return Objects.hash(field, geom, distErr, distErrPct, gridLevel, maxCells); + } + + @Override + protected boolean doEquals(Object obj) { + GeoHeatmapAggregationBuilder other = (GeoHeatmapAggregationBuilder) obj; + return Objects.equals(field, other.field) + && Objects.equals(geom, other.geom) + && Objects.equals(distErr, other.distErr) + && Objects.equals(distErrPct, other.distErrPct) + && Objects.equals(gridLevel, other.gridLevel) + && Objects.equals(maxCells, other.maxCells); + } +} diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregator.java new file mode 100644 index 000000000000..632710ca83a9 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregator.java @@ -0,0 +1,127 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.apache.lucene.index.IndexReaderContext; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.spatial.prefix.HeatmapFacetCounter; +import org.apache.lucene.spatial.prefix.HeatmapFacetCounter.Heatmap; +import org.apache.lucene.spatial.prefix.PrefixTreeStrategy; +import org.apache.lucene.util.SparseFixedBitSet; +import org.elasticsearch.search.aggregations.Aggregator; +import org.elasticsearch.search.aggregations.InternalAggregation; +import org.elasticsearch.search.aggregations.LeafBucketCollector; +import org.elasticsearch.search.aggregations.LeafBucketCollectorBase; +import org.elasticsearch.search.aggregations.metrics.MetricsAggregator; +import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator; +import org.elasticsearch.search.aggregations.support.AggregationContext; +import org.locationtech.spatial4j.shape.Shape; + +/** + * Aggregates hits on geo_shape fields as counts on a grid + */ +public class GeoHeatmapAggregator extends MetricsAggregator { + + private final int maxCells; + private final int gridLevel; + private final Shape inputShape; + private final PrefixTreeStrategy strategy; + private IndexReaderContext parentReaderContext; + private Map buckets = new HashMap<>(); + + /** + * Used by the {@link GeoHeatmapAggregatorFactory} to create a new instance + * + * @param name + * the name of this heatmap aggregator + * @param inputShape + * indexed shapes must intersect inputShape to be counted; if + * null the world rectangle is used + * @param strategy + * a geo strategy for searching the field + * @param maxCells + * the maximum number of cells (grid squares) that could possibly + * be returned from the heatmap + * @param gridLevel + * manually set the granularity of the grid + * @param aggregationContext + * the context of this aggregation + * @param parent + * the parent aggregation + * @param pipelineAggregators + * any pipeline aggregations attached + * @param metaData + * aggregation metadata + * @throws IOException + * when parsing fails + */ + public GeoHeatmapAggregator(String name, Shape inputShape, PrefixTreeStrategy strategy, int maxCells, int gridLevel, + AggregationContext aggregationContext, Aggregator parent, List pipelineAggregators, + Map metaData) throws IOException { + super(name, aggregationContext, parent, pipelineAggregators, metaData); + this.inputShape = inputShape; + this.strategy = strategy; + this.maxCells = maxCells; + this.gridLevel = gridLevel; + } + + @Override + public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException { + if (parentReaderContext == null) { + parentReaderContext = ctx.parent; + } else { + assert ctx.parent == parentReaderContext; + } + return new LeafBucketCollectorBase(sub, null) { + @Override + public void collect(int doc, long bucket) throws IOException { + SparseFixedBitSet bits = buckets.get(bucket); + if (bits == null) { + bits = new SparseFixedBitSet(parentReaderContext.reader().maxDoc()); + buckets.put(bucket, bits); + } + bits.set(ctx.docBase + doc); + } + }; + } + + @Override + public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOException { + SparseFixedBitSet matchingDocs = buckets.get(owningBucketOrdinal); + if (matchingDocs == null) + return buildEmptyAggregation(); + + Heatmap heatmap = HeatmapFacetCounter.calcFacets(strategy, parentReaderContext, matchingDocs, inputShape, gridLevel, maxCells); + + return new InternalGeoHeatmap(name, gridLevel, heatmap.rows, heatmap.columns, heatmap.region.getMinX(), heatmap.region.getMinY(), + heatmap.region.getMaxX(), heatmap.region.getMaxY(), heatmap.counts, pipelineAggregators(), metaData()); + } + + @Override + public InternalAggregation buildEmptyAggregation() { + return new InternalGeoHeatmap(name, gridLevel, 0, 0, 0, 0, 0, 0, new int[0], pipelineAggregators(), metaData()); + } + +} diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregatorFactory.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregatorFactory.java new file mode 100644 index 000000000000..52537b9ad339 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregatorFactory.java @@ -0,0 +1,134 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; + +import org.apache.lucene.spatial.prefix.PrefixTreeStrategy; +import org.apache.lucene.spatial.query.SpatialArgs; +import org.apache.lucene.spatial.query.SpatialOperation; +import org.elasticsearch.common.geo.SpatialStrategy; +import org.elasticsearch.common.unit.DistanceUnit; +import org.elasticsearch.index.mapper.GeoShapeFieldMapper; +import org.elasticsearch.index.mapper.GeoShapeFieldMapper.GeoShapeFieldType; +import org.elasticsearch.index.mapper.MappedFieldType; +import org.elasticsearch.search.aggregations.AggregationInitializationException; +import org.elasticsearch.search.aggregations.Aggregator; +import org.elasticsearch.search.aggregations.AggregatorFactories; +import org.elasticsearch.search.aggregations.AggregatorFactory; +import org.elasticsearch.search.aggregations.InternalAggregation.Type; +import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator; +import org.elasticsearch.search.aggregations.support.AggregationContext; +import org.locationtech.spatial4j.shape.Shape; + +/** + */ +public class GeoHeatmapAggregatorFactory extends AggregatorFactory { + + public static final double DEFAULT_DIST_ERR_PCT = 0.15; + public static final int DEFAULT_MAX_CELLS = 100_000; + + private final int maxCells; + private final int gridLevel; + private final Shape inputShape; + private final PrefixTreeStrategy strategy; + + /** + * Called by the {@link GeoHeatmapAggregationBuilder} + * + * @param name + * the name of this heatmap aggregator + * @param type + * passed in from the Builder + * @param field + * the indexed field on which to create the heatmap; must be a + * geo_shape + * @param inputShape + * indexed shapes must intersect inputShape to be counted; if + * null the world rectangle is used + * @param maxCells + * the maximum number of cells (grid squares) that could possibly + * be returned from the heatmap + * @param distErr + * the maximum error distance allowable between the indexed shape and the cells + * @param distErrPct + * the maximum error ratio between the indexed shape and the heatmap shape + * @param gridLevel + * manually set the granularity of the grid + * @param context + * the context of this aggregation + * @param parent + * the parent aggregation + * @param subFactoriesBuilder + * (not used) + * @param metaData + * aggregation metadata + * @throws IOException + * if an error occurs creating the factory + */ + public GeoHeatmapAggregatorFactory(String name, Type type, String field, Optional inputShape, Optional maxCells, + Optional distErr, Optional distErrPct, Optional gridLevel, + AggregationContext context, AggregatorFactory parent, AggregatorFactories.Builder subFactoriesBuilder, + Map metaData) throws IOException { + + super(name, type, context, parent, subFactoriesBuilder, metaData); + MappedFieldType fieldType = context.searchContext().mapperService().fullName(field); + if (fieldType.typeName().equals(GeoShapeFieldMapper.CONTENT_TYPE)) { + GeoShapeFieldType geoFieldType = (GeoShapeFieldType) fieldType; + this.strategy = geoFieldType.resolveStrategy(SpatialStrategy.RECURSIVE); + } else { + throw new AggregationInitializationException(String.format(Locale.ROOT, + "Field [%s] is a %s instead of a %s type", + field, fieldType.typeName(), GeoShapeFieldMapper.CONTENT_TYPE)); + } + this.inputShape = inputShape.orElse(strategy.getSpatialContext().getWorldBounds()); + this.maxCells = maxCells.orElse(DEFAULT_MAX_CELLS); + this.gridLevel = gridLevel.orElse(resolveGridLevel(distErr, distErrPct)); + } + + @Override + public Aggregator createInternal(Aggregator parent, boolean collectsFromSingleBucket, List pipelineAggregators, + Map metaData) throws IOException { + return new GeoHeatmapAggregator(name, inputShape, strategy, maxCells, gridLevel, context, parent, pipelineAggregators, metaData); + } + + private Integer resolveGridLevel(Optional distErrOp, Optional distErrPctOp) { + SpatialArgs spatialArgs = new SpatialArgs(SpatialOperation.Intersects, this.inputShape); + if (distErrOp.isPresent()) { + spatialArgs.setDistErr(distErrOp.get() * DistanceUnit.DEFAULT.getDistancePerDegree()); + } + spatialArgs.setDistErrPct(distErrPctOp.orElse(DEFAULT_DIST_ERR_PCT)); + double distErr = spatialArgs.resolveDistErr(strategy.getSpatialContext(), DEFAULT_DIST_ERR_PCT); + if (distErr <= 0) { + throw new AggregationInitializationException(String.format(Locale.ROOT, + "%s or %s should be > 0 or instead provide %s=%s for absolute maximum detail", + GeoHeatmapAggregationBuilder.DIST_ERR_PCT_FIELD, + GeoHeatmapAggregationBuilder.DIST_ERR_FIELD, + GeoHeatmapAggregationBuilder.GRID_LEVEL_FIELD, + strategy.getGrid().getMaxLevels())); + } + return strategy.getGrid().getLevelForDistance(distErr); + } + +} diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/InternalGeoHeatmap.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/InternalGeoHeatmap.java new file mode 100644 index 000000000000..55ec4ed55970 --- /dev/null +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/InternalGeoHeatmap.java @@ -0,0 +1,240 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import java.io.IOException; +import java.util.AbstractList; +import java.util.List; +import java.util.Map; + +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.search.aggregations.InternalAggregation; +import org.elasticsearch.search.aggregations.metrics.InternalMetricsAggregation; +import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator; + +/** + * + */ +public class InternalGeoHeatmap extends InternalMetricsAggregation implements GeoHeatmap { + + private int gridLevel; + private int rows; + private int columns; + private double minX; + private double minY; + private double maxX; + private double maxY; + private int[] counts; + + InternalGeoHeatmap(String name, int gridLevel, int rows, int columns, double minX, double minY, double maxX, double maxY, int[] counts, + List pipelineAggregators, Map metaData) { + super(name, pipelineAggregators, metaData); + this.gridLevel = gridLevel; + this.rows = rows; + this.columns = columns; + this.minX = minX; + this.minY = minY; + this.maxX = maxX; + this.maxY = maxY; + this.counts = counts; + } + + private InternalGeoHeatmap newInternalGeoHeatmap() { + return new InternalGeoHeatmap(name, gridLevel, rows, columns, minX, minY, maxX, maxY, new int[counts.length], pipelineAggregators(), + metaData); + } + + /** + * Read from a stream. + */ + public InternalGeoHeatmap(StreamInput in) throws IOException { + super(in); + this.gridLevel = in.readVInt(); + this.rows = in.readVInt(); + this.columns = in.readVInt(); + this.minX = in.readDouble(); + this.minY = in.readDouble(); + this.maxX = in.readDouble(); + this.maxY = in.readDouble(); + this.counts = in.readVIntArray(); + } + + @Override + protected void doWriteTo(StreamOutput out) throws IOException { + out.writeVInt(gridLevel); + out.writeVInt(rows); + out.writeVInt(columns); + out.writeDouble(minX); + out.writeDouble(minY); + out.writeDouble(maxX); + out.writeDouble(maxY); + out.writeVIntArray(counts); + } + + @Override + public String getWriteableName() { + return GeoHeatmapAggregationBuilder.NAME; + } + + @Override + public InternalAggregation doReduce(List aggregations, ReduceContext reduceContext) { + + InternalGeoHeatmap reduced = newInternalGeoHeatmap(); + for (InternalAggregation aggregation : aggregations) { + + assert aggregation.getName().equals(getName()); + assert aggregation instanceof InternalGeoHeatmap; + InternalGeoHeatmap toMerge = (InternalGeoHeatmap) aggregation; + if (toMerge.columns == 0) { + continue; + } else if (reduced.columns == 0) { + reduced.gridLevel = toMerge.gridLevel; + reduced.rows = toMerge.rows; + reduced.columns = toMerge.columns; + reduced.minX = toMerge.minX; + reduced.minY = toMerge.minY; + reduced.maxX = toMerge.maxX; + reduced.maxY = toMerge.maxY; + reduced.counts = toMerge.counts; + } else { + assert toMerge.gridLevel == reduced.gridLevel; + assert toMerge.columns == reduced.columns; + assert toMerge.maxX == reduced.maxX; + assert toMerge.counts.length == reduced.counts.length; + for (int i = 0; i < toMerge.counts.length; i++) { + reduced.counts[i] += toMerge.counts[i]; + } + } + } + return reduced; + } + + @Override + public Object getProperty(List path) { + if (path.isEmpty()) { + return this; + } else { + throw new IllegalArgumentException("path not supported for [" + getName() + "]: " + path); + } + } + + @Override + public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException { + List> countLists = asInts2D(columns, rows, counts); + builder.field("grid_level", gridLevel).field("rows", rows).field("columns", columns).field("min_x", minX).field("min_y", minY) + .field("max_x", maxX).field("max_y", maxY).startArray("counts"); + for (List row : countLists) { + builder.startArray(); + if (row != null) { + for (Integer i : row) { + builder.value(i); + } + } + builder.endArray(); + } + builder.endArray(); + return builder; + } + + // {@see org.apache.solr.handler.component.SpatialHeatmapFacets#asInts2D} + static List> asInts2D(final int columns, final int rows, final int[] counts) { + // Returns a view versus returning a copy. This saves memory. + // The data is oriented naturally for human/developer viewing: one row + // at a time top-down + return new AbstractList>() { + @Override + public List get(final int rowIdx) {// top-down remember; + // the heatmap.counts is + // bottom up + // check if all zeroes and return null if so + boolean hasNonZero = false; + int y = rows - rowIdx - 1;// flip direction for 'y' + for (int c = 0; c < columns; c++) { + if (counts[c * rows + y] > 0) { + hasNonZero = true; + break; + } + } + if (!hasNonZero) { + return null; + } + + return new AbstractList() { + @Override + public Integer get(int columnIdx) { + return counts[columnIdx * rows + y]; + } + + @Override + public int size() { + return columns; + } + }; + } + + @Override + public int size() { + return rows; + } + }; + } + + @Override + public int getGridLevel() { + return gridLevel; + } + + @Override + public int getColumns() { + return columns; + } + + @Override + public int getRows() { + return rows; + } + + @Override + public int[] getCounts() { + return counts; + } + + @Override + public double getMinX() { + return minX; + } + + @Override + public double getMinY() { + return minY; + } + + @Override + public double getMaxX() { + return maxX; + } + + @Override + public double getMaxY() { + return maxY; + } +} diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/BaseGeoHeatmapTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/BaseGeoHeatmapTests.java new file mode 100644 index 000000000000..08efb7cb5c28 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/BaseGeoHeatmapTests.java @@ -0,0 +1,57 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.GeoShapeQueryBuilderTests; +import org.elasticsearch.search.aggregations.BaseAggregationTestCase; +import org.elasticsearch.search.aggregations.metrics.geoheatmap.GeoHeatmapAggregationBuilder; + +/** + * Randomized test for construction and serialization + */ +public class BaseGeoHeatmapTests extends BaseAggregationTestCase { + + @Override + protected GeoHeatmapAggregationBuilder createTestAggregatorBuilder() { + + String name = randomAsciiOfLengthBetween(3, 20); + GeoHeatmapAggregationBuilder factory = new GeoHeatmapAggregationBuilder(name); + if (randomBoolean()) { + factory.gridLevel(randomIntBetween(1, 20)); + } + if (randomBoolean()) { + factory.maxCells(randomIntBetween(2, 5000) * 2); + } + + InternalBuilderTests it = new InternalBuilderTests(); + factory.geom(it.getRandomShapeQuery()); + + String field = randomNumericField(); + factory.field(field); + + return factory; + } + + class InternalBuilderTests extends GeoShapeQueryBuilderTests { + public GeoShapeQueryBuilder getRandomShapeQuery() { + return doCreateTestQueryBuilder(); + } + } +} diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationIT.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationIT.java new file mode 100644 index 000000000000..3a2891964f83 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapAggregationIT.java @@ -0,0 +1,142 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import com.vividsolutions.jts.geom.Coordinate; + +import org.elasticsearch.action.index.IndexRequestBuilder; +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.common.geo.builders.EnvelopeBuilder; +import org.elasticsearch.common.geo.builders.GeometryCollectionBuilder; +import org.elasticsearch.common.xcontent.ToXContent; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.QueryBuilders; +import org.elasticsearch.search.aggregations.metrics.geoheatmap.GeoHeatmap; +import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.test.geo.RandomShapeGenerator; +import org.hamcrest.Matchers; +import org.locationtech.spatial4j.shape.Rectangle; + +import java.util.ArrayList; +import java.util.List; + +import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder; +import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery; +import static org.elasticsearch.search.aggregations.metrics.geoheatmap.GeoHeatmapAggregationBuilder.heatmap; +import static org.elasticsearch.test.geo.RandomShapeGenerator.xRandomPoint; +import static org.elasticsearch.test.geo.RandomShapeGenerator.xRandomRectangle; +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.lessThanOrEqualTo; +import static org.hamcrest.core.IsNull.notNullValue; + +/** + * Indexes a few random docs with geo_shapes and performs basic checks on the + * heatmap over them + */ +@ESIntegTestCase.SuiteScopeTestCase +public class GeoHeatmapAggregationIT extends ESIntegTestCase { + + static int numDocs, numTag1Docs; + static Rectangle mbr; + + @Override + public void setupSuiteScopeCluster() throws Exception { + mbr = xRandomRectangle(random(), xRandomPoint(random()), true); + createIndex("idx2"); + numDocs = randomIntBetween(5, 20); + numTag1Docs = randomIntBetween(1, numDocs - 1); + GeometryCollectionBuilder gcb = RandomShapeGenerator.createGeometryCollectionWithin(random(), mbr, numTag1Docs); + List builders = new ArrayList<>(); + + prepareCreate("idx").addMapping("type", "location", "type=geo_shape,tree=quadtree").execute().actionGet(); + ensureGreen(); + + for (int i = 0; i < numTag1Docs; i++) { + builders.add(client().prepareIndex("idx", "type", "" + i).setSource(jsonBuilder().startObject().field("value", i + 1) + .field("tag", "tag1").field("location", gcb.getShapeAt(i)).endObject())); + } + for (int i = numTag1Docs; i < numDocs; i++) { + XContentBuilder source = jsonBuilder().startObject().field("value", i).field("tag", "tag2").field("name", "name" + i) + .endObject(); + builders.add(client().prepareIndex("idx", "type", "" + i).setSource(source)); + if (randomBoolean()) { + // randomly index the document twice so that we have deleted + // docs that match the filter + builders.add(client().prepareIndex("idx", "type", "" + i).setSource(source)); + } + } + prepareCreate("empty_bucket_idx").addMapping("type", "value", "type=integer").execute().actionGet(); + for (int i = 0; i < 2; i++) { + builders.add(client().prepareIndex("empty_bucket_idx", "type", "" + i) + .setSource(jsonBuilder().startObject().field("value", i * 2).endObject())); + } + indexRandom(true, builders); + ensureSearchable(); + } + + /** + * Create a simple heatmap over the indexed docs and verify that no cell has + * more than that number of docs + */ + public void testSimple() throws Exception { + + EnvelopeBuilder env = new EnvelopeBuilder(new Coordinate(mbr.getMinX(), mbr.getMaxY()), + new Coordinate(mbr.getMaxX(), mbr.getMinY())); + GeoShapeQueryBuilder geo = QueryBuilders.geoShapeQuery("location", env).relation(ShapeRelation.WITHIN); + + SearchResponse response2 = client().prepareSearch("idx") + .addAggregation(heatmap("heatmap1").geom(geo).field("location").gridLevel(7).maxCells(100)).execute().actionGet(); + + assertSearchResponse(response2); + + GeoHeatmap filter2 = response2.getAggregations().get("heatmap1"); + assertThat(filter2, notNullValue()); + assertThat(filter2.getName(), equalTo("heatmap1")); + + int maxHeatmapValue = 0; + for (int i = 0; i < filter2.getCounts().length; i++) { + maxHeatmapValue = Math.max(maxHeatmapValue, filter2.getCounts()[i]); + } + assertTrue(maxHeatmapValue <= numTag1Docs); + + XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); + response2.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + logger.info("Full heatmap Response Content:\n{ {} }", builder.string()); + } + + /** + * Test that the number of cells generated is not greater than maxCells + */ + public void testMaxCells() throws Exception { + int maxCells = randomIntBetween(1, 50_000) * 2; + SearchResponse searchResponse = client().prepareSearch("idx").setQuery(matchAllQuery()) + .addAggregation(heatmap("heatmap1").field("location").gridLevel(1).maxCells(maxCells)).execute().actionGet(); + + assertThat(searchResponse.getHits().getTotalHits(), equalTo(new Long(numDocs))); + GeoHeatmap heatmap = searchResponse.getAggregations().get("heatmap1"); + assertThat(heatmap, Matchers.notNullValue()); + assertThat(heatmap.getRows() * heatmap.getColumns(), lessThanOrEqualTo(maxCells)); + } + +} diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapIT.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapIT.java new file mode 100644 index 000000000000..386da91c6277 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapIT.java @@ -0,0 +1,269 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import com.vividsolutions.jts.geom.Coordinate; + +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.common.geo.builders.GeometryCollectionBuilder; +import org.elasticsearch.common.geo.builders.ShapeBuilder; +import org.elasticsearch.common.geo.builders.ShapeBuilders; +import org.elasticsearch.common.xcontent.ToXContent; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.QueryBuilders; +import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.test.geo.RandomShapeGenerator; + +import java.io.IOException; +import java.util.Arrays; + +import static org.elasticsearch.action.support.WriteRequest.RefreshPolicy.IMMEDIATE; +import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder; +import static org.elasticsearch.search.aggregations.metrics.geoheatmap.GeoHeatmapAggregationBuilder.heatmap; +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount; +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse; +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.core.IsNull.notNullValue; + +public class GeoHeatmapIT extends ESIntegTestCase { + + /** + * Indexes a random shape, builds a random heatmap with that geometry, and + * makes sure there are '1' entries in the heatmap counts + */ + public void testShapeFilterWithRandomGeoCollection() throws IOException { + String name = randomAsciiOfLengthBetween(3, 20); + GeometryCollectionBuilder gcb = RandomShapeGenerator.createGeometryCollection(random()); + logger.info("Created Random GeometryCollection containing {} shapes", gcb.numShapes()); + + client().admin().indices().prepareCreate("test").addMapping("type", "location", "type=geo_shape,tree=quadtree").execute() + .actionGet(); + + XContentBuilder docSource = gcb.toXContent(jsonBuilder().startObject().field("location"), null).endObject(); + client().prepareIndex("test", "type", "1").setSource(docSource).setRefreshPolicy(IMMEDIATE).get(); + + ShapeBuilder filterShape = (gcb.getShapeAt(randomIntBetween(0, gcb.numShapes() - 1))); + + GeoShapeQueryBuilder geom; + try { + geom = QueryBuilders.geoShapeQuery("location", filterShape); + } catch (IOException e) { + throw new RuntimeException(e); + } + geom.relation(ShapeRelation.INTERSECTS); + + GeoHeatmapAggregationBuilder factory = new GeoHeatmapAggregationBuilder(name); + if (randomBoolean()) { + factory.geom(geom); + } + if (randomBoolean()) { + int gridLevel = randomIntBetween(1, 12); + factory.gridLevel(gridLevel); + } else { + if (randomBoolean()) { + factory.distErr(randomDoubleBetween(0.0, 0.5, false)); + } + factory.distErrPct(randomDoubleBetween(0.0, 0.5, false)); + } + if (randomBoolean()) { + factory.maxCells(randomIntBetween(1, Integer.MAX_VALUE)); + } + factory.field("location"); + + SearchResponse result = client().prepareSearch("test").setTypes("type").setQuery(QueryBuilders.matchAllQuery()).setPostFilter(geom) + .get(); + assertSearchResponse(result); + assertHitCount(result, 1); + + result = client().prepareSearch("test").setTypes("type").setQuery(QueryBuilders.matchAllQuery()).addAggregation(factory).get(); + assertSearchResponse(result); + assertHitCount(result, 1); + + GeoHeatmap heatmap = result.getAggregations().get(name); + assertThat(heatmap, notNullValue()); + + int maxHeatmapValue = 0; + for (int i = 0; i < heatmap.getCounts().length; i++) { + maxHeatmapValue = Math.max(maxHeatmapValue, heatmap.getCounts()[i]); + } + assertEquals(1, maxHeatmapValue); + } + + // @see org.apache.solr.handler.component.SpatialHeatmapFacetsTest + public void testSpecificShapes() throws Exception { + createPrecalculatedIndex(); + ShapeBuilder query = ShapeBuilders.newEnvelope(new Coordinate(50, 90), new Coordinate(180, 20)); + GeoShapeQueryBuilder geo = QueryBuilders.geoShapeQuery("location", query).relation(ShapeRelation.WITHIN); + + String expected = "\"aggregations\":{\"heatmap1\":{\"grid_level\":4,\"rows\":7,\"columns\":6,\"min_x\":45.0,\"min_y\":11.25,"+ + "\"max_x\":180.0,\"max_y\":90.0,\"counts\":[[0,0,2,1,0,0],[0,0,1,1,0,0],[0,1,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[],[]]}}}"; + + assertHeatmapContents(geo, expected, 4, "test"); + } + + /** + * Check to make sure that the heatmap can be selectively built from multiple indexes + * @throws Exception + */ + public void testMultipleIndexes() throws Exception { + for (String index : Arrays.asList("test1", "test2")) { + client().admin().indices().prepareCreate(index) + .addMapping("type1", "location", "type=geo_shape,tree=quadtree").execute() + .actionGet(); + } + + // on right side + client().prepareIndex("test1", "type1", "1") + .setSource(jsonBuilder().startObject().field("name", "Document 1").startObject("location").field("type", "envelope") + .startArray("coordinates").startArray().value(100).value(80).endArray().startArray().value(120).value(40).endArray() + .endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + // just left of BOX 0 + client().prepareIndex("test2", "type1", "2") + .setSource(jsonBuilder().startObject().field("name", "Document 2").startObject("location").field("type", "point") + .startArray("coordinates").value(70).value(60).endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + ShapeBuilder query = ShapeBuilders.newEnvelope(new Coordinate(50, 90), new Coordinate(180, 20)); + GeoShapeQueryBuilder geo = QueryBuilders.geoShapeQuery("location", query).relation(ShapeRelation.WITHIN); + + // check the first index + String expected = "\"counts\":[[0,0,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[],[]]}}}"; + assertHeatmapContents(geo, expected, 1, "test1"); + + // check the second index + expected = "\"counts\":[[],[],[0,1,0,0,0,0],[],[],[],[]]}}}"; + assertHeatmapContents(geo, expected, 1, "test2"); + + // check both indexes + expected = "\"counts\":[[0,0,1,1,0,0],[0,0,1,1,0,0],[0,1,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[],[]]}}}"; + assertHeatmapContents(geo, expected, 2, "test1", "test2"); + + // check all indexes + expected = "\"counts\":[[0,0,1,1,0,0],[0,0,1,1,0,0],[0,1,1,1,0,0],[0,0,1,1,0,0],[0,0,1,1,0,0],[],[]]}}}"; + assertHeatmapContents(geo, expected, 2); + } + + /** + * Tests the various ways grid_level is calculated + */ + // @see org.apache.solr.handler.component.SpatialHeatmapFacetsTest + public void testGridLevelCalc() throws Exception { + createPrecalculatedIndex(); + ShapeBuilder query = ShapeBuilders.newEnvelope(new Coordinate(50, 90), new Coordinate(180, 20)); + GeoShapeQueryBuilder geo = QueryBuilders.geoShapeQuery("location", query).relation(ShapeRelation.WITHIN); + + SearchResponse searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location").gridLevel(4).maxCells(100_000)).execute().actionGet(); + + assertSearchResponse(searchResponse); + assertThat(searchResponse.getHits().getTotalHits(), equalTo(4L)); + assertThat(searchResponse.getHits().hits().length, equalTo(4)); + + // Default + searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location")).execute().actionGet(); + assertGridLevel("heatmap1", 7, searchResponse); + + // Explicit grid_level + searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location").gridLevel(3)).execute().actionGet(); + assertGridLevel("heatmap1", 3, searchResponse); + + // Just dist_err + searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location").distErr(100.0)).execute().actionGet(); + assertGridLevel("heatmap1", 1, searchResponse); + + // Just dist_err_pct + searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location").distErrPct(0.05)).execute().actionGet(); + assertGridLevel("heatmap1", 8, searchResponse); + + // dist_err_pct with default geom + searchResponse = client().prepareSearch("test").setTypes("type1").setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").field("location").distErrPct(0.1).maxCells(100_000)).execute().actionGet(); + assertGridLevel("heatmap1", 6, searchResponse); + + } + + private void createPrecalculatedIndex() throws Exception { + client().admin().indices().prepareCreate("test") + .addMapping("type1", "location", "type=geo_shape,tree=quadtree").execute() + .actionGet(); + + // on right side + client().prepareIndex("test", "type1", "1") + .setSource(jsonBuilder().startObject().field("name", "Document 1").startObject("location").field("type", "envelope") + .startArray("coordinates").startArray().value(100).value(80).endArray().startArray().value(120).value(40).endArray() + .endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + // on left side (outside heatmap) + client().prepareIndex("test", "type1", "2") + .setSource(jsonBuilder().startObject().field("name", "Document 2").startObject("location").field("type", "envelope") + .startArray("coordinates").startArray().value(-120).value(80).endArray().startArray().value(-110).value(20) + .endArray().endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + // just left of BOX 0 + client().prepareIndex("test", "type1", "3") + .setSource(jsonBuilder().startObject().field("name", "Document 3").startObject("location").field("type", "point") + .startArray("coordinates").value(70).value(60).endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + // just outside box 0 (above it) near pole, + client().prepareIndex("test", "type1", "4") + .setSource(jsonBuilder().startObject().field("name", "Document 4").startObject("location").field("type", "point") + .startArray("coordinates").value(91).value(89).endArray().endObject().endObject()) + .setRefreshPolicy(IMMEDIATE).get(); + + + } + + private void assertGridLevel(String aggName, int expected, SearchResponse actual) { + GeoHeatmap heatmap = actual.getAggregations().get(aggName); + assertEquals(expected, heatmap.getGridLevel()); + } + + private void assertHeatmapContents(GeoShapeQueryBuilder geo, String expected, int count, String... indices) throws IOException { + SearchResponse searchResponse = client().prepareSearch(indices).setTypes("type1") + .setQuery(QueryBuilders.matchAllQuery()) + .addAggregation(heatmap("heatmap1").geom(geo).field("location").gridLevel(4).maxCells(100_000)) + .execute().actionGet(); + + assertSearchResponse(searchResponse); + assertThat(searchResponse.getHits().getTotalHits(), equalTo(new Long(count))); + assertThat(searchResponse.getHits().hits().length, equalTo(count)); + + XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); + searchResponse.toXContent(builder, ToXContent.EMPTY_PARAMS); + builder.endObject(); + String responseString = builder.string(); + + assertThat(responseString, containsString(expected)); + } + +} diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapParserTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapParserTests.java new file mode 100644 index 000000000000..e4a66c4956d6 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/geoheatmap/GeoHeatmapParserTests.java @@ -0,0 +1,92 @@ +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.elasticsearch.search.aggregations.metrics.geoheatmap; + +import org.elasticsearch.common.ParseFieldMatcher; +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.common.xcontent.json.JsonXContent; +import org.elasticsearch.index.query.GeoShapeQueryBuilder; +import org.elasticsearch.index.query.QueryParseContext; +import org.elasticsearch.index.query.QueryParser; +import org.elasticsearch.indices.query.IndicesQueriesRegistry; +import org.elasticsearch.test.ESTestCase; + +/** + * Tests the construction of the aggregator from JSON + */ +public class GeoHeatmapParserTests extends ESTestCase { + + /** + * Randomly verifies possible field values are able to parse, except the geo_shape query + * parsing which has its own tests + */ + public void testParsing() throws Exception { + StringBuilder sb = new StringBuilder(); + sb.append("{"); + sb.append("\"field\": \"my_loc\""); + if (randomBoolean()) { + appendRandomNumericOrString(sb, "grid_level", ""+randomInt()+""); + } else { + if (randomBoolean()) { + sb.append(", \"dist_err\": \""+randomDouble()+" "+randomUnits()+ "\""); + } + appendRandomNumericOrString(sb, "dist_err_pct", ""+randomDouble()+""); + } + if (randomBoolean()) { + appendRandomNumericOrString(sb, "max_cells", ""+randomInt()+""); + } + if (randomBoolean()) { + sb.append( + ", \"geom\":{" + + " \"geo_shape\": {" + + " \"location\": {" + + " \"shape\": {" + + " \"type\": \"envelope\"," + + " \"coordinates\" : [[13.0, 53.0], [14.0, 52.0]]" + + " }," + + " \"relation\": \"within\"}}}"); + } + sb.append("}"); + XContentParser stParser = JsonXContent.jsonXContent.createParser(sb.toString()); + QueryParser parser = GeoShapeQueryBuilder::fromXContent; + IndicesQueriesRegistry mockRegistry = new IndicesQueriesRegistry(); + mockRegistry.register(parser, "geo_shape"); + QueryParseContext parseContext = new QueryParseContext(mockRegistry, stParser, ParseFieldMatcher.STRICT); + XContentParser.Token token = stParser.nextToken(); + assertSame(XContentParser.Token.START_OBJECT, token); + GeoHeatmapAggregationBuilder builder = GeoHeatmapAggregationBuilder.parse("geo_heatmap", parseContext); + assertNotNull(builder); + } + + private void appendRandomNumericOrString(StringBuilder sb, String field, String value) { + if (randomBoolean()) { + sb.append(", \""+field+"\": \""+value+"\""); + } else { + sb.append(", \"grid_level\": "+value); + } + } + + private String randomUnits() { + // Do this the hard way because not all names for the units are visible in DistanceUnit + return randomFrom("in", "inch", "yd", "yard", "mi", "miles", "km", "kilometers", "m", "meters", + "cm", "centimeters", "mm", "millimeters"); + } + +}