Make InternalComposite key comparable

Keys are compared in BucketSortPipelineAggregation so making key type (ArrayMap) implement Comparable. Maps are compared using the entry set's iterator so ordered maps order is maintain. For each entry first comparing key then value. Assuming all keys are strings. When comparing entries' values if type is not identical and\or type not implementing Comparable, throwing exception. Not implementing equals() and hashCode() functions as parent's ones are sufficient. Tests included.

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/BinaryValuesSource.java

@@ -119,7 +119,7 @@ class BinaryValuesSource extends SingleDimensionValuesSource<BytesRef> {
     }
 
     @Override
-    void setAfter(Comparable<?> value) {
+    void setAfter(Comparable value) {
         if (missingBucket && value == null) {
             afterValue = null;
         } else if (value.getClass() == String.class) {
@@ -155,7 +155,7 @@ class BinaryValuesSource extends SingleDimensionValuesSource<BytesRef> {
     }
 
     @Override
-    LeafBucketCollector getLeafCollector(Comparable<?> value, LeafReaderContext context, LeafBucketCollector next) {
+    LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) {
         if (value.getClass() != BytesRef.class) {
             throw new IllegalArgumentException("Expected BytesRef, got " + value.getClass());
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregationBuilder.java

@@ -170,7 +170,7 @@ public class CompositeAggregationBuilder extends AbstractAggregationBuilder<Comp
                 throw new IllegalArgumentException("[after] has " + after.size() +
                     " value(s) but [sources] has " + sources.size());
             }
-            Comparable<?>[] values = new Comparable<?>[sources.size()];
+            Comparable[] values = new Comparable[sources.size()];
             for (int i = 0; i < sources.size(); i++) {
                 String sourceName = sources.get(i).name();
                 if (after.containsKey(sourceName) == false) {
@@ -180,7 +180,7 @@ public class CompositeAggregationBuilder extends AbstractAggregationBuilder<Comp
                 if (configs[i].missingBucket() && obj == null) {
                     values[i] = null;
                 } else if (obj instanceof Comparable) {
-                    values[i] = (Comparable<?>) obj;
+                    values[i] = (Comparable) obj;
                 } else {
                     throw new IllegalArgumentException("Invalid value for [after." + sources.get(i).name() +
                         "], expected comparable, got [" + (obj == null ? "null" :  obj.getClass().getSimpleName()) + "]");

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeKey.java

@@ -30,16 +30,16 @@ import java.util.Arrays;
  * A key that is composed of multiple {@link Comparable} values.
  */
 class CompositeKey implements Writeable {
-    private final Comparable<?>[] values;
+    private final Comparable[] values;
 
-    CompositeKey(Comparable<?>... values) {
+    CompositeKey(Comparable... values) {
         this.values = values;
     }
 
     CompositeKey(StreamInput in) throws IOException {
-        values = new Comparable<?>[in.readVInt()];
+        values = new Comparable[in.readVInt()];
         for (int i = 0; i < values.length; i++) {
-            values[i] = (Comparable<?>) in.readGenericValue();
+            values[i] = (Comparable) in.readGenericValue();
         }
     }
 
@@ -51,7 +51,7 @@ class CompositeKey implements Writeable {
         }
     }
 
-    Comparable<?>[] values() {
+    Comparable[] values() {
         return values;
     }
 
@@ -59,7 +59,7 @@ class CompositeKey implements Writeable {
         return values.length;
     }
 
-    Comparable<?> get(int pos) {
+    Comparable get(int pos) {
         assert pos < values.length;
         return values[pos];
     }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeValuesCollectorQueue.java

@@ -98,14 +98,14 @@ final class CompositeValuesCollectorQueue implements Releasable {
     /**
      * Returns the lowest value (exclusive) of the leading source.
      */
-    Comparable<?> getLowerValueLeadSource() {
+    Comparable getLowerValueLeadSource() {
         return afterKeyIsSet ? arrays[0].getAfter() : null;
     }
 
     /**
      * Returns the upper value (inclusive) of the leading source.
      */
-    Comparable<?> getUpperValueLeadSource() throws IOException {
+    Comparable getUpperValueLeadSource() throws IOException {
         return size() >= maxSize ? arrays[0].toComparable(keys.lastKey()) : null;
     }
     /**
@@ -158,7 +158,7 @@ final class CompositeValuesCollectorQueue implements Releasable {
      */
     CompositeKey toCompositeKey(int slot) throws IOException {
         assert slot < maxSize;
-        Comparable<?>[] values = new Comparable<?>[arrays.length];
+        Comparable[] values = new Comparable[arrays.length];
         for (int i = 0; i < values.length; i++) {
             values[i] = arrays[i].toComparable(slot);
         }
@@ -178,7 +178,7 @@ final class CompositeValuesCollectorQueue implements Releasable {
      * for each document.
      * The provided collector <code>in</code> is called on each composite bucket.
      */
-    LeafBucketCollector getLeafCollector(Comparable<?> forceLeadSourceValue,
+    LeafBucketCollector getLeafCollector(Comparable forceLeadSourceValue,
                                          LeafReaderContext context, LeafBucketCollector in) throws IOException {
         int last = arrays.length - 1;
         LeafBucketCollector collector = in;

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/DoubleValuesSource.java

@@ -108,7 +108,7 @@ class DoubleValuesSource extends SingleDimensionValuesSource<Double> {
     }
 
     @Override
-    void setAfter(Comparable<?> value) {
+    void setAfter(Comparable value) {
         if (missingBucket && value == null) {
             afterValue = null;
         } else if (value instanceof Number) {
@@ -151,7 +151,7 @@ class DoubleValuesSource extends SingleDimensionValuesSource<Double> {
     }
 
     @Override
-    LeafBucketCollector getLeafCollector(Comparable<?> value, LeafReaderContext context, LeafBucketCollector next) {
+    LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) {
         if (value.getClass() != Double.class) {
             throw new IllegalArgumentException("Expected Double, got " + value.getClass());
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/GlobalOrdinalValuesSource.java

@@ -89,7 +89,7 @@ class GlobalOrdinalValuesSource extends SingleDimensionValuesSource<BytesRef> {
     }
 
     @Override
-    void setAfter(Comparable<?> value) {
+    void setAfter(Comparable value) {
         if (missingBucket && value == null) {
             afterValue = null;
             afterValueGlobalOrd = -1L;
@@ -138,7 +138,7 @@ class GlobalOrdinalValuesSource extends SingleDimensionValuesSource<BytesRef> {
     }
 
     @Override
-    LeafBucketCollector getLeafCollector(Comparable<?> value, LeafReaderContext context, LeafBucketCollector next) throws IOException {
+    LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) throws IOException {
         if (value.getClass() != BytesRef.class) {
             throw new IllegalArgumentException("Expected BytesRef, got " + value.getClass());
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java

@@ -227,7 +227,7 @@ public class InternalComposite
     }
 
     static class InternalBucket extends InternalMultiBucketAggregation.InternalBucket
-            implements CompositeAggregation.Bucket, KeyComparable<InternalBucket> {
+        implements CompositeAggregation.Bucket, KeyComparable<InternalBucket> {
 
         private final CompositeKey key;
         private final long docCount;
@@ -341,7 +341,7 @@ public class InternalComposite
                 }
                 assert key.get(i).getClass() == other.key.get(i).getClass();
                 @SuppressWarnings("unchecked")
-                int cmp = ((Comparable) key.get(i)).compareTo(other.key.get(i)) * reverseMuls[i];
+                int cmp = key.get(i).compareTo(other.key.get(i)) * reverseMuls[i];
                 if (cmp != 0) {
                     return cmp;
                 }
@@ -392,12 +392,12 @@ public class InternalComposite
         return obj;
     }
 
-    private static class ArrayMap extends AbstractMap<String, Object> {
+    static class ArrayMap extends AbstractMap<String, Object> implements Comparable<ArrayMap> {
         final List<String> keys;
+        final Comparable[] values;
         final List<DocValueFormat> formats;
-        final Object[] values;
 
-        ArrayMap(List<String> keys, List<DocValueFormat> formats, Object[] values) {
+        ArrayMap(List<String> keys, List<DocValueFormat> formats, Comparable[] values) {
             assert keys.size() == values.length && keys.size() == formats.size();
             this.keys = keys;
             this.formats = formats;
@@ -447,5 +447,45 @@ public class InternalComposite
                 }
             };
         }
+
+        @Override
+        public int compareTo(ArrayMap that) {
+            if (that == this) {
+                return 0;
+            }
+
+            int idx = 0;
+            int max = Math.min(this.keys.size(), that.keys.size());
+            while (idx < max) {
+                int compare = compareNullables(keys.get(idx), that.keys.get(idx));
+                if (compare == 0) {
+                    compare = compareNullables(values[idx], that.values[idx]);
+                }
+                if (compare != 0) {
+                    return compare;
+                }
+                idx++;
+            }
+            if (idx < keys.size()) {
+                return 1;
+            }
+            if (idx < that.keys.size()) {
+                return -1;
+            }
+            return 0;
+        }
+    }
+
+    private static int compareNullables(Comparable a, Comparable b) {
+        if (a == b) {
+            return 0;
+        }
+        if (a == null) {
+            return -1;
+        }
+        if (b == null) {
+            return 1;
+        }
+        return a.compareTo(b);
     }
 }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/LongValuesSource.java

@@ -125,7 +125,7 @@ class LongValuesSource extends SingleDimensionValuesSource<Long> {
     }
 
     @Override
-    void setAfter(Comparable<?> value) {
+    void setAfter(Comparable value) {
         if (missingBucket && value == null) {
             afterValue = null;
         } else if (value instanceof Number) {
@@ -169,7 +169,7 @@ class LongValuesSource extends SingleDimensionValuesSource<Long> {
     }
 
     @Override
-    LeafBucketCollector getLeafCollector(Comparable<?> value, LeafReaderContext context, LeafBucketCollector next) {
+    LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) {
         if (value.getClass() != Long.class) {
             throw new IllegalArgumentException("Expected Long, got " + value.getClass());
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/PointsSortedDocsProducer.java

@@ -54,7 +54,7 @@ class PointsSortedDocsProducer extends SortedDocsProducer {
             return DocIdSet.EMPTY;
         }
         long lowerBucket = Long.MIN_VALUE;
-        Comparable<?> lowerValue = queue.getLowerValueLeadSource();
+        Comparable lowerValue = queue.getLowerValueLeadSource();
         if (lowerValue != null) {
             if (lowerValue.getClass() != Long.class) {
                 throw new IllegalStateException("expected Long, got " + lowerValue.getClass());
@@ -63,7 +63,7 @@ class PointsSortedDocsProducer extends SortedDocsProducer {
         }
 
         long upperBucket = Long.MAX_VALUE;
-        Comparable<?> upperValue = queue.getUpperValueLeadSource();
+        Comparable upperValue = queue.getUpperValueLeadSource();
         if (upperValue != null) {
             if (upperValue.getClass() != Long.class) {
                 throw new IllegalStateException("expected Long, got " + upperValue.getClass());

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/SingleDimensionValuesSource.java

@@ -102,7 +102,7 @@ abstract class SingleDimensionValuesSource<T extends Comparable<T>> implements R
     /**
      * Sets the after value for this source. Values that compares smaller are filtered.
      */
-    abstract void setAfter(Comparable<?> value);
+    abstract void setAfter(Comparable value);
 
     /**
      * Returns the after value set for this source.
@@ -129,7 +129,7 @@ abstract class SingleDimensionValuesSource<T extends Comparable<T>> implements R
      * Creates a {@link LeafBucketCollector} that sets the current value for each document to the provided
      * <code>value</code> and invokes {@link LeafBucketCollector#collect} on the provided <code>next</code> collector.
      */
-    abstract LeafBucketCollector getLeafCollector(Comparable<?> value,
+    abstract LeafBucketCollector getLeafCollector(Comparable value,
                                                   LeafReaderContext context, LeafBucketCollector next) throws IOException;
 
     /**

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/SortedDocsProducer.java

@@ -51,7 +51,7 @@ abstract class SortedDocsProducer {
      * composite buckets.
      */
     protected boolean processBucket(CompositeValuesCollectorQueue queue, LeafReaderContext context, DocIdSetIterator iterator,
-                                    Comparable<?> leadSourceBucket, @Nullable DocIdSetBuilder builder) throws IOException {
+                                    Comparable leadSourceBucket, @Nullable DocIdSetBuilder builder) throws IOException {
         final int[] topCompositeCollected = new int[1];
         final boolean[] hasCollected = new boolean[1];
         final LeafBucketCollector queueCollector = new LeafBucketCollector() {

server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/InternalCompositeTests.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.aggregations.bucket.composite;
 
+import com.google.common.collect.Lists;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.time.DateFormatter;
@@ -42,9 +43,14 @@ import java.util.List;
 import java.util.Map;
 import java.util.TreeSet;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLengthBetween;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 public class InternalCompositeTests extends InternalMultiBucketAggregationTestCase<InternalComposite> {
@@ -240,4 +246,137 @@ public class InternalCompositeTests extends InternalMultiBucketAggregationTestCa
             assertThat(bucket.getDocCount(), equalTo(expectedBucket.getDocCount()*numSame));
         }
     }
+
+    public void testCompareCompositeKeyBiggerFieldName() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field3", "field2"),
+            new Comparable[]{1, 2}
+        );
+        assertThat(key1.compareTo(key2), lessThan(0));
+    }
+
+    public void testCompareCompositeKeySmallerFieldName() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field3", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        assertThat(key1.compareTo(key2), greaterThan(0));
+    }
+
+    public void testCompareCompositeKeyBiggerValue() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field3", "field2"),
+            new Comparable[]{2, 3}
+        );
+        assertThat(key1.compareTo(key2), lessThan(0));
+    }
+
+    public void testCompareCompositeKeySmallerValue() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field3", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{2, 3}
+        );
+        assertThat(key1.compareTo(key2), greaterThan(0));
+    }
+
+    public void testCompareCompositeKeyNullValueIsSmaller1() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{null, 20}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        assertThat(key1.compareTo(key2), lessThan(0));
+    }
+
+    public void testCompareCompositeKeyNullValueIsSmaller2() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{null, 20}
+        );
+        assertThat(key1.compareTo(key2), greaterThan(0));
+    }
+
+    public void testCompareCompositeKeyMoreFieldsIsGreater() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(Lists.newArrayList("field1", "field2", "field3"),new Comparable[]{1, 2, null});
+        assertThat(key1.compareTo(key2), lessThan(0));
+    }
+
+    public void testCompareCompositeKeyLessFieldsIsLesser() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2", "field3"),
+            new Comparable[]{1, 2, null}
+        );
+        InternalComposite.ArrayMap key2 = createMap(Lists.newArrayList("field1", "field2"),new Comparable[]{1, 2});
+        assertThat(key1.compareTo(key2), greaterThan(0));
+    }
+
+    public void testCompareCompositeKeyEqual() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2", "field3"),
+            new Comparable[]{null, 1, 2}
+        );
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2", "field3"),
+            new Comparable[]{null, 1, 2}
+        );
+        assertThat(key1.compareTo(key1), equalTo(0));
+        assertThat(key1.equals(key1), is(true));
+
+        assertThat(key1.compareTo(key2), equalTo(0));
+        assertThat(key1.equals(key2), is(true));
+        assertThat(key2.equals(key1), is(true));
+    }
+
+    public void testCompareCompositeKeyValuesHaveDifferentTypes() {
+        InternalComposite.ArrayMap key1 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{1, 2}
+        );
+
+        InternalComposite.ArrayMap key2 = createMap(
+            Lists.newArrayList("field1", "field2"),
+            new Comparable[]{"1", 2}
+        );
+
+        ClassCastException exception = expectThrows(ClassCastException.class, () -> key1.compareTo(key2));
+        assertThat(exception.getMessage(),
+            containsString("java.lang.String cannot be cast to"));
+    }
+
+    private InternalComposite.ArrayMap createMap(List<String> fields, Comparable[] values) {
+        List<DocValueFormat> formats = IntStream.range(0, fields.size())
+            .mapToObj(i -> DocValueFormat.RAW).collect(Collectors.toList());
+        return new InternalComposite.ArrayMap(
+            fields,
+            formats,
+            values
+        );
+    }
 }