Updated bucket sorting when isKeyOrder is true

Signed-off-by: Vinay Krishna Pudyodu <[email protected]>

Author: vinaykpud

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

@@ -262,7 +262,7 @@ private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws
                 Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningBucketOrds[ordIdx]);
 
                 // When request size is smaller than 20% of total buckets, use priority queue to get topN buckets
-                if ((size < 0.2 * bucketsInOrd) || isKeyOrder(order)) {
+                if ((size < 0.2 * bucketsInOrd)) {
                     PriorityQueue<B> ordered = buildPriorityQueue(size);
                     while (ordsEnum.next()) {
                         long docCount = bucketDocCount(ordsEnum.ord());
@@ -277,21 +277,11 @@ private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws
                         spare = ordered.insertWithOverflow(spare);
                     }
                     // Get the top buckets
-                    B[] bucketsForOrd = buildBuckets(ordered.size());
-                    topBucketsPerOrd[ordIdx] = bucketsForOrd;
-                    if (isKeyOrder(order)) {
-                        for (int b = ordered.size() - 1; b >= 0; --b) {
-                            topBucketsPerOrd[ordIdx][b] = ordered.pop();
-                            otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][b].getDocCount();
-                        }
-                    } else {
-                        // sorted buckets not needed as they will be sorted by key in buildResult() which is different from
-                        // order in priority queue ordered
-                        Iterator<B> itr = ordered.iterator();
-                        for (int b = ordered.size() - 1; b >= 0; --b) {
-                            topBucketsPerOrd[ordIdx][b] = itr.next();
-                            otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][b].getDocCount();
-                        }
+                    topBucketsPerOrd[ordIdx] =  buildBuckets(ordered.size());;
+                    Iterator<B> itr = ordered.iterator();
+                    for (int b = ordered.size() - 1; b >= 0; --b) {
+                        topBucketsPerOrd[ordIdx][b] = itr.next();
+                        otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][b].getDocCount();
                     }
                 } else {
                     B[] bucketsForOrd = buildBuckets((int) bucketsInOrd);
@@ -511,10 +501,10 @@ LongTerms buildResult(long owningBucketOrd, long otherDocCount, LongTerms.Bucket
             final BucketOrder reduceOrder;
             if (isKeyOrder(order) == false) {
                 reduceOrder = InternalOrder.key(true);
-                Arrays.sort(topBuckets, reduceOrder.comparator());
             } else {
                 reduceOrder = order;
             }
+            Arrays.sort(topBuckets, reduceOrder.comparator());
             return new LongTerms(
                 name,
                 reduceOrder,
@@ -591,10 +581,10 @@ DoubleTerms buildResult(long owningBucketOrd, long otherDocCount, DoubleTerms.Bu
             final BucketOrder reduceOrder;
             if (isKeyOrder(order) == false) {
                 reduceOrder = InternalOrder.key(true);
-                Arrays.sort(topBuckets, reduceOrder.comparator());
             } else {
                 reduceOrder = order;
             }
+            Arrays.sort(topBuckets, reduceOrder.comparator());
             return new DoubleTerms(
                 name,
                 reduceOrder,
@@ -670,10 +660,10 @@ UnsignedLongTerms buildResult(long owningBucketOrd, long otherDocCount, Unsigned
             final BucketOrder reduceOrder;
             if (isKeyOrder(order) == false) {
                 reduceOrder = InternalOrder.key(true);
-                Arrays.sort(topBuckets, reduceOrder.comparator());
             } else {
                 reduceOrder = order;
             }
+            Arrays.sort(topBuckets, reduceOrder.comparator());
             return new UnsignedLongTerms(
                 name,
                 reduceOrder,

server/src/test/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregatorTests.java

@@ -196,12 +196,13 @@ private void testSearchCase(
         }
     }
 
-    public void testBuildAggregationsForHeapSort() throws IOException {
+    public void testBuildAggregationAlgorithmSelection() throws IOException {
         List<Long> dataSet = new ArrayList<>();
-        for (long i = 0; i < 150; i++) {
+        for (long i = 0; i < 100; i++) {
             dataSet.add(i);
         }
 
+        // HeapSort
         testSearchCase(new MatchAllDocsQuery(), dataSet, aggregation -> aggregation.field(LONG_FIELD).size(2), agg -> {
             assertEquals(2, agg.getBuckets().size());
             for (int i = 0; i < 2; i++) {
@@ -210,59 +211,15 @@ public void testBuildAggregationsForHeapSort() throws IOException {
                 assertThat(bucket.getDocCount(), equalTo(1L));
             }
         }, ValueType.NUMERIC);
-    }
-
-    public void testBuildAggregationsForQuickSelect() throws IOException {
-        List<Long> dataSet = new ArrayList<>();
-        for (long i = 0; i < 150; i++) {
-            dataSet.add(i);
-        }
 
-        testSearchCase(new MatchAllDocsQuery(), dataSet, aggregation -> aggregation.field(LONG_FIELD).size(50), agg -> {
-            assertEquals(50, agg.getBuckets().size());
+        // QuickSelect
+        testSearchCase(new MatchAllDocsQuery(), dataSet, aggregation -> aggregation.field(LONG_FIELD).size(20), agg -> {
+            assertEquals(20, agg.getBuckets().size());
             for (int i = 0; i < agg.getBuckets().size(); i++) {
                 LongTerms.Bucket bucket = (LongTerms.Bucket) agg.getBuckets().get(i);
                 assertThat(bucket.getKey(), equalTo((long) i));
                 assertThat(bucket.getDocCount(), equalTo(1L));
             }
         }, ValueType.NUMERIC);
     }
-
-    public void testKeyOrderingAlgorithmSelection() throws IOException {
-        List<Long> dataset = new ArrayList<>();
-        for (long i = 100; i > 0; i--) {
-            dataset.add(i);
-        }
-
-        testSearchCase(
-            new MatchAllDocsQuery(),
-            dataset,
-            aggregation -> aggregation.field(LONG_FIELD).size(50).order(org.opensearch.search.aggregations.BucketOrder.key(true)),
-            agg -> {
-                assertEquals(50, agg.getBuckets().size());
-                for (int i = 0; i < 50; i++) {
-                    LongTerms.Bucket bucket = (LongTerms.Bucket) agg.getBuckets().get(i);
-                    assertThat(bucket.getKey(), equalTo(i + 1L));
-                    assertThat(bucket.getDocCount(), equalTo(1L));
-                }
-            },
-            ValueType.NUMERIC
-        );
-
-        testSearchCase(
-            new MatchAllDocsQuery(),
-            dataset,
-            aggregation -> aggregation.field(LONG_FIELD).size(5).order(org.opensearch.search.aggregations.BucketOrder.key(false)),
-            agg -> {
-                assertEquals(5, agg.getBuckets().size());
-                for (int i = 0; i < agg.getBuckets().size(); i++) {
-                    LongTerms.Bucket bucket = (LongTerms.Bucket) agg.getBuckets().get(i);
-                    assertThat(bucket.getKey(), equalTo(100L - i));
-                    assertThat(bucket.getDocCount(), equalTo(1L));
-                }
-            },
-            ValueType.NUMERIC
-        );
-    }
-
 }