Add Runtime Fields Contexts to Painless Execute API (#71374)

jdconrad · jdconrad · commit d009b9638f09 · 2021-04-14T09:04:42.000-07:00
This change adds support for the 7 different runtime fields contexts to the Painless Execute API. Each context can accept the standard script input (source and params) along with a user-defined document and an index name to pull mappings from. The results depend on the output of the runtime field type. Closes #70467
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
@@ -51,6 +51,7 @@
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexService;
+import org.elasticsearch.index.mapper.DateFieldMapper;
 import org.elasticsearch.index.mapper.ParsedDocument;
 import org.elasticsearch.index.mapper.SourceToParse;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
@@ -61,12 +62,19 @@
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.action.RestToXContentListener;
+import org.elasticsearch.script.BooleanFieldScript;
+import org.elasticsearch.script.DateFieldScript;
+import org.elasticsearch.script.DoubleFieldScript;
 import org.elasticsearch.script.FilterScript;
+import org.elasticsearch.script.GeoPointFieldScript;
+import org.elasticsearch.script.IpFieldScript;
+import org.elasticsearch.script.LongFieldScript;
 import org.elasticsearch.script.ScoreScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.script.ScriptType;
+import org.elasticsearch.script.StringFieldScript;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 
@@ -114,6 +122,13 @@ public static class Request extends SingleShardRequest<Request> implements ToXCo
             supportedContexts.put("painless_test", PainlessTestScript.CONTEXT);
             supportedContexts.put("filter", FilterScript.CONTEXT);
             supportedContexts.put("score", ScoreScript.CONTEXT);
+            supportedContexts.put("boolean_field", BooleanFieldScript.CONTEXT);
+            supportedContexts.put("date_field", DateFieldScript.CONTEXT);
+            supportedContexts.put("double_field", DoubleFieldScript.CONTEXT);
+            supportedContexts.put("geo_point_field", GeoPointFieldScript.CONTEXT);
+            supportedContexts.put("ip_field", IpFieldScript.CONTEXT);
+            supportedContexts.put("long_field", LongFieldScript.CONTEXT);
+            supportedContexts.put("string_field", StringFieldScript.CONTEXT);
             SUPPORTED_CONTEXTS = Collections.unmodifiableMap(supportedContexts);
         }
 
@@ -542,6 +557,69 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                     double result = scoreScript.execute(null);
                     return new Response(result);
                 }, indexService);
