Add double script fields

server/src/main/java/org/elasticsearch/index/fielddata/plain/LeafDoubleFieldData.java

@@ -34,11 +34,11 @@
 /**
  * Specialization of {@link LeafNumericFieldData} for floating-point numerics.
  */
-abstract class LeafDoubleFieldData implements LeafNumericFieldData {
+public abstract class LeafDoubleFieldData implements LeafNumericFieldData {
 
     private final long ramBytesUsed;
 
-    LeafDoubleFieldData(long ramBytesUsed) {
+    protected LeafDoubleFieldData(long ramBytesUsed) {
         this.ramBytesUsed = ramBytesUsed;
     }
 

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

@@ -416,28 +416,16 @@ 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) {
-                double l = Double.NEGATIVE_INFINITY;
-                double u = Double.POSITIVE_INFINITY;
-                if (lowerTerm != null) {
-                    l = parse(lowerTerm, false);
-                    if (includeLower == false) {
-                        l = DoublePoint.nextUp(l);
-                    }
-                }
-                if (upperTerm != null) {
-                    u = parse(upperTerm, false);
-                    if (includeUpper == false) {
-                        u = DoublePoint.nextDown(u);
+                return doubleRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, (l, u) -> {
+                    Query query = DoublePoint.newRangeQuery(field, l, u);
+                    if (hasDocValues) {
+                        Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field,
+                                NumericUtils.doubleToSortableLong(l),
+                                NumericUtils.doubleToSortableLong(u));
+                        query = new IndexOrDocValuesQuery(query, dvQuery);
                     }
-                }
-                Query query = DoublePoint.newRangeQuery(field, l, u);
-                if (hasDocValues) {
-                    Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(field,
-                            NumericUtils.doubleToSortableLong(l),
-                            NumericUtils.doubleToSortableLong(u));
-                    query = new IndexOrDocValuesQuery(query, dvQuery);
-                }
-                return query;
+                    return query;
+                });
             }
 
             @Override
@@ -844,7 +832,7 @@ static double signum(Object value) {
         /**
          * Converts an Object to a double by checking it against known types first
          */
-        private static double objectToDouble(Object value) {
+        public static double objectToDouble(Object value) {
             double doubleValue;
 
             if (value instanceof Number) {
@@ -882,6 +870,30 @@ public static long objectToLong(Object value, boolean coerce) {
             return Numbers.toLong(stringValue, coerce);
         }
 
+        public static Query doubleRangeQuery(
+            Object lowerTerm,
+            Object upperTerm,
+            boolean includeLower,
+            boolean includeUpper,
+            BiFunction<Double, Double, Query> builder
+        ) {
+            double l = Double.NEGATIVE_INFINITY;
+            double u = Double.POSITIVE_INFINITY;
+            if (lowerTerm != null) {
+                l = objectToDouble(lowerTerm);
+                if (includeLower == false) {
+                    l = DoublePoint.nextUp(l);
+                }
+            }
+            if (upperTerm != null) {
+                u = objectToDouble(upperTerm);
+                if (includeUpper == false) {
+                    u = DoublePoint.nextDown(u);
+                }
+            }
+            return builder.apply(l, u);
+        }
+
         /**
          * Processes query bounds into {@code long}s and delegates the
          * provided {@code builder} to build a range query.

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

@@ -7,6 +7,7 @@
 package org.elasticsearch.xpack.runtimefields;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.util.ArrayUtil;
 import org.elasticsearch.painless.spi.Whitelist;
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.ScriptContext;
@@ -16,10 +17,9 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
-import java.util.function.DoubleConsumer;
 
 public abstract class DoubleScriptFieldScript extends AbstractScriptFieldScript {
-    static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("double_script_field", Factory.class);
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("double_script_field", Factory.class);
 
     static List<Whitelist> whitelist() {
         return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "double_whitelist.txt"));
@@ -32,14 +32,47 @@ public interface Factory extends ScriptFactory {
     }
 
     public interface LeafFactory {
-        DoubleScriptFieldScript newInstance(LeafReaderContext ctx, DoubleConsumer sync) throws IOException;
+        DoubleScriptFieldScript newInstance(LeafReaderContext ctx) throws IOException;
     }
 
-    private final DoubleConsumer sync;
+    private double[] values = new double[1];
+    private int count;
 
-    public DoubleScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx, DoubleConsumer sync) {
+    public DoubleScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
         super(params, searchLookup, ctx);
-        this.sync = sync;
+    }
+
+    /**
+     * Execute the script for the provided {@code docId}.
+     */
+    public final void runForDoc(int docId) {
+        count = 0;
+        setDocument(docId);
+        execute();
+    }
+
+    /**
+     * Values from the last time {@link #runForDoc(int)} was called. This array
+     * is mutable and will change with the next call of {@link #runForDoc(int)}.
+     * It is also oversized and will contain garbage at all indices at and
+     * above {@link #count()}.
+     */
+    public final double[] values() {
+        return values;
+    }
+
+    /**
+     * The number of results produced the last time {@link #runForDoc(int)} was called.
+     */
+    public final int count() {
+        return count;
+    }
+
+    private void collectValue(double v) {
+        if (values.length < count + 1) {
+            values = ArrayUtil.grow(values, count + 1);
+        }
+        values[count++] = v;
     }
 
     public static class Value {
@@ -50,7 +83,7 @@ public Value(DoubleScriptFieldScript script) {
         }
 
         public void value(double v) {
-            script.sync.accept(v);
+            script.collectValue(v);
         }
     }
 }

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptDoubleDocValues.java

@@ -0,0 +1,43 @@
+/*
+ * 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.fielddata;
+
+import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
+import org.elasticsearch.xpack.runtimefields.DoubleScriptFieldScript;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+public final class ScriptDoubleDocValues extends SortedNumericDoubleValues {
+    private final DoubleScriptFieldScript script;
+    private int cursor;
+
+    ScriptDoubleDocValues(DoubleScriptFieldScript script) {
+        this.script = script;
+    }
+
+    @Override
+    public boolean advanceExact(int docId) {
+        script.runForDoc(docId);
+        if (script.count() == 0) {
+            return false;
+        }
+        Arrays.sort(script.values(), 0, script.count());
+        cursor = 0;
+        return true;
+    }
+
+    @Override
+    public double nextValue() throws IOException {
+        return script.values()[cursor++];
+    }
+
+    @Override
+    public int docValueCount() {
+        return script.count();
+    }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/fielddata/ScriptDoubleFieldData.java

@@ -0,0 +1,116 @@
+/*
+ * 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.fielddata;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.fielddata.IndexFieldDataCache;
+import org.elasticsearch.index.fielddata.IndexNumericFieldData;
+import org.elasticsearch.index.fielddata.SearchLookupAware;
+import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
+import org.elasticsearch.index.fielddata.plain.LeafDoubleFieldData;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.script.Script;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
+import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.xpack.runtimefields.DoubleScriptFieldScript;
+
+import java.io.IOException;
+
+public final class ScriptDoubleFieldData extends IndexNumericFieldData implements SearchLookupAware {
+
+    public static class Builder implements IndexFieldData.Builder {
+        private final String name;
+        private final Script script;
+        private final DoubleScriptFieldScript.Factory scriptFactory;
+
+        public Builder(String name, Script script, DoubleScriptFieldScript.Factory scriptFactory) {
+            this.name = name;
+            this.script = script;
+            this.scriptFactory = scriptFactory;
+        }
+
+        @Override
+        public ScriptDoubleFieldData build(IndexFieldDataCache cache, CircuitBreakerService breakerService, MapperService mapperService) {
+            return new ScriptDoubleFieldData(name, script, scriptFactory);
+        }
+    }
+
+    private final String fieldName;
+    private final Script script;
+    private final DoubleScriptFieldScript.Factory scriptFactory;
+    private final SetOnce<DoubleScriptFieldScript.LeafFactory> leafFactory = new SetOnce<>();
+
+    private ScriptDoubleFieldData(String fieldName, Script script, DoubleScriptFieldScript.Factory scriptFactory) {
+        this.fieldName = fieldName;
+        this.script = script;
+        this.scriptFactory = scriptFactory;
+    }
+
+    @Override
+    public void setSearchLookup(SearchLookup searchLookup) {
+        this.leafFactory.set(scriptFactory.newFactory(script.getParams(), searchLookup));
+    }
+
+    @Override
+    public String getFieldName() {
+        return fieldName;
+    }
+
+    @Override
+    public ValuesSourceType getValuesSourceType() {
+        return CoreValuesSourceType.NUMERIC;
+    }
+
+    @Override
+    public ScriptDoubleLeafFieldData load(LeafReaderContext context) {
+        try {
+            return loadDirect(context);
+        } catch (Exception e) {
+            throw ExceptionsHelper.convertToElastic(e);
+        }
+    }
+
+    @Override
+    public ScriptDoubleLeafFieldData loadDirect(LeafReaderContext context) throws IOException {
+        return new ScriptDoubleLeafFieldData(new ScriptDoubleDocValues(leafFactory.get().newInstance(context)));
+    }
+
+    @Override
+    public NumericType getNumericType() {
+        return NumericType.DOUBLE;
+    }
+
+    @Override
+    protected boolean sortRequiresCustomComparator() {
+        return true;
+    }
+
+    @Override
+    public void clear() {}
+
+    public static class ScriptDoubleLeafFieldData extends LeafDoubleFieldData {
+        private final ScriptDoubleDocValues scriptDoubleDocValues;
+
+        ScriptDoubleLeafFieldData(ScriptDoubleDocValues scriptDoubleDocValues) {
+            super(0);
+            this.scriptDoubleDocValues = scriptDoubleDocValues;
+        }
+
+        @Override
+        public SortedNumericDoubleValues getDoubleValues() {
+            return scriptDoubleDocValues;
+        }
+
+        @Override
+        public void close() {}
+    }
+}