Fix merging of terms aggregation with compound order

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalTerms.java

@@ -45,6 +45,7 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.search.aggregations.InternalOrder.isKeyAsc;
 import static org.elasticsearch.search.aggregations.InternalOrder.isKeyOrder;
 
 public abstract class InternalTerms<A extends InternalTerms<A, B>, B extends InternalTerms.Bucket<B>>
@@ -257,9 +258,9 @@ private long getDocCountError(InternalTerms<?, ?> terms) {
     }
 
     private List<B> reduceMergeSort(List<InternalAggregation> aggregations,
-                                    BucketOrder reduceOrder, ReduceContext reduceContext) {
-        assert isKeyOrder(reduceOrder);
-        final Comparator<MultiBucketsAggregation.Bucket> cmp = reduceOrder.comparator();
+                                    BucketOrder thisReduceOrder, ReduceContext reduceContext) {
+        assert isKeyOrder(thisReduceOrder);
+        final Comparator<MultiBucketsAggregation.Bucket> cmp = thisReduceOrder.comparator();
         final PriorityQueue<IteratorAndCurrent<B>> pq = new PriorityQueue<>(aggregations.size()) {
             @Override
             protected boolean lessThan(IteratorAndCurrent<B> a, IteratorAndCurrent<B> b) {
@@ -369,15 +370,22 @@ public InternalAggregation reduce(List<InternalAggregation> aggregations, Reduce
                 bucket.docCountError -= thisAggDocCountError;
             }
         }
+
+        final List<B> reducedBuckets;
         /**
          * Buckets returned by a partial reduce or a shard response are sorted by key since {@link Version#V_7_10_0}.
          * That allows to perform a merge sort when reducing multiple aggregations together.
          * For backward compatibility, we disable the merge sort and use ({@link InternalTerms#reduceLegacy} if any of
          * the provided aggregations use a different {@link InternalTerms#reduceOrder}.
          */
         BucketOrder thisReduceOrder = getReduceOrder(aggregations);
-        List<B> reducedBuckets = isKeyOrder(thisReduceOrder) ?
-            reduceMergeSort(aggregations, thisReduceOrder, reduceContext) : reduceLegacy(aggregations, reduceContext);
+        if (isKeyOrder(thisReduceOrder)) {
+            // extract the primary sort in case this is a compound order.
+            thisReduceOrder = InternalOrder.key(isKeyAsc(thisReduceOrder) ? true : false);
+            reducedBuckets = reduceMergeSort(aggregations, thisReduceOrder, reduceContext);
+        } else {
+            reducedBuckets = reduceLegacy(aggregations, reduceContext);
+        }
         final B[] list;
         if (reduceContext.isFinalReduce()) {
             final int size = Math.min(requiredSize, reducedBuckets.size());

server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/StringTermsTests.java

@@ -34,13 +34,26 @@
 import java.util.Set;
 
 public class StringTermsTests extends InternalTermsTestCase {
+    private BytesRef[] dict;
+
+    public synchronized void generateRandomDict() {
+        if (dict == null) {
+            Set<BytesRef> terms = new HashSet<>();
+            int numTerms = randomIntBetween(2, 100);
+            for (int i = 0; i < numTerms; i++) {
+                terms.add(new BytesRef(randomAlphaOfLength(10)));
+            }
+            dict = terms.stream().toArray(BytesRef[]::new);
+        }
+    }
 
     @Override
     protected InternalTerms<?, ?> createTestInstance(String name,
                                                      Map<String, Object> metadata,
                                                      InternalAggregations aggregations,
                                                      boolean showTermDocCountError,
                                                      long docCountError) {
+        generateRandomDict();
         BucketOrder order = BucketOrder.count(false);
         long minDocCount = 1;
         int requiredSize = 3;
@@ -51,11 +64,14 @@ public class StringTermsTests extends InternalTermsTestCase {
         final int numBuckets = randomNumberOfBuckets();
         Set<BytesRef> terms = new HashSet<>();
         for (int i = 0; i < numBuckets; ++i) {
-            BytesRef term = randomValueOtherThanMany(b -> terms.add(b) == false, () -> new BytesRef(randomAlphaOfLength(10)));
-            int docCount = randomIntBetween(1, 100);
-            buckets.add(new StringTerms.Bucket(term, docCount, aggregations, showTermDocCountError, docCountError, format));
+            BytesRef term = dict[randomIntBetween(0, dict.length-1)];
+            if (terms.add(term)) {
+                int docCount = randomIntBetween(1, 100);
+                buckets.add(new StringTerms.Bucket(term, docCount, aggregations, showTermDocCountError, docCountError, format));
+            }
         }
-        BucketOrder reduceOrder = rarely() ? order : BucketOrder.key(true);
+        BucketOrder reduceOrder = randomBoolean() ?
+            BucketOrder.compound(BucketOrder.key(true), BucketOrder.count(false)) : BucketOrder.key(true);
         Collections.sort(buckets, reduceOrder.comparator());
         return new StringTerms(name, reduceOrder, order, requiredSize, minDocCount,
                 metadata, format, shardSize, showTermDocCountError, otherDocCount, buckets, docCountError);