Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor Geohash Grid Aggregation #14138

Merged
merged 1 commit into from
Nov 16, 2015
Merged

Refactor Geohash Grid Aggregation #14138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
189 changes: 128 additions & 61 deletions ...src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoHashGridParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,108 +21,118 @@
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SortedNumericDocValues;
import org.apache.lucene.util.GeoHashUtils;
import org.elasticsearch.common.ParseField;
import org.elasticsearch.common.ParseFieldMatcher;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.XContentParser.Token;
import org.elasticsearch.index.fielddata.MultiGeoPointValues;
import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
import org.elasticsearch.index.fielddata.SortingNumericDocValues;
import org.elasticsearch.index.query.GeoBoundingBoxQueryBuilder;
import org.elasticsearch.search.SearchParseException;
import org.elasticsearch.search.aggregations.Aggregator;
import org.elasticsearch.search.aggregations.AggregatorFactory;
import org.elasticsearch.search.aggregations.InternalAggregation;
import org.elasticsearch.search.aggregations.NonCollectingAggregator;
import org.elasticsearch.search.aggregations.bucket.BucketUtils;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
import org.elasticsearch.search.aggregations.support.AbstractValuesSourceParser.GeoPointValuesSourceParser;
import org.elasticsearch.search.aggregations.support.AggregationContext;
import org.elasticsearch.search.aggregations.support.ValueType;
import org.elasticsearch.search.aggregations.support.ValuesSource;
import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
import org.elasticsearch.search.aggregations.support.ValuesSourceParser;
import org.elasticsearch.search.internal.SearchContext;
import org.elasticsearch.search.aggregations.support.ValuesSourceType;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
* Aggregates Geo information into cells determined by geohashes of a given precision.
* WARNING - for high-precision geohashes it may prove necessary to use a {@link GeoBoundingBoxQueryBuilder}
* aggregation to focus in on a smaller area to avoid generating too many buckets and using too much RAM
*/
public class GeoHashGridParser implements Aggregator.Parser {
public class GeoHashGridParser extends GeoPointValuesSourceParser {

public static final int DEFAULT_PRECISION = 5;
public static final int DEFAULT_MAX_NUM_CELLS = 10000;

public GeoHashGridParser() {
super(false, false);
}

@Override
public String type() {
return InternalGeoHashGrid.TYPE.name();
}
@Override
public AggregatorFactory getFactoryPrototype() {
return new GeoGridFactory(null);
}

@Override
public AggregatorFactory parse(String aggregationName, XContentParser parser, SearchContext context) throws IOException {

ValuesSourceParser<ValuesSource.GeoPoint> vsParser = ValuesSourceParser
.geoPoint(aggregationName, InternalGeoHashGrid.TYPE, context).build();

int precision = GeoHashGridParams.DEFAULT_PRECISION;
int requiredSize = GeoHashGridParams.DEFAULT_MAX_NUM_CELLS;
int shardSize = -1;

XContentParser.Token token;
String currentFieldName = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentFieldName = parser.currentName();
} else if (vsParser.token(currentFieldName, token, parser)) {
continue;
} else if (token == XContentParser.Token.VALUE_NUMBER ||
token == XContentParser.Token.VALUE_STRING) { //Be lenient and also allow numbers enclosed in quotes
if (context.parseFieldMatcher().match(currentFieldName, GeoHashGridParams.FIELD_PRECISION)) {
precision = GeoHashGridParams.checkPrecision(parser.intValue());
} else if (context.parseFieldMatcher().match(currentFieldName, GeoHashGridParams.FIELD_SIZE)) {
requiredSize = parser.intValue();
} else if (context.parseFieldMatcher().match(currentFieldName, GeoHashGridParams.FIELD_SHARD_SIZE)) {
shardSize = parser.intValue();
protected ValuesSourceAggregatorFactory<org.elasticsearch.search.aggregations.support.ValuesSource.GeoPoint> createFactory(
String aggregationName, ValuesSourceType valuesSourceType,
ValueType targetValueType, Map<ParseField, Object> otherOptions) {
GeoGridFactory factory = new GeoGridFactory(aggregationName);
Integer precision = (Integer) otherOptions.get(GeoHashGridParams.FIELD_PRECISION);
if (precision != null) {
factory.precision(precision);
}
} else if (token != XContentParser.Token.START_OBJECT) {
throw new SearchParseException(context, "Unexpected token " + token + " in [" + aggregationName + "].",
parser.getTokenLocation());
}
Integer size = (Integer) otherOptions.get(GeoHashGridParams.FIELD_SIZE);
if (size != null) {
factory.size(size);
}

if (shardSize == 0) {
shardSize = Integer.MAX_VALUE;
Integer shardSize = (Integer) otherOptions.get(GeoHashGridParams.FIELD_SHARD_SIZE);
if (shardSize != null) {
factory.shardSize(shardSize);
}
return factory;
}

if (requiredSize == 0) {
requiredSize = Integer.MAX_VALUE;
@Override
protected boolean token(String aggregationName, String currentFieldName, Token token, XContentParser parser,
ParseFieldMatcher parseFieldMatcher, Map<ParseField, Object> otherOptions) throws IOException {
if (token == XContentParser.Token.VALUE_NUMBER || token == XContentParser.Token.VALUE_STRING) {
if (parseFieldMatcher.match(currentFieldName, GeoHashGridParams.FIELD_PRECISION)) {
otherOptions.put(GeoHashGridParams.FIELD_PRECISION, parser.intValue());
return true;
} else if (parseFieldMatcher.match(currentFieldName, GeoHashGridParams.FIELD_SIZE)) {
otherOptions.put(GeoHashGridParams.FIELD_SIZE, parser.intValue());
return true;
} else if (parseFieldMatcher.match(currentFieldName, GeoHashGridParams.FIELD_SHARD_SIZE)) {
otherOptions.put(GeoHashGridParams.FIELD_SHARD_SIZE, parser.intValue());
return true;
}

if (shardSize < 0) {
//Use default heuristic to avoid any wrong-ranking caused by distributed counting
shardSize = BucketUtils.suggestShardSideQueueSize(requiredSize, context.numberOfShards());
}

if (shardSize < requiredSize) {
shardSize = requiredSize;
return false;
}

return new GeoGridFactory(aggregationName, vsParser.input(), precision, requiredSize, shardSize);
public static class GeoGridFactory extends ValuesSourceAggregatorFactory<ValuesSource.GeoPoint> {

}
private int precision = DEFAULT_PRECISION;
private int requiredSize = DEFAULT_MAX_NUM_CELLS;
private int shardSize = -1;

public GeoGridFactory(String name) {
super(name, InternalGeoHashGrid.TYPE.name(), ValuesSourceType.GEOPOINT, ValueType.GEOPOINT);
}

static class GeoGridFactory extends ValuesSourceAggregatorFactory<ValuesSource.GeoPoint> {
public void precision(int precision) {
this.precision = GeoHashGridParams.checkPrecision(precision);
}

private final int precision;
private final int requiredSize;
private final int shardSize;
public void size(int size) {
this.requiredSize = size;
}

public GeoGridFactory(String name, ValuesSourceParser.Input<ValuesSource.GeoPoint> input, int precision, int requiredSize,
int shardSize) {
super(name, InternalGeoHashGrid.TYPE.name(), input);
this.precision = precision;
this.requiredSize = requiredSize;
public void shardSize(int shardSize) {
this.shardSize = shardSize;
}

Expand All @@ -143,6 +153,23 @@ public InternalAggregation buildEmptyAggregation() {
protected Aggregator doCreateInternal(final ValuesSource.GeoPoint valuesSource, AggregationContext aggregationContext,
Aggregator parent, boolean collectsFromSingleBucket, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData)
throws IOException {
if (shardSize == 0) {
shardSize = Integer.MAX_VALUE;
}

if (requiredSize == 0) {
requiredSize = Integer.MAX_VALUE;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was curious why this is set to Integer.MAX_VALUE? Would we not want to reset it to DEFAULT_MAX_NUM_CELLS?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So we have a loose convention that size: 0 means get all buckets. I personally don't like it but we have to keep it for now I guess as this refactor should be backwards compatible

}

if (shardSize < 0) {
// Use default heuristic to avoid any wrong-ranking caused by
// distributed counting
shardSize = BucketUtils.suggestShardSideQueueSize(requiredSize, aggregationContext.searchContext().numberOfShards());
}

if (shardSize < requiredSize) {
shardSize = requiredSize;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so if shardSize == -1, and requiredSize == 0, then shardSize will be reset here to Integer.MAX_VALUE?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes

if (collectsFromSingleBucket == false) {
return asMultiBucketAggregator(this, aggregationContext, parent);
}
Expand All @@ -152,6 +179,52 @@ protected Aggregator doCreateInternal(final ValuesSource.GeoPoint valuesSource,

}

@Override
protected ValuesSourceAggregatorFactory<org.elasticsearch.search.aggregations.support.ValuesSource.GeoPoint> innerReadFrom(
String name, ValuesSourceType valuesSourceType,
ValueType targetValueType, StreamInput in) throws IOException {
GeoGridFactory factory = new GeoGridFactory(name);
factory.precision = in.readVInt();
factory.requiredSize = in.readVInt();
factory.shardSize = in.readVInt();
return factory;
}

@Override
protected void innerWriteTo(StreamOutput out) throws IOException {
out.writeVInt(precision);
out.writeVInt(requiredSize);
out.writeVInt(shardSize);
}

@Override
protected XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
builder.field(GeoHashGridParams.FIELD_PRECISION.getPreferredName(), precision);
builder.field(GeoHashGridParams.FIELD_SIZE.getPreferredName(), requiredSize);
builder.field(GeoHashGridParams.FIELD_SHARD_SIZE.getPreferredName(), shardSize);
return builder;
}

@Override
protected boolean innerEquals(Object obj) {
GeoGridFactory other = (GeoGridFactory) obj;
if (precision != other.precision) {
return false;
}
if (requiredSize != other.requiredSize) {
return false;
}
if (shardSize != other.shardSize) {
return false;
}
return true;
}

@Override
protected int innerHashCode() {
return Objects.hash(precision, requiredSize, shardSize);
}

private static class CellValues extends SortingNumericDocValues {
private MultiGeoPointValues geoValues;
private int precision;
Expand Down Expand Up @@ -209,10 +282,4 @@ public SortedBinaryDocValues bytesValues(LeafReaderContext ctx) {

}
}
// NORELEASE implement this method when refactoring this aggregation
@Override
public AggregatorFactory getFactoryPrototype() {
return null;
}

}
67 changes: 67 additions & 0 deletions core/src/test/java/org/elasticsearch/search/aggregations/bucket/GeoHashGridTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
/*
* 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.bucket;

import org.elasticsearch.search.aggregations.BaseAggregationTestCase;
import org.elasticsearch.search.aggregations.bucket.geogrid.GeoHashGridParser.GeoGridFactory;

public class GeoHashGridTests extends BaseAggregationTestCase<GeoGridFactory> {

@Override
protected GeoGridFactory createTestAggregatorFactory() {
String name = randomAsciiOfLengthBetween(3, 20);
GeoGridFactory factory = new GeoGridFactory(name);
if (randomBoolean()) {
int precision = randomIntBetween(1, 12);
factory.precision(precision);
}
if (randomBoolean()) {
int size = randomInt(5);
switch (size) {
case 0:
break;
case 1:
case 2:
case 3:
case 4:
size = randomInt();
break;
}
factory.size(size);

}
if (randomBoolean()) {
int shardSize = randomInt(5);
switch (shardSize) {
case 0:
break;
case 1:
case 2:
case 3:
case 4:
shardSize = randomInt();
break;
}
factory.shardSize(shardSize);
}
return factory;
}

}
12 changes: 11 additions & 1 deletion ...est/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoHashGridParserTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,8 @@ public void testParseValidFromInts() throws Exception {
int precision = randomIntBetween(1, 12);
XContentParser stParser = JsonXContent.jsonXContent.createParser(
"{\"field\":\"my_loc\", \"precision\":" + precision + ", \"size\": 500, \"shard_size\": 550}");
XContentParser.Token token = stParser.nextToken();
assertSame(XContentParser.Token.START_OBJECT, token);
GeoHashGridParser parser = new GeoHashGridParser();
// can create a factory
assertNotNull(parser.parse("geohash_grid", stParser, searchContext));
Expand All @@ -41,6 +43,8 @@ public void testParseValidFromStrings() throws Exception {
int precision = randomIntBetween(1, 12);
XContentParser stParser = JsonXContent.jsonXContent.createParser(
"{\"field\":\"my_loc\", \"precision\":\"" + precision + "\", \"size\": \"500\", \"shard_size\": \"550\"}");
XContentParser.Token token = stParser.nextToken();
assertSame(XContentParser.Token.START_OBJECT, token);
GeoHashGridParser parser = new GeoHashGridParser();
// can create a factory
assertNotNull(parser.parse("geohash_grid", stParser, searchContext));
Expand All @@ -49,6 +53,8 @@ public void testParseValidFromStrings() throws Exception {
public void testParseErrorOnNonIntPrecision() throws Exception {
SearchContext searchContext = new TestSearchContext();
XContentParser stParser = JsonXContent.jsonXContent.createParser("{\"field\":\"my_loc\", \"precision\":\"2.0\"}");
XContentParser.Token token = stParser.nextToken();
assertSame(XContentParser.Token.START_OBJECT, token);
GeoHashGridParser parser = new GeoHashGridParser();
try {
parser.parse("geohash_grid", stParser, searchContext);
Expand All @@ -61,18 +67,22 @@ public void testParseErrorOnNonIntPrecision() throws Exception {
public void testParseErrorOnBooleanPrecision() throws Exception {
SearchContext searchContext = new TestSearchContext();
XContentParser stParser = JsonXContent.jsonXContent.createParser("{\"field\":\"my_loc\", \"precision\":false}");
XContentParser.Token token = stParser.nextToken();
assertSame(XContentParser.Token.START_OBJECT, token);
GeoHashGridParser parser = new GeoHashGridParser();
try {
parser.parse("geohash_grid", stParser, searchContext);
fail();
} catch (SearchParseException ex) {
assertEquals("Unexpected token VALUE_BOOLEAN in [geohash_grid].", ex.getMessage());
assertEquals("Unexpected token VALUE_BOOLEAN [precision] in [geohash_grid].", ex.getMessage());
}
}

public void testParseErrorOnPrecisionOutOfRange() throws Exception {
SearchContext searchContext = new TestSearchContext();
XContentParser stParser = JsonXContent.jsonXContent.createParser("{\"field\":\"my_loc\", \"precision\":\"13\"}");
XContentParser.Token token = stParser.nextToken();
assertSame(XContentParser.Token.START_OBJECT, token);
GeoHashGridParser parser = new GeoHashGridParser();
try {
parser.parse("geohash_grid", stParser, searchContext);
Expand Down