Skip to content

Commit 1d7fc6b

Browse files
colings86colings86
committed
Aggregations: Pipeline Aggregation to filter buckets based on a script
This pipeline aggregation runs a script on each bucket in the parent aggregation to determine whether the bucket is kept in the final aggregation tree. If the script returns true the bucket is retained, if it returns false the bucket is dropped
1 parent b612cab commit 1d7fc6b

File tree

10 files changed

+945
-1
lines changed

10 files changed

+945
-1
lines changed

‎core/src/main/java/org/elasticsearch/search/aggregations/AggregationModule.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
import org.elasticsearch.search.aggregations.pipeline.cumulativesum.CumulativeSumParser;
6565
import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptParser;
6666
import org.elasticsearch.search.aggregations.pipeline.derivative.DerivativeParser;
67+
import org.elasticsearch.search.aggregations.pipeline.having.BucketSelectorParser;
6768
import org.elasticsearch.search.aggregations.pipeline.movavg.MovAvgParser;
6869
import org.elasticsearch.search.aggregations.pipeline.movavg.models.MovAvgModelModule;
6970

@@ -118,6 +119,7 @@ public AggregationModule() {
118119
pipelineAggParsers.add(MovAvgParser.class);
119120
pipelineAggParsers.add(CumulativeSumParser.class);
120121
pipelineAggParsers.add(BucketScriptParser.class);
122+
pipelineAggParsers.add(BucketSelectorParser.class);
121123
}
122124

123125
/**

‎core/src/main/java/org/elasticsearch/search/aggregations/TransportAggregationModule.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptPipelineAggregator;
7070
import org.elasticsearch.search.aggregations.pipeline.derivative.DerivativePipelineAggregator;
7171
import org.elasticsearch.search.aggregations.pipeline.derivative.InternalDerivative;
72+
import org.elasticsearch.search.aggregations.pipeline.having.BucketSelectorPipelineAggregator;
7273
import org.elasticsearch.search.aggregations.pipeline.movavg.MovAvgPipelineAggregator;
7374
import org.elasticsearch.search.aggregations.pipeline.movavg.models.TransportMovAvgModelModule;
7475

@@ -131,6 +132,7 @@ protected void configure() {
131132
MovAvgPipelineAggregator.registerStreams();
132133
CumulativeSumPipelineAggregator.registerStreams();
133134
BucketScriptPipelineAggregator.registerStreams();
135+
BucketSelectorPipelineAggregator.registerStreams();
134136
}
135137

136138
@Override

‎core/src/main/java/org/elasticsearch/search/aggregations/pipeline/PipelineAggregatorBuilders.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.elasticsearch.search.aggregations.pipeline.cumulativesum.CumulativeSumBuilder;
2727
import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptBuilder;
2828
import org.elasticsearch.search.aggregations.pipeline.derivative.DerivativeBuilder;
29+
import org.elasticsearch.search.aggregations.pipeline.having.BucketSelectorBuilder;
2930
import org.elasticsearch.search.aggregations.pipeline.movavg.MovAvgBuilder;
3031

3132
public final class PipelineAggregatorBuilders {
@@ -61,6 +62,10 @@ public static final BucketScriptBuilder bucketScript(String name) {
6162
return new BucketScriptBuilder(name);
6263
}
6364

65+
public static final BucketSelectorBuilder having(String name) {
66+
return new BucketSelectorBuilder(name);
67+
}
68+
6469
public static final CumulativeSumBuilder cumulativeSum(String name) {
6570
return new CumulativeSumBuilder(name);
6671
}

‎core/src/main/java/org/elasticsearch/search/aggregations/pipeline/having/BucketSelectorBuilder.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.search.aggregations.pipeline.having;
21+
22+
import org.elasticsearch.common.xcontent.XContentBuilder;
23+
import org.elasticsearch.script.Script;
24+
import org.elasticsearch.script.Script.ScriptField;
25+
import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy;
26+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
27+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorBuilder;
28+
29+
import java.io.IOException;
30+
import java.util.Map;
31+
32+
public class BucketSelectorBuilder extends PipelineAggregatorBuilder<BucketSelectorBuilder> {
33+
34+
private GapPolicy gapPolicy;
35+
private Script script;
36+
private Map<String, String> bucketsPathsMap;
37+
38+
public BucketSelectorBuilder(String name) {
39+
super(name, BucketSelectorPipelineAggregator.TYPE.name());
40+
}
41+
42+
public BucketSelectorBuilder script(Script script) {
43+
this.script = script;
44+
return this;
45+
}
46+
47+
public BucketSelectorBuilder gapPolicy(GapPolicy gapPolicy) {
48+
this.gapPolicy = gapPolicy;
49+
return this;
50+
}
51+
52+
/**
53+
* Sets the paths to the buckets to use for this pipeline aggregator. The
54+
* map given to this method must contain script variable name as keys with
55+
* bucket paths values to the metrics to use for each variable.
56+
*/
57+
public BucketSelectorBuilder setBucketsPathsMap(Map<String, String> bucketsPathsMap) {
58+
this.bucketsPathsMap = bucketsPathsMap;
59+
return this;
60+
}
61+
62+
@Override
63+
protected XContentBuilder internalXContent(XContentBuilder builder, Params builderParams) throws IOException {
64+
if (script != null) {
65+
builder.field(ScriptField.SCRIPT.getPreferredName(), script);
66+
}
67+
if (gapPolicy != null) {
68+
builder.field(BucketSelectorParser.GAP_POLICY.getPreferredName(), gapPolicy.getName());
69+
}
70+
if (bucketsPathsMap != null) {
71+
builder.field(PipelineAggregator.Parser.BUCKETS_PATH.getPreferredName(), bucketsPathsMap);
72+
}
73+
return builder;
74+
}
75+
76+
}

