Aggregations Refactor: Refactor Bucket Selector Aggregation

This commit is contained in:
Colin Goodheart-Smithe 2015-12-01 12:58:41 +00:00
parent 1b89c44cb5
commit 80e58e32a4
3 changed files with 141 additions and 6 deletions

View File

@ -31,6 +31,7 @@ import org.elasticsearch.search.internal.SearchContext;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@ -52,7 +53,7 @@ public class BucketSelectorParser implements PipelineAggregator.Parser {
Script script = null;
String currentFieldName = null;
Map<String, String> bucketsPathsMap = null;
GapPolicy gapPolicy = GapPolicy.SKIP;
GapPolicy gapPolicy = null;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
@ -113,13 +114,18 @@ public class BucketSelectorParser implements PipelineAggregator.Parser {
+ "] for bucket_selector aggregation [" + reducerName + "]", parser.getTokenLocation());
}
return new BucketSelectorPipelineAggregator.Factory(reducerName, bucketsPathsMap, script, gapPolicy);
BucketSelectorPipelineAggregator.Factory factory = new BucketSelectorPipelineAggregator.Factory(reducerName, bucketsPathsMap,
script);
if (gapPolicy != null) {
factory.gapPolicy(gapPolicy);
}
return factory;
}
// NORELEASE implement this method when refactoring this aggregation
@Override
public PipelineAggregatorFactory getFactoryPrototype() {
return null;
return new BucketSelectorPipelineAggregator.Factory(null, Collections.emptyMap(), null);
}
}

View File

@ -22,9 +22,11 @@ package org.elasticsearch.search.aggregations.pipeline.having;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.script.CompiledScript;
import org.elasticsearch.script.ExecutableScript;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.Script.ScriptField;
import org.elasticsearch.script.ScriptContext;
import org.elasticsearch.search.aggregations.InternalAggregation;
import org.elasticsearch.search.aggregations.InternalAggregation.ReduceContext;
@ -35,6 +37,7 @@ import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorFactory;
import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorStreams;
import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptParser;
import java.io.IOException;
import java.util.ArrayList;
@ -42,6 +45,8 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import static org.elasticsearch.search.aggregations.pipeline.BucketHelpers.resolveBucketValue;
@ -137,20 +142,83 @@ public class BucketSelectorPipelineAggregator extends PipelineAggregator {
public static class Factory extends PipelineAggregatorFactory {
private Script script;
private GapPolicy gapPolicy;
private GapPolicy gapPolicy = GapPolicy.SKIP;
private Map<String, String> bucketsPathsMap;
public Factory(String name, Map<String, String> bucketsPathsMap, Script script, GapPolicy gapPolicy) {
public Factory(String name, Map<String, String> bucketsPathsMap, Script script) {
super(name, TYPE.name(), bucketsPathsMap.values().toArray(new String[bucketsPathsMap.size()]));
this.bucketsPathsMap = bucketsPathsMap;
this.script = script;
}
/**
* Sets the gap policy to use for this aggregation.
*/
public void gapPolicy(GapPolicy gapPolicy) {
this.gapPolicy = gapPolicy;
}
/**
* Gets the gap policy to use for this aggregation.
*/
public GapPolicy gapPolicy() {
return gapPolicy;
}
@Override
protected PipelineAggregator createInternal(Map<String, Object> metaData) throws IOException {
return new BucketSelectorPipelineAggregator(name, bucketsPathsMap, script, gapPolicy, metaData);
}
@Override
protected XContentBuilder internalXContent(XContentBuilder builder, Params params) throws IOException {
builder.field(BucketScriptParser.BUCKETS_PATH.getPreferredName(), bucketsPathsMap);
builder.field(ScriptField.SCRIPT.getPreferredName(), script);
builder.field(BucketScriptParser.GAP_POLICY.getPreferredName(), gapPolicy.getName());
return builder;
}
@Override
protected boolean overrideBucketsPath() {
return true;
}
@Override
protected PipelineAggregatorFactory doReadFrom(String name, String[] bucketsPaths, StreamInput in) throws IOException {
Map<String, String> bucketsPathsMap = new HashMap<String, String>();
int mapSize = in.readVInt();
for (int i = 0; i < mapSize; i++) {
bucketsPathsMap.put(in.readString(), in.readString());
}
Script script = Script.readScript(in);
Factory factory = new Factory(name, bucketsPathsMap, script);
factory.gapPolicy = GapPolicy.readFrom(in);
return factory;
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
out.writeVInt(bucketsPathsMap.size());
for (Entry<String, String> e : bucketsPathsMap.entrySet()) {
out.writeString(e.getKey());
out.writeString(e.getValue());
}
script.writeTo(out);
gapPolicy.writeTo(out);
}
@Override
protected int doHashCode() {
return Objects.hash(bucketsPathsMap, script, gapPolicy);
}
@Override
protected boolean doEquals(Object obj) {
Factory other = (Factory) obj;
return Objects.equals(bucketsPathsMap, other.bucketsPathsMap) && Objects.equals(script, other.script)
&& Objects.equals(gapPolicy, other.gapPolicy);
}
}
}

View File

@ -0,0 +1,61 @@
/*
* 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.pipeline;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptService.ScriptType;
import org.elasticsearch.search.aggregations.BasePipelineAggregationTestCase;
import org.elasticsearch.search.aggregations.pipeline.BucketHelpers.GapPolicy;
import org.elasticsearch.search.aggregations.pipeline.having.BucketSelectorPipelineAggregator;
import org.elasticsearch.search.aggregations.pipeline.having.BucketSelectorPipelineAggregator.Factory;
import java.util.HashMap;
import java.util.Map;
public class BucketSelectorTests extends BasePipelineAggregationTestCase<BucketSelectorPipelineAggregator.Factory> {
@Override
protected Factory createTestAggregatorFactory() {
String name = randomAsciiOfLengthBetween(3, 20);
Map<String, String> bucketsPaths = new HashMap<>();
int numBucketPaths = randomIntBetween(1, 10);
for (int i = 0; i < numBucketPaths; i++) {
bucketsPaths.put(randomAsciiOfLengthBetween(1, 20), randomAsciiOfLengthBetween(1, 40));
}
Script script;
if (randomBoolean()) {
script = new Script("script");
} else {
Map<String, Object> params = null;
if (randomBoolean()) {
params = new HashMap<String, Object>();
params.put("foo", "bar");
}
script = new Script("script", randomFrom(ScriptType.values()), randomFrom("my_lang", null), params);
}
Factory factory = new Factory(name, bucketsPaths, script);
if (randomBoolean()) {
factory.gapPolicy(randomFrom(GapPolicy.values()));
}
return factory;
}
}