Add date fields to the scripting fields api (#81272)

This change adds support for date millisecond and date nanoseconds fields to the scripting fields api.

docs/changelog/81272.yaml

@@ -0,0 +1,5 @@
+pr: 81272
+summary: Add date fields to the scripting fields api
+area: Infra/Scripting
+type: enhancement
+issues: []

modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.txt

@@ -68,6 +68,16 @@ class org.elasticsearch.script.field.ShortDocValuesField @dynamic_type {
   short get(int, int)
 }
 
+class org.elasticsearch.script.field.DateMillisDocValuesField @dynamic_type {
+  ZonedDateTime get(ZonedDateTime)
+  ZonedDateTime get(int, ZonedDateTime)
+}
+
+class org.elasticsearch.script.field.DateNanosDocValuesField @dynamic_type {
+  ZonedDateTime get(ZonedDateTime)
+  ZonedDateTime get(int, ZonedDateTime)
+}
+
 class org.elasticsearch.script.field.KeywordDocValuesField @dynamic_type {
   String get(String)
   String get(int, String)

modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/50_script_doc_values.yml

@@ -11,6 +11,8 @@ setup:
                             type: boolean
                         date:
                             type: date
+                        nanos:
+                            type: date_nanos
                         geo_point:
                             type: geo_point
                         ip:
@@ -49,6 +51,7 @@ setup:
                 rank: 1
                 boolean: true
                 date: 2017-01-01T12:11:12
+                nanos: 2015-01-01T12:10:30.123456789Z
                 geo_point: 41.12,-71.34
                 ip: 192.168.0.1
                 keyword: not split at all
@@ -76,6 +79,8 @@ setup:
           body:
               rank: 3
               boolean: [true, false, true]
+              date: [2017-01-01T12:11:12, 2018-01-01T12:11:12]
+              nanos: [2015-01-01T12:10:30.123456789Z, 2015-01-01T12:10:30.987654321Z]
               keyword: ["one string", "another string"]
               long: [1152921504606846976, 576460752303423488]
               integer: [5, 17, 29]
@@ -228,6 +233,193 @@ setup:
                             source: "doc.date.value"
     - match: { hits.hits.0.fields.field.0: '2017-01-01T12:11:12.000Z' }
 
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('date').get(null)"
+    - match: { hits.hits.0.fields.field.0: '2017-01-01T12:11:12.000Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid yaml stash */ $('date', null)"
+    - match: { hits.hits.0.fields.field.0: '2017-01-01T12:11:12.000Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 2 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('date').get(ZonedDateTime.parse('2018-01-01T12:11:12.000Z'))"
+    - match: { hits.hits.0.fields.field.0: '2018-01-01T12:11:12.000Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 2 } }
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid yaml stash */ $('date', ZonedDateTime.parse('2018-01-01T12:11:12.000Z'))"
+    - match: { hits.hits.0.fields.field.0: '2018-01-01T12:11:12.000Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "doc['nanos'].value"
+    - match: { hits.hits.0.fields.field.0: '2015-01-01T12:10:30.123456789Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('nanos').get(null)"
+    - match: { hits.hits.0.fields.field.0: '2015-01-01T12:10:30.123456789Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid yaml stash */ $('nanos', null)"
+    - match: { hits.hits.0.fields.field.0: '2015-01-01T12:10:30.123456789Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 2 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('nanos').get(ZonedDateTime.parse('2016-01-01T12:10:30.123Z'))"
+    - match: { hits.hits.0.fields.field.0: '2016-01-01T12:10:30.123Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 2 } }
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid yaml stash */ $('nanos', ZonedDateTime.parse('2016-01-01T12:10:30.123Z'))"
+    - match: { hits.hits.0.fields.field.0: '2016-01-01T12:10:30.123Z' }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "doc['nanos'].value.getNano()"
+    - match: { hits.hits.0.fields.field.0: 123456789 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('nanos').get(null).getNano()"
+    - match: { hits.hits.0.fields.field.0: 123456789 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 1 } }
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid yaml stash */ $('nanos', null).getNano()"
+    - match: { hits.hits.0.fields.field.0: 123456789 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 2 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('nanos').get(ZonedDateTime.parse('2016-01-01T12:10:30.123Z')).getNano()"
+    - match: { hits.hits.0.fields.field.0: 123000000 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 3 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('date').get(1, null)"
+    - match: { hits.hits.0.fields.field.0: "2018-01-01T12:11:12.000Z" }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 3 } }
+            script_fields:
+              field:
+                script:
+                  source: "field('nanos').get(1, null)"
+    - match: { hits.hits.0.fields.field.0: "2015-01-01T12:10:30.987654321Z" }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 3 } }
+            script_fields:
+              field:
+                script:
+                  source: "List times = new ArrayList(); for (ZonedDateTime zdt : field('date')) times.add(zdt); times"
+    - match: { hits.hits.0.fields.field: ["2017-01-01T12:11:12.000Z", "2018-01-01T12:11:12.000Z"] }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query: { term: { _id: 3 } }
+            script_fields:
+              field:
+                script:
+                  source: "List times = new ArrayList(); for (ZonedDateTime zdt : field('nanos')) times.add(zdt); times"
+    - match: { hits.hits.0.fields.field: ["2015-01-01T12:10:30.123456789Z", "2015-01-01T12:10:30.987654321Z"] }
+
 ---
 "geo_point":
     - do:

server/src/main/java/org/elasticsearch/index/fielddata/ScriptDocValues.java

@@ -15,13 +15,10 @@ import org.apache.lucene.util.BytesRefBuilder;
 import org.elasticsearch.common.geo.GeoBoundingBox;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.GeoUtils;
-import org.elasticsearch.common.time.DateUtils;
 import org.elasticsearch.geometry.utils.Geohash;
 import org.elasticsearch.script.field.DocValuesField;
 
 import java.io.IOException;
-import java.time.Instant;
-import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.util.AbstractList;
 import java.util.Comparator;
@@ -161,63 +158,6 @@ public abstract class ScriptDocValues<T> extends AbstractList<T> {
         }
     }
 
-    public static class DatesSupplier implements Supplier<ZonedDateTime> {
-
-        private final SortedNumericDocValues in;
-        private final boolean isNanos;
-
-        /**
-         * Values wrapped in {@link java.time.ZonedDateTime} objects.
-         */
-        private ZonedDateTime[] dates;
-        private int count;
-
-        public DatesSupplier(SortedNumericDocValues in, boolean isNanos) {
-            this.in = in;
-            this.isNanos = isNanos;
-        }
-
-        @Override
-        public ZonedDateTime getInternal(int index) {
-            return dates[index];
-        }
-
-        @Override
-        public int size() {
-            return count;
-        }
-
-        @Override
-        public void setNextDocId(int docId) throws IOException {
-            if (in.advanceExact(docId)) {
-                count = in.docValueCount();
-            } else {
-                count = 0;
-            }
-            refreshArray();
-        }
-
-        /**
-         * Refresh the backing array. Package private so it can be called when {@link Longs} loads dates.
-         */
-        private void refreshArray() throws IOException {
-            if (count == 0) {
-                return;
-            }
-            if (dates == null || count > dates.length) {
-                // Happens for the document. We delay allocating dates so we can allocate it with a reasonable size.
-                dates = new ZonedDateTime[count];
-            }
-            for (int i = 0; i < count; ++i) {
-                if (isNanos) {
-                    dates[i] = ZonedDateTime.ofInstant(DateUtils.toInstant(in.nextValue()), ZoneOffset.UTC);
-                } else {
-                    dates[i] = ZonedDateTime.ofInstant(Instant.ofEpochMilli(in.nextValue()), ZoneOffset.UTC);
-                }
-            }
-        }
-    }
-
     public static class Dates extends ScriptDocValues<ZonedDateTime> {
 
         public Dates(Supplier<ZonedDateTime> supplier) {

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

@@ -33,8 +33,6 @@ import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData.NumericType;
-import org.elasticsearch.index.fielddata.ScriptDocValues.Dates;
-import org.elasticsearch.index.fielddata.ScriptDocValues.DatesSupplier;
 import org.elasticsearch.index.fielddata.plain.SortedNumericIndexFieldData;
 import org.elasticsearch.index.query.DateRangeIncludingNowQuery;
 import org.elasticsearch.index.query.QueryRewriteContext;
@@ -42,7 +40,8 @@ import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.script.DateFieldScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptCompiler;
-import org.elasticsearch.script.field.DelegateDocValuesField;
+import org.elasticsearch.script.field.DateMillisDocValuesField;
+import org.elasticsearch.script.field.DateNanosDocValuesField;
 import org.elasticsearch.script.field.ToScriptField;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.lookup.FieldValues;
@@ -81,7 +80,7 @@ public final class DateFieldMapper extends FieldMapper {
     private static final DateMathParser EPOCH_MILLIS_PARSER = DateFormatter.forPattern("epoch_millis").toDateMathParser();
 
     public enum Resolution {
-        MILLISECONDS(CONTENT_TYPE, NumericType.DATE, (dv, n) -> new DelegateDocValuesField(new Dates(new DatesSupplier(dv, false)), n)) {
+        MILLISECONDS(CONTENT_TYPE, NumericType.DATE, DateMillisDocValuesField::new) {
             @Override
             public long convert(Instant instant) {
                 return instant.toEpochMilli();
@@ -112,11 +111,7 @@ public final class DateFieldMapper extends FieldMapper {
                 return LongPoint.newDistanceFeatureQuery(field, boost, origin, pivot.getMillis());
             }
         },
-        NANOSECONDS(
-            DATE_NANOS_CONTENT_TYPE,
-            NumericType.DATE_NANOSECONDS,
-            (dv, n) -> new DelegateDocValuesField(new Dates(new DatesSupplier(dv, true)), n)
-        ) {
+        NANOSECONDS(DATE_NANOS_CONTENT_TYPE, NumericType.DATE_NANOSECONDS, DateNanosDocValuesField::new) {
             @Override
             public long convert(Instant instant) {
                 return toLong(instant);

server/src/main/java/org/elasticsearch/script/field/DateMillisDocValuesField.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
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.script.field;
+
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.util.ArrayUtil;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
+
+import java.io.IOException;
+import java.time.Instant;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+
+public class DateMillisDocValuesField implements DocValuesField<ZonedDateTime>, ScriptDocValues.Supplier<ZonedDateTime> {
+
+    protected final SortedNumericDocValues input;
+    protected final String name;
+
+    protected ZonedDateTime[] values = new ZonedDateTime[0];
+    protected int count;
+
+    private ScriptDocValues.Dates dates = null;
+
+    public DateMillisDocValuesField(SortedNumericDocValues input, String name) {
+        this.input = input;
+        this.name = name;
+    }
+
+    @Override
+    public void setNextDocId(int docId) throws IOException {
+        if (input.advanceExact(docId)) {
+            resize(input.docValueCount());
+            for (int i = 0; i < count; i++) {
+                values[i] = ZonedDateTime.ofInstant(Instant.ofEpochMilli(input.nextValue()), ZoneOffset.UTC);
+            }
+        } else {
+            resize(0);
+        }
+    }
+
+    protected void resize(int newSize) {
+        count = newSize;
+
+        assert count >= 0 : "size must be positive (got " + count + "): likely integer overflow?";
+        values = ArrayUtil.grow(values, count);
+    }
+
+    @Override
+    public ScriptDocValues<ZonedDateTime> getScriptDocValues() {
+        if (dates == null) {
+            dates = new ScriptDocValues.Dates(this);
+        }
+
+        return dates;
+    }
+
+    @Override
+    public ZonedDateTime getInternal(int index) {
+        return values[index];
+    }
+
+    @Override
+    public String getName() {
+        return name;
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return count == 0;
+    }
+
+    @Override
+    public int size() {
+        return count;
+    }
+
+    public ZonedDateTime get(ZonedDateTime defaultValue) {
+        return get(0, defaultValue);
+    }
+
+    public ZonedDateTime get(int index, ZonedDateTime defaultValue) {
+        if (isEmpty() || index < 0 || index >= count) {
+            return defaultValue;
+        }
+
+        return values[index];
+    }
+
+    @Override
+    public Iterator<ZonedDateTime> iterator() {
+        return new Iterator<ZonedDateTime>() {
+            private int index = 0;
+
+            @Override
+            public boolean hasNext() {
+                return index < count;
+            }
+
+            @Override
+            public ZonedDateTime next() {
+                if (hasNext() == false) {
+                    throw new NoSuchElementException();
+                }
+                return values[index++];
+            }
+        };
+    }
+}

server/src/main/java/org/elasticsearch/script/field/DateNanosDocValuesField.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
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.script.field;
+
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.util.ArrayUtil;
+import org.elasticsearch.common.time.DateUtils;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
+
+import java.io.IOException;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+
+public class DateNanosDocValuesField implements DocValuesField<ZonedDateTime>, ScriptDocValues.Supplier<ZonedDateTime> {
+
+    protected final SortedNumericDocValues input;
+    protected final String name;
+
+    protected ZonedDateTime[] values = new ZonedDateTime[0];
+    protected int count;
+
+    private ScriptDocValues.Dates dates = null;
+
+    public DateNanosDocValuesField(SortedNumericDocValues input, String name) {
+        this.input = input;
+        this.name = name;
+    }
+
+    @Override
+    public void setNextDocId(int docId) throws IOException {
+        if (input.advanceExact(docId)) {
+            resize(input.docValueCount());
+            for (int i = 0; i < count; i++) {
+                values[i] = ZonedDateTime.ofInstant(DateUtils.toInstant(input.nextValue()), ZoneOffset.UTC);
+            }
+        } else {
+            resize(0);
+        }
+    }
+
+    protected void resize(int newSize) {
+        count = newSize;
+
+        assert count >= 0 : "size must be positive (got " + count + "): likely integer overflow?";
+        values = ArrayUtil.grow(values, count);
+    }
+
+    @Override
+    public ScriptDocValues<ZonedDateTime> getScriptDocValues() {
+        if (dates == null) {
+            dates = new ScriptDocValues.Dates(this);
+        }
+
+        return dates;
+    }
+
+    @Override
+    public ZonedDateTime getInternal(int index) {
+        return values[index];
+    }
+
+    @Override
+    public String getName() {
+        return name;
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return count == 0;
+    }
+
+    @Override
+    public int size() {
+        return count;
+    }
+
+    public ZonedDateTime get(ZonedDateTime defaultValue) {
+        return get(0, defaultValue);
+    }
+
+    public ZonedDateTime get(int index, ZonedDateTime defaultValue) {
+        if (isEmpty() || index < 0 || index >= count) {
+            return defaultValue;
+        }
+
+        return values[index];
+    }
+
+    @Override
+    public Iterator<ZonedDateTime> iterator() {
+        return new Iterator<ZonedDateTime>() {
+            private int index = 0;
+
+            @Override
+            public boolean hasNext() {
+                return index < count;
+            }
+
+            @Override
+            public ZonedDateTime next() {
+                if (hasNext() == false) {
+                    throw new NoSuchElementException();
+                }
+                return values[index++];
+            }
+        };
+    }
+}

server/src/test/java/org/elasticsearch/index/mapper/DateFieldTypeTests.java

@@ -31,8 +31,6 @@ import org.elasticsearch.core.internal.io.IOUtils;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
 import org.elasticsearch.index.fielddata.LeafNumericFieldData;
-import org.elasticsearch.index.fielddata.ScriptDocValues.Dates;
-import org.elasticsearch.index.fielddata.ScriptDocValues.DatesSupplier;
 import org.elasticsearch.index.fielddata.plain.SortedNumericIndexFieldData;
 import org.elasticsearch.index.mapper.DateFieldMapper.DateFieldType;
 import org.elasticsearch.index.mapper.DateFieldMapper.Resolution;
@@ -40,7 +38,7 @@ import org.elasticsearch.index.mapper.MappedFieldType.Relation;
 import org.elasticsearch.index.query.DateRangeIncludingNowQuery;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.SearchExecutionContext;
-import org.elasticsearch.script.field.DelegateDocValuesField;
+import org.elasticsearch.script.field.DateNanosDocValuesField;
 
 import java.io.IOException;
 import java.time.Instant;
@@ -339,7 +337,7 @@ public class DateFieldTypeTests extends FieldTypeTestCase {
         SortedNumericIndexFieldData fieldData = new SortedNumericIndexFieldData(
             "my_date",
             IndexNumericFieldData.NumericType.DATE_NANOSECONDS,
-            (dv, n) -> new DelegateDocValuesField(new Dates(new DatesSupplier(dv, true)), n)
+            DateNanosDocValuesField::new
         );
         // Read index and check the doc values
         DirectoryReader reader = DirectoryReader.open(w);

x-pack/plugin/mapper-unsigned-long/src/main/java/org/elasticsearch/xpack/unsignedlong/UnsignedLongDocValuesField.java

@@ -113,7 +113,7 @@ public class UnsignedLongDocValuesField implements DocValuesField<Long>, ScriptD
 
     /** Returns the 0th index value as an {@code long} if it exists, otherwise {@code defaultValue}. */
     public long get(long defaultValue) {
-        return getValue(0, defaultValue);
+        return get(0, defaultValue);
     }
 
     /** Returns the value at {@code index} as an {@code long} if it exists, otherwise {@code defaultValue}. */