Aggregations Refactor: Refactor Bucket Script Aggregation

core/src/main/java/org/elasticsearch/search/aggregations/pipeline/PipelineAggregatorFactory.java

@@ -143,7 +143,7 @@ public abstract class PipelineAggregatorFactory extends ToXContentToBytes implem
         }
         builder.startObject(type);
 
-        if (bucketsPaths != null) {
+        if (!overrideBucketsPath() && bucketsPaths != null) {
             builder.startArray(PipelineAggregator.Parser.BUCKETS_PATH.getPreferredName());
             for (String path : bucketsPaths) {
                 builder.value(path);
@@ -158,6 +158,14 @@ public abstract class PipelineAggregatorFactory extends ToXContentToBytes implem
         return builder.endObject();
     }
 
+    /**
+     * @return <code>true</code> if the {@link PipelineAggregatorFactory}
+     *         overrides the XContent rendering of the bucketPath option.
+     */
+    protected boolean overrideBucketsPath() {
+        return false;
+    }
+
     // NORELEASE make this method abstract when agg refactor complete
     protected XContentBuilder internalXContent(XContentBuilder builder, Params params) throws IOException {
         return builder;

core/src/main/java/org/elasticsearch/search/aggregations/pipeline/bucketscript/BucketScriptParser.java

@@ -27,12 +27,11 @@ import org.elasticsearch.search.SearchParseException;
 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.support.format.ValueFormat;
-import org.elasticsearch.search.aggregations.support.format.ValueFormatter;
 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;
@@ -55,7 +54,7 @@ public class BucketScriptParser implements PipelineAggregator.Parser {
         String currentFieldName = null;
         Map<String, String> bucketsPathsMap = null;
         String format = null;
-        GapPolicy gapPolicy = GapPolicy.SKIP;
+        GapPolicy gapPolicy = null;
 
         while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
             if (token == XContentParser.Token.FIELD_NAME) {
@@ -118,20 +117,19 @@ public class BucketScriptParser implements PipelineAggregator.Parser {
                     + "] for series_arithmetic aggregation [" + reducerName + "]", parser.getTokenLocation());
         }
 
-        ValueFormatter formatter = null;
+        BucketScriptPipelineAggregator.Factory factory = new BucketScriptPipelineAggregator.Factory(reducerName, bucketsPathsMap, script);
         if (format != null) {
-            formatter = ValueFormat.Patternable.Number.format(format).formatter();
+            factory.format(format);
-        } else {
-            formatter = ValueFormatter.RAW;
         }
-
+        if (gapPolicy != null) {
-        return new BucketScriptPipelineAggregator.Factory(reducerName, bucketsPathsMap, script, formatter, gapPolicy);
+            factory.gapPolicy(gapPolicy);
+        }
+        return factory;
     }
 
-    // NORELEASE implement this method when refactoring this aggregation
     @Override
     public PipelineAggregatorFactory getFactoryPrototype() {
-        return null;
+        return new BucketScriptPipelineAggregator.Factory(null, Collections.emptyMap(), null);
     }
 
 }

core/src/main/java/org/elasticsearch/search/aggregations/pipeline/bucketscript/BucketScriptPipelineAggregator.java

@@ -21,9 +21,11 @@ package org.elasticsearch.search.aggregations.pipeline.bucketscript;
 
 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.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -37,15 +39,17 @@ import org.elasticsearch.search.aggregations.pipeline.InternalSimpleValue;
 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.support.format.ValueFormat;
 import org.elasticsearch.search.aggregations.support.format.ValueFormatter;
 import org.elasticsearch.search.aggregations.support.format.ValueFormatterStreams;
 
 import java.io.IOException;
 import java.util.ArrayList;
-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 java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 
@@ -157,22 +161,110 @@ public class BucketScriptPipelineAggregator extends PipelineAggregator {
 
     public static class Factory extends PipelineAggregatorFactory {
 
-        private Script script;
+        private final Script script;
-        private final ValueFormatter formatter;
+        private final Map<String, String> bucketsPathsMap;
-        private GapPolicy gapPolicy;
+        private String format = null;
-        private Map<String, String> bucketsPathsMap;
+        private GapPolicy gapPolicy = GapPolicy.SKIP;
 
-        public Factory(String name, Map<String, String> bucketsPathsMap, Script script, ValueFormatter formatter, 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;
-            this.formatter = formatter;
+        }
+
+        /**
+         * Sets the format to use on the output of this aggregation.
+         */
+        public void format(String format) {
+            this.format = format;
+        }
+
+        /**
+         * Gets the format to use on the output of this aggregation.
+         */
+        public String format() {
+            return format;
+        }
+
+        protected ValueFormatter formatter() {
+            if (format != null) {
+                return ValueFormat.Patternable.Number.format(format).formatter();
+            } else {
+                return ValueFormatter.RAW;
+            }
+        }
+
+        /**
+         * 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 BucketScriptPipelineAggregator(name, bucketsPathsMap, script, formatter, gapPolicy, metaData);
+            return new BucketScriptPipelineAggregator(name, bucketsPathsMap, script, formatter(), 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);
+            if (format != null) {
+                builder.field(BucketScriptParser.FORMAT.getPreferredName(), format);
+            }
+            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.format = in.readOptionalString();
+            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);
+            out.writeOptionalString(format);
+            gapPolicy.writeTo(out);
+        }
+
+        @Override
+        protected int doHashCode() {
+            return Objects.hash(bucketsPathsMap, script, format, gapPolicy);
+        }
+
+        @Override
+        protected boolean doEquals(Object obj) {
+            Factory other = (Factory) obj;
+            return Objects.equals(bucketsPathsMap, other.bucketsPathsMap) && Objects.equals(script, other.script)
+                    && Objects.equals(format, other.format) && Objects.equals(gapPolicy, other.gapPolicy);
         }
     }
 

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

@@ -0,0 +1,63 @@
+/*
+ * 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.bucketscript.BucketScriptPipelineAggregator;
+import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptPipelineAggregator.Factory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class BucketScriptTests extends BasePipelineAggregationTestCase<BucketScriptPipelineAggregator.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.format(randomAsciiOfLengthBetween(1, 10));
+        }
+        if (randomBoolean()) {
+            factory.gapPolicy(randomFrom(GapPolicy.values()));
+        }
+        return factory;
+    }
+
+}