Sorting keys for merge sort to work

Signed-off-by: expani <anijainc@amazon.com>

Author: expani

server/src/main/java/org/opensearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -771,6 +771,7 @@ public List<InternalAggregation> convert(Map<String, Object[]> shardResult, Sear
             ));
             row++;
         }
+        buckets.sort(InternalComposite.InternalBucket::compareKey);
         CompositeKey lastBucket = buckets.isEmpty() ? null : buckets.getLast().getRawKey();
         return List.of(
             new InternalComposite(

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/AbstractStringTermsAggregator.java

@@ -51,6 +51,8 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -116,34 +118,38 @@ protected SignificantStringTerms buildEmptySignificantTermsAggregation(long subs
     }
 
     @Override
-    public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row, SearchContext searchContext) {
-        String termKey = (String) searchContext.convertToComparable(shardResult.get(name)[row]);
-        long docCount = 1;
-
-        List<InternalAggregation> subAggs = new ArrayList<>();
-        for (Aggregator aggregator : subAggregators) {
-            if (aggregator instanceof ShardResultConvertor convertor) {
-                InternalAggregation subAgg = convertor.convertRow(shardResult, row, searchContext);
-                if (aggregator instanceof ValueCountAggregator) {
-                    docCount = ((InternalValueCount) subAgg).getValue();
+    public List<InternalAggregation> convert(Map<String, Object[]> shardResult, SearchContext searchContext) {
+        int rowCount = shardResult.get(shardResult.keySet().stream().findFirst().get()).length;
+        List<StringTerms.Bucket> buckets = new ArrayList<>(rowCount);
+        for (int row = 0; row < rowCount; row++) {
+            String termKey = (String) searchContext.convertToComparable(shardResult.get(name)[row]);
+            List<InternalAggregation> subAggs = new ArrayList<>();
+            long docCount = 1;
+            for (Aggregator aggregator : subAggregators) {
+                if (aggregator instanceof ShardResultConvertor convertor) {
+                    InternalAggregation subAgg = convertor.convertRow(shardResult, row, searchContext);
+                    if (aggregator instanceof ValueCountAggregator) {
+                        docCount = ((InternalValueCount) subAgg).getValue();
+                    }
+                    subAggs.add(subAgg);
                 }
-                subAggs.add(subAgg);
             }
-        }
 
+            buckets.add(new StringTerms.Bucket(
+                new BytesRef(termKey),
+                docCount,
+                InternalAggregations.from(subAggs),
+                showTermDocCountError,
+                0,
+                format
+            ));
+        }
         BucketOrder reduceOrder = order;
         if (isKeyOrder(order) == false) {
             reduceOrder = InternalOrder.key(true);
+            buckets.sort(reduceOrder.comparator());
         }
-        StringTerms.Bucket bucket = new StringTerms.Bucket(
-            new BytesRef(termKey),
-            docCount,
-            InternalAggregations.from(subAggs),
-            showTermDocCountError,
-            0,
-            format
-        );
-        return new StringTerms(
+        return List.of(new StringTerms(
             name,
             reduceOrder,
             order,
@@ -152,10 +158,10 @@ public InternalAggregation convertRow(Map<String, Object[]> shardResult, int row
             bucketCountThresholds.getShardSize(),
             showTermDocCountError,
             0,
-            List.of(bucket),
+            buckets,
             0,
             bucketCountThresholds
-        );
+        ));
     }
 
 }

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/MultiTermsAggregator.java

@@ -709,7 +709,7 @@ static InternalValuesSource doubleValueSource(ValuesSource.Numeric valuesSource,
     @Override
     public List<InternalAggregation> convert(Map<String, Object[]> shardResult, SearchContext searchContext) {
         int rowCount = shardResult.isEmpty() ? 0 : shardResult.get(fields.getFirst()).length ;
-        List<InternalMultiTerms.Bucket> buckets = new ArrayList<>();
+        List<InternalMultiTerms.Bucket> buckets = new ArrayList<>(rowCount);
         for (int i = 0; i < rowCount; i++) {
             final int j = i;
             List<Object> key = fields.stream().map(fieldName -> (Object) searchContext.convertToComparable(shardResult.get(fieldName)[j])).toList();
@@ -725,13 +725,10 @@ public List<InternalAggregation> convert(Map<String, Object[]> shardResult, Sear
             }
             buckets.add(new InternalMultiTerms.Bucket(key, docCount, InternalAggregations.from(subAggs), showTermDocCountError, 0, formats));
         }
-        // TODO : Not reducing using Priority Queue into top buckets as depending on Substrait plan.
-        BucketOrder reduceOrder;
+        BucketOrder reduceOrder = order;
         if (isKeyOrder(order) == false) {
             reduceOrder = InternalOrder.key(true);
-            //buckets.sort(reduceOrder.comparator());
-        } else {
-            reduceOrder = order;
+            buckets.sort(reduceOrder.comparator());
         }
         return Collections.singletonList(new InternalMultiTerms(
             name,