construct suggestion context from query context

22 files changed, 515 insertions, 453 deletions

diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
index ce45966200..cd2b6fb79c 100644
--- a/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
+++ b/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
@@ -717,13 +717,6 @@ public abstract class StreamInput extends InputStream {
     }
 
     /**
-     * Reads a completion {@link QueryContext} from the current stream
-     */
-    public QueryContext readCompletionSuggestionQueryContext() throws IOException {
-        return readNamedWriteable(QueryContext.class);
-    }
-
-    /**
      * Reads a {@link org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder} from the current stream
      */
     public ScoreFunctionBuilder<?> readScoreFunction() throws IOException {
diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
index 95ec0fec29..a4385dce19 100644
--- a/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
+++ b/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
@@ -740,10 +740,4 @@ public abstract class StreamOutput extends OutputStream {
         writeNamedWriteable(suggestion);
     }
 
-    /**
-     * Writes a completion {@link QueryContext} to the current stream
-     */
-    public void writeCompletionSuggestionQueryContext(QueryContext queryContext) throws IOException {
-        writeNamedWriteable(queryContext);
-    }
 }
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 04f63042d4..e5b70db699 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
@@ -21,7 +21,6 @@ package org.elasticsearch.search.suggest.completion;
 import org.apache.lucene.analysis.Analyzer;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
@@ -34,14 +33,8 @@ import org.elasticsearch.index.query.RegexpFlag;
 import org.elasticsearch.search.suggest.SuggestContextParser;
 import org.elasticsearch.search.suggest.SuggestUtils.Fields;
 import org.elasticsearch.search.suggest.SuggestionSearchContext;
-import org.elasticsearch.search.suggest.completion.context.ContextMapping;
-import org.elasticsearch.search.suggest.completion.context.ContextMappings;
 
 import java.io.IOException;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
 
 /**
  * Parses query options for {@link CompletionSuggester}
@@ -74,27 +67,20 @@ import java.util.Map;
 public class CompletionSuggestParser implements SuggestContextParser {
 
     private static ObjectParser<CompletionSuggestionContext, ContextAndSuggest> TLP_PARSER = new ObjectParser<>(CompletionSuggestionBuilder.SUGGESTION_NAME, null);
-    private static ObjectParser<RegexOptions.Builder, ContextAndSuggest> REGEXP_PARSER = new ObjectParser<>(RegexOptions.REGEX_OPTIONS.getPreferredName(), RegexOptions.Builder::new);
-    private static ObjectParser<FuzzyOptions.Builder, ContextAndSuggest> FUZZY_PARSER = new ObjectParser<>(FuzzyOptions.FUZZY_OPTIONS.getPreferredName(), FuzzyOptions.Builder::new);
     static {
-        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyMinLength, FuzzyOptions.MIN_LENGTH_FIELD);
-        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setMaxDeterminizedStates, FuzzyOptions.MAX_DETERMINIZED_STATES_FIELD);
-        FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setUnicodeAware, FuzzyOptions.UNICODE_AWARE_FIELD);
-        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyPrefixLength, FuzzyOptions.PREFIX_LENGTH_FIELD);
-        FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setTranspositions, FuzzyOptions.TRANSPOSITION_FIELD);
-        FUZZY_PARSER.declareValue((a, b) -> {
-            try {
-                a.setFuzziness(Fuzziness.parse(b).asDistance());
-            } catch (IOException e) {
-                throw new ElasticsearchException(e);
-            }
-        }, Fuzziness.FIELD);
-        REGEXP_PARSER.declareInt(RegexOptions.Builder::setMaxDeterminizedStates, RegexOptions.MAX_DETERMINIZED_STATES);
-        REGEXP_PARSER.declareStringOrNull(RegexOptions.Builder::setFlags, RegexOptions.FLAGS_VALUE);
-
         TLP_PARSER.declareStringArray(CompletionSuggestionContext::setPayloadFields, CompletionSuggestionBuilder.PAYLOAD_FIELD);
-        TLP_PARSER.declareObjectOrDefault(CompletionSuggestionContext::setFuzzyOptionsBuilder, FUZZY_PARSER, FuzzyOptions.Builder::new, FuzzyOptions.FUZZY_OPTIONS);
-        TLP_PARSER.declareObject(CompletionSuggestionContext::setRegexOptionsBuilder, REGEXP_PARSER, RegexOptions.REGEX_OPTIONS);
+        TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> {
+                if (parser.currentToken() == XContentParser.Token.VALUE_BOOLEAN) {
+                    if (parser.booleanValue()) {
+                        completionSuggestionContext.setFuzzyOptions(new FuzzyOptions.Builder().build());
+                    }
+                } else {
+                    completionSuggestionContext.setFuzzyOptions(FuzzyOptions.parse(parser));
+                }
+            },
+            FuzzyOptions.FUZZY_OPTIONS, ObjectParser.ValueType.OBJECT_OR_BOOLEAN);
+        TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> completionSuggestionContext.setRegexOptions(RegexOptions.parse(parser)),
+            RegexOptions.REGEX_OPTIONS, ObjectParser.ValueType.OBJECT);
         TLP_PARSER.declareString(SuggestionSearchContext.SuggestionContext::setField, Fields.FIELD);
         TLP_PARSER.declareField((p, v, c) -> {
             String analyzerName = p.text();
@@ -132,7 +118,7 @@ public class CompletionSuggestParser implements SuggestContextParser {
     }
 
     @Override
-    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser,  QueryShardContext shardContext) throws IOException {
+    public SuggestionSearchContext.SuggestionContext parse(XContentParser parser, QueryShardContext shardContext) throws IOException {
         MapperService mapperService = shardContext.getMapperService();
         final CompletionSuggestionContext suggestion = new CompletionSuggestionContext(shardContext);
         final ContextAndSuggest contextAndSuggest = new ContextAndSuggest(mapperService);
@@ -146,28 +132,12 @@ public class CompletionSuggestParser implements SuggestContextParser {
             if (type.hasContextMappings() == false && contextParser != null) {
                 throw new IllegalArgumentException("suggester [" + type.name() + "] doesn't expect any context");
             }
-            Map<String, List<ContextMapping.QueryContext>> queryContexts = Collections.emptyMap();
-            if (type.hasContextMappings() && contextParser != null) {
-                ContextMappings contextMappings = type.getContextMappings();
-                contextParser.nextToken();
-                queryContexts = new HashMap<>(contextMappings.size());
-                assert contextParser.currentToken() == XContentParser.Token.START_OBJECT;
-                XContentParser.Token currentToken;
-                String currentFieldName;
-                while ((currentToken = contextParser.nextToken()) != XContentParser.Token.END_OBJECT) {
-                    if (currentToken == XContentParser.Token.FIELD_NAME) {
-                        currentFieldName = contextParser.currentName();
-                        final ContextMapping mapping = contextMappings.get(currentFieldName);
-                        queryContexts.put(currentFieldName, mapping.parseQueryContext(contextParser));
-                    }
-                }
-                contextParser.close();
-            }
+            suggestion.setQueryContexts(CompletionSuggestionBuilder.parseQueryContexts(contextParser, type));
             suggestion.setFieldType(type);
-            suggestion.setQueryContexts(queryContexts);
             return suggestion;
         } else {
             throw new IllegalArgumentException("Field [" + suggestion.getField() + "] is not a completion suggest field");
         }
     }
+
 }
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 f2540e4862..38242a29ae 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
@@ -18,29 +18,37 @@
  */
 package org.elasticsearch.search.suggest.completion;
 
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.unit.Fuzziness;
+import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentFactory;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
 import org.elasticsearch.index.query.QueryParseContext;
 import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.search.suggest.SuggestUtils;
 import org.elasticsearch.search.suggest.SuggestionBuilder;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
-import org.elasticsearch.search.suggest.completion.context.CategoryQueryContext;
-import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
+import org.elasticsearch.search.suggest.completion.context.ContextMapping;
+import org.elasticsearch.search.suggest.completion.context.ContextMappings;
 import org.elasticsearch.search.suggest.completion.context.QueryContext;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Defines a suggest command based on a prefix, typically to provide "auto-complete" functionality
@@ -55,10 +63,40 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
     static final ParseField PAYLOAD_FIELD = new ParseField("payload");
     static final ParseField CONTEXTS_FIELD = new ParseField("contexts", "context");
 
