From 05ec970723534940f6d18a1a568687b543c8f720 Mon Sep 17 00:00:00 2001 From: Ignacio Vera Date: Tue, 23 Jul 2019 09:05:47 +0200 Subject: [PATCH] Support BucketScript paths of type string and array. (#44694) (#44731) --- ...ucketScriptPipelineAggregationBuilder.java | 27 ++- .../aggregations/pipeline/BucketScriptIT.java | 161 ++++++++++++++++++ .../pipeline/BucketScriptTests.java | 46 +++++ 3 files changed, 233 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregationBuilder.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregationBuilder.java index 97898b13cc2..ba2114c1152 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregationBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregationBuilder.java @@ -30,6 +30,7 @@ import org.elasticsearch.search.DocValueFormat; import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy; import java.io.IOException; +import java.util.Collections; import java.util.HashMap; import java.util.Locale; import java.util.Map; @@ -59,7 +60,10 @@ public class BucketScriptPipelineAggregationBuilder extends AbstractPipelineAggr false, o -> new BucketScriptPipelineAggregationBuilder(name, (Map) o[0], (Script) o[1])); - parser.declareObject(ConstructingObjectParser.constructorArg(), (p, c) -> p.mapStrings(), BUCKETS_PATH_FIELD); + parser.declareField(ConstructingObjectParser.constructorArg() + , BucketScriptPipelineAggregationBuilder::extractBucketPath + , BUCKETS_PATH_FIELD + , ObjectParser.ValueType.OBJECT_ARRAY_OR_STRING); parser.declareField(ConstructingObjectParser.constructorArg(), (p, c) -> Script.parse(p), Script.SCRIPT_PARSE_FIELD, ObjectParser.ValueType.OBJECT_OR_STRING); @@ -112,6 +116,27 @@ public class BucketScriptPipelineAggregationBuilder extends AbstractPipelineAggr gapPolicy.writeTo(out); } + private static Map extractBucketPath(XContentParser parser) throws IOException { + XContentParser.Token token = parser.currentToken(); + if (token == XContentParser.Token.VALUE_STRING) { + // input is a string, name of the path set to '_value'. + // This is a bit odd as there is not constructor for it + return Collections.singletonMap("_value", parser.text()); + } else if (token == XContentParser.Token.START_ARRAY) { + // input is an array, name of the path set to '_value' + position + Map bucketsPathsMap = new HashMap<>(); + int i =0; + while ((parser.nextToken()) != XContentParser.Token.END_ARRAY) { + String path = parser.text(); + bucketsPathsMap.put("_value" + i++, path); + } + return bucketsPathsMap; + } else { + // input is an object, it should contain name / value pairs + return parser.mapStrings(); + } + } + private static Map convertToBucketsPathMap(String[] bucketsPaths) { Map bucketsPathsMap = new HashMap<>(); for (int i = 0; i < bucketsPaths.length; i++) { diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptIT.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptIT.java index 3677da5df9b..88aea1532f3 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptIT.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptIT.java @@ -23,6 +23,7 @@ import org.elasticsearch.action.index.IndexRequestBuilder; import org.elasticsearch.action.search.SearchResponse; import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; import org.elasticsearch.common.xcontent.XContentType; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.script.MockScriptPlugin; @@ -117,6 +118,11 @@ public class BucketScriptIT extends ESIntegTestCase { return value0 + value1 + value2; }); + scripts.put("single_input", vars -> { + double value = (double) vars.get("_value"); + return value; + }); + scripts.put("return null", vars -> null); return scripts; @@ -628,4 +634,159 @@ public class BucketScriptIT extends ESIntegTestCase { } } } + + public void testSingleBucketPathAgg() throws Exception { + XContentBuilder content = XContentFactory.jsonBuilder() + .startObject() + .field("buckets_path", "field2Sum") + .startObject("script") + .field("source", "single_input") + .field("lang", CustomScriptPlugin.NAME) + .endObject() + .endObject(); + BucketScriptPipelineAggregationBuilder bucketScriptAgg = + BucketScriptPipelineAggregationBuilder.parse("seriesArithmetic", createParser(content)); + + SearchResponse response = client() + .prepareSearch("idx", "idx_unmapped") + .addAggregation( + histogram("histo") + .field(FIELD_1_NAME) + .interval(interval) + .subAggregation(sum("field2Sum").field(FIELD_2_NAME)) + .subAggregation(bucketScriptAgg)).get(); + + assertSearchResponse(response); + + Histogram histo = response.getAggregations().get("histo"); + assertThat(histo, notNullValue()); + assertThat(histo.getName(), equalTo("histo")); + List buckets = histo.getBuckets(); + + for (int i = 0; i < buckets.size(); ++i) { + Histogram.Bucket bucket = buckets.get(i); + if (bucket.getDocCount() == 0) { + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, nullValue()); + } else { + Sum field2Sum = bucket.getAggregations().get("field2Sum"); + assertThat(field2Sum, notNullValue()); + double field2SumValue = field2Sum.getValue(); + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, notNullValue()); + double seriesArithmeticValue = seriesArithmetic.value(); + assertThat(seriesArithmeticValue, equalTo(field2SumValue)); + } + } + } + + public void testArrayBucketPathAgg() throws Exception { + XContentBuilder content = XContentFactory.jsonBuilder() + .startObject() + .array("buckets_path", "field2Sum", "field3Sum", "field4Sum") + .startObject("script") + .field("source", "_value0 + _value1 + _value2") + .field("lang", CustomScriptPlugin.NAME) + .endObject() + .endObject(); + BucketScriptPipelineAggregationBuilder bucketScriptAgg = + BucketScriptPipelineAggregationBuilder.parse("seriesArithmetic", createParser(content)); + + SearchResponse response = client() + .prepareSearch("idx", "idx_unmapped") + .addAggregation( + histogram("histo") + .field(FIELD_1_NAME) + .interval(interval) + .subAggregation(sum("field2Sum").field(FIELD_2_NAME)) + .subAggregation(sum("field3Sum").field(FIELD_3_NAME)) + .subAggregation(sum("field4Sum").field(FIELD_4_NAME)) + .subAggregation(bucketScriptAgg)).get(); + + assertSearchResponse(response); + + Histogram histo = response.getAggregations().get("histo"); + assertThat(histo, notNullValue()); + assertThat(histo.getName(), equalTo("histo")); + List buckets = histo.getBuckets(); + + for (int i = 0; i < buckets.size(); ++i) { + Histogram.Bucket bucket = buckets.get(i); + if (bucket.getDocCount() == 0) { + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, nullValue()); + } else { + Sum field2Sum = bucket.getAggregations().get("field2Sum"); + assertThat(field2Sum, notNullValue()); + double field2SumValue = field2Sum.getValue(); + Sum field3Sum = bucket.getAggregations().get("field3Sum"); + assertThat(field3Sum, notNullValue()); + double field3SumValue = field3Sum.getValue(); + Sum field4Sum = bucket.getAggregations().get("field4Sum"); + assertThat(field4Sum, notNullValue()); + double field4SumValue = field4Sum.getValue(); + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, notNullValue()); + double seriesArithmeticValue = seriesArithmetic.value(); + assertThat(seriesArithmeticValue, equalTo(field2SumValue + field3SumValue + field4SumValue)); + } + } + } + + public void testObjectBucketPathAgg() throws Exception { + XContentBuilder content = XContentFactory.jsonBuilder() + .startObject() + .startObject("buckets_path") + .field("_value0", "field2Sum") + .field("_value1", "field3Sum") + .field("_value2", "field4Sum") + .endObject() + .startObject("script") + .field("source", "_value0 + _value1 + _value2") + .field("lang", CustomScriptPlugin.NAME) + .endObject() + .endObject(); + BucketScriptPipelineAggregationBuilder bucketScriptAgg = + BucketScriptPipelineAggregationBuilder.parse("seriesArithmetic", createParser(content)); + + SearchResponse response = client() + .prepareSearch("idx", "idx_unmapped") + .addAggregation( + histogram("histo") + .field(FIELD_1_NAME) + .interval(interval) + .subAggregation(sum("field2Sum").field(FIELD_2_NAME)) + .subAggregation(sum("field3Sum").field(FIELD_3_NAME)) + .subAggregation(sum("field4Sum").field(FIELD_4_NAME)) + .subAggregation(bucketScriptAgg)).get(); + + assertSearchResponse(response); + + Histogram histo = response.getAggregations().get("histo"); + assertThat(histo, notNullValue()); + assertThat(histo.getName(), equalTo("histo")); + List buckets = histo.getBuckets(); + + for (int i = 0; i < buckets.size(); ++i) { + Histogram.Bucket bucket = buckets.get(i); + if (bucket.getDocCount() == 0) { + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, nullValue()); + } else { + Sum field2Sum = bucket.getAggregations().get("field2Sum"); + assertThat(field2Sum, notNullValue()); + double field2SumValue = field2Sum.getValue(); + Sum field3Sum = bucket.getAggregations().get("field3Sum"); + assertThat(field3Sum, notNullValue()); + double field3SumValue = field3Sum.getValue(); + Sum field4Sum = bucket.getAggregations().get("field4Sum"); + assertThat(field4Sum, notNullValue()); + double field4SumValue = field4Sum.getValue(); + SimpleValue seriesArithmetic = bucket.getAggregations().get("seriesArithmetic"); + assertThat(seriesArithmetic, notNullValue()); + double seriesArithmeticValue = seriesArithmetic.value(); + assertThat(seriesArithmeticValue, equalTo(field2SumValue + field3SumValue + field4SumValue)); + } + } + } } diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptTests.java index 20684b6383f..0746fa1782f 100644 --- a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptTests.java +++ b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptTests.java @@ -19,11 +19,14 @@ package org.elasticsearch.search.aggregations.pipeline; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptType; import org.elasticsearch.search.aggregations.BasePipelineAggregationTestCase; import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy; +import java.io.IOException; import java.util.HashMap; import java.util.Map; @@ -59,4 +62,47 @@ public class BucketScriptTests extends BasePipelineAggregationTestCase