Added enum for script sort type

6 files changed, 85 insertions, 28 deletions

diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
index a767faa3e0..6254d5b1e4 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java
@@ -24,6 +24,7 @@ import org.elasticsearch.common.ParseFieldMatcher;
 import org.elasticsearch.common.io.stream.NamedWriteable;
 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.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -35,6 +36,7 @@ import org.elasticsearch.script.ScriptParameterParser.ScriptParameterValue;
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 
@@ -45,7 +47,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
     SortElementParserTemp<ScriptSortBuilder> {
 
     private static final String NAME = "_script";
-    static final ScriptSortBuilder PROTOTYPE = new ScriptSortBuilder(new Script("_na_"), "_na_");
+    static final ScriptSortBuilder PROTOTYPE = new ScriptSortBuilder(new Script("_na_"), ScriptSortType.STRING);
     public static final ParseField TYPE_FIELD = new ParseField("type");
     public static final ParseField SCRIPT_FIELD = new ParseField("script");
     public static final ParseField SORTMODE_FIELD = new ParseField("mode");
@@ -55,8 +57,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
 
     private final Script script;
 
-    // TODO make this an enum
-    private final String type;
+    private ScriptSortType type;
 
     // TODO make this an enum
     private String sortMode;
@@ -74,7 +75,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
      *            The type of the script, can be either {@link ScriptSortParser#STRING_SORT_TYPE} or
      *            {@link ScriptSortParser#NUMBER_SORT_TYPE}
      */
-    public ScriptSortBuilder(Script script, String type) {
+    public ScriptSortBuilder(Script script, ScriptSortType type) {
         Objects.requireNonNull(script, "script cannot be null");
         Objects.requireNonNull(type, "type cannot be null");
         this.script = script;
@@ -100,7 +101,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
     /**
      * Get the type used in this sort.
      */
-    public String type() {
+    public ScriptSortType type() {
         return this.type;
     }
 
@@ -177,7 +178,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
         XContentParser parser = context.parser();
         ParseFieldMatcher parseField = context.parseFieldMatcher();
         Script script = null;
-        String type = null;
+        ScriptSortType type = null;
         String sortMode = null;
         SortOrder order = null;
         QueryBuilder<?> nestedFilter = null;
@@ -203,7 +204,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
                 } else if (scriptParameterParser.token(currentName, token, parser, parseField)) {
                     // Do Nothing (handled by ScriptParameterParser
                 } else if (parseField.match(currentName, TYPE_FIELD)) {
-                    type = parser.text();
+                    type = ScriptSortType.fromString(parser.text());
                 } else if (parseField.match(currentName, SORTMODE_FIELD)) {
                     sortMode = parser.text();
                 } else if (parseField.match(currentName, NESTED_PATH_FIELD)) {
@@ -263,7 +264,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         script.writeTo(out);
-        out.writeString(type);
+        type.writeTo(out);
         order.writeTo(out);
         out.writeOptionalString(sortMode);
         out.writeOptionalString(nestedPath);
@@ -276,7 +277,7 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
 
     @Override
     public ScriptSortBuilder readFrom(StreamInput in) throws IOException {
-        ScriptSortBuilder builder = new ScriptSortBuilder(Script.readScript(in), in.readString());
+        ScriptSortBuilder builder = new ScriptSortBuilder(Script.readScript(in), ScriptSortType.PROTOTYPE.readFrom(in));
         builder.order(SortOrder.readOrderFrom(in));
         builder.sortMode = in.readOptionalString();
         builder.nestedPath = in.readOptionalString();
@@ -290,4 +291,44 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> implements
     public String getWriteableName() {
         return NAME;
     }
+
+    public enum ScriptSortType implements Writeable<ScriptSortType> {
+        /** script sort for a string value **/
+        STRING,
+        /** script sort for a numeric value **/
+        NUMBER;
+
+        static ScriptSortType PROTOTYPE = STRING;
+
+        @Override
+        public void writeTo(final StreamOutput out) throws IOException {
+            out.writeVInt(ordinal());
+        }
+
+        @Override
+        public ScriptSortType readFrom(final StreamInput in) throws IOException {
+            int ordinal = in.readVInt();
+            if (ordinal < 0 || ordinal >= values().length) {
+                throw new IOException("Unknown ScriptSortType ordinal [" + ordinal + "]");
+            }
+            return values()[ordinal];
+        }
+
+        public static ScriptSortType fromString(final String str) {
+            Objects.requireNonNull(str, "input string is null");
+            switch (str.toLowerCase(Locale.ROOT)) {
+                case ("string"):
+                    return ScriptSortType.STRING;
+                case ("number"):
+                    return ScriptSortType.NUMBER;
+                default:
+                    throw new IllegalArgumentException("Unknown ScriptSortType [" + str + "]");
+            }
+        }
+
+        @Override
+        public String toString() {
+            return name().toLowerCase(Locale.ROOT);
+        }
+    }
 }
diff --git a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
index 9bf4dde711..c238ad6cca 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/ScriptSortParser.java
@@ -48,6 +48,7 @@ import org.elasticsearch.script.ScriptParameterParser;
 import org.elasticsearch.script.ScriptParameterParser.ScriptParameterValue;
 import org.elasticsearch.script.SearchScript;
 import org.elasticsearch.search.MultiValueMode;
+import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -59,9 +60,6 @@ import java.util.Map;
  */
 public class ScriptSortParser implements SortParser {
 
-    public static final String STRING_SORT_TYPE = "string";
-    public static final String NUMBER_SORT_TYPE = "number";
-
     @Override
     public String[] names() {
         return new String[]{"_script"};
@@ -71,7 +69,7 @@ public class ScriptSortParser implements SortParser {
     public SortField parse(XContentParser parser, QueryShardContext context) throws Exception {
         ScriptParameterParser scriptParameterParser = new ScriptParameterParser();
         Script script = null;
-        String type = null;
+        ScriptSortType type = null;
         Map<String, Object> params = null;
         boolean reverse = false;
         MultiValueMode sortMode = null;
@@ -101,7 +99,7 @@ public class ScriptSortParser implements SortParser {
                 } else if (scriptParameterParser.token(currentName, token, parser, context.parseFieldMatcher())) {
                     // Do Nothing (handled by ScriptParameterParser
                 } else if ("type".equals(currentName)) {
-                    type = parser.text();
+                    type = ScriptSortType.fromString(parser.text());
                 } else if ("mode".equals(currentName)) {
                     sortMode = MultiValueMode.fromString(parser.text());
                 } else if ("nested_path".equals(currentName) || "nestedPath".equals(currentName)) {
@@ -134,7 +132,7 @@ public class ScriptSortParser implements SortParser {
         final SearchScript searchScript = context.getScriptService().search(
                 context.lookup(), script, ScriptContext.Standard.SEARCH, Collections.emptyMap());
 
-        if (STRING_SORT_TYPE.equals(type) && (sortMode == MultiValueMode.SUM || sortMode == MultiValueMode.AVG)) {
+        if (ScriptSortType.STRING.equals(type) && (sortMode == MultiValueMode.SUM || sortMode == MultiValueMode.AVG)) {
             throw new ParsingException(parser.getTokenLocation(), "type [string] doesn't support mode [" + sortMode + "]");
         }
 
@@ -160,7 +158,7 @@ public class ScriptSortParser implements SortParser {
 
         final IndexFieldData.XFieldComparatorSource fieldComparatorSource;
         switch (type) {
-            case STRING_SORT_TYPE:
+            case STRING:
                 fieldComparatorSource = new BytesRefFieldComparatorSource(null, null, sortMode, nested) {
                     LeafSearchScript leafScript;
                     @Override
@@ -183,7 +181,7 @@ public class ScriptSortParser implements SortParser {
                     }
                 };
                 break;
-            case NUMBER_SORT_TYPE:
+            case NUMBER:
                 // TODO: should we rather sort missing values last?
                 fieldComparatorSource = new DoubleValuesComparatorSource(null, Double.MAX_VALUE, sortMode, nested) {
                     LeafSearchScript leafScript;
diff --git a/core/src/main/java/org/elasticsearch/search/sort/SortBuilders.java b/core/src/main/java/org/elasticsearch/search/sort/SortBuilders.java
index f326fee383..3eae9b8d01 100644
--- a/core/src/main/java/org/elasticsearch/search/sort/SortBuilders.java
+++ b/core/src/main/java/org/elasticsearch/search/sort/SortBuilders.java
@@ -21,8 +21,7 @@ package org.elasticsearch.search.sort;
 
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.script.Script;
-
-import java.util.Arrays;
+import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
 
 /**
  * A set of static factory methods for {@link SortBuilder}s.
@@ -53,7 +52,7 @@ public class SortBuilders {
      * @param script The script to use.
      * @param type   The type, can either be "string" or "number".
      */
-    public static ScriptSortBuilder scriptSort(Script script, String type) {
+    public static ScriptSortBuilder scriptSort(Script script, ScriptSortType type) {
         return new ScriptSortBuilder(script, type);
     }
 
@@ -63,12 +62,12 @@ public class SortBuilders {
      * @param fieldName The geo point like field name.
      * @param lat Latitude of the point to create the range distance facets from.
      * @param lon Longitude of the point to create the range distance facets from.
-     * 
+     *
      */
     public static GeoDistanceSortBuilder geoDistanceSort(String fieldName, double lat, double lon) {
         return new GeoDistanceSortBuilder(fieldName, lat, lon);
     }
-    
+
     /**
      * Constructs a new distance based sort on a geo point like field.
      *
@@ -87,5 +86,5 @@ public class SortBuilders {
      */
     public static GeoDistanceSortBuilder geoDistanceSort(String fieldName, String ... geohashes) {
         return new GeoDistanceSortBuilder(fieldName, geohashes);
-    }    
+    }
 }
diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/TopHitsTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/TopHitsTests.java
index cccac925a1..05ea6148a5 100644
--- a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/TopHitsTests.java
+++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/TopHitsTests.java
@@ -29,6 +29,7 @@ import org.elasticsearch.search.aggregations.BaseAggregationTestCase;
 import org.elasticsearch.search.aggregations.metrics.tophits.TopHitsAggregatorBuilder;
 import org.elasticsearch.search.fetch.source.FetchSourceContext;
 import org.elasticsearch.search.highlight.HighlightBuilderTests;
+import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
 import org.elasticsearch.search.sort.SortBuilders;
 import org.elasticsearch.search.sort.SortOrder;
 
@@ -132,7 +133,7 @@ public class TopHitsTests extends BaseAggregationTestCase<TopHitsAggregatorBuild
                     factory.sort(SortBuilders.scoreSort().order(randomFrom(SortOrder.values())));
                     break;
                 case 3:
-                    factory.sort(SortBuilders.scriptSort(new Script("foo"), "number").order(randomFrom(SortOrder.values())));
+                    factory.sort(SortBuilders.scriptSort(new Script("foo"), ScriptSortType.NUMBER).order(randomFrom(SortOrder.values())));
                     break;
                 case 4:
                     factory.sort(randomAsciiOfLengthBetween(5, 20));
diff --git a/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java b/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java
index 2cb4554ed0..60537d4554 100644
--- a/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java
+++ b/core/src/test/java/org/elasticsearch/search/builder/SearchSourceBuilderTests.java
@@ -74,6 +74,7 @@ import org.elasticsearch.search.fetch.source.FetchSourceContext;
 import org.elasticsearch.search.highlight.HighlightBuilderTests;
 import org.elasticsearch.search.rescore.QueryRescoreBuilderTests;
 import org.elasticsearch.search.searchafter.SearchAfterBuilder;
+import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
 import org.elasticsearch.search.sort.SortBuilders;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.search.suggest.SuggestBuilder;
@@ -344,7 +345,7 @@ public class SearchSourceBuilderTests extends ESTestCase {
                     builder.sort(SortBuilders.scoreSort().order(randomFrom(SortOrder.values())));
                     break;
                 case 3:
-                    builder.sort(SortBuilders.scriptSort(new Script("foo"), "number").order(randomFrom(SortOrder.values())));
+                    builder.sort(SortBuilders.scriptSort(new Script("foo"), ScriptSortType.NUMBER).order(randomFrom(SortOrder.values())));
                     break;
                 case 4:
                     builder.sort(randomAsciiOfLengthBetween(5, 20));
diff --git a/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java b/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java
index 2b46858b4c..40798cbcb8 100644
--- a/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java
+++ b/core/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java
@@ -21,6 +21,7 @@ package org.elasticsearch.search.sort;
 
 
 import org.elasticsearch.script.Script;
+import org.elasticsearch.search.sort.ScriptSortBuilder.ScriptSortType;
 
 import java.io.IOException;
 
@@ -29,7 +30,7 @@ public class ScriptSortBuilderTests extends AbstractSortTestCase<ScriptSortBuild
     @Override
     protected ScriptSortBuilder createTestItem() {
         ScriptSortBuilder builder = new ScriptSortBuilder(new Script(randomAsciiOfLengthBetween(5, 10)),
-                randomBoolean() ? ScriptSortParser.NUMBER_SORT_TYPE : ScriptSortParser.STRING_SORT_TYPE);
+                randomBoolean() ? ScriptSortType.NUMBER : ScriptSortType.STRING);
         if (randomBoolean()) {
             builder.order(RandomSortDataGenerator.order(builder.order()));
         }
@@ -51,11 +52,11 @@ public class ScriptSortBuilderTests extends AbstractSortTestCase<ScriptSortBuild
         if (randomBoolean()) {
             // change one of the constructor args, copy the rest over
             Script script = original.script();
-            String type = original.type();
+            ScriptSortType type = original.type();
             if (randomBoolean()) {
                 result = new ScriptSortBuilder(new Script(script.getScript() + "_suffix"), type);
             } else {
-                result = new ScriptSortBuilder(script, type + "_suffix");
+                result = new ScriptSortBuilder(script, type.equals(ScriptSortType.NUMBER) ? ScriptSortType.STRING : ScriptSortType.NUMBER);
             }
             result.order(original.order());
             result.sortMode(original.sortMode());
@@ -84,4 +85,20 @@ public class ScriptSortBuilderTests extends AbstractSortTestCase<ScriptSortBuild
         }
         return result;
     }
+
+    public void testScriptSortType() {
+        // we rely on these ordinals in serialization, so changing them breaks bwc.
+        assertEquals(0, ScriptSortType.STRING.ordinal());
+        assertEquals(1, ScriptSortType.NUMBER.ordinal());
+
+        assertEquals("string", ScriptSortType.STRING.toString());
+        assertEquals("number", ScriptSortType.NUMBER.toString());
+
+        assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("string"));
+        assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("String"));
+        assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("STRING"));
+        assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("number"));
+        assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("Number"));
+        assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("NUMBER"));
+    }
 }