Allow different data types for category in Context suggester (#23491)

The "category" in context suggester could be String, Number or Boolean. However with the changes in version 5 this is failing and only accepting String. This will have problem for existing users of Elasticsearch if they choose to migrate to higher version; as their existing Mapping and query will fail as mentioned in a bug #22358

This PR fixes the above mentioned issue and allows user to migrate seamlessly.

Closes #22358

4 files changed, 24 insertions, 17 deletions

diff --git a/core/src/main/java/org/elasticsearch/index/mapper/CompletionFieldMapper.java b/core/src/main/java/org/elasticsearch/index/mapper/CompletionFieldMapper.java
index 3855489efe..1ab84eda63 100644
--- a/core/src/main/java/org/elasticsearch/index/mapper/CompletionFieldMapper.java
+++ b/core/src/main/java/org/elasticsearch/index/mapper/CompletionFieldMapper.java
@@ -528,14 +528,10 @@ public class CompletionFieldMapper extends FieldMapper implements ArrayValueMapp
                                 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 {
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 150b7bf4f9..38e31ec92a 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
@@ -107,21 +107,24 @@ public class CategoryContextMapping extends ContextMapping<CategoryQueryContext>
      *  </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;
     }
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 59f59075bd..51b740a352 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,6 +21,7 @@ package org.elasticsearch.search.suggest.completion.context;
 
 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 final class CategoryQueryContext implements ToXContent {
 
     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 final class CategoryQueryContext implements ToXContent {
         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();
     }
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 f41273662a..273138bbb7 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
@@ -109,13 +109,14 @@ public abstract class ContextMapping<T extends ToXContent> implements ToXContent
         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);
     }