+    private static ObjectParser<CompletionSuggestionBuilder, Void> TLP_PARSER =
+        new ObjectParser<>(CompletionSuggestionBuilder.SUGGESTION_NAME, null);
+    static {
+        TLP_PARSER.declareStringArray(CompletionSuggestionBuilder::payload, CompletionSuggestionBuilder.PAYLOAD_FIELD);
+        TLP_PARSER.declareField((parser, completionSuggestionContext, context) -> {
+                if (parser.currentToken() == XContentParser.Token.VALUE_BOOLEAN) {
+                    if (parser.booleanValue()) {
+                        completionSuggestionContext.fuzzyOptions = new FuzzyOptions.Builder().build();
+                    }
+                } else {
+                    completionSuggestionContext.fuzzyOptions = FuzzyOptions.parse(parser);
+                }
+            },
+            FuzzyOptions.FUZZY_OPTIONS, ObjectParser.ValueType.OBJECT_OR_BOOLEAN);
+        TLP_PARSER.declareField((parser, completionSuggestionContext, context) ->
+            completionSuggestionContext.regexOptions = RegexOptions.parse(parser),
+            RegexOptions.REGEX_OPTIONS, ObjectParser.ValueType.OBJECT);
+        TLP_PARSER.declareString(CompletionSuggestionBuilder::field, SuggestUtils.Fields.FIELD);
+        TLP_PARSER.declareString(CompletionSuggestionBuilder::analyzer, SuggestUtils.Fields.ANALYZER);
+        TLP_PARSER.declareInt(CompletionSuggestionBuilder::size, SuggestUtils.Fields.SIZE);
+        TLP_PARSER.declareInt(CompletionSuggestionBuilder::shardSize, SuggestUtils.Fields.SHARD_SIZE);
+        TLP_PARSER.declareField((p, v, c) -> {
+            // Copy the current structure. We will parse, once the mapping is provided
+            XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+            builder.copyCurrentStructure(p);
+            v.contextBytes = builder.bytes();
+            p.skipChildren();
+        }, CompletionSuggestionBuilder.CONTEXTS_FIELD, ObjectParser.ValueType.OBJECT); // context is deprecated
+    }
+
     private FuzzyOptions fuzzyOptions;
     private RegexOptions regexOptions;
