From 4be7743fa3f9684260a2b43bf3bbcb75395f676b Mon Sep 17 00:00:00 2001 From: Luca Cavanna Date: Mon, 19 Dec 2022 17:02:20 +0100 Subject: [PATCH] Share common runForDoc impl in AbstractFieldScript (#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. --- .../action/PainlessExecuteAction.java | 3 ++- .../fielddata/StringScriptDocValues.java | 3 ++- .../script/AbstractFieldScript.java | 11 ++++++++++ .../script/AbstractLongFieldScript.java | 8 ++----- .../script/BooleanFieldScript.java | 8 ++----- .../script/CompositeFieldScript.java | 10 +++++---- .../script/DoubleFieldScript.java | 8 ++----- .../script/GeoPointFieldScript.java | 8 ++----- .../elasticsearch/script/IpFieldScript.java | 8 ++----- .../script/StringFieldScript.java | 22 +++++++++---------- .../AbstractStringScriptFieldQuery.java | 3 ++- .../index/mapper/StringFieldScriptTests.java | 8 +++---- 12 files changed, 48 insertions(+), 52 deletions(-) diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java index 6a3d54e734975..a65d18a762ec5 100644 --- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java +++ b/modules/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 index e42659743521f..f5e236e9e8640 100644 --- 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 results = script.resultsForDoc(docId); + script.runForDoc(docId); + List 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 index c31e5d3db352b..cda67064240ea 100644 --- 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 index 10cb90607b5a1..101598b233cdc 100644 --- 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 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 index 5591ff6d796fd..5c6e4828471ff 100644 --- 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 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 consumer) { diff --git a/server/src/main/java/org/elasticsearch/script/CompositeFieldScript.java b/server/src/main/java/org/elasticsearch/script/CompositeFieldScript.java index 253d83fd2596a..c02968059ae06 100644 --- 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 params, Search */ public final List 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 values = fieldValues.get(field); fieldValues.clear(); // don't hold on to values unnecessarily return values; } - public final Map> runForDoc(int doc) { - setDocument(doc); + @Override + protected void prepareExecute() { fieldValues.clear(); - execute(); + } + + public final Map> 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 index f59759e65bdd9..d259e9a42d891 100644 --- 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 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 index d89cf9d4d0699..c4276fc32b248 100644 --- 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 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 index 665eaee24cedb..c8cb1f0a2278f 100644 --- 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 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 consumer) { diff --git a/server/src/main/java/org/elasticsearch/script/StringFieldScript.java b/server/src/main/java/org/elasticsearch/script/StringFieldScript.java index 7e366f4b72b18..b996293f1c512 100644 --- 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 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 resultsForDoc(int docId) { + @Override + protected void prepareExecute() { results.clear(); chars = 0; - setDocument(docId); - execute(); - return results; } public final void runForDoc(int docId, Consumer 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 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 index bf2be13aab774..dde67d2143d73 100644 --- 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 null, (ft, lookup, fdt) -> null, new SourceLookup.ReaderSourceProvider()) ); StringFieldScript stringFieldScript = leafFactory.newInstance(reader.leaves().get(0)); - List 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 results = stringFieldScript.resultsForDoc(0); - assertEquals(5, results.size()); + stringFieldScript.runForDoc(0); + assertEquals(5, stringFieldScript.getValues().size()); } } }