Share common runForDoc impl in AbstractFieldScript (elastic#92450)

Each runtime field script class implements its own slighly different variant of the runForDoc method, sometimes called resultsForDoc in case it also returns the results. For simplicity, we can unify runForDoc in the base class and introduce additional getValues methods where missing. This should also simplify unifying error handling down the line introducing the error handling logic directly in runForDoc.
droberts195 · Dec 19, 2022 · 4be7743 · 4be7743
1 parent 2d2b82b
commit 4be7743
Show file tree

Hide file tree

Showing 12 changed files with 48 additions and 52 deletions.
diff --git a/.../lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java b/.../lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java
@@ -663,7 +663,8 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                         context.lookup()
                     );
                     CompositeFieldScript compositeFieldScript = leafFactory.newInstance(leafReaderContext);
-                    return new Response(compositeFieldScript.runForDoc(0));
+                    compositeFieldScript.runForDoc(0);
+                    return new Response(compositeFieldScript.getFieldValues());
                 }, indexService);
             } else {
                 throw new UnsupportedOperationException("unsupported context [" + scriptContext.name + "]");

diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/StringScriptDocValues.java b/server/src/main/java/org/elasticsearch/index/fielddata/StringScriptDocValues.java
@@ -21,7 +21,8 @@ public final class StringScriptDocValues extends SortingBinaryDocValues {
 
     @Override
     public boolean advanceExact(int docId) {
-        List<String> results = script.resultsForDoc(docId);
+        script.runForDoc(docId);
+        List<String> results = script.getValues();
         count = results.size();
         if (count == 0) {
             return false;

diff --git a/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java b/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java
@@ -133,5 +133,16 @@ protected final void checkMaxSize(int currentSize) {
         }
     }
 
+    protected abstract void prepareExecute();
+
+    /**
+     * Execute the script for the provided {@code docId}.
+     */
+    public final void runForDoc(int docId) {
+        prepareExecute();
+        setDocument(docId);
+        execute();
+    }
+
     public abstract void execute();
 }
diff --git a/server/src/main/java/org/elasticsearch/script/AbstractLongFieldScript.java b/server/src/main/java/org/elasticsearch/script/AbstractLongFieldScript.java
@@ -26,13 +26,9 @@ public AbstractLongFieldScript(String fieldName, Map<String, Object> params, Sea
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     */
-    public final void runForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         count = 0;
-        setDocument(docId);
-        execute();
     }
 
     /**

diff --git a/server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java b/server/src/main/java/org/elasticsearch/script/BooleanFieldScript.java
@@ -75,14 +75,10 @@ public BooleanFieldScript(String fieldName, Map<String, Object> params, SearchLo
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     */
-    public final void runForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         trues = 0;
         falses = 0;
-        setDocument(docId);
-        execute();
     }
 
     public final void runForDoc(int docId, Consumer<Boolean> consumer) {

diff --git a/server/src/main/java/org/elasticsearch/script/CompositeFieldScript.java b/server/src/main/java/org/elasticsearch/script/CompositeFieldScript.java
@@ -47,17 +47,19 @@ public CompositeFieldScript(String fieldName, Map<String, Object> params, Search
      */
     public final List<Object> getValues(String field) {
         // TODO for now we re-run the script every time a leaf field is accessed, but we could cache the values?
-        fieldValues.clear();
+        prepareExecute();
         execute();
         List<Object> values = fieldValues.get(field);
         fieldValues.clear();    // don't hold on to values unnecessarily
         return values;
     }
 
-    public final Map<String, List<Object>> runForDoc(int doc) {
-        setDocument(doc);
+    @Override
+    protected void prepareExecute() {
         fieldValues.clear();
-        execute();
+    }
+
+    public final Map<String, List<Object>> getFieldValues() {
         return fieldValues;
     }
 

diff --git a/server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java b/server/src/main/java/org/elasticsearch/script/DoubleFieldScript.java
@@ -74,13 +74,9 @@ public DoubleFieldScript(String fieldName, Map<String, Object> params, SearchLoo
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     */
-    public final void runForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         count = 0;
-        setDocument(docId);
-        execute();
     }
 
     /**

diff --git a/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java b/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java
@@ -85,13 +85,9 @@ public GeoPointFieldScript(String fieldName, Map<String, Object> params, SearchL
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     */
-    public final void runForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         count = 0;
-        setDocument(docId);
-        execute();
     }
 
     /**

diff --git a/server/src/main/java/org/elasticsearch/script/IpFieldScript.java b/server/src/main/java/org/elasticsearch/script/IpFieldScript.java
@@ -95,13 +95,9 @@ public IpFieldScript(String fieldName, Map<String, Object> params, SearchLookup
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     */
-    public final void runForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         count = 0;
-        setDocument(docId);
-        execute();
     }
 
     public final void runForDoc(int docId, Consumer<InetAddress> consumer) {

diff --git a/server/src/main/java/org/elasticsearch/script/StringFieldScript.java b/server/src/main/java/org/elasticsearch/script/StringFieldScript.java
@@ -81,22 +81,22 @@ public StringFieldScript(String fieldName, Map<String, Object> params, SearchLoo
         super(fieldName, params, searchLookup, ctx);
     }
 
-    /**
-     * Execute the script for the provided {@code docId}.
-     *
-     * @return a mutable {@link List} that contains the results of the script
-     * and will be modified the next time you call {@linkplain #resultsForDoc}.
-     */
-    public final List<String> resultsForDoc(int docId) {
+    @Override
+    protected void prepareExecute() {
         results.clear();
         chars = 0;
-        setDocument(docId);
-        execute();
-        return results;
     }
 
     public final void runForDoc(int docId, Consumer<String> consumer) {
-        resultsForDoc(docId).forEach(consumer);
+        runForDoc(docId);
+        results.forEach(consumer);
+    }
+
+    /**
+     * Values from the last time runForDoc(int) was called. This list is mutable and will change with the next call of runForDoc(int).
+     */
+    public List<String> getValues() {
+        return results;
     }
 
     @Override

diff --git a/server/src/main/java/org/elasticsearch/search/runtime/AbstractStringScriptFieldQuery.java b/server/src/main/java/org/elasticsearch/search/runtime/AbstractStringScriptFieldQuery.java
@@ -24,7 +24,8 @@ abstract class AbstractStringScriptFieldQuery extends AbstractScriptFieldQuery<S
 
     @Override
     protected final boolean matches(StringFieldScript scriptContext, int docId) {
-        return matches(scriptContext.resultsForDoc(docId));
+        scriptContext.runForDoc(docId);
+        return matches(scriptContext.getValues());
     }
 
     /**

diff --git a/server/src/test/java/org/elasticsearch/index/mapper/StringFieldScriptTests.java b/server/src/test/java/org/elasticsearch/index/mapper/StringFieldScriptTests.java
@@ -133,8 +133,8 @@ public final void testFromSourceDoesNotEnforceValuesLimit() throws IOException {
                     new SearchLookup(field -> null, (ft, lookup, fdt) -> null, new SourceLookup.ReaderSourceProvider())
                 );
                 StringFieldScript stringFieldScript = leafFactory.newInstance(reader.leaves().get(0));
-                List<String> results = stringFieldScript.resultsForDoc(0);
-                assertEquals(numValues, results.size());
+                stringFieldScript.runForDoc(0);
+                assertEquals(numValues, stringFieldScript.getValues().size());
             }
         }
     }
@@ -162,8 +162,8 @@ public final void testFromSourceDoesNotEnforceCharsLimit() throws IOException {
                     new SearchLookup(field -> null, (ft, lookup, fdt) -> null, new SourceLookup.ReaderSourceProvider())
                 );
                 StringFieldScript stringFieldScript = leafFactory.newInstance(reader.leaves().get(0));
-                List<String> results = stringFieldScript.resultsForDoc(0);
-                assertEquals(5, results.size());
+                stringFieldScript.runForDoc(0);
+                assertEquals(5, stringFieldScript.getValues().size());
             }
         }
     }