-    private final Map<String, List<QueryContext>> queryContexts = new HashMap<>();
-    private final Set<String> payloadFields = new HashSet<>();
+    private BytesReference contextBytes = null;
+    private List<String> payloadFields = Collections.emptyList();
 
     public CompletionSuggestionBuilder(String fieldname) {
         super(fieldname);
@@ -117,36 +155,33 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
      * Note: Only doc values enabled fields are supported
      */
     public CompletionSuggestionBuilder payload(List<String> fields) {
-        this.payloadFields.addAll(fields);
+        this.payloadFields = fields;
         return this;
     }
 
     /**
-     * Sets query contexts for a category context
-     * @param name of the category context to execute on
-     * @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}
+     * Sets query contexts for completion
+     * @param queryContexts named query contexts
+     *                      see {@link org.elasticsearch.search.suggest.completion.context.CategoryQueryContext}
+     *                      and {@link org.elasticsearch.search.suggest.completion.context.GeoQueryContext}
      */
-    public CompletionSuggestionBuilder geoContexts(String name, GeoQueryContext... queryContexts) {
-        return contexts(name, queryContexts);
-    }
-
-    private CompletionSuggestionBuilder contexts(String name, QueryContext... queryContexts) {
-        List<QueryContext> contexts = this.queryContexts.get(name);
-        if (contexts == null) {
-            contexts = new ArrayList<>(2);
-            this.queryContexts.put(name, contexts);
+    public CompletionSuggestionBuilder contexts(Map<String, List<? extends QueryContext>> queryContexts) {
+        try {
+            XContentBuilder contentBuilder = XContentFactory.jsonBuilder();
+            contentBuilder.startObject();
+            for (Map.Entry<String, List<? extends QueryContext>> contextEntry : queryContexts.entrySet()) {
+                contentBuilder.startArray(contextEntry.getKey());
+                for (ToXContent queryContext : contextEntry.getValue()) {
+                    queryContext.toXContent(contentBuilder, EMPTY_PARAMS);
+                }
+                contentBuilder.endArray();
+            }
+            contentBuilder.endObject();
+            contextBytes = contentBuilder.bytes();
+            return this;
+        } catch (IOException e) {
+            throw new IllegalArgumentException(e);
         }
-        Collections.addAll(contexts, queryContexts);
-        return this;
     }
 
     @Override
@@ -164,33 +199,44 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
         if (regexOptions != null) {
             regexOptions.toXContent(builder, params);
         }
-        if (queryContexts.isEmpty() == false) {
-            builder.startObject(CONTEXTS_FIELD.getPreferredName());
-            for (Map.Entry<String, List<QueryContext>> entry : this.queryContexts.entrySet()) {
-                builder.startArray(entry.getKey());
-                for (ToXContent queryContext : entry.getValue()) {
-                    queryContext.toXContent(builder, params);
-                }
-                builder.endArray();
-            }
-            builder.endObject();
+        if (contextBytes != null) {
+            XContentParser contextParser = XContentFactory.xContent(XContentType.JSON).createParser(contextBytes);
+            builder.field(CONTEXTS_FIELD.getPreferredName());
+            builder.copyCurrentStructure(contextParser);
         }
         return builder;
     }
 
     @Override
     protected CompletionSuggestionBuilder innerFromXContent(QueryParseContext parseContext) throws IOException {
-        // NORELEASE implement parsing logic
-        throw new UnsupportedOperationException();
+        CompletionSuggestionBuilder builder = new CompletionSuggestionBuilder();
+        TLP_PARSER.parse(parseContext.parser(), builder);
+        return builder;
     }
 
     @Override
     protected SuggestionContext innerBuild(QueryShardContext context) throws IOException {
         CompletionSuggestionContext suggestionContext = new CompletionSuggestionContext(context);
         // copy over common settings to each suggestion builder
-        populateCommonFields(context.getMapperService(), suggestionContext);
-        // NORELEASE
-        // still need to populate CompletionSuggestionContext's specific settings
+        final MapperService mapperService = context.getMapperService();
+        populateCommonFields(mapperService, suggestionContext);
+        suggestionContext.setPayloadFields(payloadFields);
+        suggestionContext.setFuzzyOptions(fuzzyOptions);
+        suggestionContext.setRegexOptions(regexOptions);
+        MappedFieldType mappedFieldType = mapperService.fullName(suggestionContext.getField());
+        if (mappedFieldType == null) {
+            throw new ElasticsearchException("Field [" + suggestionContext.getField() + "] is not a completion suggest field");
+        } else if (mappedFieldType instanceof CompletionFieldMapper.CompletionFieldType) {
+            CompletionFieldMapper.CompletionFieldType type = (CompletionFieldMapper.CompletionFieldType) mappedFieldType;
+            if (type.hasContextMappings() && contextBytes != null) {
+                XContentParser contextParser = XContentFactory.xContent(contextBytes).createParser(contextBytes);
+                suggestionContext.setQueryContexts(parseQueryContexts(contextParser, type));
+            } else if (contextBytes != null) {
+                throw new IllegalArgumentException("suggester [" + type.name() + "] doesn't expect any context");
+            }
+        } else {
+            throw new IllegalArgumentException("Field [" + suggestionContext.getField() + "] is not a completion suggest field");
+        }
         return suggestionContext;
     }
 
@@ -217,18 +263,10 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
         if (regexOptions != null) {
             regexOptions.writeTo(out);
         }
-        boolean queryContextsExists = queryContexts.isEmpty() == false;
+        boolean queryContextsExists = contextBytes != null;
         out.writeBoolean(queryContextsExists);
         if (queryContextsExists) {
-            out.writeVInt(queryContexts.size());
-            for (Map.Entry<String, List<QueryContext>> namedQueryContexts : queryContexts.entrySet()) {
-                out.writeString(namedQueryContexts.getKey());
-                List<QueryContext> queryContexts = namedQueryContexts.getValue();
-                out.writeVInt(queryContexts.size());
-                for (QueryContext queryContext : queryContexts) {
-                    out.writeCompletionSuggestionQueryContext(queryContext);
-                }
-            }
+            out.writeBytesReference(contextBytes);
         }
     }
 
@@ -237,9 +275,11 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
         CompletionSuggestionBuilder completionSuggestionBuilder = new CompletionSuggestionBuilder(fieldname);
         if (in.readBoolean()) {
             int numPayloadField = in.readVInt();
+            List<String> payloadFields = new ArrayList<>(numPayloadField);
             for (int i = 0; i < numPayloadField; i++) {
-                completionSuggestionBuilder.payloadFields.add(in.readString());
+                payloadFields.add(in.readString());
             }
+            completionSuggestionBuilder.payloadFields = payloadFields;
         }
         if (in.readBoolean()) {
             completionSuggestionBuilder.fuzzyOptions = FuzzyOptions.readFuzzyOptions(in);
@@ -248,30 +288,43 @@ public class CompletionSuggestionBuilder extends SuggestionBuilder<CompletionSug
             completionSuggestionBuilder.regexOptions = RegexOptions.readRegexOptions(in);
         }
         if (in.readBoolean()) {
-            int numNamedQueryContexts = in.readVInt();
-            for (int i = 0; i < numNamedQueryContexts; i++) {
-                String queryContextName = in.readString();
-                int numQueryContexts = in.readVInt();
-                List<QueryContext> queryContexts = new ArrayList<>(numQueryContexts);
-                for (int j = 0; j < numQueryContexts; j++) {
-                    queryContexts.add(in.readCompletionSuggestionQueryContext());
-                }
-                completionSuggestionBuilder.queryContexts.put(queryContextName, queryContexts);
-            }
+            completionSuggestionBuilder.contextBytes = in.readBytesReference();
         }
         return completionSuggestionBuilder;
     }
 
     @Override
     protected boolean doEquals(CompletionSuggestionBuilder other) {
-        return Objects.equals(payloadFields, other.payloadFields) &&
+         return Objects.equals(payloadFields, other.payloadFields) &&
             Objects.equals(fuzzyOptions, other.fuzzyOptions) &&
             Objects.equals(regexOptions, other.regexOptions) &&
-            Objects.equals(queryContexts, other.queryContexts);
+            Objects.equals(contextBytes, other.contextBytes);
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(payloadFields, fuzzyOptions, regexOptions, queryContexts);
+        return Objects.hash(payloadFields, fuzzyOptions, regexOptions, contextBytes);
+    }
+
+    static Map<String, List<ContextMapping.InternalQueryContext>> parseQueryContexts(
+        XContentParser contextParser, CompletionFieldMapper.CompletionFieldType type) throws IOException {
+        Map<String, List<ContextMapping.InternalQueryContext>> queryContexts = Collections.emptyMap();
+        if (type.hasContextMappings() && contextParser != null) {
+            ContextMappings contextMappings = type.getContextMappings();
+            contextParser.nextToken();
+            queryContexts = new HashMap<>(contextMappings.size());
+            assert contextParser.currentToken() == XContentParser.Token.START_OBJECT;
+            XContentParser.Token currentToken;
+            String currentFieldName;
+            while ((currentToken = contextParser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                if (currentToken == XContentParser.Token.FIELD_NAME) {
+                    currentFieldName = contextParser.currentName();
+                    final ContextMapping mapping = contextMappings.get(currentFieldName);
+                    queryContexts.put(currentFieldName, mapping.parseQueryContext(contextParser));
+                }
+            }
+            contextParser.close();
+        }
+        return queryContexts;
     }
 }
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 b20b9a5aee..c681455117 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
@@ -44,7 +44,7 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
     private CompletionFieldMapper.CompletionFieldType fieldType;
     private FuzzyOptions fuzzyOptions;
     private RegexOptions regexOptions;
-    private Map<String, List<ContextMapping.QueryContext>> queryContexts = Collections.emptyMap();
+    private Map<String, List<ContextMapping.InternalQueryContext>> queryContexts = Collections.emptyMap();
     private Set<String> payloadFields = Collections.emptySet();
 
     CompletionFieldMapper.CompletionFieldType getFieldType() {
@@ -55,15 +55,15 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
         this.fieldType = fieldType;
     }
 
-    void setRegexOptionsBuilder(RegexOptions.Builder regexOptionsBuilder) {
-        this.regexOptions = regexOptionsBuilder.build();
+    void setRegexOptions(RegexOptions regexOptions) {
+        this.regexOptions = regexOptions;
     }
 
-    void setFuzzyOptionsBuilder(FuzzyOptions.Builder fuzzyOptionsBuilder) {
-        this.fuzzyOptions = fuzzyOptionsBuilder.build();
+    void setFuzzyOptions(FuzzyOptions fuzzyOptions) {
+        this.fuzzyOptions = fuzzyOptions;
     }
 
-    void setQueryContexts(Map<String, List<ContextMapping.QueryContext>> queryContexts) {
+    void setQueryContexts(Map<String, List<ContextMapping.InternalQueryContext>> queryContexts) {
         this.queryContexts = queryContexts;
     }
 
@@ -79,6 +79,18 @@ public class CompletionSuggestionContext extends SuggestionSearchContext.Suggest
         return payloadFields;
     }
 
+    public FuzzyOptions getFuzzyOptions() {
+        return fuzzyOptions;
+    }
+
+    public RegexOptions getRegexOptions() {
+        return regexOptions;
+    }
+
+    public Map<String, List<ContextMapping.InternalQueryContext>> getQueryContexts() {
+        return queryContexts;
+    }
+
     CompletionQuery toQuery() {
         CompletionFieldMapper.CompletionFieldType fieldType = getFieldType();
         final CompletionQuery query;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/FuzzyOptions.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/FuzzyOptions.java
index 317ac049d6..aac58d7cb3 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/FuzzyOptions.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/FuzzyOptions.java
@@ -21,13 +21,16 @@ package org.elasticsearch.search.suggest.completion;
 
 import org.apache.lucene.search.suggest.document.FuzzyCompletionQuery;
 import org.apache.lucene.util.automaton.Operations;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.unit.Fuzziness;
+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.Objects;
@@ -43,6 +46,22 @@ public class FuzzyOptions implements ToXContent, Writeable<FuzzyOptions> {
     static final ParseField UNICODE_AWARE_FIELD = new ParseField("unicode_aware");
     static final ParseField MAX_DETERMINIZED_STATES_FIELD = new ParseField("max_determinized_states");
 
+    static ObjectParser<FuzzyOptions.Builder, Void> FUZZY_PARSER = new ObjectParser<>(FUZZY_OPTIONS.getPreferredName(), Builder::new);
+    static {
+        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyMinLength, MIN_LENGTH_FIELD);
+        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setMaxDeterminizedStates, MAX_DETERMINIZED_STATES_FIELD);
+        FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setUnicodeAware, UNICODE_AWARE_FIELD);
+        FUZZY_PARSER.declareInt(FuzzyOptions.Builder::setFuzzyPrefixLength, PREFIX_LENGTH_FIELD);
+        FUZZY_PARSER.declareBoolean(FuzzyOptions.Builder::setTranspositions, TRANSPOSITION_FIELD);
+        FUZZY_PARSER.declareValue((a, b) -> {
+            try {
+                a.setFuzziness(Fuzziness.parse(b).asDistance());
+            } catch (IOException e) {
+                throw new ElasticsearchException(e);
+            }
+        }, Fuzziness.FIELD);
+    }
+
     private int editDistance;
     private boolean transpositions;
     private int fuzzyMinLength;
@@ -63,6 +82,10 @@ public class FuzzyOptions implements ToXContent, Writeable<FuzzyOptions> {
     private FuzzyOptions() {
     }
 
+    public static FuzzyOptions parse(XContentParser parser) throws IOException {
+        return FUZZY_PARSER.parse(parser).build();
+    }
+
     public static Builder builder() {
         return new Builder();
     }
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/RegexOptions.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/RegexOptions.java
index fc183cdb1c..a1dcec2d55 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/RegexOptions.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/RegexOptions.java
@@ -21,12 +21,15 @@ package org.elasticsearch.search.suggest.completion;
 
 import org.apache.lucene.util.automaton.Operations;
 import org.apache.lucene.util.automaton.RegExp;
+import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+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 org.elasticsearch.index.query.RegexpFlag;
 
 import java.io.IOException;
@@ -39,6 +42,25 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
     static final ParseField REGEX_OPTIONS = new ParseField(NAME);
     static final ParseField FLAGS_VALUE = new ParseField("flags", "flags_value");
     static final ParseField MAX_DETERMINIZED_STATES = new ParseField("max_determinized_states");
+
+
+    private static ObjectParser<RegexOptions.Builder, Void> REGEXP_PARSER =
+        new ObjectParser<>(REGEX_OPTIONS.getPreferredName(), RegexOptions.Builder::new);
+    static {
+        REGEXP_PARSER.declareInt(RegexOptions.Builder::setMaxDeterminizedStates, MAX_DETERMINIZED_STATES);
+        REGEXP_PARSER.declareField((parser, builder, aVoid) -> {
+            if (parser.currentToken() == XContentParser.Token.VALUE_STRING) {
+                builder.setFlags(parser.text());
+            } else if (parser.currentToken() == XContentParser.Token.VALUE_NUMBER) {
+                builder.setFlagsValue(parser.intValue());
+            } else {
+                throw new ElasticsearchParseException(REGEX_OPTIONS.getPreferredName()
+                    + " " + FLAGS_VALUE.getPreferredName() + " supports string or number");
+            }
+        }, FLAGS_VALUE, ObjectParser.ValueType.VALUE);
+        REGEXP_PARSER.declareStringOrNull(RegexOptions.Builder::setFlags, FLAGS_VALUE);
+    }
+
     private int flagsValue;
     private int maxDeterminizedStates;
 
@@ -69,6 +91,10 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
         return new Builder();
     }
 
+    public static RegexOptions parse(XContentParser parser) throws IOException {
+        return REGEXP_PARSER.parse(parser).build();
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
@@ -135,6 +161,11 @@ public class RegexOptions implements ToXContent, Writeable<RegexOptions> {
             return this;
         }
 
+        private Builder setFlagsValue(int flagsValue) {
+            this.flagsValue = flagsValue;
+            return this;
+        }
+
         /**
          * Sets the maximum automaton states allowed for the regular expression expansion
          */
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 10ac3935cc..c9cb165aef 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
@@ -36,6 +36,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * A {@link ContextMapping} that uses a simple string as a criteria
@@ -44,7 +45,7 @@ import java.util.Set;
  * {@link CategoryQueryContext} defines options for constructing
  * a unit of query context for this context type
  */
-public class CategoryContextMapping extends ContextMapping {
+public class CategoryContextMapping extends ContextMapping<CategoryQueryContext> {
 
     private static final String FIELD_FIELDNAME = "path";
 
@@ -137,6 +138,11 @@ public class CategoryContextMapping extends ContextMapping {
         return (values == null) ? Collections.<CharSequence>emptySet() : values;
     }
 
+    @Override
+    protected CategoryQueryContext prototype() {
+        return CategoryQueryContext.PROTOTYPE;
+    }
+
     /**
      * Parse a list of {@link CategoryQueryContext}
      * using <code>parser</code>. A QueryContexts accepts one of the following forms:
@@ -154,19 +160,13 @@ public class CategoryContextMapping extends ContextMapping {
      *  </ul>
      */
     @Override
-    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) {
-            CategoryQueryContext parse = CategoryQueryContext.PROTOTYPE.fromXContext(parser);
-            queryContexts.add(new QueryContext(parse.getCategory(), parse.getBoost(), parse.isPrefix()));
-        } else if (token == Token.START_ARRAY) {
-            while (parser.nextToken() != Token.END_ARRAY) {
-                CategoryQueryContext parse = CategoryQueryContext.PROTOTYPE.fromXContext(parser);
-                queryContexts.add(new QueryContext(parse.getCategory(), parse.getBoost(), parse.isPrefix()));
-            }
-        }
-        return queryContexts;
+    public List<InternalQueryContext> toInternalQueryContexts(List<CategoryQueryContext> queryContexts) {
+        List<InternalQueryContext> internalInternalQueryContexts = new ArrayList<>(queryContexts.size());
+        internalInternalQueryContexts.addAll(
+            queryContexts.stream()
+                .map(queryContext -> new InternalQueryContext(queryContext.getCategory(), queryContext.getBoost(), queryContext.isPrefix()))
+                .collect(Collectors.toList()));
+        return internalInternalQueryContexts;
     }
 
     @Override
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 8db9afe5ae..a164faff8b 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
@@ -21,14 +21,11 @@ package org.elasticsearch.search.suggest.completion.context;
 
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.ParseField;
-import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.Objects;
 
 import static org.elasticsearch.search.suggest.completion.context.CategoryContextMapping.CONTEXT_BOOST;
@@ -98,11 +95,6 @@ public final class CategoryQueryContext implements QueryContext {
         return result;
     }
 
-    @Override
-    public String getWriteableName() {
-        return NAME;
-    }
-
     private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>(NAME, null);
     static {
         CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField(CONTEXT_VALUE));
@@ -134,22 +126,6 @@ public final class CategoryQueryContext implements QueryContext {
         return builder;
     }
 
-    @Override
-    public void writeTo(StreamOutput out) throws IOException {
-        out.writeBoolean(isPrefix);
-        out.writeVInt(boost);
-        out.writeString(category);
-    }
-
-    @Override
-    public QueryContext readFrom(StreamInput in) throws IOException {
-        Builder builder = new Builder();
-        builder.isPrefix = in.readBoolean();
-        builder.boost = in.readVInt();
-        builder.category = in.readString();
-        return builder.build();
-    }
-
     public static class Builder {
         private String category;
         private boolean isPrefix = false;
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 42e5cc0a15..501f4d153d 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
@@ -23,11 +23,13 @@ import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentParser.Token;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.index.mapper.ParseContext;
 import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -38,7 +40,7 @@ import java.util.Set;
  *
  * Implementations have to define how contexts are parsed at query/index time
  */
-public abstract class ContextMapping implements ToXContent {
+public abstract class ContextMapping<T extends QueryContext> implements ToXContent {
 
     public static final String FIELD_TYPE = "type";
     public static final String FIELD_NAME = "name";
@@ -94,10 +96,25 @@ public abstract class ContextMapping implements ToXContent {
      */
     protected abstract Set<CharSequence> parseContext(ParseContext.Document document);
 
+    protected abstract T prototype();
+
     /**
      * Parses query contexts for this mapper
      */
-    public abstract List<QueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException;
+    public List<InternalQueryContext> parseQueryContext(XContentParser parser) throws IOException, ElasticsearchParseException {
+        List<T> queryContexts = new ArrayList<>();
+        Token token = parser.nextToken();
+        if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
+            queryContexts.add((T) prototype().fromXContext(parser));
+        } else if (token == Token.START_ARRAY) {
+            while (parser.nextToken() != Token.END_ARRAY) {
+                queryContexts.add((T) prototype().fromXContext(parser));
+            }
+        }
+        return toInternalQueryContexts(queryContexts);
+    }
+
+    protected abstract List<InternalQueryContext> toInternalQueryContexts(List<T> queryContexts);
 
     /**
      * Implementations should add specific configurations
@@ -136,18 +153,39 @@ public abstract class ContextMapping implements ToXContent {
         }
     }
 
-    public static class QueryContext {
+    public static class InternalQueryContext {
         public final String context;
         public final int boost;
         public final boolean isPrefix;
 
-        public QueryContext(String context, int boost, boolean isPrefix) {
+        public InternalQueryContext(String context, int boost, boolean isPrefix) {
             this.context = context;
             this.boost = boost;
             this.isPrefix = isPrefix;
         }
 
         @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+
+            InternalQueryContext that = (InternalQueryContext) o;
+
+            if (boost != that.boost) return false;
+            if (isPrefix != that.isPrefix) return false;
+            return context != null ? context.equals(that.context) : that.context == null;
+
+        }
+
+        @Override
+        public int hashCode() {
+            int result = context != null ? context.hashCode() : 0;
+            result = 31 * result + boost;
+            result = 31 * result + (isPrefix ? 1 : 0);
+            return result;
+        }
+
+        @Override
         public String toString() {
             return "QueryContext{" +
                     "context='" + context + '\'' +
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 ccd4b2d584..ff550a8d34 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
@@ -152,7 +152,7 @@ public class ContextMappings implements ToXContent {
      * @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<ContextMapping.QueryContext>> queryContexts) {
+    public ContextQuery toContextQuery(CompletionQuery query, Map<String, List<ContextMapping.InternalQueryContext>> queryContexts) {
         ContextQuery typedContextQuery = new ContextQuery(query);
         if (queryContexts.isEmpty() == false) {
             CharsRefBuilder scratch = new CharsRefBuilder();
@@ -161,9 +161,9 @@ public class ContextMappings implements ToXContent {
                 scratch.setCharAt(0, (char) typeId);
                 scratch.setLength(1);
                 ContextMapping mapping = contextMappings.get(typeId);
-                List<ContextMapping.QueryContext> queryContext = queryContexts.get(mapping.name());
-                if (queryContext != null) {
-                    for (ContextMapping.QueryContext context : queryContext) {
+                List<ContextMapping.InternalQueryContext> internalQueryContext = queryContexts.get(mapping.name());
+                if (internalQueryContext != null) {
+                    for (ContextMapping.InternalQueryContext context : internalQueryContext) {
                         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 870ec3bd12..41d78e7535 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
@@ -42,6 +42,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.apache.lucene.spatial.util.GeoHashUtils.addNeighbors;
 import static org.apache.lucene.spatial.util.GeoHashUtils.stringEncode;
@@ -56,7 +57,7 @@ import static org.apache.lucene.spatial.util.GeoHashUtils.stringEncode;
  * {@link GeoQueryContext} defines the options for constructing
  * a unit of query context for this context type
  */
-public class GeoContextMapping extends ContextMapping {
+public class GeoContextMapping extends ContextMapping<GeoQueryContext> {
 
     public static final String FIELD_PRECISION = "precision";
     public static final String FIELD_FIELDNAME = "path";
@@ -221,6 +222,11 @@ public class GeoContextMapping extends ContextMapping {
         return locations;
     }
 
+    @Override
+    protected GeoQueryContext prototype() {
+        return GeoQueryContext.PROTOTYPE;
+    }
+
     /**
      * Parse a list of {@link GeoQueryContext}
      * using <code>parser</code>. A QueryContexts accepts one of the following forms:
@@ -245,17 +251,8 @@ public class GeoContextMapping extends ContextMapping {
      * see {@link GeoUtils#parseGeoPoint(String, GeoPoint)} for GEO POINT
      */
     @Override
-    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(GeoQueryContext.PROTOTYPE.fromXContext(parser));
-        } else if (token == Token.START_ARRAY) {
-            while (parser.nextToken() != Token.END_ARRAY) {
-                queryContexts.add(GeoQueryContext.PROTOTYPE.fromXContext(parser));
-            }
-        }
-        List<QueryContext> queryContextList = new ArrayList<>();
+    public List<InternalQueryContext> toInternalQueryContexts(List<GeoQueryContext> queryContexts) {
+        List<InternalQueryContext> internalQueryContextList = new ArrayList<>();
         for (GeoQueryContext queryContext : queryContexts) {
             int minPrecision = Math.min(this.precision, queryContext.getPrecision());
             GeoPoint point = queryContext.getGeoPoint();
@@ -265,19 +262,20 @@ public class GeoContextMapping extends ContextMapping {
             if (queryContext.getNeighbours().isEmpty() && geoHash.length() == this.precision) {
                 addNeighbors(geoHash, locations);
             } else if (queryContext.getNeighbours().isEmpty() == false) {
-                for (Integer neighbourPrecision : queryContext.getNeighbours()) {
-                    if (neighbourPrecision < geoHash.length()) {
+                queryContext.getNeighbours().stream()
+                    .filter(neighbourPrecision -> neighbourPrecision < geoHash.length())
+                    .forEach(neighbourPrecision -> {
                         String truncatedGeoHash = geoHash.substring(0, neighbourPrecision);
                         locations.add(truncatedGeoHash);
                         addNeighbors(truncatedGeoHash, locations);
-                    }
-                }
-            }
-            for (String location : locations) {
-                queryContextList.add(new QueryContext(location, queryContext.getBoost(), location.length() < this.precision));
+                    });
             }
+            internalQueryContextList.addAll(
+                locations.stream()
+                    .map(location -> new InternalQueryContext(location, queryContext.getBoost(), location.length() < this.precision))
+                    .collect(Collectors.toList()));
         }
-        return queryContextList;
+        return internalQueryContextList;
     }
 
     @Override
@@ -301,7 +299,7 @@ public class GeoContextMapping extends ContextMapping {
         private int precision = DEFAULT_PRECISION;
         private String fieldName = null;
 
-        protected Builder(String name) {
+        public Builder(String name) {
             super(name);
         }
 
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 5b406abc1d..913702c18d 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
@@ -23,14 +23,11 @@ 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.io.stream.StreamInput;
-import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -114,11 +111,6 @@ public final class GeoQueryContext implements QueryContext {
         return new Builder();
     }
 
-    @Override
-    public String getWriteableName() {
-        return NAME;
-    }
-
     private static ObjectParser<GeoQueryContext.Builder, Void> GEO_CONTEXT_PARSER = new ObjectParser<>(NAME, null);
     static {
         GEO_CONTEXT_PARSER.declareField((parser, geoQueryContext, geoContextMapping) -> geoQueryContext.setGeoPoint(GeoUtils.parseGeoPoint(parser)), new ParseField(CONTEXT_VALUE), ObjectParser.ValueType.OBJECT);
@@ -159,33 +151,6 @@ public final class GeoQueryContext implements QueryContext {
         return builder;
     }
 
-    @Override
-    public void writeTo(StreamOutput out) throws IOException {
-        out.writeGeoPoint(geoPoint);
-        out.writeVInt(boost);
-        out.writeInt(precision);
-        out.writeVInt(neighbours.size());
-        for (Integer neighbour : neighbours) {
-            out.writeVInt(neighbour);
-        }
-    }
-
-    @Override
-    public QueryContext readFrom(StreamInput in) throws IOException {
-        Builder builder = new Builder();
-        builder.geoPoint = in.readGeoPoint();
-        builder.boost = in.readVInt();
-        builder.precision = in.readInt();
-        int nNeighbour = in.readVInt();
-        if (nNeighbour != 0) {
-            builder.neighbours = new ArrayList<>(nNeighbour);
-            for (int i = 0; i < nNeighbour; i++) {
-                builder.neighbours.add(in.readVInt());
-            }
-        }
-        return builder.build();
-    }
-
     public static class Builder {
         private GeoPoint geoPoint;
         private int boost = 1;
diff --git a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/QueryContext.java b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/QueryContext.java
index ccfd4a8d3d..9d96bf8144 100644
--- a/core/src/main/java/org/elasticsearch/search/suggest/completion/context/QueryContext.java
+++ b/core/src/main/java/org/elasticsearch/search/suggest/completion/context/QueryContext.java
@@ -19,7 +19,6 @@
 
 package org.elasticsearch.search.suggest.completion.context;
 
-import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentParser;
 
@@ -28,7 +27,7 @@ import java.io.IOException;
 /**
  * Interface for serializing/de-serializing completion query context
  */
-public interface QueryContext extends ToXContent, NamedWriteable<QueryContext> {
+public interface QueryContext extends ToXContent {
 
     QueryContext fromXContext(XContentParser parser) throws IOException;
 }
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java b/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
index f31b57d851..2c310b0487 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/AbstractSuggestionBuilderTestCase.java
@@ -21,6 +21,7 @@ package org.elasticsearch.search.suggest;
 
 import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
 import org.elasticsearch.common.ParseFieldMatcher;
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
@@ -160,7 +161,8 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
             assertTrue("suggestion builder is not equal to self", firstBuilder.equals(firstBuilder));
             assertThat("same suggestion builder's hashcode returns different values if called multiple times", firstBuilder.hashCode(),
                     equalTo(firstBuilder.hashCode()));
-            assertThat("different suggestion builders should not be equal", mutate(firstBuilder), not(equalTo(firstBuilder)));
+        final SB mutate = mutate(firstBuilder);
+        assertThat("different suggestion builders should not be equal", mutate, not(equalTo(firstBuilder)));
 
             SB secondBuilder = serializedCopy(firstBuilder);
             assertTrue("suggestion builder is not equal to self", secondBuilder.equals(secondBuilder));
@@ -211,6 +213,22 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
         }
     }
 
+    protected Tuple<MapperService, SB> mockMapperServiceAndSuggestionBuilder(
+        IndexSettings idxSettings, AnalysisService mockAnalysisService, SB suggestBuilder) {
+        final MapperService mapperService = new MapperService(idxSettings, mockAnalysisService, null,
+            new IndicesModule().getMapperRegistry(), null) {
+            @Override
+            public MappedFieldType fullName(String fullName) {
+                StringFieldType type = new StringFieldType();
+                if (randomBoolean()) {
+                    type.setSearchAnalyzer(new NamedAnalyzer("foo", new WhitespaceAnalyzer()));
+                }
+                return type;
+            }
+        };
+        return new Tuple<>(mapperService, suggestBuilder);
+    }
+
     /**
      * parses random suggestion builder via old parseElement method and via
      * build, comparing the results for equality
@@ -226,30 +244,21 @@ public abstract class AbstractSuggestionBuilderTestCase<SB extends SuggestionBui
             }
         };
 
-        MapperService mockMapperService = new MapperService(idxSettings, mockAnalysisService, null,
-                                                               new IndicesModule().getMapperRegistry(), null) {
-            @Override
-            public MappedFieldType fullName(String fullName) {
-                StringFieldType type = new StringFieldType();
-                if (randomBoolean()) {
-                    type.setSearchAnalyzer(new NamedAnalyzer("foo", new WhitespaceAnalyzer()));
-                }
-                return type;
-            }
-        };
-
-        QueryShardContext mockShardContext = new QueryShardContext(idxSettings, null, null, mockMapperService, null, scriptService, null) {
-            @Override
-            public MappedFieldType fieldMapper(String name) {
-                StringFieldMapper.Builder builder = new StringFieldMapper.Builder(name);
-                return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
-            }
-        };
-        mockShardContext.setMapUnmappedFieldAsString(true);
-
         for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
             SuggestBuilder suggestBuilder = new SuggestBuilder();
             SB suggestionBuilder = randomTestBuilder();
+            Tuple<MapperService, SB> mapperServiceSBTuple =
+                mockMapperServiceAndSuggestionBuilder(idxSettings, mockAnalysisService, suggestionBuilder);
+            suggestionBuilder = mapperServiceSBTuple.v2();
+            QueryShardContext mockShardContext = new QueryShardContext(idxSettings,
+                null, null, mapperServiceSBTuple.v1(), null, scriptService, null) {
+                @Override
+                public MappedFieldType fieldMapper(String name) {
+                    StringFieldMapper.Builder builder = new StringFieldMapper.Builder(name);
+                    return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
+                }
+            };
+            mockShardContext.setMapUnmappedFieldAsString(true);
             suggestBuilder.addSuggestion(randomAsciiOfLength(10), suggestionBuilder);
 
             if (suggestionBuilder.text() == null) {
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java b/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
index fa552172b6..58458c9d24 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/ContextCompletionSuggestSearchIT.java
@@ -36,12 +36,14 @@ import org.elasticsearch.search.suggest.completion.context.ContextBuilder;
 import org.elasticsearch.search.suggest.completion.context.ContextMapping;
 import org.elasticsearch.search.suggest.completion.context.GeoContextMapping;
 import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
+import org.elasticsearch.search.suggest.completion.context.QueryContext;
 import org.elasticsearch.test.ESIntegTestCase;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Locale;
@@ -180,7 +182,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         indexRandom(true, indexRequestBuilders);
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
+                .contexts(Collections.singletonMap("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat0").build())));
 
         assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
     }
@@ -207,9 +209,9 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         indexRandom(true, indexRequestBuilders);
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .categoryContexts("cat",
-                        CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
-                        CategoryQueryContext.builder().setCategory("cat1").build()
+                .contexts(Collections.singletonMap("cat",
+                        Arrays.asList(CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                        CategoryQueryContext.builder().setCategory("cat1").build()))
                 );
         assertSuggestions("foo", prefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
     }
@@ -267,24 +269,21 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
 
         // filter only on context cat
         CompletionSuggestionBuilder catFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
-        catFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat0").build());
+        catFilterSuggest.contexts(Collections.singletonMap("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat0").build())));
         assertSuggestions("foo", catFilterSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
 
         // filter only on context type
         CompletionSuggestionBuilder typeFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
-        typeFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build(),
-                CategoryQueryContext.builder().setCategory("type1").build());
+        typeFilterSuggest.contexts(Collections.singletonMap("type", Arrays.asList(CategoryQueryContext.builder().setCategory("type2").build(),
+                CategoryQueryContext.builder().setCategory("type1").build())));
         assertSuggestions("foo", typeFilterSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
 
         CompletionSuggestionBuilder multiContextFilterSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
         // query context order should never matter
-        if (randomBoolean()) {
-            multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
-            multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
-        } else {
-            multiContextFilterSuggest.categoryContexts("cat", CategoryQueryContext.builder().setCategory("cat2").build());
-            multiContextFilterSuggest.categoryContexts("type", CategoryQueryContext.builder().setCategory("type2").build());
-        }
+        Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
+        contextMap.put("type", Collections.singletonList(CategoryQueryContext.builder().setCategory("type2").build()));
+        contextMap.put("cat", Collections.singletonList(CategoryQueryContext.builder().setCategory("cat2").build()));
+        multiContextFilterSuggest.contexts(contextMap);
         assertSuggestions("foo", multiContextFilterSuggest, "suggestion6", "suggestion2");
     }
 
@@ -315,36 +314,33 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
 
         // boost only on context cat
         CompletionSuggestionBuilder catBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
-        catBoostSuggest.categoryContexts("cat",
-                CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
-                CategoryQueryContext.builder().setCategory("cat1").build());
+        catBoostSuggest.contexts(Collections.singletonMap("cat",
+                Arrays.asList(
+                    CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+                    CategoryQueryContext.builder().setCategory("cat1").build())));
         assertSuggestions("foo", catBoostSuggest, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion2");
 
         // boost only on context type
         CompletionSuggestionBuilder typeBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
-        typeBoostSuggest.categoryContexts("type",
-                CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
-                CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
+        typeBoostSuggest.contexts(Collections.singletonMap("type",
+                Arrays.asList(
+                    CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
+                    CategoryQueryContext.builder().setCategory("type1").setBoost(4).build())));
         assertSuggestions("foo", typeBoostSuggest, "suggestion9", "suggestion5", "suggestion6", "suggestion1", "suggestion2");
 
         // boost on both contexts
         CompletionSuggestionBuilder multiContextBoostSuggest = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg");
         // query context order should never matter
-        if (randomBoolean()) {
-            multiContextBoostSuggest.categoryContexts("type",
-                    CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
-                    CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
-            multiContextBoostSuggest.categoryContexts("cat",
-                    CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
-                    CategoryQueryContext.builder().setCategory("cat1").build());
-        } else {
-            multiContextBoostSuggest.categoryContexts("cat",
-                    CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
-                    CategoryQueryContext.builder().setCategory("cat1").build());
-            multiContextBoostSuggest.categoryContexts("type",
-                    CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
-                    CategoryQueryContext.builder().setCategory("type1").setBoost(4).build());
-        }
+        Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
+        contextMap.put("type", Arrays.asList(
+            CategoryQueryContext.builder().setCategory("type2").setBoost(2).build(),
+            CategoryQueryContext.builder().setCategory("type1").setBoost(4).build())
+        );
+        contextMap.put("cat", Arrays.asList(
+            CategoryQueryContext.builder().setCategory("cat0").setBoost(3).build(),
+            CategoryQueryContext.builder().setCategory("cat1").build())
+        );
+        multiContextBoostSuggest.contexts(contextMap);
         assertSuggestions("foo", multiContextBoostSuggest, "suggestion9", "suggestion6", "suggestion5", "suggestion2", "suggestion1");
     }
 
@@ -463,7 +459,8 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
 
         CompletionSuggestionBuilder geoFilteringPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(geoPoints[0])).build());
+                .contexts(Collections.singletonMap("geo", Collections.singletonList(
+                    GeoQueryContext.builder().setGeoPoint(new GeoPoint(geoPoints[0])).build())));
 
         assertSuggestions("foo", geoFilteringPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion2", "suggestion0");
     }
@@ -497,7 +494,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         GeoQueryContext context1 = GeoQueryContext.builder().setGeoPoint(geoPoints[0]).setBoost(2).build();
         GeoQueryContext context2 = GeoQueryContext.builder().setGeoPoint(geoPoints[1]).build();
         CompletionSuggestionBuilder geoBoostingPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .geoContexts("geo", context1, context2);
+                .contexts(Collections.singletonMap("geo", Arrays.asList(context1, context2)));
 
         assertSuggestions("foo", geoBoostingPrefix, "suggestion8", "suggestion6", "suggestion4", "suggestion9", "suggestion7");
     }
@@ -528,7 +525,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         indexRandom(true, indexRequestBuilders);
         ensureYellow(INDEX);
         CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.2263, 4.543)).build());
+                .contexts(Collections.singletonMap("geo", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.2263, 4.543)).build())));
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
     }
 
@@ -569,7 +566,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
         assertSuggestions("foo", prefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
 
         CompletionSuggestionBuilder geoNeighbourPrefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg")
-                .geoContexts("geo", GeoQueryContext.builder().setGeoPoint(GeoPoint.fromGeohash(geohash)).build());
+                .contexts(Collections.singletonMap("geo", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(GeoPoint.fromGeohash(geohash)).build())));
 
         assertSuggestions("foo", geoNeighbourPrefix, "suggestion9", "suggestion8", "suggestion7", "suggestion6", "suggestion5");
     }
@@ -626,7 +623,7 @@ public class ContextCompletionSuggestSearchIT extends ESIntegTestCase {
 
         String suggestionName = randomAsciiOfLength(10);
         CompletionSuggestionBuilder context = SuggestBuilders.completionSuggestion(FIELD).text("h").size(10)
-                .geoContexts("st", GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.52, 13.4)).build());
+                .contexts(Collections.singletonMap("st", Collections.singletonList(GeoQueryContext.builder().setGeoPoint(new GeoPoint(52.52, 13.4)).build())));
         SuggestResponse suggestResponse = client().prepareSuggest(INDEX).addSuggestion(suggestionName, context).get();
 
         assertEquals(suggestResponse.getSuggest().size(), 1);
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryContextMappingTests.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryContextMappingTests.java
index 0d27ba04a9..e0aabe7730 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryContextMappingTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryContextMappingTests.java
@@ -190,11 +190,11 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
         XContentBuilder builder = jsonBuilder().value("context1");
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         CategoryContextMapping mapping = ContextBuilder.category("cat").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1));
-        assertThat(queryContexts.get(0).context, equalTo("context1"));
-        assertThat(queryContexts.get(0).boost, equalTo(1));
-        assertThat(queryContexts.get(0).isPrefix, equalTo(false));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1));
+        assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
+        assertThat(internalQueryContexts.get(0).boost, equalTo(1));
+        assertThat(internalQueryContexts.get(0).isPrefix, equalTo(false));
     }
 
     public void testQueryContextParsingArray() throws Exception {
@@ -204,14 +204,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
                 .endArray();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         CategoryContextMapping mapping = ContextBuilder.category("cat").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(2));
-        assertThat(queryContexts.get(0).context, equalTo("context1"));
-        assertThat(queryContexts.get(0).boost, equalTo(1));
-        assertThat(queryContexts.get(0).isPrefix, equalTo(false));
-        assertThat(queryContexts.get(1).context, equalTo("context2"));
-        assertThat(queryContexts.get(1).boost, equalTo(1));
-        assertThat(queryContexts.get(1).isPrefix, equalTo(false));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(2));
+        assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
+        assertThat(internalQueryContexts.get(0).boost, equalTo(1));
+        assertThat(internalQueryContexts.get(0).isPrefix, equalTo(false));
+        assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
+        assertThat(internalQueryContexts.get(1).boost, equalTo(1));
+        assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
     }
 
     public void testQueryContextParsingObject() throws Exception {
@@ -222,11 +222,11 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
                 .endObject();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         CategoryContextMapping mapping = ContextBuilder.category("cat").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1));
-        assertThat(queryContexts.get(0).context, equalTo("context1"));
-        assertThat(queryContexts.get(0).boost, equalTo(10));
-        assertThat(queryContexts.get(0).isPrefix, equalTo(true));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1));
+        assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
+        assertThat(internalQueryContexts.get(0).boost, equalTo(10));
+        assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
     }
 
 
@@ -245,14 +245,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
                 .endArray();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         CategoryContextMapping mapping = ContextBuilder.category("cat").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(2));
-        assertThat(queryContexts.get(0).context, equalTo("context1"));
-        assertThat(queryContexts.get(0).boost, equalTo(2));
-        assertThat(queryContexts.get(0).isPrefix, equalTo(true));
-        assertThat(queryContexts.get(1).context, equalTo("context2"));
-        assertThat(queryContexts.get(1).boost, equalTo(3));
-        assertThat(queryContexts.get(1).isPrefix, equalTo(false));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(2));
+        assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
+        assertThat(internalQueryContexts.get(0).boost, equalTo(2));
+        assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
+        assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
+        assertThat(internalQueryContexts.get(1).boost, equalTo(3));
+        assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
     }
 
     public void testQueryContextParsingMixed() throws Exception {
@@ -266,14 +266,14 @@ public class CategoryContextMappingTests extends ESSingleNodeTestCase {
                 .endArray();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         CategoryContextMapping mapping = ContextBuilder.category("cat").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(2));
-        assertThat(queryContexts.get(0).context, equalTo("context1"));
-        assertThat(queryContexts.get(0).boost, equalTo(2));
-        assertThat(queryContexts.get(0).isPrefix, equalTo(true));
-        assertThat(queryContexts.get(1).context, equalTo("context2"));
-        assertThat(queryContexts.get(1).boost, equalTo(1));
-        assertThat(queryContexts.get(1).isPrefix, equalTo(false));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(2));
+        assertThat(internalQueryContexts.get(0).context, equalTo("context1"));
+        assertThat(internalQueryContexts.get(0).boost, equalTo(2));
+        assertThat(internalQueryContexts.get(0).isPrefix, equalTo(true));
+        assertThat(internalQueryContexts.get(1).context, equalTo("context2"));
+        assertThat(internalQueryContexts.get(1).boost, equalTo(1));
+        assertThat(internalQueryContexts.get(1).isPrefix, equalTo(false));
     }
 
     public void testParsingContextFromDocument() throws Exception {
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryQueryContextTests.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryQueryContextTests.java
index b4d80ebba9..a4cfc71a3b 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryQueryContextTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/CategoryQueryContextTests.java
@@ -39,25 +39,6 @@ public class CategoryQueryContextTests extends QueryContextTestCase<CategoryQuer
     }
 
     @Override
-    protected CategoryQueryContext createMutation(CategoryQueryContext original) throws IOException {
-        final CategoryQueryContext.Builder builder = CategoryQueryContext.builder();
-        builder.setCategory(original.getCategory()).setBoost(original.getBoost()).setPrefix(original.isPrefix());
-        switch (randomIntBetween(0, 2)) {
-            case 0:
-                builder.setCategory(randomValueOtherThan(original.getCategory(), () -> randomAsciiOfLength(10)));
-                break;
-            case 1:
-                builder.setBoost(randomValueOtherThan(original.getBoost(), () -> randomIntBetween(1, 5)));
-                break;
-            case 2:
-                builder.setPrefix(!original.isPrefix());
-                break;
-
-        }
-        return builder.build();
-    }
-
-    @Override
     protected CategoryQueryContext prototype() {
         return CategoryQueryContext.PROTOTYPE;
     }
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/CompletionSuggesterBuilderTests.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/CompletionSuggesterBuilderTests.java
index c5547163d2..c92d2a1eaf 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/CompletionSuggesterBuilderTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/CompletionSuggesterBuilderTests.java
@@ -19,10 +19,21 @@
 
 package org.elasticsearch.search.suggest.completion;
 
+import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.unit.Fuzziness;
+import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.analysis.AnalysisService;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.mapper.core.CompletionFieldMapper;
+import org.elasticsearch.indices.IndicesModule;
 import org.elasticsearch.search.suggest.AbstractSuggestionBuilderTestCase;
 import org.elasticsearch.search.suggest.SuggestionSearchContext.SuggestionContext;
+import org.elasticsearch.search.suggest.completion.context.CategoryContextMapping;
 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 org.elasticsearch.search.suggest.completion.context.GeoContextMapping;
 import org.elasticsearch.search.suggest.completion.context.GeoQueryContext;
 import org.elasticsearch.search.suggest.completion.context.QueryContext;
 import org.junit.BeforeClass;
@@ -30,18 +41,28 @@ import org.junit.BeforeClass;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
 
-public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTestCase<CompletionSuggestionBuilder> {
+import static org.hamcrest.core.IsInstanceOf.instanceOf;
 
-    @BeforeClass
-    public static void initQueryContexts() {
-        namedWriteableRegistry.registerPrototype(QueryContext.class, CategoryQueryContext.PROTOTYPE);
-        namedWriteableRegistry.registerPrototype(QueryContext.class, GeoQueryContext.PROTOTYPE);
-    }
+public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTestCase<CompletionSuggestionBuilder> {
 
     @Override
     protected CompletionSuggestionBuilder randomSuggestionBuilder() {
+        return randomSuggestionBuilderWithContextInfo().builder;
+    }
+
+    private static class BuilderAndInfo {
+        CompletionSuggestionBuilder builder;
+        List<String> catContexts = new ArrayList<>();
+        List<String> geoContexts = new ArrayList<>();
+    }
+
+    private BuilderAndInfo randomSuggestionBuilderWithContextInfo() {
+        final BuilderAndInfo builderAndInfo = new BuilderAndInfo();
         CompletionSuggestionBuilder testBuilder = new CompletionSuggestionBuilder(randomAsciiOfLengthBetween(2, 20));
         switch (randomIntBetween(0, 3)) {
             case 0:
@@ -60,38 +81,66 @@ public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTe
         List<String> payloads = new ArrayList<>();
         Collections.addAll(payloads, generateRandomStringArray(5, 10, false, false));
         maybeSet(testBuilder::payload, payloads);
+        Map<String, List<? extends QueryContext>> contextMap = new HashMap<>();
         if (randomBoolean()) {
             int numContext = randomIntBetween(1, 5);
-            CategoryQueryContext[] contexts = new CategoryQueryContext[numContext];
+            List<CategoryQueryContext> contexts = new ArrayList<>(numContext);
             for (int i = 0; i < numContext; i++) {
-                contexts[i] = CategoryQueryContextTests.randomCategoryQueryContext();
+                contexts.add(CategoryQueryContextTests.randomCategoryQueryContext());
             }
-            testBuilder.categoryContexts(randomAsciiOfLength(10), contexts);
+            String name = randomAsciiOfLength(10);
+            contextMap.put(name, contexts);
+            builderAndInfo.catContexts.add(name);
         }
         if (randomBoolean()) {
             int numContext = randomIntBetween(1, 5);
-            GeoQueryContext[] contexts = new GeoQueryContext[numContext];
+            List<GeoQueryContext> contexts = new ArrayList<>(numContext);
             for (int i = 0; i < numContext; i++) {
-                contexts[i] = GeoQueryContextTests.randomGeoQueryContext();
+                contexts.add(GeoQueryContextTests.randomGeoQueryContext());
             }
-            testBuilder.geoContexts(randomAsciiOfLength(10), contexts);
+            String name = randomAsciiOfLength(10);
+            contextMap.put(name, contexts);
+            builderAndInfo.geoContexts.add(name);
         }
-        return testBuilder;
+        testBuilder.contexts(contextMap);
+        builderAndInfo.builder = testBuilder;
+        return builderAndInfo;
     }
 
     @Override
     protected void assertSuggestionContext(SuggestionContext oldSuggestion, SuggestionContext newSuggestion) {
+        assertThat(oldSuggestion, instanceOf(CompletionSuggestionContext.class));
+        assertThat(newSuggestion, instanceOf(CompletionSuggestionContext.class));
+        CompletionSuggestionContext oldCompletionSuggestion = (CompletionSuggestionContext) oldSuggestion;
+        CompletionSuggestionContext newCompletionSuggestion = (CompletionSuggestionContext) newSuggestion;
+        assertEquals(oldCompletionSuggestion.getPayloadFields(), newCompletionSuggestion.getPayloadFields());
+        assertEquals(oldCompletionSuggestion.getFuzzyOptions(), newCompletionSuggestion.getFuzzyOptions());
+        assertEquals(oldCompletionSuggestion.getRegexOptions(), newCompletionSuggestion.getRegexOptions());
+        assertEquals(oldCompletionSuggestion.getQueryContexts(), newCompletionSuggestion.getQueryContexts());
 
     }
 
     @Override
-    public void testBuild() throws IOException {
-        // skip for now
-    }
-
-    @Override
-    public void testFromXContent() throws IOException {
-        // skip for now
+    protected Tuple<MapperService, CompletionSuggestionBuilder> mockMapperServiceAndSuggestionBuilder(
+        IndexSettings idxSettings, AnalysisService mockAnalysisService, CompletionSuggestionBuilder suggestBuilder) {
+        final BuilderAndInfo builderAndInfo = randomSuggestionBuilderWithContextInfo();
+        final MapperService mapperService = new MapperService(idxSettings, mockAnalysisService, null,
+            new IndicesModule().getMapperRegistry(), null) {
+            @Override
+            public MappedFieldType fullName(String fullName) {
+                CompletionFieldMapper.CompletionFieldType type = new CompletionFieldMapper.CompletionFieldType();
+                List<ContextMapping> contextMappings = builderAndInfo.catContexts.stream()
+                    .map(catContext -> new CategoryContextMapping.Builder(catContext).build())
+                    .collect(Collectors.toList());
+                contextMappings.addAll(builderAndInfo.geoContexts.stream()
+                    .map(geoContext -> new GeoContextMapping.Builder(geoContext).build())
+                    .collect(Collectors.toList()));
+                type.setContextMappings(new ContextMappings(contextMappings));
+                return type;
+            }
+        };
+        final CompletionSuggestionBuilder builder = builderAndInfo.builder;
+        return new Tuple<>(mapperService, builder);
     }
 
     @Override
@@ -103,20 +152,20 @@ public class CompletionSuggesterBuilderTests extends AbstractSuggestionBuilderTe
                 builder.payload(payloads);
                 break;
             case 1:
-                int numCategoryContext = randomIntBetween(1, 5);
-                CategoryQueryContext[] categoryContexts = new CategoryQueryContext[numCategoryContext];
-                for (int i = 0; i < numCategoryContext; i++) {
-                    categoryContexts[i] = CategoryQueryContextTests.randomCategoryQueryContext();
+                int nCatContext = randomIntBetween(1, 5);
+                List<CategoryQueryContext> contexts = new ArrayList<>(nCatContext);
+                for (int i = 0; i < nCatContext; i++) {
+                    contexts.add(CategoryQueryContextTests.randomCategoryQueryContext());
                 }
-                builder.categoryContexts(randomAsciiOfLength(10), categoryContexts);
+                builder.contexts(Collections.singletonMap(randomAsciiOfLength(10), contexts));
                 break;
             case 2:
-                int numGeoContext = randomIntBetween(1, 5);
-                GeoQueryContext[] geoContexts = new GeoQueryContext[numGeoContext];
-                for (int i = 0; i < numGeoContext; i++) {
-                    geoContexts[i] = GeoQueryContextTests.randomGeoQueryContext();
+                int nGeoContext = randomIntBetween(1, 5);
+                List<GeoQueryContext> geoContexts = new ArrayList<>(nGeoContext);
+                for (int i = 0; i < nGeoContext; i++) {
+                    geoContexts.add(GeoQueryContextTests.randomGeoQueryContext());
                 }
-                builder.geoContexts(randomAsciiOfLength(10), geoContexts);
+                builder.contexts(Collections.singletonMap(randomAsciiOfLength(10), geoContexts));
                 break;
             case 3:
                 builder.prefix(randomAsciiOfLength(10), FuzzyOptionsTests.randomFuzzyOptions());
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoContextMappingTests.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoContextMappingTests.java
index 471de9c3e9..c16b0ce645 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoContextMappingTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoContextMappingTests.java
@@ -202,15 +202,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
         XContentBuilder builder = jsonBuilder().value("ezs42e44yx96");
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         GeoContextMapping mapping = ContextBuilder.geo("geo").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1 + 8));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1 + 8));
         Collection<String> locations = new ArrayList<>();
         locations.add("ezs42e");
         addNeighbors("ezs42e", GeoContextMapping.DEFAULT_PRECISION, locations);
-        for (ContextMapping.QueryContext queryContext : queryContexts) {
-            assertThat(queryContext.context, isIn(locations));
-            assertThat(queryContext.boost, equalTo(1));
-            assertThat(queryContext.isPrefix, equalTo(false));
+        for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
+            assertThat(internalQueryContext.context, isIn(locations));
+            assertThat(internalQueryContext.boost, equalTo(1));
+            assertThat(internalQueryContext.isPrefix, equalTo(false));
         }
     }
 
@@ -221,15 +221,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
                 .endObject();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         GeoContextMapping mapping = ContextBuilder.geo("geo").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1 + 8));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1 + 8));
         Collection<String> locations = new ArrayList<>();
         locations.add("wh0n94");
         addNeighbors("wh0n94", GeoContextMapping.DEFAULT_PRECISION, locations);
-        for (ContextMapping.QueryContext queryContext : queryContexts) {
-            assertThat(queryContext.context, isIn(locations));
-            assertThat(queryContext.boost, equalTo(1));
-            assertThat(queryContext.isPrefix, equalTo(false));
+        for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
+            assertThat(internalQueryContext.context, isIn(locations));
+            assertThat(internalQueryContext.boost, equalTo(1));
+            assertThat(internalQueryContext.isPrefix, equalTo(false));
         }
     }
 
@@ -244,8 +244,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
                 .endObject();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         GeoContextMapping mapping = ContextBuilder.geo("geo").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
         Collection<String> locations = new ArrayList<>();
         locations.add("wh0n94");
         locations.add("w");
@@ -254,10 +254,10 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
         addNeighbors("wh", 2, locations);
         locations.add("wh0");
         addNeighbors("wh0", 3, locations);
-        for (ContextMapping.QueryContext queryContext : queryContexts) {
-            assertThat(queryContext.context, isIn(locations));
-            assertThat(queryContext.boost, equalTo(10));
-            assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
+        for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
+            assertThat(internalQueryContext.context, isIn(locations));
+            assertThat(internalQueryContext.boost, equalTo(10));
+            assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
         }
     }
 
@@ -282,8 +282,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
                 .endArray();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         GeoContextMapping mapping = ContextBuilder.geo("geo").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8 + 1 + 1 + 8));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8 + 1 + 1 + 8));
         Collection<String> firstLocations = new ArrayList<>();
         firstLocations.add("wh0n94");
         firstLocations.add("w");
@@ -296,15 +296,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
         secondLocations.add("w5cx04");
         secondLocations.add("w5cx0");
         addNeighbors("w5cx0", 5, secondLocations);
-        for (ContextMapping.QueryContext queryContext : queryContexts) {
-            if (firstLocations.contains(queryContext.context)) {
-                assertThat(queryContext.boost, equalTo(10));
-            } else if (secondLocations.contains(queryContext.context)) {
-                assertThat(queryContext.boost, equalTo(2));
+        for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
+            if (firstLocations.contains(internalQueryContext.context)) {
+                assertThat(internalQueryContext.boost, equalTo(10));
+            } else if (secondLocations.contains(internalQueryContext.context)) {
+                assertThat(internalQueryContext.boost, equalTo(2));
             } else {
-                fail(queryContext.context + " was not expected");
+                fail(internalQueryContext.context + " was not expected");
             }
-            assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
+            assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
         }
     }
 
@@ -325,8 +325,8 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
                 .endArray();
         XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(builder.bytes());
         GeoContextMapping mapping = ContextBuilder.geo("geo").build();
-        List<ContextMapping.QueryContext> queryContexts = mapping.parseQueryContext(parser);
-        assertThat(queryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
+        List<ContextMapping.InternalQueryContext> internalQueryContexts = mapping.parseQueryContext(parser);
+        assertThat(internalQueryContexts.size(), equalTo(1 + 1 + 8 + 1 + 8 + 1 + 8));
         Collection<String> firstLocations = new ArrayList<>();
         firstLocations.add("wh0n94");
         firstLocations.add("w");
@@ -336,15 +336,15 @@ public class GeoContextMappingTests extends ESSingleNodeTestCase {
         Collection<String> secondLocations = new ArrayList<>();
         secondLocations.add("w5cx04");
         addNeighbors("w5cx04", 6, secondLocations);
-        for (ContextMapping.QueryContext queryContext : queryContexts) {
-            if (firstLocations.contains(queryContext.context)) {
-                assertThat(queryContext.boost, equalTo(10));
-            } else if (secondLocations.contains(queryContext.context)) {
-                assertThat(queryContext.boost, equalTo(1));
+        for (ContextMapping.InternalQueryContext internalQueryContext : internalQueryContexts) {
+            if (firstLocations.contains(internalQueryContext.context)) {
+                assertThat(internalQueryContext.boost, equalTo(10));
+            } else if (secondLocations.contains(internalQueryContext.context)) {
+                assertThat(internalQueryContext.boost, equalTo(1));
             } else {
-                fail(queryContext.context + " was not expected");
+                fail(internalQueryContext.context + " was not expected");
             }
-            assertThat(queryContext.isPrefix, equalTo(queryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
+            assertThat(internalQueryContext.isPrefix, equalTo(internalQueryContext.context.length() < GeoContextMapping.DEFAULT_PRECISION));
         }
     }
 }
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoQueryContextTests.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoQueryContextTests.java
index d26be5036e..1f72496782 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoQueryContextTests.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/GeoQueryContextTests.java
@@ -50,35 +50,6 @@ public class GeoQueryContextTests extends QueryContextTestCase<GeoQueryContext>
     }
 
     @Override
-    protected GeoQueryContext createMutation(GeoQueryContext original) throws IOException {
-        final GeoQueryContext.Builder builder = GeoQueryContext.builder();
-        builder.setGeoPoint(original.getGeoPoint()).setBoost(original.getBoost())
-            .setNeighbours(original.getNeighbours()).setPrecision(original.getPrecision());
-        switch (randomIntBetween(0, 3)) {
-            case 0:
-                builder.setGeoPoint(randomValueOtherThan(original.getGeoPoint() ,() ->
-                    new GeoPoint(randomDouble(), randomDouble())));
-                break;
-            case 1:
-                builder.setBoost(randomValueOtherThan(original.getBoost() ,() -> randomIntBetween(1, 5)));
-                break;
-            case 2:
-                builder.setPrecision(randomValueOtherThan(original.getPrecision() ,() -> randomIntBetween(1, 12)));
-                break;
-            case 3:
-                builder.setNeighbours(randomValueOtherThan(original.getNeighbours(), () -> {
-                    List<Integer> newNeighbours = new ArrayList<>();
-                    for (int i = 0; i < randomIntBetween(1, 12); i++) {
-                        newNeighbours.add(randomIntBetween(1, 12));
-                    }
-                    return newNeighbours;
-                }));
-                break;
-        }
-        return builder.build();
-    }
-
-    @Override
     protected GeoQueryContext prototype() {
         return GeoQueryContext.PROTOTYPE;
     }
diff --git a/core/src/test/java/org/elasticsearch/search/suggest/completion/QueryContextTestCase.java b/core/src/test/java/org/elasticsearch/search/suggest/completion/QueryContextTestCase.java
index b4a6a5b1da..78b73e6889 100644
--- a/core/src/test/java/org/elasticsearch/search/suggest/completion/QueryContextTestCase.java
+++ b/core/src/test/java/org/elasticsearch/search/suggest/completion/QueryContextTestCase.java
@@ -26,24 +26,27 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.search.suggest.completion.context.QueryContext;
+import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
 
+import static junit.framework.TestCase.assertEquals;
 
-public abstract class QueryContextTestCase<QC extends QueryContext> extends WritableTestCase<QC> {
+
+public abstract class QueryContextTestCase<QC extends QueryContext> extends ESTestCase {
 
     private static final int NUMBER_OF_RUNS = 20;
 
     /**
+     * create random model that is put under test
+     */
+    protected abstract QC createTestModel();
+
+    /**
      * query context prototype to read serialized format
      */
     protected abstract QC prototype();
 
-    @Override
-    protected QC readFrom(StreamInput in) throws IOException {
-        return (QC) prototype().readFrom(in);
-    }
-
     public void testToXContext() throws IOException {
         for (int i = 0; i < NUMBER_OF_RUNS; i++) {
             QueryContext toXContent = createTestModel();