From 75d55e11ade619aa64d8244b348dd0113d0bbb7e Mon Sep 17 00:00:00 2001 From: Areek Zillur Date: Tue, 3 Nov 2015 03:39:08 -0500 Subject: cut over to using ObjectParser for context mappings --- .../completion/CompletionSuggestParser.java | 3 +- .../completion/CompletionSuggestionBuilder.java | 30 +++-- .../completion/CompletionSuggestionContext.java | 6 +- .../completion/context/CategoryContextMapping.java | 63 +-------- .../completion/context/CategoryQueryContext.java | 48 ++++++- .../suggest/completion/context/ContextMapping.java | 30 +++-- .../completion/context/ContextMappings.java | 11 +- .../completion/context/GeoContextMapping.java | 147 ++------------------- .../completion/context/GeoQueryContext.java | 101 +++++++++++++- 9 files changed, 204 insertions(+), 235 deletions(-) (limited to 'core/src/main') diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java index 8b0dbd758b..9ca6c0ca8c 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestParser.java @@ -33,7 +33,6 @@ import org.elasticsearch.index.mapper.core.CompletionFieldMapper; import org.elasticsearch.index.query.RegexpFlag; import org.elasticsearch.search.suggest.SuggestContextParser; import org.elasticsearch.search.suggest.SuggestionSearchContext; -import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext; import org.elasticsearch.search.suggest.completion.context.ContextMapping; import org.elasticsearch.search.suggest.completion.context.ContextMappings; @@ -178,7 +177,7 @@ public class CompletionSuggestParser implements SuggestContextParser { if (type.hasContextMappings() == false && contextParser != null) { throw new IllegalArgumentException("suggester [" + type.names().fullName() + "] doesn't expect any context"); } - Map> queryContexts = Collections.emptyMap(); + Map> queryContexts = Collections.emptyMap(); if (type.hasContextMappings() && contextParser != null) { ContextMappings contextMappings = type.getContextMappings(); contextParser.nextToken(); diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java index 8e08b5a315..b4c88e2064 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionBuilder.java @@ -43,7 +43,7 @@ import static org.elasticsearch.search.suggest.completion.context.CategoryContex public class CompletionSuggestionBuilder extends SuggestBuilder.SuggestionBuilder { private FuzzyOptionsBuilder fuzzyOptionsBuilder; private RegexOptionsBuilder regexOptionsBuilder; - private Map> queryContexts; + private Map> queryContexts; private String[] payloadFields; public CompletionSuggestionBuilder(String name) { @@ -293,10 +293,23 @@ public class CompletionSuggestionBuilder extends SuggestBuilder.SuggestionBuilde * @param queryContexts a list of {@link CategoryQueryContext} */ public CompletionSuggestionBuilder categoryContexts(String name, CategoryQueryContext... queryContexts) { + return contexts(name, queryContexts); + } + + /** + * Sets query contexts for a geo context + * @param name of the geo context to execute on + * @param queryContexts a list of {@link GeoQueryContext} + */ + public CompletionSuggestionBuilder geoContexts(String name, GeoQueryContext... queryContexts) { + return contexts(name, queryContexts); + } + + private CompletionSuggestionBuilder contexts(String name, ToXContent... queryContexts) { if (this.queryContexts == null) { this.queryContexts = new HashMap<>(2); } - List contexts = this.queryContexts.get(name); + List contexts = this.queryContexts.get(name); if (contexts == null) { contexts = new ArrayList<>(2); this.queryContexts.put(name, contexts); @@ -305,15 +318,6 @@ public class CompletionSuggestionBuilder extends SuggestBuilder.SuggestionBuilde return this; } - /** - * Sets query contexts for a geo context - * @param name of the geo context to execute on - * @param queryContexts a list of {@link GeoQueryContext} - */ - public CompletionSuggestionBuilder geoContexts(String name, GeoQueryContext... queryContexts) { - return categoryContexts(name, queryContexts); - } - @Override protected XContentBuilder innerToXContent(XContentBuilder builder, Params params) throws IOException { if (payloadFields != null) { @@ -331,9 +335,9 @@ public class CompletionSuggestionBuilder extends SuggestBuilder.SuggestionBuilde } if (queryContexts != null) { builder.startObject("contexts"); - for (Map.Entry> entry : this.queryContexts.entrySet()) { + for (Map.Entry> entry : this.queryContexts.entrySet()) { builder.startArray(entry.getKey()); - for (CategoryQueryContext queryContext : entry.getValue()) { + for (ToXContent queryContext : entry.getValue()) { queryContext.toXContent(builder, params); } builder.endArray(); diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java index 87ea90cdda..2f4b729e1a 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/CompletionSuggestionContext.java @@ -25,7 +25,7 @@ import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.core.CompletionFieldMapper; import org.elasticsearch.search.suggest.Suggester; import org.elasticsearch.search.suggest.SuggestionSearchContext; -import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext; +import org.elasticsearch.search.suggest.completion.context.ContextMapping; import org.elasticsearch.search.suggest.completion.context.ContextMappings; import java.util.List; @@ -40,7 +40,7 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest private CompletionFieldMapper.CompletionFieldType fieldType; private CompletionSuggestionBuilder.FuzzyOptionsBuilder fuzzyOptionsBuilder; private CompletionSuggestionBuilder.RegexOptionsBuilder regexOptionsBuilder; - private Map> queryContexts; + private Map> queryContexts; private MapperService mapperService; private IndexFieldDataService fieldData; private Set payloadFields; @@ -65,7 +65,7 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest this.fuzzyOptionsBuilder = fuzzyOptionsBuilder; } - void setQueryContexts(Map> queryContexts) { + void setQueryContexts(Map> queryContexts) { this.queryContexts = queryContexts; } diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java index 554a2d341d..62fec54a33 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java @@ -148,72 +148,21 @@ public class CategoryContextMapping extends ContextMapping { * */ @Override - public List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { - List queryContexts = new ArrayList<>(); + public List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { + List queryContexts = new ArrayList<>(); Token token = parser.nextToken(); if (token == Token.START_OBJECT || token == Token.VALUE_STRING) { - queryContexts.add(innerParseQueryContext(parser)); + CategoryQueryContext parse = CategoryQueryContext.parse(parser); + queryContexts.add(new QueryContext(parse.context.toString(), parse.boost, parse.isPrefix)); } else if (token == Token.START_ARRAY) { while (parser.nextToken() != Token.END_ARRAY) { - queryContexts.add(innerParseQueryContext(parser)); + CategoryQueryContext parse = CategoryQueryContext.parse(parser); + queryContexts.add(new QueryContext(parse.context.toString(), parse.boost, parse.isPrefix)); } } return queryContexts; } - private CategoryQueryContext innerParseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { - Token token = parser.currentToken(); - if (token == Token.VALUE_STRING) { - return new CategoryQueryContext(parser.text()); - } else if (token == Token.START_OBJECT) { - String currentFieldName = null; - String context = null; - boolean isPrefix = false; - int boost = 1; - while ((token = parser.nextToken()) != Token.END_OBJECT) { - if (token == Token.FIELD_NAME) { - currentFieldName = parser.currentName(); - } else if (token == Token.VALUE_STRING) { - // context, exact - if (CONTEXT_VALUE.equals(currentFieldName)) { - context = parser.text(); - } else if (CONTEXT_PREFIX.equals(currentFieldName)) { - isPrefix = Boolean.valueOf(parser.text()); - } else if (CONTEXT_BOOST.equals(currentFieldName)) { - Number number; - try { - number = Long.parseLong(parser.text()); - } catch (NumberFormatException e) { - throw new IllegalArgumentException("boost must be a string representing a numeric value, but was [" + parser.text() + "]"); - } - boost = number.intValue(); - } - } else if (token == Token.VALUE_NUMBER) { - // boost - if (CONTEXT_BOOST.equals(currentFieldName)) { - Number number = parser.numberValue(); - if (parser.numberType() == XContentParser.NumberType.INT) { - boost = number.intValue(); - } else { - throw new ElasticsearchParseException("boost must be in the interval [0..2147483647], but was [" + number.longValue() + "]"); - } - } - } else if (token == Token.VALUE_BOOLEAN) { - // exact - if (CONTEXT_PREFIX.equals(currentFieldName)) { - isPrefix = parser.booleanValue(); - } - } - } - if (context == null) { - throw new ElasticsearchParseException("no context provided"); - } - return new CategoryQueryContext(context, boost, isPrefix); - } else { - throw new ElasticsearchParseException("contexts field expected string or object but was [" + token.name() + "]"); - } - } - @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java index da4ed802c4..977587ffed 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java @@ -19,8 +19,9 @@ package org.elasticsearch.search.suggest.completion.context; -import org.elasticsearch.common.xcontent.ToXContent; -import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.ElasticsearchParseException; +import org.elasticsearch.common.ParseField; +import org.elasticsearch.common.xcontent.*; import java.io.IOException; @@ -31,13 +32,13 @@ import static org.elasticsearch.search.suggest.completion.context.CategoryContex /** * Defines the query context for {@link CategoryContextMapping} */ -public class CategoryQueryContext implements ToXContent { +public final class CategoryQueryContext implements ToXContent { - public final CharSequence context; + public CharSequence context; - public final boolean isPrefix; + public boolean isPrefix = false; - public final int boost; + public int boost = 1; /** * Creates a query context with a provided context and a @@ -65,6 +66,41 @@ public class CategoryQueryContext implements ToXContent { this.isPrefix = isPrefix; } + private CategoryQueryContext() { + } + + void setContext(CharSequence context) { + this.context = context; + } + + void setIsPrefix(boolean isPrefix) { + this.isPrefix = isPrefix; + } + + void setBoost(int boost) { + this.boost = boost; + } + + private static ObjectParser CATEGORY_PARSER = new ObjectParser<>("category", null); + static { + CATEGORY_PARSER.declareString(CategoryQueryContext::setContext, new ParseField("context")); + CATEGORY_PARSER.declareInt(CategoryQueryContext::setBoost, new ParseField("boost")); + CATEGORY_PARSER.declareBoolean(CategoryQueryContext::setIsPrefix, new ParseField("prefix")); + } + + public static CategoryQueryContext parse(XContentParser parser) throws IOException { + XContentParser.Token token = parser.currentToken(); + CategoryQueryContext queryContext = new CategoryQueryContext(); + if (token == XContentParser.Token.START_OBJECT) { + CATEGORY_PARSER.parse(parser, queryContext); + } else if (token == XContentParser.Token.VALUE_STRING) { + queryContext.setContext(parser.text()); + } else { + throw new ElasticsearchParseException("category context must be an object or string"); + } + return queryContext; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMapping.java index 0f56de55f7..b15577d6fb 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMapping.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMapping.java @@ -95,14 +95,7 @@ public abstract class ContextMapping implements ToXContent { /** * Parses query contexts for this mapper */ - public abstract List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException; - - /** - * Adds query contexts to a completion query - */ - protected List getQueryContexts(List queryContexts) { - return queryContexts; - } + public abstract List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException; /** * Implementations should add specific configurations @@ -140,4 +133,25 @@ public abstract class ContextMapping implements ToXContent { return super.toString(); } } + + public static class QueryContext { + public final String context; + public final int boost; + public final boolean isPrefix; + + public QueryContext(String context, int boost, boolean isPrefix) { + this.context = context; + this.boost = boost; + this.isPrefix = isPrefix; + } + + @Override + public String toString() { + return "QueryContext{" + + "context='" + context + '\'' + + ", boost=" + boost + + ", isPrefix=" + isPrefix + + '}'; + } + } } diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java index 76eb5a0ffd..87b702c2ff 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMappings.java @@ -27,7 +27,6 @@ import org.elasticsearch.ElasticsearchParseException; import org.elasticsearch.Version; import org.elasticsearch.common.xcontent.ToXContent; import org.elasticsearch.common.xcontent.XContentBuilder; -import org.elasticsearch.common.xcontent.XContentParser; import org.elasticsearch.index.mapper.DocumentMapperParser; import org.elasticsearch.index.mapper.ParseContext; import org.elasticsearch.index.mapper.core.CompletionFieldMapper; @@ -138,15 +137,13 @@ public class ContextMappings implements ToXContent { } /** - * Wraps a {@link CompletionQuery} with context queries, - * individual context mappings adds query contexts using - * {@link ContextMapping#getQueryContexts(List)}s + * Wraps a {@link CompletionQuery} with context queries * * @param query base completion query to wrap * @param queryContexts a map of context mapping name and collected query contexts * @return a context-enabled query */ - public ContextQuery toContextQuery(CompletionQuery query, Map> queryContexts) { + public ContextQuery toContextQuery(CompletionQuery query, Map> queryContexts) { ContextQuery typedContextQuery = new ContextQuery(query); if (queryContexts.isEmpty() == false) { CharsRefBuilder scratch = new CharsRefBuilder(); @@ -155,9 +152,9 @@ public class ContextMappings implements ToXContent { scratch.setCharAt(0, (char) typeId); scratch.setLength(1); ContextMapping mapping = contextMappings.get(typeId); - List queryContext = queryContexts.get(mapping.name()); + List queryContext = queryContexts.get(mapping.name()); if (queryContext != null) { - for (CategoryQueryContext context : mapping.getQueryContexts(queryContext)) { + for (QueryContext context : queryContext) { scratch.append(context.context); typedContextQuery.addContext(scratch.toCharsRef(), context.boost, !context.isPrefix); scratch.setLength(1); diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java index 9a01322145..f895f280d2 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoContextMapping.java @@ -232,151 +232,30 @@ public class GeoContextMapping extends ContextMapping { * see {@link GeoUtils#parseGeoPoint(String, GeoPoint)} for GEO POINT */ @Override - public List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { - List queryContexts = new ArrayList<>(); + public List parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { + List queryContexts = new ArrayList<>(); Token token = parser.nextToken(); if (token == Token.START_OBJECT || token == Token.VALUE_STRING) { - queryContexts.add(innerParseQueryContext(parser)); + queryContexts.add(GeoQueryContext.parse(parser)); } else if (token == Token.START_ARRAY) { while (parser.nextToken() != Token.END_ARRAY) { - queryContexts.add(innerParseQueryContext(parser)); + queryContexts.add(GeoQueryContext.parse(parser)); } } - return queryContexts; - } - - private GeoQueryContext innerParseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException { - Token token = parser.currentToken(); - if (token == Token.VALUE_STRING) { - return new GeoQueryContext(GeoUtils.parseGeoPoint(parser), 1, precision, precision); - } else if (token == Token.START_OBJECT) { - String currentFieldName = null; - GeoPoint point = null; - double lat = Double.NaN; - double lon = Double.NaN; - int precision = this.precision; - List neighbours = new ArrayList<>(); - int boost = 1; - while ((token = parser.nextToken()) != Token.END_OBJECT) { - if (token == Token.FIELD_NAME) { - currentFieldName = parser.currentName(); - } else if (currentFieldName != null) { - if ("lat".equals(currentFieldName)) { - if (token == Token.VALUE_STRING || token == Token.VALUE_NUMBER) { - if (point == null) { - lat = parser.doubleValue(true); - } else { - throw new ElasticsearchParseException("context must have either lat/lon or geohash"); - } - } else { - throw new ElasticsearchParseException("lat must be a number"); - } - } else if ("lon".equals(currentFieldName)) { - if (token == Token.VALUE_STRING || token == Token.VALUE_NUMBER) { - if (point == null) { - lon = parser.doubleValue(true); - } else { - throw new ElasticsearchParseException("context must have either lat/lon or geohash"); - } - } else { - throw new ElasticsearchParseException("lon must be a number"); - } - } else if (CONTEXT_VALUE.equals(currentFieldName)) { - point = GeoUtils.parseGeoPoint(parser); - } else if (CONTEXT_BOOST.equals(currentFieldName)) { - final Number number; - if (token == Token.VALUE_STRING) { - try { - number = Long.parseLong(parser.text()); - } catch (NumberFormatException e) { - throw new IllegalArgumentException("boost must be a string representing a numeric value, but was [" + parser.text() + "]"); - } - } else if (token == Token.VALUE_NUMBER) { - XContentParser.NumberType numberType = parser.numberType(); - number = parser.numberValue(); - if (numberType != XContentParser.NumberType.INT) { - throw new ElasticsearchParseException("boost must be in the interval [0..2147483647], but was [" + number.longValue() + "]"); - } - } else { - throw new ElasticsearchParseException("boost must be an int"); - } - boost = number.intValue(); - } else if (CONTEXT_NEIGHBOURS.equals(currentFieldName)) { - if (token == Token.VALUE_STRING) { - neighbours.add(GeoUtils.geoHashLevelsForPrecision(parser.text())); - } else if (token == Token.VALUE_NUMBER) { - XContentParser.NumberType numberType = parser.numberType(); - if (numberType == XContentParser.NumberType.INT || numberType == XContentParser.NumberType.LONG) { - neighbours.add(parser.intValue()); - } else { - neighbours.add(GeoUtils.geoHashLevelsForPrecision(parser.doubleValue())); - } - } else if (token == Token.START_ARRAY) { - while ((token = parser.nextToken()) != Token.END_ARRAY) { - if (token == Token.VALUE_STRING || token == Token.VALUE_NUMBER) { - neighbours.add(parser.intValue(true)); - } else { - throw new ElasticsearchParseException("neighbours array must have only numbers"); - } - } - } else { - throw new ElasticsearchParseException("neighbours must be a number or a list of numbers"); - } - } else if (CONTEXT_PRECISION.equals(currentFieldName)) { - if (token == Token.VALUE_STRING) { - precision = GeoUtils.geoHashLevelsForPrecision(parser.text()); - } else if (token == Token.VALUE_NUMBER) { - XContentParser.NumberType numberType = parser.numberType(); - if (numberType == XContentParser.NumberType.INT || numberType == XContentParser.NumberType.LONG) { - precision = parser.intValue(); - } else { - precision = GeoUtils.geoHashLevelsForPrecision(parser.doubleValue()); - } - } else { - throw new ElasticsearchParseException("precision must be a number"); - } - } - } - } - if (point == null) { - if (Double.isNaN(lat) == false && Double.isNaN(lon) == false) { - point = new GeoPoint(lat, lon); - } else { - throw new ElasticsearchParseException("no context provided"); - } - } - - String geoHash = GeoHashUtils.stringEncode(point.getLon(), point.getLat(), precision); - if (neighbours.size() > 0) { - final int[] neighbourValues = new int[neighbours.size()]; - for (int i = 0; i < neighbours.size(); i++) { - neighbourValues[i] = neighbours.get(i); - } - return new GeoQueryContext(geoHash, boost, precision, neighbourValues); - } else { - return new GeoQueryContext(geoHash, boost, precision, precision); - } - } else { - throw new ElasticsearchParseException("contexts field expected string or object but was [" + token.name() + "]"); - } - } - - @Override - public List getQueryContexts(List queryContexts) { - List queryContextList = new ArrayList<>(); - for (CategoryQueryContext queryContext : queryContexts) { - GeoQueryContext geoQueryContext = ((GeoQueryContext) queryContext); - int precision = Math.min(this.precision, geoQueryContext.context.length()); - String truncatedGeohash = geoQueryContext.context.toString().substring(0, precision); - queryContextList.add(new CategoryQueryContext(truncatedGeohash, geoQueryContext.boost, false)); + List queryContextList = new ArrayList<>(); + for (GeoQueryContext geoQueryContext : queryContexts) { + int minPrecision = Math.min(this.precision, geoQueryContext.precision); + int precision = Math.min(minPrecision, geoQueryContext.geoHash.length()); + String truncatedGeohash = geoQueryContext.geoHash.toString().substring(0, precision); + queryContextList.add(new QueryContext(truncatedGeohash, geoQueryContext.boost, truncatedGeohash.length() < this.precision)); for (int neighboursPrecision : geoQueryContext.neighbours) { int neighbourPrecision = Math.min(neighboursPrecision, truncatedGeohash.length()); String neighbourGeohash = truncatedGeohash.substring(0, neighbourPrecision); - Collection locations = new HashSet<>(); + Collection locations = new ArrayList<>(); GeoHashUtils.addNeighbors(neighbourGeohash, neighbourPrecision, locations); - boolean isPrefix = neighbourPrecision < precision; + boolean isPrefix = neighbourPrecision < this.precision; for (String location : locations) { - queryContextList.add(new CategoryQueryContext(location, geoQueryContext.boost, isPrefix)); + queryContextList.add(new QueryContext(location, geoQueryContext.boost, isPrefix)); } } } diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java index caf1a3eeb9..19ea55d804 100644 --- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java +++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/GeoQueryContext.java @@ -19,20 +19,28 @@ package org.elasticsearch.search.suggest.completion.context; +import org.elasticsearch.ElasticsearchParseException; +import org.elasticsearch.common.ParseField; import org.elasticsearch.common.geo.GeoPoint; +import org.elasticsearch.common.geo.GeoUtils; +import org.elasticsearch.common.xcontent.ObjectParser; import org.elasticsearch.common.xcontent.ToXContent; import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentParser; import java.io.IOException; +import java.util.List; import static org.elasticsearch.search.suggest.completion.context.GeoContextMapping.*; /** * Defines the query context for {@link GeoContextMapping} */ -public class GeoQueryContext extends CategoryQueryContext { - public final int precision; - public final int[] neighbours; +public final class GeoQueryContext implements ToXContent { + public CharSequence geoHash; + public int boost = 1; + public int precision = DEFAULT_PRECISION; + public int[] neighbours; /** * Creates a query context for a given geo point with a boost of 1 @@ -81,17 +89,100 @@ public class GeoQueryContext extends CategoryQueryContext { * at specified precisions */ public GeoQueryContext(CharSequence geoHash, int boost, int precision, int... neighbours) { - super(geoHash, boost, true); + this.geoHash = geoHash; + this.boost = boost; this.precision = precision; this.neighbours = neighbours; } + private GeoQueryContext() { + } + + void setBoost(int boost) { + this.boost = boost; + } + + void setPrecision(int precision) { + this.precision = precision; + } + + void setNeighbours(List neighbours) { + int[] neighbourArray = new int[neighbours.size()]; + for (int i = 0; i < neighbours.size(); i++) { + neighbourArray[i] = neighbours.get(i); + } + this.neighbours = neighbourArray; + } + + private GeoPoint point; + void setPoint(GeoPoint point) { + this.point = point; + } + + private double lat = Double.NaN; + void setLat(double lat) { + this.lat = lat; + } + + private double lon = Double.NaN; + void setLon(double lon) { + this.lon = lon; + } + + void setGeoHash(String geoHash) { + this.geoHash = geoHash; + } + + void finish() { + if (point == null) { + if (Double.isNaN(lat) == false && Double.isNaN(lon) == false) { + point = new GeoPoint(lat, lon); + } else { + throw new ElasticsearchParseException("no geohash or geo point provided"); + } + } + this.geoHash = point.geohash(); + if (this.neighbours == null) { + this.neighbours = new int[]{precision}; + } + } + + private static ObjectParser GEO_CONTEXT_PARSER = new ObjectParser<>("geo", null); + static { + GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setPoint(GeoUtils.parseGeoPoint(parser)), new ParseField("context"), ObjectParser.ValueType.OBJECT); + GEO_CONTEXT_PARSER.declareString(GeoQueryContext::setGeoHash, new ParseField("context")); + GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setBoost, new ParseField("boost")); + // TODO : add string support for precision for GeoUtils.geoHashLevelsForPrecision() + GEO_CONTEXT_PARSER.declareInt(GeoQueryContext::setPrecision, new ParseField("precision")); + // TODO : add string array support for precision for GeoUtils.geoHashLevelsForPrecision() + GEO_CONTEXT_PARSER.declareIntArray(GeoQueryContext::setNeighbours, new ParseField("neighbours")); + GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext::setLat, new ParseField("lat")); + GEO_CONTEXT_PARSER.declareDouble(GeoQueryContext::setLon, new ParseField("lon")); + } + + public static GeoQueryContext parse(XContentParser parser) throws IOException { + XContentParser.Token token = parser.currentToken(); + GeoQueryContext queryContext = new GeoQueryContext(); + if (token == XContentParser.Token.START_OBJECT) { + GEO_CONTEXT_PARSER.parse(parser, queryContext); + } else if (token == XContentParser.Token.VALUE_STRING) { + queryContext.setPoint(GeoPoint.fromGeohash(parser.text())); + } else { + throw new ElasticsearchParseException("geo context must be an object or string"); + } + queryContext.finish(); + return queryContext; + } + @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.startObject(); - builder.field(CONTEXT_VALUE, context); + builder.startObject(CONTEXT_VALUE); + builder.field("geohash", geoHash); + builder.endObject(); builder.field(CONTEXT_BOOST, boost); builder.field(CONTEXT_NEIGHBOURS, neighbours); + builder.field(CONTEXT_PRECISION, precision); builder.endObject(); return builder; } -- cgit v1.2.3