Allow pipeline aggs to select specific buckets from multi-bucket aggs

docs/reference/aggregations/pipeline.asciidoc

@@ -35,11 +35,12 @@ parameter, which follows a specific format:
 // https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_Form
 [source,ebnf]
 --------------------------------------------------
-AGG_SEPARATOR       =  '>' ;
-METRIC_SEPARATOR    =  '.' ;
+AGG_SEPARATOR       =  `>` ;
+METRIC_SEPARATOR    =  `.` ;
 AGG_NAME            =  <the name of the aggregation> ;
 METRIC              =  <the name of the metric (in case of multi-value metrics aggregation)> ;
-PATH                =  <AGG_NAME> [ <AGG_SEPARATOR>, <AGG_NAME> ]* [ <METRIC_SEPARATOR>, <METRIC> ] ;
+MULTIBUCKET_KEY     =  `[<KEY_NAME>]`
+PATH                =  <AGG_NAME><MULTIBUCKET_KEY>? (<AGG_SEPARATOR>, <AGG_NAME> )* ( <METRIC_SEPARATOR>, <METRIC> ) ;
 --------------------------------------------------
 
 For example, the path `"my_bucket>my_stats.avg"` will path to the `avg` value in the `"my_stats"` metric, which is
@@ -110,6 +111,52 @@ POST /_search
 <1> `buckets_path` instructs this max_bucket aggregation that we want the maximum value of the `sales` aggregation in the
 `sales_per_month` date histogram.
 
