From a48c7369a2d5590eedfb5217398ba0578816054f Mon Sep 17 00:00:00 2001 From: Jack Conradson Date: Wed, 11 Aug 2021 08:32:36 -0700 Subject: [PATCH] Add Fields API to aggregation scripts and field scripts (#76325) This change updates the aggregation script, map script for aggregations, and field scripts to extend DocBasedScript to give them access to the new fields api. --- .../painless/PainlessPlugin.java | 21 +++++ ...org.elasticsearch.script.fields.aggmap.txt | 15 ++++ ...lasticsearch.script.fields.aggregation.txt | 15 ++++ .../org.elasticsearch.script.fields.field.txt | 15 ++++ .../test/painless/100_terms_agg.yml | 54 ++++++++++++ .../test/painless/130_metric_agg.yml | 85 +++++++++++++++++++ .../test/painless/20_scriptfield.yml | 75 ++++++++++++++++ .../script/AbstractFieldScript.java | 27 ++---- .../script/AggregationScript.java | 29 +------ .../org/elasticsearch/script/DocReader.java | 4 - .../script/DocValuesDocReader.java | 9 -- .../org/elasticsearch/script/FieldScript.java | 28 +----- .../script/GeoPointFieldScript.java | 2 +- .../org/elasticsearch/script/ScoreScript.java | 2 +- .../script/ScriptedMetricAggContexts.java | 26 ++---- .../script/MockScriptEngine.java | 2 +- 16 files changed, 304 insertions(+), 105 deletions(-) create mode 100644 modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggmap.txt create mode 100644 modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggregation.txt create mode 100644 modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.field.txt create mode 100644 modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/130_metric_agg.yml diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java index d83fd9fcf92c3..7d30f8d22dc5b 100644 --- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java +++ b/modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessPlugin.java @@ -38,6 +38,8 @@ import org.elasticsearch.repositories.RepositoriesService; import org.elasticsearch.rest.RestController; import org.elasticsearch.rest.RestHandler; +import org.elasticsearch.script.AggregationScript; +import org.elasticsearch.script.FieldScript; import org.elasticsearch.script.FilterScript; import org.elasticsearch.script.IngestScript; import org.elasticsearch.script.NumberSortScript; @@ -46,6 +48,7 @@ import org.elasticsearch.script.ScriptEngine; import org.elasticsearch.script.ScriptModule; import org.elasticsearch.script.ScriptService; +import org.elasticsearch.script.ScriptedMetricAggContexts; import org.elasticsearch.script.StringSortScript; import org.elasticsearch.search.aggregations.pipeline.MovingFunctionScript; import org.elasticsearch.threadpool.ThreadPool; @@ -140,6 +143,24 @@ public final class PainlessPlugin extends Plugin implements ScriptPlugin, Extens filter.add(filterWhitelist); map.put(FilterScript.CONTEXT, filter); + List aggregation = new ArrayList<>(); + Whitelist aggregationWhitelist = + WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.aggregation.txt"); + aggregation.add(aggregationWhitelist); + map.put(AggregationScript.CONTEXT, aggregation); + + List aggmap = new ArrayList<>(); + Whitelist aggmapWhitelist = + WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.aggmap.txt"); + aggmap.add(aggmapWhitelist); + map.put(ScriptedMetricAggContexts.MapScript.CONTEXT, aggmap); + + List field = new ArrayList<>(); + Whitelist fieldWhitelist = + WhitelistLoader.loadFromResourceFiles(PainlessPlugin.class, "org.elasticsearch.script.fields.field.txt"); + field.add(fieldWhitelist); + map.put(FieldScript.CONTEXT, field); + // Execute context gets everything List test = new ArrayList<>(); test.add(movFnWhitelist); diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggmap.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggmap.txt new file mode 100644 index 0000000000000..84df365c4fe2c --- /dev/null +++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggmap.txt @@ -0,0 +1,15 @@ +# +# 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. +# + +# The whitelist for the fields api + +# The scripts must be whitelisted for painless to find the classes for the field API +class org.elasticsearch.script.ScriptedMetricAggContexts$MapScript @no_import { +} +class org.elasticsearch.script.ScriptedMetricAggContexts$MapScript$Factory @no_import { +} diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggregation.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggregation.txt new file mode 100644 index 0000000000000..a09bce101a155 --- /dev/null +++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.aggregation.txt @@ -0,0 +1,15 @@ +# +# 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. +# + +# The whitelist for the fields api + +# The scripts must be whitelisted for painless to find the classes for the field API +class org.elasticsearch.script.AggregationScript @no_import { +} +class org.elasticsearch.script.AggregationScript$Factory @no_import { +} diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.field.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.field.txt new file mode 100644 index 0000000000000..9f1c4b2368bc8 --- /dev/null +++ b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.field.txt @@ -0,0 +1,15 @@ +# +# 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. +# + +# The whitelist for the fields api + +# The scripts must be whitelisted for painless to find the classes for the field API +class org.elasticsearch.script.FieldScript @no_import { +} +class org.elasticsearch.script.FieldScript @no_import { +} diff --git a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/100_terms_agg.yml b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/100_terms_agg.yml index 000e1af694d7d..824036e193037 100644 --- a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/100_terms_agg.yml +++ b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/100_terms_agg.yml @@ -66,6 +66,24 @@ setup: - is_false: aggregations.str_terms.buckets.1.key_as_string - match: { aggregations.str_terms.buckets.1.doc_count: 1 } +--- +"String Value Script with doc notation (fields api)": + + - do: + search: + rest_total_hits_as_int: true + body: { "size" : 0, "aggs" : { "str_terms" : { "terms" : { "field" : "str", "script": { "source": "return field('str').getValue('') + \"1\""} } } } } + + - match: { hits.total: 3 } + + - length: { aggregations.str_terms.buckets: 2 } + - match: { aggregations.str_terms.buckets.0.key: "abc1" } + - is_false: aggregations.str_terms.buckets.0.key_as_string + - match: { aggregations.str_terms.buckets.0.doc_count: 2 } + - match: { aggregations.str_terms.buckets.1.key: "bcd1" } + - is_false: aggregations.str_terms.buckets.1.key_as_string + - match: { aggregations.str_terms.buckets.1.doc_count: 1 } + --- "Long Value Script with doc notation": @@ -84,6 +102,24 @@ setup: - is_false: aggregations.long_terms.buckets.1.key_as_string - match: { aggregations.long_terms.buckets.1.doc_count: 1 } +--- +"Long Value Script with doc notation (fields api)": + + - do: + search: + rest_total_hits_as_int: true + body: { "size" : 0, "aggs" : { "long_terms" : { "terms" : { "field" : "number", "script": { "source": "return field('number').getValue(0L) + 1"} } } } } + + - match: { hits.total: 3 } + + - length: { aggregations.long_terms.buckets: 2 } + - match: { aggregations.long_terms.buckets.0.key: 2.0 } + - is_false: aggregations.long_terms.buckets.0.key_as_string + - match: { aggregations.long_terms.buckets.0.doc_count: 2 } + - match: { aggregations.long_terms.buckets.1.key: 3.0 } + - is_false: aggregations.long_terms.buckets.1.key_as_string + - match: { aggregations.long_terms.buckets.1.doc_count: 1 } + --- "Double Value Script with doc notation": @@ -102,6 +138,24 @@ setup: - is_false: aggregations.double_terms.buckets.1.key_as_string - match: { aggregations.double_terms.buckets.1.doc_count: 1 } +--- +"Double Value Script with doc notation (fields api)": + + - do: + search: + rest_total_hits_as_int: true + body: { "size" : 0, "aggs" : { "double_terms" : { "terms" : { "field" : "double", "script": { "source": "return field('double').getValue(0.0) + 1"} } } } } + + - match: { hits.total: 3 } + + - length: { aggregations.double_terms.buckets: 2 } + - match: { aggregations.double_terms.buckets.0.key: 2.0 } + - is_false: aggregations.double_terms.buckets.0.key_as_string + - match: { aggregations.double_terms.buckets.0.doc_count: 2 } + - match: { aggregations.double_terms.buckets.1.key: 3.0 } + - is_false: aggregations.double_terms.buckets.1.key_as_string + - match: { aggregations.double_terms.buckets.1.doc_count: 1 } + --- "Bucket script with keys": diff --git a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/130_metric_agg.yml b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/130_metric_agg.yml new file mode 100644 index 0000000000000..60cc31f623cfe --- /dev/null +++ b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/130_metric_agg.yml @@ -0,0 +1,85 @@ +setup: + - do: + indices.create: + index: test + body: + settings: + number_of_replicas: 0 + mappings: + properties: + double: + type: double + + - do: + cluster.health: + wait_for_status: green + + - do: + index: + index: test + id: 1 + body: + double: 1.0 + + - do: + index: + index: test + id: 2 + body: + double: 1.0 + + - do: + index: + index: test + id: 3 + body: + double: 2.0 + + - do: + indices.refresh: {} + +--- +"Scripted Metric Agg Total": + + - do: + search: + rest_total_hits_as_int: true + body: { + "size": 0, + "aggs": { + "total": { + "scripted_metric": { + "init_script": "state.transactions = []", + "map_script": "state.transactions.add(doc['double'].value)", + "combine_script": "double total = 0.0; for (t in state.transactions) { total += t } return total", + "reduce_script": "double total = 0; for (a in states) { total += a } return total" + } + } + } + } + + - match: { hits.total: 3 } + - match: { aggregations.total.value: 4.0 } + +--- +"Scripted Metric Agg Total (fields api)": + + - do: + search: + rest_total_hits_as_int: true + body: { + "size": 0, + "aggs": { + "total": { + "scripted_metric": { + "init_script": "state.transactions = []", + "map_script": "state.transactions.add(field('double').getValue(0.0))", + "combine_script": "double total = 0.0; for (t in state.transactions) { total += t } return total", + "reduce_script": "double total = 0; for (a in states) { total += a } return total" + } + } + } + } + + - match: { hits.total: 3 } + - match: { aggregations.total.value: 4.0 } diff --git a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/20_scriptfield.yml b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/20_scriptfield.yml index 3a3e5608e02bf..0ff7c38b77633 100644 --- a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/20_scriptfield.yml +++ b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/20_scriptfield.yml @@ -44,6 +44,21 @@ setup: - match: { hits.hits.0.fields.bar.0: "aaabbb"} +--- +"Scripted Field (fields api)": + - do: + search: + rest_total_hits_as_int: true + body: + script_fields: + bar: + script: + source: "field('foo').getValue('') + params.x;" + params: + x: "bbb" + + - match: { hits.hits.0.fields.bar.0: "aaabbb"} + --- "Scripted Field Doing Compare": - do: @@ -73,6 +88,35 @@ setup: - match: { hits.hits.0.fields.bar.0: false} +--- +"Scripted Field Doing Compare (fields api)": + - do: + search: + rest_total_hits_as_int: true + body: + script_fields: + bar: + script: + source: "boolean compare(Supplier s, def v) {return s.get() == v;} + compare(() -> { return field('foo').getValue('') }, params.x);" + params: + x: "aaa" + + - match: { hits.hits.0.fields.bar.0: true} + - do: + search: + rest_total_hits_as_int: true + body: + script_fields: + bar: + script: + source: "boolean compare(Supplier s, def v) {return s.get() == v;} + compare(() -> { return doc['foo'].value }, params.x);" + params: + x: "bbb" + + - match: { hits.hits.0.fields.bar.0: false} + --- "Scripted Field with a null safe dereference (non-null)": - do: @@ -116,6 +160,19 @@ setup: - match: { hits.hits.0.fields.bar.0: 7} +--- +"Access a date (fields api)": + - do: + search: + rest_total_hits_as_int: true + body: + script_fields: + bar: + script: + source: "field('date').getValue(new Date()).dayOfWeekEnum.value" + + - match: { hits.hits.0.fields.bar.0: 7} + --- "Access many dates": - do: @@ -134,6 +191,24 @@ setup: - match: { hits.hits.0.fields.bar.0: "7 3 3"} +--- +"Access many dates (fields api)": + - do: + search: + rest_total_hits_as_int: true + body: + script_fields: + bar: + script: + source: > + StringBuilder b = new StringBuilder(); + for (def date : field('dates').getValues()) { + b.append(" ").append(date.getDayOfWeekEnum().value); + } + return b.toString().trim() + + - match: { hits.hits.0.fields.bar.0: "7 3 3"} + --- "Scripted Field with script error": - do: diff --git a/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java b/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java index 27f653980177e..2792ce48b94f2 100644 --- a/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java +++ b/server/src/main/java/org/elasticsearch/script/AbstractFieldScript.java @@ -10,8 +10,6 @@ import org.apache.lucene.index.LeafReaderContext; import org.elasticsearch.common.xcontent.support.XContentMapValues; -import org.elasticsearch.index.fielddata.ScriptDocValues; -import org.elasticsearch.search.lookup.LeafSearchLookup; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.lookup.SourceLookup; @@ -69,27 +67,21 @@ static ScriptContext newContext(String name, Class factoryClass) { ); protected final String fieldName; + protected final SourceLookup sourceLookup; private final Map params; - protected final LeafSearchLookup leafSearchLookup; public AbstractFieldScript(String fieldName, Map params, SearchLookup searchLookup, LeafReaderContext ctx) { super(new DocValuesDocReader(searchLookup, ctx)); this.fieldName = fieldName; - this.leafSearchLookup = searchLookup.getLeafSearchLookup(ctx); + Map docAsMap = docAsMap(); + this.sourceLookup = (SourceLookup)docAsMap.get("_source"); params = new HashMap<>(params); - params.put("_source", leafSearchLookup.source()); - params.put("_fields", leafSearchLookup.fields()); + params.put("_source", sourceLookup); + params.put("_fields", docAsMap.get("_fields")); this.params = new DynamicMap(params, PARAMS_FUNCTIONS); } - /** - * Set the document to run the script against. - */ - public void setDocument(int docId) { - this.leafSearchLookup.setDocument(docId); - } - /** * Expose the {@code params} of the script to the script itself. */ @@ -97,15 +89,8 @@ public final Map getParams() { return params; } - /** - * Expose field data to the script as {@code doc}. - */ - public final Map> getDoc() { - return leafSearchLookup.doc(); - } - protected List extractFromSource(String path) { - return XContentMapValues.extractRawValues(path, leafSearchLookup.source().source()); + return XContentMapValues.extractRawValues(path, sourceLookup.source()); } protected final void emitFromCompositeScript(CompositeFieldScript compositeFieldScript) { diff --git a/server/src/main/java/org/elasticsearch/script/AggregationScript.java b/server/src/main/java/org/elasticsearch/script/AggregationScript.java index 085fc349f867b..b91d126dad2da 100644 --- a/server/src/main/java/org/elasticsearch/script/AggregationScript.java +++ b/server/src/main/java/org/elasticsearch/script/AggregationScript.java @@ -13,8 +13,6 @@ import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.common.lucene.ScorerAware; -import org.elasticsearch.index.fielddata.ScriptDocValues; -import org.elasticsearch.search.lookup.LeafSearchLookup; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.lookup.SourceLookup; @@ -23,7 +21,7 @@ import java.util.Map; import java.util.function.Function; -public abstract class AggregationScript implements ScorerAware { +public abstract class AggregationScript extends DocBasedScript implements ScorerAware { public static final String[] PARAMETERS = {}; @@ -51,11 +49,6 @@ public abstract class AggregationScript implements ScorerAware { */ private final Map params; - /** - * A leaf lookup for the bound segment this script will operate on. - */ - private final LeafSearchLookup leafLookup; - /** * A scorer that will return the score for the current document when the script is run. */ @@ -64,14 +57,14 @@ public abstract class AggregationScript implements ScorerAware { private Object value; public AggregationScript(Map params, SearchLookup lookup, LeafReaderContext leafContext) { + super(new DocValuesDocReader(lookup, leafContext)); this.params = new DynamicMap(new HashMap<>(params), PARAMS_FUNCTIONS); - this.leafLookup = lookup.getLeafSearchLookup(leafContext); - this.params.putAll(leafLookup.asMap()); + this.params.putAll(docAsMap()); } protected AggregationScript() { + super(null); params = null; - leafLookup = null; } /** @@ -81,20 +74,6 @@ public Map getParams() { return params; } - /** - * The doc lookup for the Lucene segment this script was created for. - */ - public Map> getDoc() { - return leafLookup.doc(); - } - - /** - * Set the current document to run the script on next. - */ - public void setDocument(int docid) { - leafLookup.setDocument(docid); - } - @Override public void setScorer(Scorable scorer) { this.scorer = scorer; diff --git a/server/src/main/java/org/elasticsearch/script/DocReader.java b/server/src/main/java/org/elasticsearch/script/DocReader.java index 3326cdfb559e7..ece8497b61737 100644 --- a/server/src/main/java/org/elasticsearch/script/DocReader.java +++ b/server/src/main/java/org/elasticsearch/script/DocReader.java @@ -35,8 +35,4 @@ public interface DocReader { /** Old-style doc['field'] access */ Map> doc(); - - /** Base document-id of the current reader, used as seed for RandomScore */ - // should be replaced - int getDocBase(); } diff --git a/server/src/main/java/org/elasticsearch/script/DocValuesDocReader.java b/server/src/main/java/org/elasticsearch/script/DocValuesDocReader.java index 3d4d3b591809c..a10e8219f3d3d 100644 --- a/server/src/main/java/org/elasticsearch/script/DocValuesDocReader.java +++ b/server/src/main/java/org/elasticsearch/script/DocValuesDocReader.java @@ -23,13 +23,9 @@ public class DocValuesDocReader implements DocReader, LeafReaderContextSupplier // provide access to the leaf context reader for expressions protected final LeafReaderContext leafReaderContext; - // backwards compatibility access for random score script - protected final int docBase; - public DocValuesDocReader(SearchLookup lookup, LeafReaderContext leafContext) { this.leafReaderContext = leafContext; this.leafLookup = lookup.getLeafSearchLookup(leafReaderContext); - this.docBase = leafContext.docBase; } @Override @@ -63,11 +59,6 @@ public Map> doc() { return leafLookup.doc(); } - @Override - public int getDocBase() { - return docBase; - } - @Override public LeafReaderContext getLeafReaderContext() { return leafReaderContext; diff --git a/server/src/main/java/org/elasticsearch/script/FieldScript.java b/server/src/main/java/org/elasticsearch/script/FieldScript.java index ec3adf3f45d0e..42c7a75a2a520 100644 --- a/server/src/main/java/org/elasticsearch/script/FieldScript.java +++ b/server/src/main/java/org/elasticsearch/script/FieldScript.java @@ -11,8 +11,6 @@ import org.apache.lucene.index.LeafReaderContext; import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; -import org.elasticsearch.index.fielddata.ScriptDocValues; -import org.elasticsearch.search.lookup.LeafSearchLookup; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.lookup.SourceLookup; @@ -24,7 +22,7 @@ /** * A script to produce dynamic values for return fields. */ -public abstract class FieldScript { +public abstract class FieldScript extends DocBasedScript { public static final String[] PARAMETERS = {}; @@ -48,44 +46,26 @@ public abstract class FieldScript { /** The generic runtime parameters for the script. */ private final Map params; - /** A leaf lookup for the bound segment this script will operate on. */ - private final LeafSearchLookup leafLookup; - public FieldScript(Map params, SearchLookup lookup, LeafReaderContext leafContext) { - this.leafLookup = lookup.getLeafSearchLookup(leafContext); + super(new DocValuesDocReader(lookup, leafContext)); params = new HashMap<>(params); - params.putAll(leafLookup.asMap()); + params.putAll(docAsMap()); this.params = new DynamicMap(params, PARAMS_FUNCTIONS); } // for expression engine protected FieldScript() { + super(null); params = null; - leafLookup = null; } public abstract Object execute(); - /** The leaf lookup for the Lucene segment this script was created for. */ - protected final LeafSearchLookup getLeafLookup() { - return leafLookup; - } - /** Return the parameters for this script. */ public Map getParams() { return params; } - /** The doc lookup for the Lucene segment this script was created for. */ - public final Map> getDoc() { - return leafLookup.doc(); - } - - /** Set the current document to run the script on next. */ - public void setDocument(int docid) { - leafLookup.setDocument(docid); - } - /** A factory to construct {@link FieldScript} instances. */ public interface LeafFactory { FieldScript newInstance(LeafReaderContext ctx) throws IOException; diff --git a/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java b/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java index 5773fc7bca93e..6a8aba788c957 100644 --- a/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java +++ b/server/src/main/java/org/elasticsearch/script/GeoPointFieldScript.java @@ -100,7 +100,7 @@ public void runGeoPointForDoc(int doc, Consumer consumer) { @Override protected List extractFromSource(String path) { - Object value = XContentMapValues.extractValue(path, leafSearchLookup.source().source()); + Object value = XContentMapValues.extractValue(path, sourceLookup.source()); if (value instanceof List) { @SuppressWarnings("unchecked") List list = (List) value; diff --git a/server/src/main/java/org/elasticsearch/script/ScoreScript.java b/server/src/main/java/org/elasticsearch/script/ScoreScript.java index 9b800f6e98e8e..865e49c0b2f67 100644 --- a/server/src/main/java/org/elasticsearch/script/ScoreScript.java +++ b/server/src/main/java/org/elasticsearch/script/ScoreScript.java @@ -92,7 +92,7 @@ public ScoreScript(Map params, SearchLookup searchLookup, DocRea params = new HashMap<>(params); params.putAll(docReader.docAsMap()); this.params = new DynamicMap(params, PARAMS_FUNCTIONS); - this.docBase = docReader.getDocBase(); + this.docBase = ((DocValuesDocReader)docReader).getLeafReaderContext().docBase; } } diff --git a/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java b/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java index 82b26cfb457d9..c74788776d13a 100644 --- a/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java +++ b/server/src/main/java/org/elasticsearch/script/ScriptedMetricAggContexts.java @@ -14,7 +14,6 @@ import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.index.fielddata.ScriptDocValues; -import org.elasticsearch.search.lookup.LeafSearchLookup; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.lookup.SourceLookup; @@ -53,7 +52,7 @@ public interface Factory extends ScriptFactory { public static ScriptContext CONTEXT = new ScriptContext<>("aggs_init", Factory.class); } - public abstract static class MapScript { + public abstract static class MapScript extends DocBasedScript { private static final DeprecationLogger deprecationLogger = DeprecationLogger.getLogger(DynamicMap.class); private static final Map> PARAMS_FUNCTIONS = Map.of( @@ -79,18 +78,14 @@ public abstract static class MapScript { private final Map params; private final Map state; - private final LeafSearchLookup leafLookup; private Scorable scorer; public MapScript(Map params, Map state, SearchLookup lookup, LeafReaderContext leafContext) { + super(leafContext == null ? null : new DocValuesDocReader(lookup, leafContext)); this.state = state; - this.leafLookup = leafContext == null ? null : lookup.getLeafSearchLookup(leafContext); - if (leafLookup != null) { - params = new HashMap<>(params); // copy params so we aren't modifying input - params.putAll(leafLookup.asMap()); // add lookup vars - params = new DynamicMap(params, PARAMS_FUNCTIONS); // wrap with deprecations - } - this.params = params; + params = new HashMap<>(params); // copy params so we aren't modifying input + params.putAll(docAsMap()); // add lookup vars + this.params = new DynamicMap(params, PARAMS_FUNCTIONS); // wrap with deprecations } public Map getParams() { @@ -101,16 +96,9 @@ public Map getState() { return state; } - // Return the doc as a map (instead of LeafDocLookup) in order to abide by type whitelisting rules for - // Painless scripts. + // Override this to ensure null is returned for backcompat rather than an empty map. public Map> getDoc() { - return leafLookup == null ? null : leafLookup.doc(); - } - - public void setDocument(int docId) { - if (leafLookup != null) { - leafLookup.setDocument(docId); - } + return docReader == null ? null : docReader.doc(); } public void setScorer(Scorable scorer) { diff --git a/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java b/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java index 390971e32e080..dfbff8fdfffbb 100644 --- a/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java +++ b/test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java @@ -698,7 +698,7 @@ public FieldScript.LeafFactory newFactory(Map parameters, Search @Override public Object execute() { Map vars = createVars(parameters); - vars.putAll(getLeafLookup().asMap()); + vars.putAll(docAsMap()); return script.apply(vars); }