Aggregations Refactor: Refactor Sampler Aggregation

core/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -21,7 +21,6 @@
 
 import org.elasticsearch.common.inject.AbstractModule;
 import org.elasticsearch.common.inject.multibindings.Multibinder;
-import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.query.functionscore.ScoreFunctionParser;
 import org.elasticsearch.index.query.functionscore.ScoreFunctionParserMapper;
 import org.elasticsearch.search.action.SearchServiceTransportAction;
@@ -56,6 +55,7 @@
 import org.elasticsearch.search.aggregations.bucket.range.geodistance.InternalGeoDistance;
 import org.elasticsearch.search.aggregations.bucket.range.ipv4.InternalIPv4Range;
 import org.elasticsearch.search.aggregations.bucket.range.ipv4.IpRangeParser;
+import org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedSamplerParser;
 import org.elasticsearch.search.aggregations.bucket.sampler.InternalSampler;
 import org.elasticsearch.search.aggregations.bucket.sampler.SamplerParser;
 import org.elasticsearch.search.aggregations.bucket.sampler.UnmappedSampler;
@@ -65,8 +65,11 @@
 import org.elasticsearch.search.aggregations.bucket.significant.UnmappedSignificantTerms;
 import org.elasticsearch.search.aggregations.bucket.significant.heuristics.SignificanceHeuristicParser;
 import org.elasticsearch.search.aggregations.bucket.significant.heuristics.SignificanceHeuristicParserMapper;
-import org.elasticsearch.search.aggregations.bucket.significant.heuristics.SignificanceHeuristicStreams;
-import org.elasticsearch.search.aggregations.bucket.terms.*;
+import org.elasticsearch.search.aggregations.bucket.terms.DoubleTerms;
+import org.elasticsearch.search.aggregations.bucket.terms.LongTerms;
+import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
+import org.elasticsearch.search.aggregations.bucket.terms.TermsParser;
+import org.elasticsearch.search.aggregations.bucket.terms.UnmappedTerms;
 import org.elasticsearch.search.aggregations.metrics.avg.AvgParser;
 import org.elasticsearch.search.aggregations.metrics.avg.InternalAvg;
 import org.elasticsearch.search.aggregations.metrics.cardinality.CardinalityParser;
@@ -127,7 +130,6 @@
 import org.elasticsearch.search.aggregations.pipeline.movavg.MovAvgPipelineAggregator;
 import org.elasticsearch.search.aggregations.pipeline.movavg.models.MovAvgModel;
 import org.elasticsearch.search.aggregations.pipeline.movavg.models.MovAvgModelParserMapper;
-import org.elasticsearch.search.aggregations.pipeline.movavg.models.MovAvgModelStreams;
 import org.elasticsearch.search.aggregations.pipeline.serialdiff.SerialDiffParser;
 import org.elasticsearch.search.aggregations.pipeline.serialdiff.SerialDiffPipelineAggregator;
 import org.elasticsearch.search.controller.SearchPhaseController;
