Allow different data types for category in Context suggester

core/src/main/java/org/elasticsearch/index/mapper/CompletionFieldMapper.java

@@ -528,14 +528,10 @@ private void parse(ParseContext parseContext, Token token, XContentParser parser
                                 if (currentToken == XContentParser.Token.FIELD_NAME) {
                                     fieldName = parser.currentName();
                                     contextMapping = contextMappings.get(fieldName);
-                                } else if (currentToken == XContentParser.Token.VALUE_STRING
-                                        || currentToken == XContentParser.Token.START_ARRAY
-                                        || currentToken == XContentParser.Token.START_OBJECT) {
+                                } else {
                                     assert fieldName != null;
                                     assert !contextsMap.containsKey(fieldName);
                                     contextsMap.put(fieldName, contextMapping.parseContext(parseContext, parser));
-                                } else {
-                                    throw new IllegalArgumentException("contexts must be an object or an array , but was [" + currentToken + "]");
                                 }
                             }
                         } else {

core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryContextMapping.java

@@ -107,21 +107,24 @@ protected XContentBuilder toInnerXContent(XContentBuilder builder, Params params
      *  </ul>
      */
     @Override
-    public Set<CharSequence> parseContext(ParseContext parseContext, XContentParser parser) throws IOException, ElasticsearchParseException {
+    public Set<CharSequence> parseContext(ParseContext parseContext, XContentParser parser)
+            throws IOException, ElasticsearchParseException {
         final Set<CharSequence> contexts = new HashSet<>();
         Token token = parser.currentToken();
-        if (token == Token.VALUE_STRING) {
+        if (token == Token.VALUE_STRING || token == Token.VALUE_NUMBER || token == Token.VALUE_BOOLEAN) {
             contexts.add(parser.text());
         } else if (token == Token.START_ARRAY) {
             while ((token = parser.nextToken()) != Token.END_ARRAY) {
-                if (token == Token.VALUE_STRING) {
+                if (token == Token.VALUE_STRING || token == Token.VALUE_NUMBER || token == Token.VALUE_BOOLEAN) {
                     contexts.add(parser.text());
                 } else {
-                    throw new ElasticsearchParseException("context array must have string values");
+                    throw new ElasticsearchParseException(
+                            "context array must have string, number or boolean values, but was [" + token + "]");
                 }
             }
         } else {
-            throw new ElasticsearchParseException("contexts must be a string or a list of strings");
+            throw new ElasticsearchParseException(
+                    "contexts must be a string, number or boolean or a list of string, number or boolean, but was [" + token + "]");
         }
         return contexts;
     }

core/src/main/java/org/elasticsearch/search/suggest/completion/context/CategoryQueryContext.java

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
@@ -98,7 +99,8 @@ public int hashCode() {
 
     private static ObjectParser<Builder, Void> CATEGORY_PARSER = new ObjectParser<>(NAME, null);
     static {
-        CATEGORY_PARSER.declareString(Builder::setCategory, new ParseField(CONTEXT_VALUE));
+        CATEGORY_PARSER.declareField(Builder::setCategory, XContentParser::text, new ParseField(CONTEXT_VALUE),
+                ObjectParser.ValueType.VALUE);
         CATEGORY_PARSER.declareInt(Builder::setBoost, new ParseField(CONTEXT_BOOST));
         CATEGORY_PARSER.declareBoolean(Builder::setPrefix, new ParseField(CONTEXT_PREFIX));
     }
@@ -108,11 +110,16 @@ public static CategoryQueryContext fromXContent(QueryParseContext context) throw
         XContentParser.Token token = parser.currentToken();
         Builder builder = builder();
         if (token == XContentParser.Token.START_OBJECT) {
-            CATEGORY_PARSER.parse(parser, builder, null);
-        } else if (token == XContentParser.Token.VALUE_STRING) {
+            try {
+                CATEGORY_PARSER.parse(parser, builder, null);
+            } catch(ParsingException e) {
+                throw new ElasticsearchParseException("category context must be a string, number or boolean");
+            } 
+        } else if (token == XContentParser.Token.VALUE_STRING || token == XContentParser.Token.VALUE_BOOLEAN
+                || token == XContentParser.Token.VALUE_NUMBER) {
             builder.setCategory(parser.text());
         } else {
-            throw new ElasticsearchParseException("category context must be an object or string");
+            throw new ElasticsearchParseException("category context must be an object, string, number or boolean");
         }
         return builder.build();
     }

core/src/main/java/org/elasticsearch/search/suggest/completion/context/ContextMapping.java

@@ -109,13 +109,14 @@ public final List<InternalQueryContext> parseQueryContext(QueryParseContext cont
         List<T> queryContexts = new ArrayList<>();
         XContentParser parser = context.parser();
         Token token = parser.nextToken();
-        if (token == Token.START_OBJECT || token == Token.VALUE_STRING) {
-            queryContexts.add(fromXContent(context));
-        } else if (token == Token.START_ARRAY) {
+        if (token == Token.START_ARRAY) {
             while (parser.nextToken() != Token.END_ARRAY) {
                 queryContexts.add(fromXContent(context));
             }
+        } else {
+            queryContexts.add(fromXContent(context));
         }
+
         return toInternalQueryContexts(queryContexts);
     }