Add parsing method for binary range aggregation (#24706)

Related to #23331

core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRange.java

@@ -44,6 +44,7 @@ import static java.util.Collections.unmodifiableList;
 public final class InternalBinaryRange
         extends InternalMultiBucketAggregation<InternalBinaryRange, InternalBinaryRange.Bucket>
         implements Range {
+
     public static class Bucket extends InternalMultiBucketAggregation.InternalBucket implements Range.Bucket {
 
         private final transient DocValueFormat format;

core/src/main/java/org/elasticsearch/search/aggregations/bucket/range/ParsedBinaryRange.java

@@ -0,0 +1,168 @@
+/*
+ * 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.range;
+
+import org.elasticsearch.common.xcontent.ObjectParser;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentParserUtils;
+import org.elasticsearch.search.aggregations.Aggregation;
+import org.elasticsearch.search.aggregations.Aggregations;
+import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
+import org.elasticsearch.search.aggregations.bucket.range.ip.IpRangeAggregationBuilder;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
+
+public class ParsedBinaryRange extends ParsedMultiBucketAggregation<ParsedBinaryRange.ParsedBucket> implements Range {
+
+    @Override
+    public String getType() {
+        return IpRangeAggregationBuilder.NAME;
+    }
+
+    @Override
+    public List<? extends Range.Bucket> getBuckets() {
+        return buckets;
+    }
+
+    private static ObjectParser<ParsedBinaryRange, Void> PARSER =
+            new ObjectParser<>(ParsedBinaryRange.class.getSimpleName(), true, ParsedBinaryRange::new);
+    static {
+        declareMultiBucketAggregationFields(PARSER,
+                parser -> ParsedBucket.fromXContent(parser, false),
+                parser -> ParsedBucket.fromXContent(parser, true));
+    }
+
+    public static ParsedBinaryRange fromXContent(XContentParser parser, String name) throws IOException {
+        ParsedBinaryRange aggregation = PARSER.parse(parser, null);
+        aggregation.setName(name);
+        return aggregation;
+    }
+
+    public static class ParsedBucket extends ParsedMultiBucketAggregation.ParsedBucket implements Range.Bucket {
+
+        private String key;
+        private String from;
+        private String to;
+
+        @Override
+        public Object getKey() {
+            return key;
+        }
+
+        @Override
+        public String getKeyAsString() {
+            return key;
+        }
+
+        @Override
+        public Object getFrom() {
+            return from;
+        }
+
+        @Override
+        public String getFromAsString() {
+            return from;
+        }
+
+        @Override
+        public Object getTo() {
+            return to;
+        }
+
+        @Override
+        public String getToAsString() {
+            return to;
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            if (isKeyed()) {
+                builder.startObject(key != null ? key : rangeKey(from, to));
+            } else {
+                builder.startObject();
+                if (key != null) {
+                    builder.field(CommonFields.KEY.getPreferredName(), key);
+                }
+            }
+            if (from != null) {
+                builder.field(CommonFields.FROM.getPreferredName(), getFrom());
+            }
+            if (to != null) {
+                builder.field(CommonFields.TO.getPreferredName(), getTo());
+            }
+            builder.field(CommonFields.DOC_COUNT.getPreferredName(), getDocCount());
+            getAggregations().toXContentInternal(builder, params);
+            builder.endObject();
+            return builder;
+        }
+
+        static ParsedBucket fromXContent(final XContentParser parser, final boolean keyed) throws IOException {
+            final ParsedBucket bucket = new ParsedBucket();
+            bucket.setKeyed(keyed);
+            XContentParser.Token token = parser.currentToken();
+            String currentFieldName = parser.currentName();
+
+            String rangeKey = null;
+            if (keyed) {
+                ensureExpectedToken(XContentParser.Token.FIELD_NAME, token, parser::getTokenLocation);
+                rangeKey = currentFieldName;
+                ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser::getTokenLocation);
+            }
+
+            List<Aggregation> aggregations = new ArrayList<>();
+            while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                if (token == XContentParser.Token.FIELD_NAME) {
+                    currentFieldName = parser.currentName();
+                } else if (token.isValue()) {
+                    if (CommonFields.KEY.getPreferredName().equals(currentFieldName)) {
+                        bucket.key = parser.text();
+                    } else if (CommonFields.DOC_COUNT.getPreferredName().equals(currentFieldName)) {
+                        bucket.setDocCount(parser.longValue());
+                    } else if (CommonFields.FROM.getPreferredName().equals(currentFieldName)) {
+                        bucket.from = parser.text();
+                    } else if (CommonFields.TO.getPreferredName().equals(currentFieldName)) {
+                        bucket.to = parser.text();
+                    }
+                } else if (token == XContentParser.Token.START_OBJECT) {
+                    aggregations.add(XContentParserUtils.parseTypedKeysObject(parser, Aggregation.TYPED_KEYS_DELIMITER, Aggregation.class));
+                }
+            }
+            bucket.setAggregations(new Aggregations(aggregations));
+
+            if (keyed) {
+                if (rangeKey(bucket.from, bucket.to).equals(rangeKey)) {
+                    bucket.key = null;
+                } else {
+                    bucket.key = rangeKey;
+                }
+            }
+            return bucket;
+        }
+
+        private static String rangeKey(String from, String to) {
+            return (from == null ? "*" : from) + '-' + (to == null ? "*" : to);
+        }
+    }
+}

core/src/test/java/org/elasticsearch/search/aggregations/AggregationsTests.java

@@ -36,6 +36,7 @@ import org.elasticsearch.search.aggregations.bucket.histogram.InternalHistogramT
 import org.elasticsearch.search.aggregations.bucket.missing.InternalMissingTests;
 import org.elasticsearch.search.aggregations.bucket.nested.InternalNestedTests;
 import org.elasticsearch.search.aggregations.bucket.nested.InternalReverseNestedTests;
+import org.elasticsearch.search.aggregations.bucket.range.InternalBinaryRangeTests;
 import org.elasticsearch.search.aggregations.bucket.range.InternalRangeTests;
 import org.elasticsearch.search.aggregations.bucket.range.date.InternalDateRangeTests;
 import org.elasticsearch.search.aggregations.bucket.range.geodistance.InternalGeoDistanceTests;
@@ -132,6 +133,7 @@ public class AggregationsTests extends ESTestCase {
         aggsTests.add(new SignificantLongTermsTests());
         aggsTests.add(new SignificantStringTermsTests());
         aggsTests.add(new InternalScriptedMetricTests());
+        aggsTests.add(new InternalBinaryRangeTests());
         return Collections.unmodifiableList(aggsTests);
     }
 

core/src/test/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregationTestCase.java

@@ -98,9 +98,12 @@ public abstract class InternalMultiBucketAggregationTestCase<T extends InternalA
             }
         } else {
             for (MultiBucketsAggregation.Bucket expectedBucket : expectedBuckets) {
+                final Object expectedKey = expectedBucket.getKey();
                 boolean found = false;
+
                 for (MultiBucketsAggregation.Bucket actualBucket : actualBuckets) {
-                    if (actualBucket.getKey().equals(expectedBucket.getKey())) {
+                    final Object actualKey = actualBucket.getKey();
+                    if ((actualKey != null && actualKey.equals(expectedKey)) || (actualKey == null && expectedKey == null)) {
                         found = true;
                         assertBucket(expectedBucket, actualBucket, false);
                         break;

core/src/test/java/org/elasticsearch/search/aggregations/bucket/range/InternalBinaryRangeTests.java

@@ -23,47 +23,61 @@ import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
+import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
-import org.elasticsearch.test.InternalAggregationTestCase;
-import org.junit.Before;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
-public class InternalBinaryRangeTests extends InternalAggregationTestCase<InternalBinaryRange> {
+public class InternalBinaryRangeTests extends InternalRangeTestCase<InternalBinaryRange> {
-    private Tuple<BytesRef, BytesRef>[] RANGES;
+
+    private List<Tuple<BytesRef, BytesRef>> ranges;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+
+        final int numRanges = randomIntBetween(1, 10);
+        ranges = new ArrayList<>(numRanges);
 
-    @Before
-    public void randomSortedRanges() {
-        int numRanges = randomIntBetween(1, 10);
-        Tuple<BytesRef, BytesRef>[] ranges = new Tuple[numRanges];
         for (int i = 0; i < numRanges; i++) {
             BytesRef[] values = new BytesRef[2];
             values[0] = new BytesRef(randomAlphaOfLength(15));
             values[1] = new BytesRef(randomAlphaOfLength(15));
             Arrays.sort(values);
-            ranges[i] = new Tuple(values[0], values[1]);
+            ranges.add(Tuple.tuple(values[0], values[1]));
+        }
+
+        if (randomBoolean()) {
+            ranges.add(Tuple.tuple(null, new BytesRef(randomAlphaOfLength(15))));
+        }
+        if (randomBoolean()) {
+            ranges.add(Tuple.tuple(new BytesRef(randomAlphaOfLength(15)), null));
+        }
+        if (randomBoolean()) {
+            ranges.add(Tuple.tuple(null, null));
         }
-        Arrays.sort(ranges, (t1, t2) -> t1.v1().compareTo(t2.v1()));
-        RANGES = ranges;
     }
 
-
     @Override
-    protected InternalBinaryRange createTestInstance(String name, List<PipelineAggregator> pipelineAggregators,
+    protected InternalBinaryRange createTestInstance(String name,
-                                                     Map<String, Object> metaData) {
+                                                     List<PipelineAggregator> pipelineAggregators,
-        boolean keyed = randomBoolean();
+                                                     Map<String, Object> metaData,
+                                                     InternalAggregations aggregations,
+                                                     boolean keyed) {
         DocValueFormat format = DocValueFormat.RAW;
         List<InternalBinaryRange.Bucket> buckets = new ArrayList<>();
-        for (int i = 0; i < RANGES.length; ++i) {
+
+        int nullKey = randomBoolean() ? randomIntBetween(0, ranges.size() -1) : -1;
+        for (int i = 0; i < ranges.size(); ++i) {
             final int docCount = randomIntBetween(1, 100);
-            buckets.add(new InternalBinaryRange.Bucket(format, keyed, randomAlphaOfLength(10),
+            final String key = (i == nullKey) ? null: randomAlphaOfLength(10);
-                RANGES[i].v1(), RANGES[i].v2(), docCount, InternalAggregations.EMPTY));
+            buckets.add(new InternalBinaryRange.Bucket(format, keyed, key, ranges.get(i).v1(), ranges.get(i).v2(), docCount, aggregations));
         }
-        return new InternalBinaryRange(name, format, keyed, buckets, pipelineAggregators, Collections.emptyMap());
+        return new InternalBinaryRange(name, format, keyed, buckets, pipelineAggregators, metaData);
     }
 
     @Override
@@ -71,6 +85,11 @@ public class InternalBinaryRangeTests extends InternalAggregationTestCase<Intern
         return InternalBinaryRange::new;
     }
 
+    @Override
+    protected Class<? extends ParsedMultiBucketAggregation> implementationClass() {
+        return ParsedBinaryRange.class;
+    }
+
     @Override
     protected void assertReduced(InternalBinaryRange reduced, List<InternalBinaryRange> inputs) {
         int pos = 0;
@@ -86,4 +105,14 @@ public class InternalBinaryRangeTests extends InternalAggregationTestCase<Intern
             pos ++;
         }
     }
+
+    @Override
+    protected Class<? extends InternalMultiBucketAggregation.InternalBucket> internalRangeBucketClass() {
+        return InternalBinaryRange.Bucket.class;
+    }
+
+    @Override
+    protected Class<? extends ParsedMultiBucketAggregation.ParsedBucket> parsedRangeBucketClass() {
+        return ParsedBinaryRange.ParsedBucket.class;
+    }
 }

core/src/test/java/org/elasticsearch/search/aggregations/bucket/range/InternalRangeTestCase.java

@@ -21,7 +21,9 @@ package org.elasticsearch.search.aggregations.bucket.range;
 
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.InternalMultiBucketAggregationTestCase;
+import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.junit.Before;
@@ -73,11 +75,16 @@ public abstract class InternalRangeTestCase<T extends InternalAggregation & Rang
     }
 
     @Override
-    protected void assertBucket(MultiBucketsAggregation.Bucket expected, MultiBucketsAggregation.Bucket actual, boolean checkOrder) {
+    protected final void assertBucket(MultiBucketsAggregation.Bucket expected, MultiBucketsAggregation.Bucket actual, boolean checkOrder) {
         super.assertBucket(expected, actual, checkOrder);
 
-        assertTrue(expected instanceof InternalRange.Bucket);
+        Class<?> internalBucketClass = internalRangeBucketClass();
-        assertTrue(actual instanceof ParsedRange.ParsedBucket);
+        assertNotNull("Internal bucket class must not be null", internalBucketClass);
+        assertTrue(internalBucketClass.isInstance(expected));
+
+        Class<?> parsedBucketClass = parsedRangeBucketClass();
+        assertNotNull("Parsed bucket class must not be null", parsedBucketClass);
+        assertTrue(parsedBucketClass.isInstance(actual));
 
         Range.Bucket expectedRange = (Range.Bucket) expected;
         Range.Bucket actualRange = (Range.Bucket) actual;
@@ -87,4 +94,8 @@ public abstract class InternalRangeTestCase<T extends InternalAggregation & Rang
         assertEquals(expectedRange.getTo(), actualRange.getTo());
         assertEquals(expectedRange.getToAsString(), actualRange.getToAsString());
     }
+
+    protected abstract Class<? extends InternalMultiBucketAggregation.InternalBucket> internalRangeBucketClass();
+
+    protected abstract Class<? extends ParsedMultiBucketAggregation.ParsedBucket> parsedRangeBucketClass();
 }

core/src/test/java/org/elasticsearch/search/aggregations/bucket/range/InternalRangeTests.java

@@ -23,6 +23,7 @@ import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.junit.Before;
@@ -97,4 +98,14 @@ public class InternalRangeTests extends InternalRangeTestCase<InternalRange> {
     protected Class<? extends ParsedMultiBucketAggregation> implementationClass() {
         return ParsedRange.class;
     }
+
+    @Override
+    protected Class<? extends InternalMultiBucketAggregation.InternalBucket> internalRangeBucketClass() {
+        return InternalRange.Bucket.class;
+    }
+
+    @Override
+    protected Class<? extends ParsedMultiBucketAggregation.ParsedBucket> parsedRangeBucketClass() {
+        return ParsedRange.ParsedBucket.class;
+    }
 }

core/src/test/java/org/elasticsearch/search/aggregations/bucket/range/date/InternalDateRangeTests.java

@@ -23,6 +23,7 @@ import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.range.InternalRangeTestCase;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
@@ -101,4 +102,14 @@ public class InternalDateRangeTests extends InternalRangeTestCase<InternalDateRa
     protected Class<? extends ParsedMultiBucketAggregation> implementationClass() {
         return ParsedDateRange.class;
     }
+
+    @Override
+    protected Class<? extends InternalMultiBucketAggregation.InternalBucket> internalRangeBucketClass() {
+        return InternalDateRange.Bucket.class;
+    }
+
+    @Override
+    protected Class<? extends ParsedMultiBucketAggregation.ParsedBucket> parsedRangeBucketClass() {
+        return ParsedDateRange.ParsedBucket.class;
+    }
 }

core/src/test/java/org/elasticsearch/search/aggregations/bucket/range/geodistance/InternalGeoDistanceTests.java

@@ -22,6 +22,7 @@ package org.elasticsearch.search.aggregations.bucket.range.geodistance;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.search.aggregations.InternalAggregations;
+import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.ParsedMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.range.InternalRangeTestCase;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
@@ -86,4 +87,14 @@ public class InternalGeoDistanceTests extends InternalRangeTestCase<InternalGeoD
     protected Class<? extends ParsedMultiBucketAggregation> implementationClass() {
         return ParsedGeoDistance.class;
     }
+
+    @Override
+    protected Class<? extends InternalMultiBucketAggregation.InternalBucket> internalRangeBucketClass() {
+        return InternalGeoDistance.Bucket.class;
+    }
+
+    @Override
+    protected Class<? extends ParsedMultiBucketAggregation.ParsedBucket> parsedRangeBucketClass() {
+        return ParsedGeoDistance.ParsedBucket.class;
+    }
 }

test/framework/src/main/java/org/elasticsearch/test/InternalAggregationTestCase.java

@@ -58,12 +58,14 @@ import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuil
 import org.elasticsearch.search.aggregations.bucket.nested.ParsedNested;
 import org.elasticsearch.search.aggregations.bucket.nested.ParsedReverseNested;
 import org.elasticsearch.search.aggregations.bucket.nested.ReverseNestedAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.range.ParsedBinaryRange;
 import org.elasticsearch.search.aggregations.bucket.range.ParsedRange;
 import org.elasticsearch.search.aggregations.bucket.range.RangeAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.range.date.DateRangeAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.range.date.ParsedDateRange;
 import org.elasticsearch.search.aggregations.bucket.range.geodistance.GeoDistanceAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.range.geodistance.ParsedGeoDistance;
+import org.elasticsearch.search.aggregations.bucket.range.ip.IpRangeAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.sampler.InternalSampler;
 import org.elasticsearch.search.aggregations.bucket.sampler.ParsedSampler;
 import org.elasticsearch.search.aggregations.bucket.significant.ParsedSignificantLongTerms;
@@ -161,8 +163,7 @@ public abstract class InternalAggregationTestCase<T extends InternalAggregation>
         map.put(StatsAggregationBuilder.NAME, (p, c) -> ParsedStats.fromXContent(p, (String) c));
         map.put(StatsBucketPipelineAggregationBuilder.NAME, (p, c) -> ParsedStatsBucket.fromXContent(p, (String) c));
         map.put(ExtendedStatsAggregationBuilder.NAME, (p, c) -> ParsedExtendedStats.fromXContent(p, (String) c));
-        map.put(ExtendedStatsBucketPipelineAggregationBuilder.NAME,
+        map.put(ExtendedStatsBucketPipelineAggregationBuilder.NAME, (p, c) -> ParsedExtendedStatsBucket.fromXContent(p, (String) c));
-                (p, c) -> ParsedExtendedStatsBucket.fromXContent(p, (String) c));
         map.put(GeoBoundsAggregationBuilder.NAME, (p, c) -> ParsedGeoBounds.fromXContent(p, (String) c));
         map.put(GeoCentroidAggregationBuilder.NAME, (p, c) -> ParsedGeoCentroid.fromXContent(p, (String) c));
         map.put(HistogramAggregationBuilder.NAME, (p, c) -> ParsedHistogram.fromXContent(p, (String) c));
@@ -185,6 +186,7 @@ public abstract class InternalAggregationTestCase<T extends InternalAggregation>
         map.put(SignificantLongTerms.NAME, (p, c) -> ParsedSignificantLongTerms.fromXContent(p, (String) c));
         map.put(SignificantStringTerms.NAME, (p, c) -> ParsedSignificantStringTerms.fromXContent(p, (String) c));
         map.put(ScriptedMetricAggregationBuilder.NAME, (p, c) -> ParsedScriptedMetric.fromXContent(p, (String) c));
+        map.put(IpRangeAggregationBuilder.NAME, (p, c) -> ParsedBinaryRange.fromXContent(p, (String) c));
 
         namedXContents = map.entrySet().stream()
                 .map(entry -> new NamedXContentRegistry.Entry(Aggregation.class, new ParseField(entry.getKey()), entry.getValue()))