+            } else if (scriptContext == BooleanFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    BooleanFieldScript.Factory factory = scriptService.compile(request.script, BooleanFieldScript.CONTEXT);
+                    BooleanFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("boolean_field", request.getScript().getParams(), context.lookup());
+                    BooleanFieldScript booleanFieldScript = leafFactory.newInstance(leafReaderContext);
+                    booleanFieldScript.runForDoc(0);
+                    return new Response(booleanFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == DateFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    DateFieldScript.Factory factory = scriptService.compile(request.script, DateFieldScript.CONTEXT);
+                    DateFieldScript.LeafFactory leafFactory = factory.newFactory("date_field",
+                            request.getScript().getParams(), context.lookup(), DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER);
+                    DateFieldScript dateFieldScript = leafFactory.newInstance(leafReaderContext);
+                    dateFieldScript.runForDoc(0);
+                    return new Response(dateFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == DoubleFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    DoubleFieldScript.Factory factory = scriptService.compile(request.script, DoubleFieldScript.CONTEXT);
+                    DoubleFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("double_field", request.getScript().getParams(), context.lookup());
+                    DoubleFieldScript doubleFieldScript = leafFactory.newInstance(leafReaderContext);
+                    doubleFieldScript.runForDoc(0);
+                    return new Response(doubleFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == GeoPointFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    GeoPointFieldScript.Factory factory = scriptService.compile(request.script, GeoPointFieldScript.CONTEXT);
+                    GeoPointFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("geo_point_field", request.getScript().getParams(), context.lookup());
+                    GeoPointFieldScript geoPointFieldScript = leafFactory.newInstance(leafReaderContext);
+                    geoPointFieldScript.runForDoc(0);
+                    return new Response(geoPointFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == IpFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    IpFieldScript.Factory factory = scriptService.compile(request.script, IpFieldScript.CONTEXT);
+                    IpFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("ip_field", request.getScript().getParams(), context.lookup());
+                    IpFieldScript ipFieldScript = leafFactory.newInstance(leafReaderContext);
+                    ipFieldScript.runForDoc(0);
+                    return new Response(ipFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == LongFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    LongFieldScript.Factory factory = scriptService.compile(request.script, LongFieldScript.CONTEXT);
+                    LongFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("long_field", request.getScript().getParams(), context.lookup());
+                    LongFieldScript longFieldScript = leafFactory.newInstance(leafReaderContext);
+                    longFieldScript.runForDoc(0);
+                    return new Response(longFieldScript.asDocValues());
+                }, indexService);
+            } else if (scriptContext == StringFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    StringFieldScript.Factory factory = scriptService.compile(request.script, StringFieldScript.CONTEXT);
+                    StringFieldScript.LeafFactory leafFactory =
+                            factory.newFactory("string_field", request.getScript().getParams(), context.lookup());
+                    StringFieldScript stringFieldScript = leafFactory.newInstance(leafReaderContext);
+                    stringFieldScript.resultsForDoc(0);
+                    return new Response(stringFieldScript.asDocValues());
+                }, indexService);
             } else {
                 throw new UnsupportedOperationException("unsupported context [" + scriptContext.name + "]");
             }
diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java
@@ -7,10 +7,12 @@
  */
 package org.elasticsearch.painless.action;
 
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.IndexService;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.painless.PainlessPlugin;
 import org.elasticsearch.painless.action.PainlessExecuteAction.Request;
@@ -28,6 +30,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.painless.action.PainlessExecuteAction.TransportAction.innerShardOperation;
 import static org.hamcrest.Matchers.equalTo;
@@ -100,6 +103,168 @@ public void testScoreExecutionContext() throws IOException {
         assertThat(response.getResult(), equalTo(0.93D));
     }
 
+    public void testBooleanFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "rank", "type=long", "text", "type=text");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"rank\": 4.0, \"text\": \"quick brown fox\"}"), new MatchQueryBuilder("text", "fox"));
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['rank'].value < params.max_rank)", singletonMap("max_rank", 5.0)), "boolean_field",
+                contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((boolean[])response.getResult(), new boolean[] {true});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(false); emit(true); emit (false);", emptyMap()), "boolean_field",
+                contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((boolean[])response.getResult(), new boolean[] {false, false, true});
+    }
+
+    public void testDateFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_date", "type=date");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"test_date\":\"2015-01-01T12:10:30Z\"}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['test_date'].value.toInstant().toEpochMilli())", emptyMap()), "date_field",
+                contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {1420114230000L});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(ZonedDateTime.parse(\"2021-01-01T00:00:00Z\").toInstant().toEpochMilli());\n" +
+                "emit(ZonedDateTime.parse(\"1942-05-31T15:16:17Z\").toInstant().toEpochMilli());\n" +
+                "emit(ZonedDateTime.parse(\"2035-10-13T10:54:19Z\").toInstant().toEpochMilli());",
+                emptyMap()), "date_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {-870597823000L, 1609459200000L, 2075885659000L});
+    }
+
+    public void testDoubleFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "rank", "type=long", "text", "type=text");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"rank\": 4.0, \"text\": \"quick brown fox\"}"), new MatchQueryBuilder("text", "fox"));
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['rank'].value); emit(Math.log(doc['rank'].value))", emptyMap()), "double_field",
+                contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((double[])response.getResult(), new double[] {Math.log(4.0), 4.0}, 0.00001);
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(3.1); emit(2.29); emit(-12.47); emit(-12.46); emit(Double.MAX_VALUE); emit(0.0);",
+                emptyMap()), "double_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((double[])response.getResult(), new double[] {-12.47, -12.46, 0.0, 2.29, 3.1, Double.MAX_VALUE}, 0.00001);
+    }
+
+    public void testGeoPointFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_point", "type=geo_point");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"test_point\":\"30.0,40.0\"}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['test_point'].value.lat + 1.0, doc['test_point'].value.lon - 1.0)", emptyMap()),
+                "geo_point_field", contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {3176939252927413179L});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(78.96, 12.12); emit(13.45, 56.78);",
+                emptyMap()), "geo_point_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {1378381707499043786L, 8091971733044486384L});
+    }
+
+    public void testIpFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_ip", "type=ip");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"test_ip\":\"192.168.1.254\"}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['test_ip'].value);", emptyMap()),
+                "ip_field", contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((BytesRef[])response.getResult(),
+                new BytesRef[] {new BytesRef(new byte[] {0,0,0,0,0,0,0,0,0,0,(byte)255,(byte)255,(byte)192,(byte)168,1,(byte)254})});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(\"192.168.0.1\"); emit(\"127.0.0.1\"); emit(\"255.255.255.255\"); emit(\"0.0.0.0\");",
+                emptyMap()), "ip_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((BytesRef[])response.getResult(), 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})});
+    }
+
+    public void testLongFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_point", "type=geo_point");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"test_point\":\"30.2,40.2\"}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit((long)doc['test_point'].value.lat); emit((long)doc['test_point'].value.lon);", emptyMap()),
+                "long_field", contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {30, 40});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(3L); emit(1L); emit(20000000000L); emit(10L); emit(-1000L); emit(0L);",
+                emptyMap()), "long_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((long[])response.getResult(), new long[] {-1000L, 0L, 1L, 3L, 10L, 20000000000L});
+    }
+
+    public void testStringFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_point", "type=geo_point");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup("index",
+                new BytesArray("{\"test_point\":\"30.2,40.2\"}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(doc['test_point'].value.lat.toString().substring(0, 5)); " +
+                        "emit(doc['test_point'].value.lon.toString().substring(0, 5));", emptyMap()),
+                "string_field", contextSetup);
+        Response response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((String[])response.getResult(), new String[] {"30.19", "40.19"});
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(new Script(ScriptType.INLINE, "painless",
+                "emit(\"test\"); emit(\"baz was not here\"); emit(\"Data\"); emit(\"-10\"); emit(\"20\"); emit(\"9\");",
+                emptyMap()), "string_field", contextSetup);
+        response = innerShardOperation(request, scriptService, indexService);
+        assertArrayEquals((String[])response.getResult(), new String[] {"-10", "20", "9", "Data", "baz was not here", "test"});
+    }
+
     public void testContextWhitelists() throws IOException {
         ScriptService scriptService = getInstanceFromNode(ScriptService.class);
         // score
diff --git a/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/70_execute_painless_scripts.yml b/modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/70_execute_painless_scripts.yml
