Merge pull request #18262 from rmuir/painless_doc_access

Painless doc access

modules/lang-painless/src/main/java/org/elasticsearch/painless/Analyzer.java

@@ -93,9 +93,22 @@ class Analyzer extends PainlessParserBaseVisitor<Void> {
 
         utility.incrementScope();
         utility.addVariable(null, "#this", definition.execType);
+        //
+        // reserved words parameters.
+        //
+        // input map of variables passed to the script. TODO: rename to 'params' since that will be its use
         metadata.inputValueSlot = utility.addVariable(null, "input", definition.smapType).slot;
+        // scorer parameter passed to the script. internal use only.
         metadata.scorerValueSlot = utility.addVariable(null, "#scorer", definition.objectType).slot;
+        // doc parameter passed to the script.
+        // TODO: currently working as a def type, should be smapType...
+        metadata.docValueSlot = utility.addVariable(null, "doc", definition.defType).slot;
+        //
+        // reserved words implemented as local variables
+        //
+        // loop counter to catch runaway scripts. internal use only.
         metadata.loopCounterSlot = utility.addVariable(null, "#loop", definition.intType).slot;
+        // document's score as a read-only float.
         metadata.scoreValueSlot = utility.addVariable(null, "_score", definition.floatType).slot;
 
         metadata.createStatementMetadata(metadata.root);

modules/lang-painless/src/main/java/org/elasticsearch/painless/AnalyzerExternal.java

@@ -448,6 +448,15 @@ class AnalyzerExternal {
                 throw new IllegalArgumentException(AnalyzerUtility.error(ctx) + "Unknown variable [" + id + "].");
             }
 
+            // special cases: reserved words
+            if ("_score".equals(id) || "doc".equals(id)) {
+                // read-only: don't allow stores
+                if (parentemd.storeExpr != null) {
+                    throw new IllegalArgumentException(AnalyzerUtility.error(ctx) + "Variable [" + id + "] is read-only.");
+                }
+            }
+
+            // track if the _score value is ever used, we will invoke Scorer.score() only once if so.
             if ("_score".equals(id)) {
                 metadata.scoreValueUsed = true;
             }

modules/lang-painless/src/main/java/org/elasticsearch/painless/Executable.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 import org.apache.lucene.search.Scorer;
+import org.elasticsearch.search.lookup.LeafDocLookup;
 
 import java.util.Map;
 
@@ -48,5 +49,5 @@ public abstract class Executable {
         return definition;
     }
 
-    public abstract Object execute(Map<String, Object> input, Scorer scorer);
+    public abstract Object execute(Map<String, Object> input, Scorer scorer, LeafDocLookup doc);
 }

modules/lang-painless/src/main/java/org/elasticsearch/painless/Metadata.java

