Updated the final reduce part to select topN buckets

Signed-off-by: Vinay Krishna Pudyodu <vinkrish.neo@gmail.com>

Author: vinaykpud

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

@@ -459,68 +459,80 @@ For backward compatibility, we disable the merge sort and use ({@link InternalTe
         final B[] list;
         if (reduceContext.isFinalReduce() || reduceContext.isSliceLevel()) {
             final int size = Math.min(localBucketCountThresholds.getRequiredSize(), reducedBuckets.size());
-            if (size < reducedBuckets.size()) {
-                Comparator<MultiBucketsAggregation.Bucket> cmp = order.comparator();
-                B[] reducedBucketsArr  = createBucketsArray(reducedBuckets.size());;
-                for (int i = 0; i < reducedBuckets.size(); i++) {
-                    reducedBucketsArr[i] = reducedBuckets.get(i);
-                }
-                ArrayUtil.select(
-                    reducedBucketsArr,
-                    0,
-                    reducedBuckets.size(),
-                    size,
-                    cmp
-                );
-                int sz = 0;
-                for (B bucket : reducedBucketsArr) {
+
+            if (size <= 1000) {
+                final BucketPriorityQueue<B> ordered = new BucketPriorityQueue<>(size, order.comparator());
+                for (B bucket : reducedBuckets) {
                     if (sumDocCountError == -1) {
                         bucket.setDocCountError(-1);
                     } else {
                         final long finalSumDocCountError = sumDocCountError;
                         bucket.setDocCountError(docCountError -> docCountError + finalSumDocCountError);
                     }
                     if (bucket.getDocCount() >= localBucketCountThresholds.getMinDocCount()) {
-                        B removed = ((sz == size) ? bucket : null);
+                        B removed = ordered.insertWithOverflow(bucket);
                         if (removed != null) {
                             otherDocCount += removed.getDocCount();
                             reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(removed));
                         } else {
-                            sz++;
                             reduceContext.consumeBucketsAndMaybeBreak(1);
                         }
                     } else {
                         reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(bucket));
                     }
                 }
-                list = createBucketsArray(sz);
-                if (sz >= 0) System.arraycopy(reducedBucketsArr, 0, list, 0, sz);
-                Arrays.sort(list, cmp);
+                list = createBucketsArray(ordered.size());
+                for (int i = ordered.size() - 1; i >= 0; i--) {
+                    list[i] = ordered.pop();
+                }
             } else {
-                final BucketPriorityQueue<B> ordered = new BucketPriorityQueue<>(size, order.comparator());
+                final Comparator<MultiBucketsAggregation.Bucket> cmp = order.comparator();
+                int validBucketCount = 0;
+                // Process buckets and update doc count errors and count valid buckets
                 for (B bucket : reducedBuckets) {
                     if (sumDocCountError == -1) {
                         bucket.setDocCountError(-1);
                     } else {
                         final long finalSumDocCountError = sumDocCountError;
                         bucket.setDocCountError(docCountError -> docCountError + finalSumDocCountError);
                     }
+
+                    if (bucket.getDocCount() < localBucketCountThresholds.getMinDocCount()) {
+                        reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(bucket));
+                    } else {
+                        validBucketCount++;
+                    }
+                }
+                // Create array and populate with valid buckets
+                B[] validBuckets = createBucketsArray(validBucketCount);
+                int arrayIndex = 0;
+                for (B bucket : reducedBuckets) {
                     if (bucket.getDocCount() >= localBucketCountThresholds.getMinDocCount()) {
-                        B removed = ordered.insertWithOverflow(bucket);
-                        if (removed != null) {
-                            otherDocCount += removed.getDocCount();
-                            reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(removed));
-                        } else {
-                            reduceContext.consumeBucketsAndMaybeBreak(1);
-                        }
+                        validBuckets[arrayIndex++] = bucket;
+                    }
+                }
+                // Select top buckets if needed
+                if (size < validBuckets.length) {
+                    ArrayUtil.select(validBuckets, 0, validBuckets.length, size, cmp);
+                }
+                // Process selected buckets and calculate otherDocCount
+                int finalSize = Math.min(size, validBuckets.length);
+                for (int i = 0; i < validBuckets.length; i++) {
+                    B bucket = validBuckets[i];
+                    if (i < finalSize) {
+                        reduceContext.consumeBucketsAndMaybeBreak(1);
                     } else {
+                        otherDocCount += bucket.getDocCount();
                         reduceContext.consumeBucketsAndMaybeBreak(-countInnerBucket(bucket));
                     }
                 }
-                list = createBucketsArray(ordered.size());
-                for (int i = ordered.size() - 1; i >= 0; i--) {
-                    list[i] = ordered.pop();
+                if (finalSize != validBuckets.length) {
+                    list = createBucketsArray(finalSize);
+                    System.arraycopy(validBuckets, 0, list, 0, finalSize);
+                } else {
+                    list = validBuckets;
                 }
+                Arrays.sort(list, cmp);
             }
         } else {
             // we can prune the list on partial reduce if the aggregation is ordered by key