Add utility methods to return runtime field values in doc value order (#71599) (#71604)

This adds utility methods to each type of runtime field to return the results of a document in an ordered array based on the same order that doc values are ordered in. This is useful for supporting execute api in this #71374.

Author: jdconrad

server/src/main/java/org/elasticsearch/script/AbstractLongFieldScript.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.elasticsearch.search.lookup.SearchLookup;
 
+import java.util.Arrays;
 import java.util.Map;
 import java.util.function.LongConsumer;
 
@@ -55,6 +56,17 @@ public final long[] values() {
         return values;
     }
 
+    /**
+     * Reorders the values from the last time {@link #values()} was called to
+     * how this would appear in doc-values order. Truncates garbage values
+     * based on {@link #count()}.
+     */
+    public final long[] asDocValues() {
+        long[] truncated = Arrays.copyOf(values, count());
+        Arrays.sort(truncated);
+        return truncated;
+    }
+
     /**
      * The number of results produced the last time {@link #runForDoc(int)} was called.
      */

server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java

@@ -11,6 +11,7 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.search.lookup.SearchLookup;
 
+import java.util.Arrays;
 import java.util.Map;
 import java.util.function.Consumer;
 
@@ -68,6 +69,17 @@ public final int falses() {
         return falses;
     }
 
+    /**
+     * Reorders the values from the last time {@link #runForDoc(int)} was called to
+     * how this would appear in doc-values order.
+     */
+    public final boolean[] asDocValues() {
+        boolean[] values = new boolean[falses + trues];
+        Arrays.fill(values, 0, falses, false);
+        Arrays.fill(values, falses, falses + trues, true);
+        return values;
+    }
+
     public final void emit(boolean v) {
         if (v) {
             trues++;

server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.util.ArrayUtil;
 import org.elasticsearch.search.lookup.SearchLookup;
 
+import java.util.Arrays;
 import java.util.Map;
 import java.util.function.DoubleConsumer;
 
@@ -65,6 +66,17 @@ public final double[] values() {
         return values;
     }
 
+    /**
+     * Reorders the values from the last time {@link #values()} was called to
+     * how this would appear in doc-values order. Truncates garbage values
+     * based on {@link #count()}.
+     */
+    public final double[] asDocValues() {
+        double[] truncated = Arrays.copyOf(values, count());
+        Arrays.sort(truncated);
+        return truncated;
+    }
+
     /**
      * The number of results produced the last time {@link #runForDoc(int)} was called.
      */

server/src/main/java/org/elasticsearch/script/IpFieldScript.java

@@ -19,6 +19,7 @@
 import java.net.Inet4Address;
 import java.net.Inet6Address;
 import java.net.InetAddress;
+import java.util.Arrays;
 import java.util.Map;
 import java.util.function.Consumer;
 
@@ -86,6 +87,17 @@ public final BytesRef[] values() {
         return values;
     }
 
+    /**
+     * Reorders the values from the last time {@link #values()} was called to
+     * how this would appear in doc-values order. Truncates garbage values
+     * based on {@link #count()}.
+     */
+    public final BytesRef[] asDocValues() {
+        BytesRef[] truncated = Arrays.copyOf(values, count());
+        Arrays.sort(truncated);
+        return truncated;
+    }
+
     /**
      * The number of results produced the last time {@link #runForDoc(int)} was called.
      */

server/src/main/java/org/elasticsearch/script/StringFieldScript.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -61,6 +62,17 @@ public final void runForDoc(int docId, Consumer<String> consumer) {
         resultsForDoc(docId).forEach(consumer);
     }
 
+    /**
+     * Reorders the values from the last time {@link #resultsForDoc(int)} was called to
+     * how this would appear in doc-values order.
+     */
+    public final String[] asDocValues() {
+        String[] values = new String[results.size()];
+        results.toArray(values);
+        Arrays.sort(values);
+        return values;
+    }
+
     public final void emit(String v) {
         checkMaxSize(results.size());
         chars += v.length();

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 public class BooleanFieldScriptTests extends FieldScriptTestCase<BooleanFieldScript.Factory> {
     public static final BooleanFieldScript.Factory DUMMY = (fieldName, params, lookup) -> ctx -> new BooleanFieldScript(
@@ -43,6 +44,27 @@ protected BooleanFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        BooleanFieldScript script = new BooleanFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit(true);
+                emit(false);
+                emit(true);
+                emit(true);
+                emit(false);
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new boolean[] {false, false, true, true, true}, script.asDocValues());
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -20,6 +20,8 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.time.ZonedDateTime;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -47,6 +49,26 @@ protected DateFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        DateFieldScript script = new DateFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                DateFormatter.forPattern("YYYY-MM-DD 'T' HH:MM:SSZ"),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit(ZonedDateTime.parse("2021-01-01T00:00:00Z").toInstant().toEpochMilli());
+                emit(ZonedDateTime.parse("1942-05-31T15:16:17Z").toInstant().toEpochMilli());
+                emit(ZonedDateTime.parse("2035-10-13T10:54:19Z").toInstant().toEpochMilli());
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new long[] {-870597823000L, 1609459200000L, 2075885659000L}, script.asDocValues());
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,6 +46,28 @@ protected DoubleFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        DoubleFieldScript script = new DoubleFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit(3.1);
+                emit(2.29);
+                emit(-12.47);
+                emit(-12.46);
+                emit(Double.MAX_VALUE);
+                emit(0.0);
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new double[] {-12.47, -12.46, 0.0, 2.29, 3.1, Double.MAX_VALUE}, script.asDocValues(), 0.000000001);
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,6 +46,24 @@ protected GeoPointFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        GeoPointFieldScript script = new GeoPointFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit(78.96, 12.12);
+                emit(13.45, 56.78);
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new long[] {1378381707499043786L, 8091971733044486384L}, script.asDocValues());
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,6 +46,32 @@ protected IpFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        IpFieldScript script = new IpFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit("192.168.0.1");
+                emit("127.0.0.1");
+                emit("255.255.255.255");
+                emit("0.0.0.0");
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new BytesRef[] {
+                new BytesRef(new byte[] {0,0,0,0,0,0,0,0,0,0,-1,-1,0,0,0,0}),
+                new BytesRef(new byte[] {0,0,0,0,0,0,0,0,0,0,-1,-1,127,0,0,1}),
+                new BytesRef(new byte[] {0,0,0,0,0,0,0,0,0,0,-1,-1,-64,-88,0,1}),
+                new BytesRef(new byte[] {0,0,0,0,0,0,0,0,0,0,-1,-1,-1,-1,-1,-1})},
+                script.asDocValues()
+        );
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,6 +46,28 @@ protected LongFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        LongFieldScript script = new LongFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit(3L);
+                emit(1L);
+                emit(20000000000L);
+                emit(10L);
+                emit(-1000L);
+                emit(0L);
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new long[] {-1000L, 0L, 1L, 3L, 10L, 20000000000L}, script.asDocValues());
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
+import java.util.Collections;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -45,6 +46,28 @@ protected StringFieldScript.Factory dummyScript() {
         return DUMMY;
     }
 
+    public void testAsDocValues() {
+        StringFieldScript script = new StringFieldScript(
+                "test",
+                Collections.emptyMap(),
+                new SearchLookup(field -> null, (ft, lookup) -> null),
+                null
+        ) {
+            @Override
+            public void execute() {
+                emit("test");
+                emit("baz was not here");
+                emit("Data");
+                emit("-10");
+                emit("20");
+                emit("9");
+            }
+        };
+        script.execute();
+
+        assertArrayEquals(new String[] {"-10", "20", "9", "Data", "baz was not here", "test"}, script.asDocValues());
+    }
+
     public void testTooManyValues() throws IOException {
         try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
             iw.addDocument(org.elasticsearch.common.collect.List.of(new StoredField("_source", new BytesRef("{}"))));