@@ -433,6 +433,12 @@ class Metadata {
      * variable slots at the completion of analysis if _score is not used.
      */
     boolean scoreValueUsed = false;
+    
+    /**
+     * Used to determine what slot the doc variable is stored in.  This is used in the {@link Writer} whenever
+     * the doc variable is accessed.
+     */
+    int docValueSlot = -1;
 
     /**
      * Maps the relevant ANTLR node to its metadata.

modules/lang-painless/src/main/java/org/elasticsearch/painless/ScriptImpl.java

@@ -22,6 +22,7 @@ package org.elasticsearch.painless;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.script.ExecutableScript;
 import org.elasticsearch.script.LeafSearchScript;
+import org.elasticsearch.search.lookup.LeafDocLookup;
 import org.elasticsearch.search.lookup.LeafSearchLookup;
 
 import java.util.HashMap;
@@ -46,6 +47,11 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      * The lookup is used to access search field values at run-time.
      */
     private final LeafSearchLookup lookup;
+    
+    /**
+     * the 'doc' object accessed by the script, if available.
+     */
+    private final LeafDocLookup doc;
 
     /**
      * Current scorer being used
@@ -70,6 +76,9 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
 
         if (lookup != null) {
             variables.putAll(lookup.asMap());
+            doc = lookup.doc();
+        } else {
+            doc = null;
         }
     }
 
@@ -89,7 +98,7 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     @Override
     public Object run() {
-        return executable.execute(variables, scorer);
+        return executable.execute(variables, scorer, doc);
     }
 
     /**

modules/lang-painless/src/main/java/org/elasticsearch/painless/WriterConstants.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.painless;
 
 import org.apache.lucene.search.Scorer;
+import org.elasticsearch.search.lookup.LeafDocLookup;
 import org.objectweb.asm.Handle;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
@@ -37,7 +38,7 @@ class WriterConstants {
     final static Type CLASS_TYPE        = Type.getType("L" + CLASS_NAME.replace(".", "/") + ";");
 
     final static Method CONSTRUCTOR = getAsmMethod(void.class, "<init>", Definition.class, String.class, String.class);
-    final static Method EXECUTE     = getAsmMethod(Object.class, "execute", Map.class, Scorer.class);
+    final static Method EXECUTE     = getAsmMethod(Object.class, "execute", Map.class, Scorer.class, LeafDocLookup.class);
 
     final static Type PAINLESS_ERROR_TYPE = Type.getType(PainlessError.class);
 

modules/lang-painless/src/test/java/org/elasticsearch/painless/NeedsScoreTests.java

@@ -46,7 +46,7 @@ public class NeedsScoreTests extends ESSingleNodeTestCase {
                                          lookup, Collections.<String, Object>emptyMap());
         assertFalse(ss.needsScores());
 
-        compiled = service.compile("input.doc['d'].value", Collections.emptyMap());
+        compiled = service.compile("doc['d'].value", Collections.emptyMap());
         ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "painless", compiled), 
                             lookup, Collections.<String, Object>emptyMap());
         assertFalse(ss.needsScores());
@@ -56,7 +56,7 @@ public class NeedsScoreTests extends ESSingleNodeTestCase {
                             lookup, Collections.<String, Object>emptyMap());
         assertTrue(ss.needsScores());
 
-        compiled = service.compile("input.doc['d'].value * _score", Collections.emptyMap());
+        compiled = service.compile("doc['d'].value * _score", Collections.emptyMap());
         ss = service.search(new CompiledScript(ScriptType.INLINE, "randomName", "painless", compiled), 
                             lookup, Collections.<String, Object>emptyMap());
         assertTrue(ss.needsScores());

modules/lang-painless/src/test/java/org/elasticsearch/painless/ReservedWordTests.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless;
+
+/** Tests for special reserved words such as _score */
+public class ReservedWordTests extends ScriptTestCase {
+    
+    /** check that we can't declare a variable of _score, its really reserved! */
+    public void testScoreVar() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("int _score = 5; return _score;");
+        });
+        assertTrue(expected.getMessage().contains("Variable name [_score] already defined"));
+    }
+    
+    /** check that we can't write to _score, its read-only! */
+    public void testScoreStore() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("_score = 5; return _score;");
+        });
+        assertTrue(expected.getMessage().contains("Variable [_score] is read-only"));
+    }
+    
+    /** check that we can't declare a variable of doc, its really reserved! */
+    public void testDocVar() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("int doc = 5; return doc;");
+        });
+        assertTrue(expected.getMessage().contains("Variable name [doc] already defined"));
+    }
+    
+    /** check that we can't write to _score, its read-only! */
+    public void testDocStore() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("doc = 5; return doc;");
+        });
+        assertTrue(expected.getMessage().contains("Variable [doc] is read-only"));
+    }
+}

modules/lang-painless/src/test/resources/rest-api-spec/test/plan_a/16_update2.yaml

@@ -5,7 +5,7 @@
       put_script:
         id: "1"
         lang: "painless"
-        body: { "script":  "_score * input.doc[\"myParent.weight\"].value" }
+        body: { "script":  "_score * doc[\"myParent.weight\"].value" }
   - match: { acknowledged: true }
 
   - do:
@@ -15,7 +15,7 @@
   - match: { found: true }
   - match: { lang: painless }
   - match: { _id: "1" }
-  - match: { "script":  "_score * input.doc[\"myParent.weight\"].value" }
+  - match: { "script":  "_score * doc[\"myParent.weight\"].value" }
 
   - do:
      catch: missing
