Aggregations Refactor: Refactor Geo Centroid Aggregation

Author: colings86

core/src/main/java/org/elasticsearch/search/aggregations/metrics/geocentroid/GeoCentroidAggregator.java

@@ -22,21 +22,24 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.util.GeoUtils;
 import org.elasticsearch.common.geo.GeoPoint;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.index.fielddata.MultiGeoPointValues;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.LeafBucketCollector;
 import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
 import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
-import org.elasticsearch.search.aggregations.metrics.geobounds.InternalGeoBounds;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
+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.ValuesSourceParser;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 
 import java.io.IOException;
 import java.util.List;
@@ -123,8 +126,9 @@ public void doClose() {
     }
 
     public static class Factory extends ValuesSourceAggregatorFactory.LeafOnly<ValuesSource.GeoPoint> {
-        protected Factory(String name, ValuesSourceParser.Input<ValuesSource.GeoPoint> config) {
-            super(name, InternalGeoBounds.TYPE, config);
+
+        public Factory(String name) {
+            super(name, InternalGeoCentroid.TYPE, ValuesSourceType.GEOPOINT, ValueType.GEOPOINT);
         }
 
         @Override
@@ -139,5 +143,31 @@ protected Aggregator doCreateInternal(ValuesSource.GeoPoint valuesSource, Aggreg
                 throws IOException {
             return new GeoCentroidAggregator(name, aggregationContext, parent, valuesSource, pipelineAggregators, metaData);
         }
+
+        @Override
+        protected ValuesSourceAggregatorFactory<ValuesSource.GeoPoint> innerReadFrom(String name, ValuesSourceType valuesSourceType,
+                ValueType targetValueType, StreamInput in) throws IOException {
+            return new Factory(name);
+        }
+
+        @Override
+        protected void innerWriteTo(StreamOutput out) {
+            // Do nothing, no extra state to write to stream
+        }
+
+        @Override
+        public XContentBuilder doXContentBody(XContentBuilder builder, Params params) throws IOException {
+            return builder;
+        }
+
+        @Override
+        protected int innerHashCode() {
+            return 0;
+        }
+
+        @Override
+        protected boolean innerEquals(Object obj) {
+            return true;
+        }
     }
 }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/geocentroid/GeoCentroidParser.java

@@ -19,51 +19,48 @@
 
 package org.elasticsearch.search.aggregations.metrics.geocentroid;
 
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.ParseFieldMatcher;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.search.SearchParseException;
-import org.elasticsearch.search.aggregations.Aggregator;
+import org.elasticsearch.common.xcontent.XContentParser.Token;
 import org.elasticsearch.search.aggregations.AggregatorFactory;
+import org.elasticsearch.search.aggregations.support.AbstractValuesSourceParser.GeoPointValuesSourceParser;
 import org.elasticsearch.search.aggregations.support.ValueType;
-import org.elasticsearch.search.aggregations.support.ValuesSource;
-import org.elasticsearch.search.aggregations.support.ValuesSourceParser;
-import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.aggregations.support.ValuesSource.GeoPoint;
+import org.elasticsearch.search.aggregations.support.ValuesSourceAggregatorFactory;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  * Parser class for {@link org.elasticsearch.search.aggregations.metrics.geocentroid.GeoCentroidAggregator}
  */
-public class GeoCentroidParser implements Aggregator.Parser {
+public class GeoCentroidParser extends GeoPointValuesSourceParser {
+
+    public GeoCentroidParser() {
+        super(true, false);
+    }
 
     @Override
     public String type() {
         return InternalGeoCentroid.TYPE.name();
     }
 
     @Override
-    public AggregatorFactory parse(String aggregationName, XContentParser parser, SearchContext context) throws IOException {
-        ValuesSourceParser<ValuesSource.GeoPoint> vsParser = ValuesSourceParser.geoPoint(aggregationName, InternalGeoCentroid.TYPE, context)
-                .targetValueType(ValueType.GEOPOINT)
-                .formattable(true)
-                .build();
-        XContentParser.Token token;
-        String currentFieldName = null;
-        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
-            if (token == XContentParser.Token.FIELD_NAME) {
-                currentFieldName = parser.currentName();
-            } else if (vsParser.token(currentFieldName, token, parser)) {
-                continue;
-            } else {
-                throw new SearchParseException(context, "Unknown key for a " + token + " in aggregation [" + aggregationName + "]: ["
-                        + currentFieldName + "].", parser.getTokenLocation());
-            }
-        }
-        return new GeoCentroidAggregator.Factory(aggregationName, vsParser.input());
+    protected boolean token(String aggregationName, String currentFieldName, Token token, XContentParser parser,
+            ParseFieldMatcher parseFieldMatcher, Map<ParseField, Object> otherOptions) throws IOException {
+        return false;
+    }
+
+    @Override
+    protected ValuesSourceAggregatorFactory<GeoPoint> createFactory(String aggregationName, ValuesSourceType valuesSourceType,
+            ValueType targetValueType, Map<ParseField, Object> otherOptions) {
+        return new GeoCentroidAggregator.Factory(aggregationName);
     }
 
-    // NORELEASE implement this method when refactoring this aggregation
     @Override
     public AggregatorFactory[] getFactoryPrototypes() {
-        return null;
+        return new AggregatorFactory[] { new GeoCentroidAggregator.Factory(null) };
     }
 }

core/src/main/java/org/elasticsearch/search/aggregations/support/AbstractValuesSourceParser.java

@@ -155,9 +155,52 @@ public final AggregatorFactory parse(String aggregationName, XContentParser pars
         return factory;
     }
 
+    /**
+     * Creates a {@link ValuesSourceAggregatorFactory} from the information
+     * gathered by the subclass. Options parsed in
+     * {@link AbstractValuesSourceParser} itself will be added to the factory
+     * after it has been returned by this method.
+     *
+     * @param aggregationName
+     *            the name of the aggregation
+     * @param valuesSourceType
+     *            the type of the {@link ValuesSource}
+     * @param targetValueType
+     *            the target type of the final value output by the aggregation
+     * @param otherOptions
+     *            a {@link Map} containing the extra options parsed by the
+     *            {@link #token(String, String, org.elasticsearch.common.xcontent.XContentParser.Token, XContentParser, ParseFieldMatcher, Map)}
+     *            method
+     * @return the created factory
+     */
     protected abstract ValuesSourceAggregatorFactory<VS> createFactory(String aggregationName, ValuesSourceType valuesSourceType,
             ValueType targetValueType, Map<ParseField, Object> otherOptions);
 
+    /**
+     * Allows subclasses of {@link AbstractValuesSourceParser} to parse extra
+     * parameters and store them in a {@link Map} which will later be passed to
+     * {@link #createFactory(String, ValuesSourceType, ValueType, Map)}.
+     *
+     * @param aggregationName
+     *            the name of the aggregation
+     * @param currentFieldName
+     *            the name of the current field being parsed
+     * @param token
+     *            the current token for the parser
+     * @param parser
+     *            the parser
+     * @param parseFieldMatcher
+     *            the {@link ParseFieldMatcher} to use to match field names
+     * @param otherOptions
+     *            a {@link Map} of options to be populated by successive calls
+     *            to this method which will then be passed to the
+     *            {@link #createFactory(String, ValuesSourceType, ValueType, Map)}
+     *            method
+     * @return <code>true</code> if the current token was correctly parsed,
+     *         <code>false</code> otherwise
+     * @throws IOException
+     *             if an error occurs whilst parsing
+     */
     protected abstract boolean token(String aggregationName, String currentFieldName, XContentParser.Token token, XContentParser parser,
             ParseFieldMatcher parseFieldMatcher, Map<ParseField, Object> otherOptions) throws IOException;
 }

core/src/test/java/org/elasticsearch/search/aggregations/metrics/GeoCentroidTests.java

@@ -0,0 +1,52 @@
+/*
+ * 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.metrics;
+
+import org.elasticsearch.script.Script;
+import org.elasticsearch.search.aggregations.BaseAggregationTestCase;
+import org.elasticsearch.search.aggregations.metrics.geocentroid.GeoCentroidAggregator;
+import org.elasticsearch.search.aggregations.metrics.geocentroid.GeoCentroidAggregator.Factory;
+
+public class GeoCentroidTests extends BaseAggregationTestCase<GeoCentroidAggregator.Factory> {
+
+    @Override
+    protected Factory createTestAggregatorFactory() {
+        Factory factory = new Factory(randomAsciiOfLengthBetween(1, 20));
+        String field = randomNumericField();
+        int randomFieldBranch = randomInt(3);
+        switch (randomFieldBranch) {
+        case 0:
+            factory.field(field);
+            break;
+        case 1:
+            factory.field(field);
+            factory.script(new Script("_value + 1"));
+            break;
+        case 2:
+            factory.script(new Script("doc[" + field + "] + 1"));
+            break;
+        }
+        if (randomBoolean()) {
+            factory.missing("0,0");
+        }
+        return factory;
+    }
+
+}