construct suggestion context from query context

1 files changed, 14 insertions, 14 deletions

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