@@ -44,11 +44,11 @@
       put_script:
         id: "1"
         lang: "painless"
-        body: { "script":  "_score * foo bar + input.doc[\"myParent.weight\"].value" }
+        body: { "script":  "_score * foo bar + doc[\"myParent.weight\"].value" }
 
   - do:
       catch: /Unable.to.parse.*/
       put_script:
         id: "1"
         lang: "painless"
-        body: { "script":  "_score * foo bar + input.doc[\"myParent.weight\"].value" }
+        body: { "script":  "_score * foo bar + doc[\"myParent.weight\"].value" }

modules/lang-painless/src/test/resources/rest-api-spec/test/plan_a/20_scriptfield.yaml

@@ -28,7 +28,7 @@ setup:
                 script_fields:
                     bar:
                         script: 
-                            inline: "input.doc['foo'].value + input.x;"
+                            inline: "doc['foo'].value + input.x;"
                             lang: painless
                             params:
                                 x: "bbb"

modules/lang-painless/src/test/resources/rest-api-spec/test/plan_a/30_search.yaml

@@ -29,12 +29,12 @@
                 query:
                     script:
                         script:
-                            inline: "input.doc['num1'].value > 1;"
+                            inline: "doc['num1'].value > 1;"
                             lang: painless
                 script_fields:
                     sNum1:
                         script: 
-                            inline: "input.doc['num1'].value;"
+                            inline: "doc['num1'].value;"
                             lang: painless
                 sort:
                     num1:
@@ -51,7 +51,7 @@
                 query:
                     script:
                         script:
-                            inline: "input.doc['num1'].value > input.param1;"
+                            inline: "doc['num1'].value > input.param1;"
                             lang: painless
                             params:
                                 param1: 1
@@ -59,7 +59,7 @@
                 script_fields:
                     sNum1:
                         script:
-                            inline: "return input.doc['num1'].value;"
+                            inline: "return doc['num1'].value;"
                             lang: painless
                 sort:
                     num1:
@@ -76,7 +76,7 @@
                 query:
                     script:
                         script:
-                            inline: "input.doc['num1'].value > input.param1;"
+                            inline: "doc['num1'].value > input.param1;"
                             lang: painless
                             params:
                                 param1: -1
@@ -84,7 +84,7 @@
                 script_fields:
                     sNum1:
                         script: 
-                            inline: "input.doc['num1'].value;"
+                            inline: "doc['num1'].value;"
                             lang: painless
                 sort:
                     num1:
@@ -126,7 +126,7 @@
                             "script_score": {
                                 "script": {
                                     "lang": "painless",
-                                    "inline": "input.doc['num1'].value"
+                                    "inline": "doc['num1'].value"
                                 }
                             }
                         }]
@@ -148,7 +148,7 @@
                             "script_score": {
                                 "script": {
                                     "lang": "painless",
-                                    "inline": "-input.doc['num1'].value"
+                                    "inline": "-doc['num1'].value"
                                 }
                             }
                         }]
@@ -170,7 +170,7 @@
                             "script_score": {
                                 "script": {
                                     "lang": "painless",
-                                    "inline": "Math.pow(input.doc['num1'].value, 2)"
+                                    "inline": "Math.pow(doc['num1'].value, 2)"
                                 }
                             }
                         }]
@@ -192,7 +192,7 @@
                             "script_score": {
                                 "script": {
                                     "lang": "painless",
-                                    "inline": "Math.max(input.doc['num1'].value, 1)"
+                                    "inline": "Math.max(doc['num1'].value, 1)"
                                 }
                             }
                         }]
@@ -214,7 +214,7 @@
                             "script_score": {
                                 "script": {
                                     "lang": "painless",
-                                    "inline": "input.doc['num1'].value * _score"
+                                    "inline": "doc['num1'].value * _score"
                                 }
                             }
                         }]
@@ -357,7 +357,7 @@
                 script_fields:
                     foobar:
                         script: 
-                            inline: "input.doc['f'].values.size()"
+                            inline: "doc['f'].values.size()"
                             lang: painless