From 85723153ac2d59546a70b2f3f4bb69687356724d Mon Sep 17 00:00:00 2001 From: Nikita_Glashchenko Date: Thu, 18 Jul 2019 21:06:22 +0400 Subject: [PATCH] Allow BucketAggregationScript to return strings --- .../expression/ExpressionScriptEngine.java | 2 +- .../script/BucketAggregationScript.java | 2 +- .../BucketScriptPipelineAggregator.java | 13 ++-- .../pipeline/InternalStringValue.java | 61 +++++++++++++++++++ 4 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/search/aggregations/pipeline/InternalStringValue.java diff --git a/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java index abd8738b0e4c1..b996b61135b8c 100644 --- a/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java +++ b/modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionScriptEngine.java @@ -109,7 +109,7 @@ protected Class loadClass(String name, boolean resolve) throws ClassNotFoundE BucketAggregationSelectorScript.Factory wrappedFactory = parameters -> new BucketAggregationSelectorScript(parameters) { @Override public boolean execute() { - return factory.newInstance(getParams()).execute().doubleValue() == 1.0; + return ((Number)factory.newInstance(getParams()).execute()).doubleValue() == 1.0; } }; return context.factoryClazz.cast(wrappedFactory); diff --git a/server/src/main/java/org/elasticsearch/script/BucketAggregationScript.java b/server/src/main/java/org/elasticsearch/script/BucketAggregationScript.java index 76ff776353ef2..bdfc339ae0f4f 100644 --- a/server/src/main/java/org/elasticsearch/script/BucketAggregationScript.java +++ b/server/src/main/java/org/elasticsearch/script/BucketAggregationScript.java @@ -46,7 +46,7 @@ public Map getParams() { return params; } - public abstract Number execute(); + public abstract Object execute(); public interface Factory { BucketAggregationScript newInstance(Map params); diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java index 996c9b2f0e12f..f3a303e1a2787 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java @@ -108,16 +108,21 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext if (skipBucket) { newBuckets.add(bucket); } else { - Number returned = factory.newInstance(vars).execute(); + Object returned = factory.newInstance(vars).execute(); if (returned == null) { newBuckets.add(bucket); } else { final List aggs = StreamSupport.stream(bucket.getAggregations().spliterator(), false).map( (p) -> (InternalAggregation) p).collect(Collectors.toList()); - InternalSimpleValue simpleValue = new InternalSimpleValue(name(), returned.doubleValue(), - formatter, new ArrayList<>(), metaData()); - aggs.add(simpleValue); + InternalAggregation value; + if (returned instanceof Number) { + value = new InternalSimpleValue(name(), ((Number)returned).doubleValue(), + formatter, new ArrayList<>(), metaData()); + } else { + value = new InternalStringValue(name(), returned.toString(), new ArrayList<>(), metaData()); + } + aggs.add(value); InternalMultiBucketAggregation.InternalBucket newBucket = originalAgg.createBucket(new InternalAggregations(aggs), bucket); newBuckets.add(newBucket); diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/InternalStringValue.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/InternalStringValue.java new file mode 100644 index 0000000000000..0cee91e8e272f --- /dev/null +++ b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/InternalStringValue.java @@ -0,0 +1,61 @@ +package org.elasticsearch.search.aggregations.pipeline; + +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.search.aggregations.InternalAggregation; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +public class InternalStringValue extends InternalAggregation { + public static final String NAME = "string_value"; + + private final String text; + + protected InternalStringValue(String name, String text, List pipelineAggregators, Map metaData) { + super(name, pipelineAggregators, metaData); + this.text = text; + } + + @Override + protected void doWriteTo(StreamOutput out) throws IOException { + throw new UnsupportedOperationException("not supported yet"); + } + + @Override + public InternalAggregation doReduce(List aggregations, ReduceContext reduceContext) { + throw new UnsupportedOperationException("Not supported"); + } + + @Override + public Object getProperty(List path) { + if (path.isEmpty()) { + return this; + } else if (path.size() == 1 && "value".equals(path.get(0))) { + return text(); + } else { + throw new IllegalArgumentException("path not supported for [" + getName() + "]: " + path); + } + } + + @Override + public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException { + boolean hasValue = text != null; + builder.field(CommonFields.VALUE.getPreferredName(), hasValue ? text : null); + if (hasValue) { + builder.field(CommonFields.VALUE_AS_STRING.getPreferredName(), text); + } + return builder; + } + + @Override + public String getWriteableName() { + return NAME; + } + + public String text() { + return text; + } +}