Aggregations Refactor: Refactor Bucket Selector Aggregation

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

@@ -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);
     }
 
 }

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

@@ -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);
+        }
+
     }
 
 }

core/src/test/java/org/elasticsearch/search/aggregations/pipeline/BucketSelectorTests.java

@@ -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;
+    }
+
+
+}