Remaining queries for long script field

server/src/main/java/org/elasticsearch/index/mapper/NumberFieldMapper.java

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.exc.InputCoercionException;
+
 import org.apache.lucene.document.DoublePoint;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.document.FieldType;
@@ -66,6 +67,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 /** A {@link FieldMapper} for numeric types: byte, short, int, long, float and double. */
@@ -748,44 +750,17 @@ public Query termsQuery(String field, List<Object> values) {
             public Query rangeQuery(String field, Object lowerTerm, Object upperTerm,
                                     boolean includeLower, boolean includeUpper,
                                     boolean hasDocValues, QueryShardContext context) {
-                long l = Long.MIN_VALUE;
-                long u = Long.MAX_VALUE;
-                if (lowerTerm != null) {
-                    l = parse(lowerTerm, true);
-                    // if the lower bound is decimal:
-                    // - if the bound is positive then we increment it:
-                    //      if lowerTerm=1.5 then the (inclusive) bound becomes 2
-                    // - if the bound is negative then we leave it as is:
-                    //      if lowerTerm=-1.5 then the (inclusive) bound becomes -1 due to the call to longValue
-                    boolean lowerTermHasDecimalPart = hasDecimalPart(lowerTerm);
-                    if ((lowerTermHasDecimalPart == false && includeLower == false) ||
-                            (lowerTermHasDecimalPart && signum(lowerTerm) > 0)) {
-                        if (l == Long.MAX_VALUE) {
-                            return new MatchNoDocsQuery();
-                        }
-                        ++l;
-                    }
-                }
-                if (upperTerm != null) {
-                    u = parse(upperTerm, true);
-                    boolean upperTermHasDecimalPart = hasDecimalPart(upperTerm);
-                    if ((upperTermHasDecimalPart == false && includeUpper == false) ||
-                            (upperTermHasDecimalPart && signum(upperTerm) < 0)) {
-                        if (u == Long.MIN_VALUE) {
-                            return new MatchNoDocsQuery();
+                return longRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, (l, u) -> {
+                    Query query = LongPoint.newRangeQuery(field, l, u); 
+                    if (hasDocValues) {
+                        Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u);
+                        query = new IndexOrDocValuesQuery(query, dvQuery);
+                        if (context.indexSortedOnField(field)) {
+                            query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query);
                         }
-                        --u;
                     }
-                }
-                Query query = LongPoint.newRangeQuery(field, l, u);
-                if (hasDocValues) {
-                    Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field, l, u);
-                    query = new IndexOrDocValuesQuery(query, dvQuery);
-                    if (context.indexSortedOnField(field)) {
-                        query = new IndexSortSortedNumericDocValuesRangeQuery(field, l, u, query);
-                    }
-                }
-                return query;
+                    return query;
+                });
             }
 
             @Override