+If a Sibling pipeline agg references a multi-bucket aggregation, such as a `terms` agg, it also has the option to
+select specific keys from the multi-bucket.  For example, a `bucket_script` could select two specific buckets (via
+their bucket keys) to perform the calculation:
+
+[source,js]
+--------------------------------------------------
+POST /_search
+{
+    "aggs" : {
+        "sales_per_month" : {
+            "date_histogram" : {
+                "field" : "date",
+                "calendar_interval" : "month"
+            },
+            "aggs": {
+                "sale_type": {
+                    "terms": {
+                        "field": "type"
+                    },
+                    "aggs": {
+                        "sales": {
+                            "sum": {
+                                "field": "price"
+                            }
+                        }
+                    }
+                },
+                "hat_vs_bag_ratio": {
+                    "bucket_script": {
+                        "buckets_path": {
+                            "hats": "sale_type['hat']>sales", <1>
+                            "bags": "sale_type['bag']>sales"  <1>
+                        },
+                        "script": "params.hats / params.bags"
+                    }
+                }
+            }
+        }
+    }
+}
+--------------------------------------------------
+// CONSOLE
+// TEST[setup:sales]
+<1> `buckets_path` selects the hats and bags buckets (via `['hat']`/`['bag']``) to use in the script specifically,
+instead of fetching all the buckets from `sale_type` aggregation
+
 [float]
 === Special Paths
 

modules/lang-painless/src/test/resources/rest-api-spec/test/painless/100_terms_agg.yml

@@ -102,3 +102,41 @@ setup:
   - is_false: aggregations.double_terms.buckets.1.key_as_string
   - match: { aggregations.double_terms.buckets.1.doc_count: 1 }
 
+---
+"Bucket script with keys":
+
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        body:
+          size: 0
+          aggs:
+            placeholder:
+              filters:
+                filters:
+                 - match_all: {}
+              aggs:
+                str_terms:
+                  terms:
+                    field: "str"
+                  aggs:
+                    the_avg:
+                      avg:
+                        field: "number"
+                the_bucket_script:
+                  bucket_script:
+                    buckets_path:
+                      foo: "str_terms['bcd']>the_avg.value"
+                    script: "params.foo"
+
+  - match: { hits.total: 3 }
+
+  - length: { aggregations.placeholder.buckets.0.str_terms.buckets: 2 }
+  - match: { aggregations.placeholder.buckets.0.str_terms.buckets.0.key: "abc" }
+  - is_false: aggregations.placeholder.buckets.0.str_terms.buckets.0.key_as_string
+  - match: { aggregations.placeholder.buckets.0.str_terms.buckets.0.doc_count: 2 }
+  - match: { aggregations.placeholder.buckets.0.str_terms.buckets.1.key: "bcd" }
+  - is_false: aggregations.placeholder.buckets.0.str_terms.buckets.1.key_as_string
+  - match: { aggregations.placeholder.buckets.0.str_terms.buckets.1.doc_count: 1 }
+  - match: { aggregations.placeholder.buckets.0.the_bucket_script.value: 2.0 }
+

server/src/main/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregation.java

@@ -73,16 +73,33 @@ protected InternalMultiBucketAggregation(StreamInput in) throws IOException {
     public Object getProperty(List<String> path) {
         if (path.isEmpty()) {
             return this;
-        } else if (path.get(0).equals("_bucket_count")) {
-            return getBuckets().size();
-        } else {
-            List<? extends InternalBucket> buckets = getBuckets();
-            Object[] propertyArray = new Object[buckets.size()];
-            for (int i = 0; i < buckets.size(); i++) {
-                propertyArray[i] = buckets.get(i).getProperty(getName(), path);
+        }
+        return resolvePropertyFromPath(path, getBuckets(), getName());
+    }
+
+    static Object resolvePropertyFromPath(List<String> path, List<? extends InternalBucket> buckets, String name) {
+        String aggName = path.get(0);
+        if (aggName.equals("_bucket_count")) {
+            return buckets.size();
+        }
+
+        // This is a bucket key, look through our buckets and see if we can find a match
+        if (aggName.startsWith("'") && aggName.endsWith("'")) {
+            for (InternalBucket bucket : buckets) {
+                if (bucket.getKeyAsString().equals(aggName.substring(1, aggName.length() - 1))) {
+                    return bucket.getProperty(name, path.subList(1, path.size()));
+                }
             }
-            return propertyArray;
+            // No key match, time to give up
+            throw new InvalidAggregationPathException("Cannot find an key [" + aggName + "] in [" + name + "]");
+        }
+
+        Object[] propertyArray = new Object[buckets.size()];
+        for (int i = 0; i < buckets.size(); i++) {
+            propertyArray[i] = buckets.get(i).getProperty(name, path);
         }
+        return propertyArray;
+
     }
 
     /**

server/src/test/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregationTests.java

@@ -0,0 +1,183 @@
+/*
+ * 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.search.aggregations;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.bucket.terms.InternalTerms;
+import org.elasticsearch.search.aggregations.bucket.terms.LongTerms;
+import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
+import org.elasticsearch.search.aggregations.metrics.InternalAvg;
+import org.elasticsearch.search.aggregations.support.AggregationPath;
+import org.elasticsearch.test.ESTestCase;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import static org.elasticsearch.search.aggregations.InternalMultiBucketAggregation.resolvePropertyFromPath;
+import static org.hamcrest.Matchers.equalTo;
+
+public class InternalMultiBucketAggregationTests extends ESTestCase {
+
+    public void testResolveToAgg() {
+        AggregationPath path = AggregationPath.parse("the_avg");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object[] value = (Object[]) resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value[0], equalTo(agg));
+    }
+
+    public void testResolveToAggValue() {
+        AggregationPath path = AggregationPath.parse("the_avg.value");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object[] value = (Object[]) resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value[0], equalTo(2.0));
+    }
+
+    public void testResolveToNothing() {
+        AggregationPath path = AggregationPath.parse("foo.value");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        InvalidAggregationPathException e = expectThrows(InvalidAggregationPathException.class,
+            () -> resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms"));
+        assertThat(e.getMessage(), equalTo("Cannot find an aggregation named [foo] in [the_long_terms]"));
+    }
+
+    public void testResolveToUnknown() {
+        AggregationPath path = AggregationPath.parse("the_avg.unknown");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        IllegalArgumentException e = expectThrows(IllegalArgumentException.class,
+            () -> resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms"));
+        assertThat(e.getMessage(), equalTo("path not supported for [the_avg]: [unknown]"));
+    }
+
+    public void testResolveToBucketCount() {
+        AggregationPath path = AggregationPath.parse("_bucket_count");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object value = resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value, equalTo(1));
+    }
+
+    public void testResolveToCount() {
+        AggregationPath path = AggregationPath.parse("_count");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(1, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object[] value = (Object[]) resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value[0], equalTo(1L));
+    }
+
+    public void testResolveToKey() {
+        AggregationPath path = AggregationPath.parse("_key");
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(agg));
+
+        LongTerms.Bucket bucket = new LongTerms.Bucket(19, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object[] value = (Object[]) resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value[0], equalTo(19L));
+    }
+
+    public void testResolveToSpecificBucket() {
+        AggregationPath path = AggregationPath.parse("string_terms['foo']>the_avg.value");
+
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalStringAggs = new InternalAggregations(Collections.singletonList(agg));
+        List<StringTerms.Bucket> stringBuckets = Collections.singletonList(new StringTerms.Bucket(
+            new BytesRef("foo".getBytes(StandardCharsets.UTF_8), 0, "foo".getBytes(StandardCharsets.UTF_8).length), 1,
+            internalStringAggs, false, 0, DocValueFormat.RAW));
+
+        InternalTerms termsAgg = new StringTerms("string_terms", BucketOrder.count(false), 1, 0, Collections.emptyList(),
+            Collections.emptyMap(), DocValueFormat.RAW, 1, false, 0, stringBuckets, 0);
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(termsAgg));
+        LongTerms.Bucket bucket = new LongTerms.Bucket(19, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        Object[] value = (Object[]) resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms");
+        assertThat(value[0], equalTo(2.0));
+    }
+
+    public void testResolveToMissingSpecificBucket() {
+        AggregationPath path = AggregationPath.parse("string_terms['bar']>the_avg.value");
+
+        List<LongTerms.Bucket> buckets = new ArrayList<>();
+        InternalAggregation agg = new InternalAvg("the_avg", 2, 1,
+            DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap());
+        InternalAggregations internalStringAggs = new InternalAggregations(Collections.singletonList(agg));
+        List<StringTerms.Bucket> stringBuckets = Collections.singletonList(new StringTerms.Bucket(
+            new BytesRef("foo".getBytes(StandardCharsets.UTF_8), 0, "foo".getBytes(StandardCharsets.UTF_8).length), 1,
+            internalStringAggs, false, 0, DocValueFormat.RAW));
+
+        InternalTerms termsAgg = new StringTerms("string_terms", BucketOrder.count(false), 1, 0, Collections.emptyList(),
+            Collections.emptyMap(), DocValueFormat.RAW, 1, false, 0, stringBuckets, 0);
+        InternalAggregations internalAggregations = new InternalAggregations(Collections.singletonList(termsAgg));
+        LongTerms.Bucket bucket = new LongTerms.Bucket(19, 1, internalAggregations, false, 0, DocValueFormat.RAW);
+        buckets.add(bucket);
+
+        InvalidAggregationPathException e = expectThrows(InvalidAggregationPathException.class,
+            () -> resolvePropertyFromPath(path.getPathElementsAsStringList(), buckets, "the_long_terms"));
+        assertThat(e.getMessage(), equalTo("Cannot find an key ['bar'] in [string_terms]"));
+    }
+}

server/src/test/java/org/elasticsearch/search/aggregations/metrics/ScriptedMetricAggregatorTests.java

@@ -173,6 +173,17 @@ public static void initMockScripts() {
         });
     }
 
+    @Override
+    protected ScriptService getMockScriptService() {
+        MockScriptEngine scriptEngine = new MockScriptEngine(MockScriptEngine.NAME,
+            SCRIPTS,
+            Collections.emptyMap());
+        Map<String, ScriptEngine> engines = Collections.singletonMap(scriptEngine.getType(), scriptEngine);
+
+        return new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
+    }
+
+
     @SuppressWarnings("unchecked")
     public void testNoDocs() throws IOException {
         try (Directory directory = newDirectory()) {
@@ -311,7 +322,7 @@ public void testAggParamsPassedToReduceScript() throws IOException {
                         .initScript(INIT_SCRIPT_PARAMS).mapScript(MAP_SCRIPT_PARAMS)
                         .combineScript(COMBINE_SCRIPT_PARAMS).reduceScript(REDUCE_SCRIPT_PARAMS);
                 ScriptedMetric scriptedMetric = searchAndReduce(
-                        newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder, 0, scriptService);
+                        newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder, 0);
 
                 // The result value depends on the script params.
                 assertEquals(4803, scriptedMetric.aggregation());

server/src/test/java/org/elasticsearch/search/aggregations/pipeline/AvgBucketAggregatorTests.java

@@ -120,9 +120,9 @@ public void testSameAggNames() throws IOException {
                 valueFieldType.setName(VALUE_FIELD);
                 valueFieldType.setHasDocValues(true);
 
-                avgResult = searchAndReduce(indexSearcher, query, avgBuilder, 10000, null,
+                avgResult = searchAndReduce(indexSearcher, query, avgBuilder, 10000,
                     new MappedFieldType[]{fieldType, valueFieldType});
-                histogramResult = searchAndReduce(indexSearcher, query, histo, 10000, null,
+                histogramResult = searchAndReduce(indexSearcher, query, histo, 10000,
                     new MappedFieldType[]{fieldType, valueFieldType});
             }