cut over to using ObjectParser for context mappings

9 files changed, 204 insertions, 235 deletions

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<String, List<CategoryQueryContext>> queryContexts = Collections.emptyMap();
+            Map<String, List<ContextMapping.QueryContext>> 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<CompletionSuggestionBuilder> {
     private FuzzyOptionsBuilder fuzzyOptionsBuilder;
     private RegexOptionsBuilder regexOptionsBuilder;
-    private Map<String, List<CategoryQueryContext>> queryContexts;
+    private Map<String, List<ToXContent>> 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<CategoryQueryContext> contexts = this.queryContexts.get(name);
+        List<ToXContent> 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<String, List<CategoryQueryContext>> entry : this.queryContexts.entrySet()) {
+            for (Map.Entry<String, List<ToXContent>> 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<String, List<CategoryQueryContext>> queryContexts;
+    private Map<String, List<ContextMapping.QueryContext>> queryContexts;
     private MapperService mapperService;
     private IndexFieldDataService fieldData;
     private Set<String> payloadFields;
@@ -65,7 +65,7 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
         this.fuzzyOptionsBuilder = fuzzyOptionsBuilder;
     }
 
-    void setQueryContexts(Map<String, List<CategoryQueryContext>> queryContexts) {
+    void setQueryContexts(Map<String, List<ContextMapping.QueryContext>> 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 {
      *  </ul>
      */
     @Override
-    public List<CategoryQueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
-        List<CategoryQueryContext> queryContexts = new ArrayList<>();
+    public List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
+        List<QueryContext> 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<CategoryQueryContext, CategoryContextMapping> 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<CategoryQueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException;
-
-    /**
-     * Adds query contexts to a completion query
-     */
-    protected List<CategoryQueryContext> getQueryContexts(List<CategoryQueryContext> queryContexts) {
-        return queryContexts;
-    }
+    public abstract List<QueryContext> 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<String, List<CategoryQueryContext>> queryContexts) {
+    public ContextQuery toContextQuery(CompletionQuery query, Map<String, List<QueryContext>> 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<CategoryQueryContext> queryContext = queryContexts.get(mapping.name());
+                List<QueryContext> 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<CategoryQueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
-        List<CategoryQueryContext> queryContexts = new ArrayList<>();
+    public List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
+        List<GeoQueryContext> 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<Integer> 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<CategoryQueryContext> getQueryContexts(List<CategoryQueryContext> queryContexts) {
-        List<CategoryQueryContext> 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<QueryContext> 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<String> locations = new HashSet<>();
+                Collection<String> 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<Integer> 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<GeoQueryContext, GeoContextMapping> 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;
     }