Script: Fields API for Filter context (#76119)

elastic · Aug 4, 2021 · bb9c91f · bb9c91f
1 parent b62ab7e
commit bb9c91f
Show file tree

Hide file tree

Showing 10 changed files with 209 additions and 149 deletions.
diff --git a/...-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java b/...-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java
@@ -25,7 +25,6 @@
 import org.elasticsearch.script.BucketAggregationScript;
 import org.elasticsearch.script.BucketAggregationSelectorScript;
 import org.elasticsearch.script.ClassPermission;
-import org.elasticsearch.script.DocValuesDocReader;
 import org.elasticsearch.script.FieldScript;
 import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.NumberSortScript;
@@ -332,9 +331,9 @@ private static FieldScript.LeafFactory newFieldScript(Expression expr, SearchLoo
      */
     private static FilterScript.LeafFactory newFilterScript(Expression expr, SearchLookup lookup, @Nullable Map<String, Object> vars) {
         ScoreScript.LeafFactory searchLeafFactory = newScoreScript(expr, lookup, vars);
-        return ctx -> {
-            ScoreScript script = searchLeafFactory.newInstance(new DocValuesDocReader(lookup, ctx));
-            return new FilterScript(vars, lookup, ctx) {
+        return docReader -> {
+            ScoreScript script = searchLeafFactory.newInstance(docReader);
+            return new FilterScript(vars, lookup, docReader) {
                 @Override
                 public boolean execute() {
                     return script.execute(null) != 0.0;

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
@@ -37,6 +37,7 @@
 import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.rest.RestController;
 import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.IngestScript;
 import org.elasticsearch.script.NumberSortScript;
 import org.elasticsearch.script.ScoreScript;
@@ -109,6 +110,11 @@ public final class PainlessPlugin extends Plugin implements ScriptPlugin, Extens
         strSort.add(strSortField);
         map.put(StringSortScript.CONTEXT, strSort);
 
+        List<Whitelist> filter = new ArrayList<>(Whitelist.BASE_WHITELISTS);
+        Whitelist filterWhitelist = WhitelistLoader.loadFromResourceFiles(Whitelist.class, "org.elasticsearch.script.fields.filter.txt");
+        filter.add(filterWhitelist);
+        map.put(FilterScript.CONTEXT, filter);
+
         // Execute context gets everything
         List<Whitelist> test = new ArrayList<>(Whitelist.BASE_WHITELISTS);
         test.add(movFnWhitelist);

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
@@ -520,9 +520,9 @@ static Response innerShardOperation(Request request, ScriptService scriptService
             } else if (scriptContext == FilterScript.CONTEXT) {
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     FilterScript.Factory factory = scriptService.compile(request.script, FilterScript.CONTEXT);
-                    FilterScript.LeafFactory leafFactory =
-                            factory.newFactory(request.getScript().getParams(), context.lookup());
-                    FilterScript filterScript = leafFactory.newInstance(leafReaderContext);
+                    SearchLookup lookup = context.lookup();
+                    FilterScript.LeafFactory leafFactory = factory.newFactory(request.getScript().getParams(), lookup);
+                    FilterScript filterScript = leafFactory.newInstance(new DocValuesDocReader(lookup, leafReaderContext));
                     filterScript.setDocument(0);
                     boolean result = filterScript.execute();
                     return new Response(result);

diff --git a/.../main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.fields.filter.txt b/.../main/resources/org/elasticsearch/painless/spi/org.elasticsearch.script.fields.filter.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.FilterScript @no_import {
+}
+class org.elasticsearch.script.FilterScript$Factory @no_import {
+}
diff --git a/...es/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/40_fields_api.yml b/...es/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/40_fields_api.yml
@@ -0,0 +1,160 @@
+# Integration tests for sort script queries using Painless
+
+setup:
+- skip:
+    version: " - 7.14.99"
+    reason: "sort script fields api was added in 7.15.0"
+
+---
+"sort script fields api":
+     - do:
+         indices.create:
+           index: test
+           body:
+             settings:
+               number_of_shards: 2
+             mappings:
+               properties:
+                 dval:
+                   type: double
+     - do:
+         index:
+           index: test
+           id: d1
+           body: {"dval": 10, "sval": "f"}
+     - do:
+         index:
+           index: test
+           id: d2
+           body: {}
+     - do:
+         index:
+           index: test
+           id: d3
+           body: {"dval": 5, "sval": "a"}
+     - do:
+         indices.refresh: {}
+     - do:
+         search:
+           rest_total_hits_as_int: true
+           index: test
+           body:
+             sort:
+               _script:
+                 type: number
+                 script:
+                   source: "field('dval').getValue(3)"
+     - match: { hits.total: 3 }
+     - match: { hits.hits.0._id: d2 }
+     - match: { hits.hits.1._id: d3 }
+     - match: { hits.hits.2._id: d1 }
+     - do:
+         search:
+           rest_total_hits_as_int: true
+           index: test
+           body:
+             sort:
+               _script:
+                 type: string
+                 script:
+                   source: "field('sval.keyword').getValue('g')"
+     - match: { hits.total: 3 }
+     - match: { hits.hits.0._id: d3 }
+     - match: { hits.hits.1._id: d1 }
+     - match: { hits.hits.2._id: d2 }
+
+---
+"script score fields api":
+    - do:
+        indices.create:
+          index: test
+          body:
+            settings:
+              number_of_shards: 2
+            mappings:
+              properties:
+                dval:
+                  type: double
+    - do:
+        index:
+          index: test
+          id: d1
+          body: {"dval": 10}
+    - do:
+        index:
+          index: test
+          id: d2
+          body: {}
+    - do:
+        index:
+          index: test
+          id: d3
+          body: {"dval": 5}
+    - do:
+        indices.refresh: {}
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          index: test
+          body:
+            query:
+              script_score:
+                query: {match_all: {} }
+                script:
+                  source: "field('dval').getValue(3)"
+    - match: { hits.total: 3 }
+    - match: { hits.hits.0._id: d1 }
+    - match: { hits.hits.1._id: d3 }
+    - match: { hits.hits.2._id: d2 }
+
+---
+"filter script fields api":
+    - do:
+        indices.create:
+          index: test
+          body:
+            settings:
+              number_of_shards: 2
+            mappings:
+              properties:
+                dval:
+                  type: double
+    - do:
+        index:
+          index: test
+          id: d1
+          body: {"dval": 10, "sval": "f"}
+    - do:
+        index:
+          index: test
+          id: d2
+          body: {}
+    - do:
+        index:
+          index: test
+          id: d3
+          body: {"dval": 5, "sval": "a"}
+    - do:
+        indices.refresh: {}
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query:
+              bool:
+                filter:
+                  script:
+                    script: "field('dval').getValue(6) > 6"
+    - match: { hits.total: 1 }
+    - match: { hits.hits.0._id: d1 }
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            query:
+              bool:
+                filter:
+                  script:
+                    script: "field('sval.keyword').getValue('b') == 'a'"
+    - match: { hits.total: 1 }
+    - match: { hits.hits.0._id: d3 }
diff --git a/...less/src/yamlRestTest/resources/rest-api-spec/test/painless/40_sort_script_fields_api.yml b/...less/src/yamlRestTest/resources/rest-api-spec/test/painless/40_sort_script_fields_api.yml
diff --git a/...ess/src/yamlRestTest/resources/rest-api-spec/test/painless/88_script_score_fields_api.yml b/...ess/src/yamlRestTest/resources/rest-api-spec/test/painless/88_script_score_fields_api.yml
diff --git a/server/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java
@@ -24,8 +24,10 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.script.DocValuesDocReader;
 import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.Script;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.Objects;
@@ -127,18 +129,21 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
                     ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to false.");
         }
         FilterScript.Factory factory = context.compile(script, FilterScript.CONTEXT);
-        FilterScript.LeafFactory filterScript = factory.newFactory(script.getParams(), context.lookup());
-        return new ScriptQuery(script, filterScript);
+        SearchLookup lookup = context.lookup();
+        FilterScript.LeafFactory filterScript = factory.newFactory(script.getParams(), lookup);
+        return new ScriptQuery(script, filterScript, lookup);
     }
 
     static class ScriptQuery extends Query {
 
         final Script script;
         final FilterScript.LeafFactory filterScript;
+        final SearchLookup lookup;
 
-        ScriptQuery(Script script, FilterScript.LeafFactory filterScript) {
+        ScriptQuery(Script script, FilterScript.LeafFactory filterScript, SearchLookup lookup) {
             this.script = script;
             this.filterScript = filterScript;
+            this.lookup = lookup;
         }
 
         @Override
@@ -172,7 +177,7 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
                 @Override
                 public Scorer scorer(LeafReaderContext context) throws IOException {
                     DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
-                    final FilterScript leafScript = filterScript.newInstance(context);
+                    final FilterScript leafScript = filterScript.newInstance(new DocValuesDocReader(lookup, context));
                     TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
 
                         @Override