@@ -263,6 +265,7 @@ protected void configureAggs() {
         multibinderAggParser.addBinding().to(FilterParser.class);
         multibinderAggParser.addBinding().to(FiltersParser.class);
         multibinderAggParser.addBinding().to(SamplerParser.class);
+        multibinderAggParser.addBinding().to(DiversifiedSamplerParser.class);
         multibinderAggParser.addBinding().to(TermsParser.class);
         multibinderAggParser.addBinding().to(SignificantTermsParser.class);
         multibinderAggParser.addBinding().to(RangeParser.class);
@@ -385,7 +388,7 @@ protected void configureSearch() {
         SumBucketPipelineAggregator.registerStreams();
         StatsBucketPipelineAggregator.registerStreams();
         ExtendedStatsBucketPipelineAggregator.registerStreams();
-        PercentilesBucketPipelineAggregator.registerStreams();        
+        PercentilesBucketPipelineAggregator.registerStreams();
         MovAvgPipelineAggregator.registerStreams();
         CumulativeSumPipelineAggregator.registerStreams();
         BucketScriptPipelineAggregator.registerStreams();

core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerAggregationBuilder.java

@@ -0,0 +1,79 @@
+/*
+ * 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.bucket.sampler;
+
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.search.aggregations.ValuesSourceAggregationBuilder;
+
+import java.io.IOException;
+
+/**
+ * Builder for the {@link Sampler} aggregation.
+ */
+public class DiversifiedSamplerAggregationBuilder extends ValuesSourceAggregationBuilder<DiversifiedSamplerAggregationBuilder> {
+
+    private int shardSize = SamplerAggregator.Factory.DEFAULT_SHARD_SAMPLE_SIZE;
+
+    int maxDocsPerValue = SamplerAggregator.DiversifiedFactory.MAX_DOCS_PER_VALUE_DEFAULT;
+    String executionHint = null;
+
+    /**
+     * Sole constructor.
+     */
+    public DiversifiedSamplerAggregationBuilder(String name) {
+        super(name, SamplerAggregator.DiversifiedFactory.TYPE.name());
+    }
+
+    /**
+     * Set the max num docs to be returned from each shard.
+     */
+    public DiversifiedSamplerAggregationBuilder shardSize(int shardSize) {
+        this.shardSize = shardSize;
+        return this;
+    }
+
+    public DiversifiedSamplerAggregationBuilder maxDocsPerValue(int maxDocsPerValue) {
+        this.maxDocsPerValue = maxDocsPerValue;
+        return this;
+    }
+
+    public DiversifiedSamplerAggregationBuilder executionHint(String executionHint) {
+        this.executionHint = executionHint;
+        return this;
+    }
+
+    @Override
+    protected XContentBuilder doInternalXContent(XContentBuilder builder, Params params) throws IOException {
+        if (shardSize != SamplerAggregator.Factory.DEFAULT_SHARD_SAMPLE_SIZE) {
+            builder.field(SamplerAggregator.SHARD_SIZE_FIELD.getPreferredName(), shardSize);
+        }
+
+        if (maxDocsPerValue != SamplerAggregator.DiversifiedFactory.MAX_DOCS_PER_VALUE_DEFAULT) {
+            builder.field(SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD.getPreferredName(), maxDocsPerValue);
+        }
+        if (executionHint != null) {
+            builder.field(SamplerAggregator.EXECUTION_HINT_FIELD.getPreferredName(), executionHint);
+        }
+
+        return builder;
+    }
+
+
+}

core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/DiversifiedSamplerParser.java

@@ -0,0 +1,97 @@
+/*
+ * 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.bucket.sampler;
+
+
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.ParseFieldMatcher;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.search.aggregations.AggregatorFactory;
+import org.elasticsearch.search.aggregations.support.AbstractValuesSourceParser.AnyValuesSourceParser;
+import org.elasticsearch.search.aggregations.support.ValueType;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ *
+ */
+public class DiversifiedSamplerParser extends AnyValuesSourceParser {
+
+    public DiversifiedSamplerParser() {
+        super(true, false);
+    }
+
+    @Override
+    public String type() {
+        return SamplerAggregator.DiversifiedFactory.TYPE.name();
+    }
+
+    @Override
+    protected ValuesSourceAggregatorFactory<ValuesSource> createFactory(String aggregationName, ValuesSourceType valuesSourceType,
+            ValueType targetValueType, Map<ParseField, Object> otherOptions) {
+        SamplerAggregator.DiversifiedFactory factory = new SamplerAggregator.DiversifiedFactory(aggregationName, valuesSourceType,
+                targetValueType);
+        Integer shardSize = (Integer) otherOptions.get(SamplerAggregator.SHARD_SIZE_FIELD);
+        if (shardSize != null) {
+            factory.shardSize(shardSize);
+        }
+        Integer maxDocsPerValue = (Integer) otherOptions.get(SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD);
+        if (maxDocsPerValue != null) {
+            factory.maxDocsPerValue(maxDocsPerValue);
+        }
+        String executionHint = (String) otherOptions.get(SamplerAggregator.EXECUTION_HINT_FIELD);
+        if (executionHint != null) {
+            factory.executionHint(executionHint);
+        }
+        return factory;
+    }
+
+    @Override
+    protected boolean token(String aggregationName, String currentFieldName, XContentParser.Token token, XContentParser parser,
+            ParseFieldMatcher parseFieldMatcher, Map<ParseField, Object> otherOptions) throws IOException {
+        if (token == XContentParser.Token.VALUE_NUMBER) {
+            if (parseFieldMatcher.match(currentFieldName, SamplerAggregator.SHARD_SIZE_FIELD)) {
+                int shardSize = parser.intValue();
+                otherOptions.put(SamplerAggregator.SHARD_SIZE_FIELD, shardSize);
+                return true;
+            } else if (parseFieldMatcher.match(currentFieldName, SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD)) {
+                int maxDocsPerValue = parser.intValue();
+                otherOptions.put(SamplerAggregator.MAX_DOCS_PER_VALUE_FIELD, maxDocsPerValue);
+                return true;
+            }
+        } else if (token == XContentParser.Token.VALUE_STRING) {
+            if (parseFieldMatcher.match(currentFieldName, SamplerAggregator.EXECUTION_HINT_FIELD)) {
+                String executionHint = parser.text();
+                otherOptions.put(SamplerAggregator.EXECUTION_HINT_FIELD, executionHint);
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public AggregatorFactory[] getFactoryPrototypes() {
+        return new AggregatorFactory[] { new SamplerAggregator.DiversifiedFactory(null, null, null) };
+    }
+
+}

core/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregationBuilder.java

@@ -29,10 +29,7 @@
  */
 public class SamplerAggregationBuilder extends ValuesSourceAggregationBuilder<SamplerAggregationBuilder> {
 
-    private int shardSize = SamplerParser.DEFAULT_SHARD_SAMPLE_SIZE;
-
-    int maxDocsPerValue = SamplerParser.MAX_DOCS_PER_VALUE_DEFAULT;
-    String executionHint = null;
+    private int shardSize = SamplerAggregator.Factory.DEFAULT_SHARD_SAMPLE_SIZE;
 
     /**
      * Sole constructor.
@@ -49,28 +46,10 @@ public SamplerAggregationBuilder shardSize(int shardSize) {
         return this;
     }
 
-    public SamplerAggregationBuilder maxDocsPerValue(int maxDocsPerValue) {
-        this.maxDocsPerValue = maxDocsPerValue;
-        return this;
-    }
-
-    public SamplerAggregationBuilder executionHint(String executionHint) {
-        this.executionHint = executionHint;
-        return this;
-    }
-
     @Override
     protected XContentBuilder doInternalXContent(XContentBuilder builder, Params params) throws IOException {
-        // builder.startObject();
-        if (shardSize != SamplerParser.DEFAULT_SHARD_SAMPLE_SIZE) {
-            builder.field(SamplerParser.SHARD_SIZE_FIELD.getPreferredName(), shardSize);
-        }
-
-        if (maxDocsPerValue != SamplerParser.MAX_DOCS_PER_VALUE_DEFAULT) {
-            builder.field(SamplerParser.MAX_DOCS_PER_VALUE_FIELD.getPreferredName(), maxDocsPerValue);
-        }
-        if (executionHint != null) {
-            builder.field(SamplerParser.EXECUTION_HINT_FIELD.getPreferredName(), executionHint);
+        if (shardSize != SamplerAggregator.Factory.DEFAULT_SHARD_SAMPLE_SIZE) {
+            builder.field(SamplerAggregator.SHARD_SIZE_FIELD.getPreferredName(), shardSize);
         }
 
         return builder;