‎core/src/main/java/org/elasticsearch/search/aggregations/pipeline/having/BucketSelectorParser.java‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.elasticsearch.search.aggregations.pipeline.having;
21+
22+
import org.elasticsearch.common.ParseField;
23+
import org.elasticsearch.common.xcontent.XContentParser;
24+
import org.elasticsearch.script.Script;
25+
import org.elasticsearch.script.Script.ScriptField;
26+
import org.elasticsearch.search.SearchParseException;
27+
import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy;
28+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
29+
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorFactory;
30+
import org.elasticsearch.search.internal.SearchContext;
31+
32+
import java.io.IOException;
33+
import java.util.ArrayList;
34+
import java.util.HashMap;
35+
import java.util.List;
36+
import java.util.Map;
37+
38+
public class BucketSelectorParser implements PipelineAggregator.Parser {
39+
40+
public static final ParseField FORMAT = new ParseField("format");
41+
public static final ParseField GAP_POLICY = new ParseField("gap_policy");
42+
public static final ParseField PARAMS_FIELD = new ParseField("params");
43+
44+
@Override
45+
public String type() {
46+
return BucketSelectorPipelineAggregator.TYPE.name();
47+
}
48+
49+
@Override
50+
public PipelineAggregatorFactory parse(String reducerName, XContentParser parser, SearchContext context) throws IOException {
51+
XContentParser.Token token;
52+
Script script = null;
53+
String currentFieldName = null;
54+
Map<String, String> bucketsPathsMap = null;
55+
GapPolicy gapPolicy = GapPolicy.SKIP;
56+
57+
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
58+
if (token == XContentParser.Token.FIELD_NAME) {
59+
currentFieldName = parser.currentName();
60+
} else if (token == XContentParser.Token.VALUE_STRING) {
61+
if (context.parseFieldMatcher().match(currentFieldName, BUCKETS_PATH)) {
62+
bucketsPathsMap = new HashMap<>();
63+
bucketsPathsMap.put("_value", parser.text());
64+
} else if (context.parseFieldMatcher().match(currentFieldName, GAP_POLICY)) {
65+
gapPolicy = GapPolicy.parse(context, parser.text(), parser.getTokenLocation());
66+
} else if (context.parseFieldMatcher().match(currentFieldName, ScriptField.SCRIPT)) {
67+
script = Script.parse(parser, context.parseFieldMatcher());
68+
} else {
69+
throw new SearchParseException(context, "Unknown key for a " + token + " in [" + reducerName + "]: ["
70+
+ currentFieldName + "].", parser.getTokenLocation());
71+
}
72+
} else if (token == XContentParser.Token.START_ARRAY) {
73+
if (context.parseFieldMatcher().match(currentFieldName, BUCKETS_PATH)) {
74+
List<String> paths = new ArrayList<>();
75+
while ((token = parser.nextToken()) != XContentParser.Token.END_ARRAY) {
76+
String path = parser.text();
77+
paths.add(path);
78+
}
79+
bucketsPathsMap = new HashMap<>();
80+
for (int i = 0; i < paths.size(); i++) {
81+
bucketsPathsMap.put("_value" + i, paths.get(i));
82+
}
83+
} else {
84+
throw new SearchParseException(context, "Unknown key for a " + token + " in [" + reducerName + "]: ["
85+
+ currentFieldName + "].", parser.getTokenLocation());
86+
}
87+
} else if (token == XContentParser.Token.START_OBJECT) {
88+
if (context.parseFieldMatcher().match(currentFieldName, ScriptField.SCRIPT)) {
89+
script = Script.parse(parser, context.parseFieldMatcher());
90+
} else if (context.parseFieldMatcher().match(currentFieldName, BUCKETS_PATH)) {
91+
Map<String, Object> map = parser.map();
92+
bucketsPathsMap = new HashMap<>();
93+
for (Map.Entry<String, Object> entry : map.entrySet()) {
94+
bucketsPathsMap.put(entry.getKey(), String.valueOf(entry.getValue()));
95+
}
96+
} else {
97+
throw new SearchParseException(context, "Unknown key for a " + token + " in [" + reducerName + "]: ["
98+
+ currentFieldName + "].", parser.getTokenLocation());
99+
}
100+
} else {
101+
throw new SearchParseException(context, "Unexpected token " + token + " in [" + reducerName + "].",
102+
parser.getTokenLocation());
103+
}
104+
}
105+
106+
if (bucketsPathsMap == null) {
107+
throw new SearchParseException(context, "Missing required field [" + BUCKETS_PATH.getPreferredName()
108+
+ "] for bucket_selector aggregation [" + reducerName + "]", parser.getTokenLocation());
109+
}
110+
111+
if (script == null) {
112+
throw new SearchParseException(context, "Missing required field [" + ScriptField.SCRIPT.getPreferredName()
113+
+ "] for bucket_selector aggregation [" + reducerName + "]", parser.getTokenLocation());
114+
}
115+
116+
return new BucketSelectorPipelineAggregator.Factory(reducerName, bucketsPathsMap, script, gapPolicy);
117+
}
118+
119+
}

0 commit comments

Comments
 (0)