@@ -855,7 +830,7 @@ public static boolean hasDecimalPart(Object number) {
         /**
          * Returns -1, 0, or 1 if the value is lower than, equal to, or greater than 0
          */
-        double signum(Object value) {
+        static double signum(Object value) {
             if (value instanceof Number) {
                 double doubleValue = ((Number) value).doubleValue();
                 return Math.signum(doubleValue);
@@ -906,6 +881,47 @@ public static long objectToLong(Object value, boolean coerce) {
             String stringValue = (value instanceof BytesRef) ? ((BytesRef) value).utf8ToString() : value.toString();
             return Numbers.toLong(stringValue, coerce);
         }
+
+        /**
+         * Processes query bounds into {@code long}s and delegates the
+         * provided {@code builder} to build a range query.
+         */
+        public static Query longRangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            BiFunction<Long, Long, Query> builder
+        ) {
+            long l = Long.MIN_VALUE;
+            long u = Long.MAX_VALUE;
+            if (lowerTerm != null) {
+                l = objectToLong(lowerTerm, true);
+                // if the lower bound is decimal:
+                // - if the bound is positive then we increment it:
+                // if lowerTerm=1.5 then the (inclusive) bound becomes 2
+                // - if the bound is negative then we leave it as is:
+                // if lowerTerm=-1.5 then the (inclusive) bound becomes -1 due to the call to longValue
+                boolean lowerTermHasDecimalPart = hasDecimalPart(lowerTerm);
+                if ((lowerTermHasDecimalPart == false && includeLower == false) || (lowerTermHasDecimalPart && signum(lowerTerm) > 0)) {
+                    if (l == Long.MAX_VALUE) {
+                        return new MatchNoDocsQuery();
+                    }
+                    ++l;
+                }
+            }
+            if (upperTerm != null) {
+                u = objectToLong(upperTerm, true);
+                boolean upperTermHasDecimalPart = hasDecimalPart(upperTerm);
+                if ((upperTermHasDecimalPart == false && includeUpper == false) || (upperTermHasDecimalPart && signum(upperTerm) < 0)) {
+                    if (u == Long.MIN_VALUE) {
+                        return new MatchNoDocsQuery();
+                    }
+                    --u;
+                }
+            }
+            return builder.apply(l, u);
+        }
     }
 
     public static final class NumberFieldType extends SimpleMappedFieldType {

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/mapper/ScriptLongMappedFieldType.java

@@ -6,16 +6,25 @@
 
 package org.elasticsearch.xpack.runtimefields.mapper;
 
+import com.carrotsearch.hppc.LongHashSet;
+import com.carrotsearch.hppc.LongSet;
+
 import org.apache.lucene.search.Query;
+import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.common.time.DateMathParser;
 import org.elasticsearch.index.mapper.NumberFieldMapper.NumberType;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.xpack.runtimefields.LongScriptFieldScript;
 import org.elasticsearch.xpack.runtimefields.fielddata.ScriptLongFieldData;
 import org.elasticsearch.xpack.runtimefields.query.LongScriptFieldExistsQuery;
+import org.elasticsearch.xpack.runtimefields.query.LongScriptFieldRangeQuery;
 import org.elasticsearch.xpack.runtimefields.query.LongScriptFieldTermQuery;
+import org.elasticsearch.xpack.runtimefields.query.LongScriptFieldTermsQuery;
 
+import java.time.ZoneId;
+import java.util.List;
 import java.util.Map;
 
 public class ScriptLongMappedFieldType extends AbstractScriptMappedFieldType {
@@ -52,6 +61,27 @@ public Query existsQuery(QueryShardContext context) {
         return new LongScriptFieldExistsQuery(script, leafFactory(context), name());
     }
 
+    @Override
+    public Query rangeQuery(
+        Object lowerTerm,
+        Object upperTerm,
+        boolean includeLower,
+        boolean includeUpper,
+        ShapeRelation relation,
+        ZoneId timeZone,
+        DateMathParser parser,
+        QueryShardContext context
+    ) {
+        checkAllowExpensiveQueries(context);
+        return NumberType.longRangeQuery(
+            lowerTerm,
+            upperTerm,
+            includeLower,
+            includeUpper,
+            (l, u) -> new LongScriptFieldRangeQuery(script, leafFactory(context), name(), l, u)
+        );
+    }
+
     @Override
     public Query termQuery(Object value, QueryShardContext context) {
         if (NumberType.hasDecimalPart(value)) {
@@ -60,4 +90,23 @@ public Query termQuery(Object value, QueryShardContext context) {
         checkAllowExpensiveQueries(context);
         return new LongScriptFieldTermQuery(script, leafFactory(context), name(), NumberType.objectToLong(value, true));
     }
+
+    @Override
+    public Query termsQuery(List<?> values, QueryShardContext context) {
+        if (values.isEmpty()) {
+            return Queries.newMatchAllQuery();
+        }
+        LongSet terms = new LongHashSet(values.size());
+        for (Object value : values) {
+            if (NumberType.hasDecimalPart(value)) {
+                continue;
+            }
+            terms.add(NumberType.objectToLong(value, true));
+        }
+        if (terms.isEmpty()) {
+            return Queries.newMatchNoDocsQuery("All values have a decimal part");
+        }
+        checkAllowExpensiveQueries(context);
+        return new LongScriptFieldTermsQuery(script, leafFactory(context), name(), terms);
+    }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldRangeQuery.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.runtimefields.query;
+
+import org.elasticsearch.script.Script;
+import org.elasticsearch.xpack.runtimefields.LongScriptFieldScript;
+
+import java.util.Objects;
+
+public class LongScriptFieldRangeQuery extends AbstractLongScriptFieldQuery {
+    private final long lowerValue;
+    private final long upperValue;
+
+    public LongScriptFieldRangeQuery(
+        Script script,
+        LongScriptFieldScript.LeafFactory leafFactory,
+        String fieldName,
+        long lowerValue,
+        long upperValue
+    ) {
+        super(script, leafFactory, fieldName);
+        this.lowerValue = lowerValue;
+        this.upperValue = upperValue;
+        assert lowerValue <= upperValue;
+    }
+
+    @Override
+    protected boolean matches(long[] values, int count) {
+        for (int i = 0; i < count; i++) {
+            if (lowerValue <= values[i] && values[i] <= upperValue) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public final String toString(String field) {
+        StringBuilder b = new StringBuilder();
+        if (false == fieldName().contentEquals(field)) {
+            b.append(fieldName()).append(':');
+        }
+        b.append('[').append(lowerValue).append(" TO ").append(upperValue).append(']');
+        return b.toString();
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), lowerValue, upperValue);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (false == super.equals(obj)) {
+            return false;
+        }
+        LongScriptFieldRangeQuery other = (LongScriptFieldRangeQuery) obj;
+        return lowerValue == other.lowerValue && upperValue == other.upperValue;
+    }
+
+    long lowerValue() {
+        return lowerValue;
+    }
+
+    long upperValue() {
+        return upperValue;
+    }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/query/LongScriptFieldTermsQuery.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.runtimefields.query;
+
+import com.carrotsearch.hppc.LongSet;
+
+import org.elasticsearch.script.Script;
+import org.elasticsearch.xpack.runtimefields.LongScriptFieldScript;
+
+import java.util.Objects;
+
+public class LongScriptFieldTermsQuery extends AbstractLongScriptFieldQuery {
+    private final LongSet terms;
+
+    public LongScriptFieldTermsQuery(Script script, LongScriptFieldScript.LeafFactory leafFactory, String fieldName, LongSet terms) {
+        super(script, leafFactory, fieldName);
+        this.terms = terms;
+    }
+
+    @Override
+    protected boolean matches(long[] values, int count) {
+        for (int i = 0; i < count; i++) {
+            if (terms.contains(values[i])) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public final String toString(String field) {
+        if (fieldName().contentEquals(field)) {
+            return terms.toString();
+        }
+        return fieldName() + ":" + terms;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), terms);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (false == super.equals(obj)) {
+            return false;
+        }
+        LongScriptFieldTermsQuery other = (LongScriptFieldTermsQuery) obj;
+        return terms.equals(other.terms);
+    }
+
+    LongSet terms() {
+        return terms;
+    }
+}