Merge pull request #18284 from rmuir/painless_value_aggregations

_value support in painess?

core/src/main/java/org/elasticsearch/script/ExecutableScript.java

@@ -24,6 +24,14 @@ package org.elasticsearch.script;
  */
 public interface ExecutableScript {
 
+    /**
+     * Sets a runtime script parameter.
+     * <p>
+     * Note that this method may be slow, involving put() and get() calls
+     * to a hashmap or similar.
+     * @param name parameter name
+     * @param value parameter value
+     */
     void setNextVar(String name, Object value);
 
     /**

core/src/main/java/org/elasticsearch/script/LeafSearchScript.java

@@ -31,6 +31,18 @@ public interface LeafSearchScript extends ScorerAware, ExecutableScript {
     void setDocument(int doc);
 
     void setSource(Map<String, Object> source);
+    
+    /**
+     * Sets per-document aggregation {@code _value}.
+     * <p>
+     * The default implementation just calls {@code setNextVar("_value", value)} but
+     * some engines might want to handle this differently for better performance.
+     * <p>
+     * @param value per-document value, typically a String, Long, or Double
+     */
+    default void setNextAggregationValue(Object value) {
+        setNextVar("_value", value);
+    }
 
     float runAsFloat();
 

core/src/main/java/org/elasticsearch/search/aggregations/support/ValuesSource.java

@@ -329,7 +329,7 @@ public abstract class ValuesSource {
                     resize(longValues.count());
                     script.setDocument(doc);
                     for (int i = 0; i < count(); ++i) {
-                        script.setNextVar("_value", longValues.valueAt(i));
+                        script.setNextAggregationValue(longValues.valueAt(i));
                         values[i] = script.runAsLong();
                     }
                     sort();
@@ -357,7 +357,7 @@ public abstract class ValuesSource {
                     resize(doubleValues.count());
                     script.setDocument(doc);
                     for (int i = 0; i < count(); ++i) {
-                        script.setNextVar("_value", doubleValues.valueAt(i));
+                        script.setNextAggregationValue(doubleValues.valueAt(i));
                         values[i] = script.runAsDouble();
                     }
                     sort();
@@ -474,7 +474,7 @@ public abstract class ValuesSource {
                 grow();
                 for (int i = 0; i < count; ++i) {
                     final BytesRef value = bytesValues.valueAt(i);
-                    script.setNextVar("_value", value.utf8ToString());
+                    script.setNextAggregationValue(value.utf8ToString());
                     values[i].copyChars(script.run().toString());
                 }
                 sort();

docs/reference/aggregations/bucket/terms-aggregation.asciidoc

@@ -443,7 +443,10 @@ Generating the terms using a script:
     "aggs" : {
         "genders" : {
             "terms" : {
-                "script" : "doc['gender'].value"
+                "script" : { 
+                    "inline": "doc['gender'].value"
+                    "lang": "painless"
+                }
             }
         }
     }
@@ -482,7 +485,10 @@ TIP: for indexed scripts replace the `file` parameter with an `id` parameter.
         "genders" : {
             "terms" : {
                 "field" : "gender",
-                "script" : "'Gender: ' +_value"
+                "script" : { 
+                    "inline" : "'Gender: ' +_value"
+                    "lang" : "painless" 
+                }
             }
         }
     }

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionSearchScript.java

@@ -111,10 +111,7 @@ class ExpressionSearchScript implements SearchScript {
             }
 
             @Override
-            public void setNextVar(String name, Object value) {
-                // this should only be used for the special "_value" variable used in aggregations
-                assert(name.equals("_value"));
-
+            public void setNextAggregationValue(Object value) {
                 // _value isn't used in script if specialValue == null
                 if (specialValue != null) {
                     if (value instanceof Number) {
@@ -124,6 +121,12 @@ class ExpressionSearchScript implements SearchScript {
                     }
                 }
             }
+
+            @Override
+            public void setNextVar(String name, Object value) {
+                // other per-document variables aren't supported yet, even if they are numbers
+                // but we shouldn't encourage this anyway.
+            }
         };
     }
 

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

@@ -103,6 +103,8 @@ class Analyzer extends PainlessParserBaseVisitor<Void> {
         // doc parameter passed to the script.
         // TODO: currently working as a Map<String,Def>, we can do better?
         metadata.docValueSlot = utility.addVariable(null, "doc", definition.smapType).slot;
+        // aggregation _value parameter passed to the script
+        metadata.aggregationValueSlot = utility.addVariable(null, "_value", definition.defType).slot;
         //
         // reserved words implemented as local variables
         //

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

@@ -449,7 +449,7 @@ class AnalyzerExternal {
             }
 
             // special cases: reserved words
-            if ("_score".equals(id) || "doc".equals(id) || "ctx".equals(id)) {
+            if ("_score".equals(id) || "_value".equals(id) || "doc".equals(id) || "ctx".equals(id)) {
                 // read-only: don't allow stores to ourself
                 if (varenmd.last && parentemd.storeExpr != null) {
                     throw new IllegalArgumentException(AnalyzerUtility.error(ctx) + "Variable [" + id + "] is read-only.");

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

@@ -49,5 +49,5 @@ public abstract class Executable {
         return definition;
     }
 
-    public abstract Object execute(Map<String, Object> input, Scorer scorer, LeafDocLookup doc);
+    public abstract Object execute(Map<String, Object> input, Scorer scorer, LeafDocLookup doc, Object value);
 }

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

@@ -415,6 +415,11 @@ class Metadata {
      * _score from it, if _score will be accessed by the script.
      */
     int scorerValueSlot = -1;
+    
+    /**
+     * Used to determine what slot the _value variable is scored in. 
+     */
+    int aggregationValueSlot = -1;
 
     /**
      * Used to determine what slot the loopCounter variable is stored in.  This is used n the {@link Writer} whenever

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

@@ -59,6 +59,12 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     private Scorer scorer;
 
+    /**
+     * Current _value for aggregation
+     * @see #setNextAggregationValue(Object)
+     */
+    private Object aggregationValue;
+
     /**
      * Creates a ScriptImpl for the a previously compiled Painless script.
      * @param executable The previously compiled Painless script.
@@ -91,6 +97,11 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
     public void setNextVar(final String name, final Object value) {
         variables.put(name, value);
     }
+    
+    @Override
+    public void setNextAggregationValue(Object value) {
+        this.aggregationValue = value;
+    }
 
     /**
      * Run the script.
@@ -98,7 +109,7 @@ final class ScriptImpl implements ExecutableScript, LeafSearchScript {
      */
     @Override
     public Object run() {
-        return executable.execute(variables, scorer, doc);
+        return executable.execute(variables, scorer, doc, aggregationValue);
     }
 
     /**

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

@@ -38,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, LeafDocLookup.class);
+    final static Method EXECUTE     = getAsmMethod(Object.class, "execute", Map.class, Scorer.class, LeafDocLookup.class, Object.class);
 
     final static Type PAINLESS_ERROR_TYPE = Type.getType(PainlessError.class);
 

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

@@ -77,4 +77,20 @@ public class ReservedWordTests extends ScriptTestCase {
     public void testCtxStoreMap() {
         assertEquals(5, exec("ctx.foo = 5; return ctx.foo;", Collections.singletonMap("ctx", new HashMap<String,Object>())));
     }
+    
+    /** check that we can't declare a variable of _value, its really reserved! */
+    public void testAggregationValueVar() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("int _value = 5; return _value;");
+        });
+        assertTrue(expected.getMessage().contains("Variable name [_value] already defined"));
+    }
+    
+    /** check that we can't write to _value, its read-only! */
+    public void testAggregationValueStore() {
+        IllegalArgumentException expected = expectThrows(IllegalArgumentException.class, () -> {
+            exec("_value = 5; return _value;");
+        });
+        assertTrue(expected.getMessage().contains("Variable [_value] is read-only"));
+    }
 }

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

@@ -363,3 +363,33 @@
 
     - match: { hits.total: 1 }
     - match: { hits.hits.0.fields.foobar.0: 1 }
+
+---
+
+"Agg _value":
+    - do:
+        index:
+            index: test
+            type: test
+            id: 1
+            body: { "dummy_field": 1 }
+    - do:
+        indices.refresh: {}
+
+
+    - do:
+        index: test
+        search:
+            body:
+                aggs:
+                    value_agg:
+                        terms:
+                            field: dummy_field
+                            script:
+                                lang: painless
+                                inline: "_value + 1"
+
+    - match: { hits.total: 1 }
+    - match: { hits.hits.0._score: 1.0 }
+    - match: { aggregations.value_agg.buckets.0.key: 2 }
+    - match: { aggregations.value_agg.buckets.0.doc_count: 1 }