diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java index 3d0839b7fb477..602a0964ee9e9 100644 --- a/core/src/main/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java +++ b/core/src/main/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrix.java @@ -35,9 +35,10 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; -public class InternalAdjacencyMatrix - extends InternalMultiBucketAggregation +public class InternalAdjacencyMatrix + extends InternalMultiBucketAggregation implements AdjacencyMatrix { public static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket implements AdjacencyMatrix.Bucket { @@ -111,12 +112,31 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.endObject(); return builder; } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + InternalBucket that = (InternalBucket) other; + return Objects.equals(key, that.key) + && Objects.equals(docCount, that.docCount) + && Objects.equals(aggregations, that.aggregations); + } + + @Override + public int hashCode() { + return Objects.hash(getClass(), key, docCount, aggregations); + } } private final List buckets; private Map bucketMap; - public InternalAdjacencyMatrix(String name, List buckets, + public InternalAdjacencyMatrix(String name, List buckets, List pipelineAggregators, Map metaData) { super(name, pipelineAggregators, metaData); this.buckets = buckets; @@ -193,13 +213,13 @@ public InternalAggregation doReduce(List aggregations, Redu ArrayList reducedBuckets = new ArrayList<>(bucketsMap.size()); for (List sameRangeList : bucketsMap.values()) { InternalBucket reducedBucket = sameRangeList.get(0).reduce(sameRangeList, reduceContext); - if(reducedBucket.docCount >= 1){ + if(reducedBucket.docCount >= 1){ reducedBuckets.add(reducedBucket); } } Collections.sort(reducedBuckets, Comparator.comparing(InternalBucket::getKey)); - - InternalAdjacencyMatrix reduced = new InternalAdjacencyMatrix(name, reducedBuckets, pipelineAggregators(), + + InternalAdjacencyMatrix reduced = new InternalAdjacencyMatrix(name, reducedBuckets, pipelineAggregators(), getMetaData()); return reduced; @@ -215,4 +235,14 @@ public XContentBuilder doXContentBody(XContentBuilder builder, Params params) th return builder; } + @Override + protected int doHashCode() { + return Objects.hash(buckets); + } + + @Override + protected boolean doEquals(Object obj) { + InternalAdjacencyMatrix that = (InternalAdjacencyMatrix) obj; + return Objects.equals(buckets, that.buckets); + } } diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrixTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrixTests.java new file mode 100644 index 0000000000000..53ae1db4544b1 --- /dev/null +++ b/core/src/test/java/org/elasticsearch/search/aggregations/bucket/adjacency/InternalAdjacencyMatrixTests.java @@ -0,0 +1,92 @@ +/* + * 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.adjacency; + +import org.elasticsearch.common.io.stream.Writeable.Reader; +import org.elasticsearch.search.aggregations.InternalAggregations; +import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator; +import org.elasticsearch.test.InternalAggregationTestCase; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; + +public class InternalAdjacencyMatrixTests extends InternalAggregationTestCase { + + private List keys; + + @Override + @Before + public void setUp() throws Exception { + super.setUp(); + keys = new ArrayList<>(); + int numFilters = randomIntBetween(2, 4); + String[] filters = new String[numFilters]; + for (int i = 0; i < numFilters; i++) { + filters[i] = randomAlphaOfLength(5); + } + for (int i = 0; i < filters.length; i++) { + keys.add(filters[i]); + for (int j = i + 1; j < filters.length; j++) { + if (filters[i].compareTo(filters[j]) <= 0) { + keys.add(filters[i] + "&" + filters[j]); + } else { + keys.add(filters[j] + "&" + filters[i]); + } + } + } + } + + @Override + protected InternalAdjacencyMatrix createTestInstance(String name, List pipelineAggregators, + Map metaData) { + final List buckets = new ArrayList<>(); + for (int i = 0; i < keys.size(); ++i) { + String key = keys.get(i); + int docCount = randomIntBetween(0, 1000); + buckets.add(new InternalAdjacencyMatrix.InternalBucket(key, docCount, InternalAggregations.EMPTY)); + } + return new InternalAdjacencyMatrix(name, buckets, pipelineAggregators, metaData); + } + + @Override + protected void assertReduced(InternalAdjacencyMatrix reduced, List inputs) { + final Map expectedCounts = new TreeMap<>(); + for (InternalAdjacencyMatrix input : inputs) { + for (InternalAdjacencyMatrix.InternalBucket bucket : input.getBuckets()) { + expectedCounts.compute(bucket.getKeyAsString(), + (key, oldValue) -> (oldValue == null ? 0 : oldValue) + bucket.getDocCount()); + } + } + final Map actualCounts = new TreeMap<>(); + for (InternalAdjacencyMatrix.InternalBucket bucket : reduced.getBuckets()) { + actualCounts.compute(bucket.getKeyAsString(), + (key, oldValue) -> (oldValue == null ? 0 : oldValue) + bucket.getDocCount()); + } + assertEquals(expectedCounts, actualCounts); + } + + @Override + protected Reader instanceReader() { + return InternalAdjacencyMatrix::new; + } +}