Replace TopBucketBuilder with a BucketPriorityQueue (#112602)

BucketPriorityQueue is a much better option nowadays

Author: iverase

server/src/main/java/org/elasticsearch/search/aggregations/BucketOrder.java

@@ -133,7 +133,7 @@ public final void validate(Aggregator aggregator) throws AggregationExecutionExc
      * The comparator might need to reduce the {@link DelayedBucket} and therefore we need to provide the
      * reducer and the reduce context.The context must be on the final reduce phase.
      */
-    abstract <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
+    public abstract <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
         BiFunction<List<B>, AggregationReduceContext, B> reduce,
         AggregationReduceContext reduceContext
     );

server/src/main/java/org/elasticsearch/search/aggregations/DelayedBucket.java

@@ -92,7 +92,7 @@ public String toString() {
      * Called to mark a bucket as non-competitive so it can release it can release
      * any sub-buckets from the breaker.
      */
-    void nonCompetitive(AggregationReduceContext reduceContext) {
+    public void nonCompetitive(AggregationReduceContext reduceContext) {
         if (reduced != null) {
             // -1 for itself, -countInnerBucket for all the sub-buckets.
             reduceContext.consumeBucketsAndMaybeBreak(-1 - InternalMultiBucketAggregation.countInnerBucket(reduced));

server/src/main/java/org/elasticsearch/search/aggregations/InternalOrder.java

@@ -81,7 +81,7 @@ public Comparator<Bucket> comparator() {
         }
 
         @Override
-        <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
+        public <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
             BiFunction<List<B>, AggregationReduceContext, B> reduce,
             AggregationReduceContext reduceContext
         ) {
@@ -216,7 +216,7 @@ public Comparator<Bucket> comparator() {
         }
 
         @Override
-        <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
+        public <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
             BiFunction<List<B>, AggregationReduceContext, B> reduce,
             AggregationReduceContext reduceContext
         ) {
@@ -284,7 +284,7 @@ public Comparator<Bucket> comparator() {
         }
 
         @Override
-        <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
+        public <B extends InternalMultiBucketAggregation.InternalBucket> Comparator<DelayedBucket<B>> delayedBucketComparator(
             BiFunction<List<B>, AggregationReduceContext, B> reduce,
             AggregationReduceContext reduceContext
         ) {

server/src/main/java/org/elasticsearch/search/aggregations/TopBucketBuilder.java

@@ -20,13 +20,13 @@
 import org.elasticsearch.search.aggregations.InternalMultiBucketAggregation;
 import org.elasticsearch.search.aggregations.InternalOrder;
 import org.elasticsearch.search.aggregations.KeyComparable;
-import org.elasticsearch.search.aggregations.TopBucketBuilder;
 import org.elasticsearch.search.aggregations.bucket.IteratorAndCurrent;
 import org.elasticsearch.search.aggregations.support.SamplingContext;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
@@ -295,19 +295,34 @@ public InternalAggregation get() {
                     }
                 });
             } else if (reduceContext.isFinalReduce()) {
-                TopBucketBuilder<B> top = TopBucketBuilder.build(
-                    getRequiredSize(),
-                    getOrder(),
-                    removed -> otherDocCount[0] += removed.getDocCount(),
+                final Comparator<DelayedBucket<B>> comparator = getOrder().delayedBucketComparator(
                     AbstractInternalTerms.this::reduceBucket,
                     reduceContext
                 );
-                thisReduceOrder = reduceBuckets(bucketsList, getThisReduceOrder(), bucket -> {
-                    if (bucket.getDocCount() >= getMinDocCount()) {
-                        top.add(bucket);
+                try (
+                    BucketPriorityQueue<DelayedBucket<B>> top = new BucketPriorityQueue<>(
+                        getRequiredSize(),
+                        reduceContext.bigArrays(),
+                        comparator
+                    )
+                ) {
+                    thisReduceOrder = reduceBuckets(bucketsList, getThisReduceOrder(), bucket -> {
+                        if (bucket.getDocCount() >= getMinDocCount()) {
+                            final DelayedBucket<B> removed = top.insertWithOverflow(bucket);
+                            if (removed != null) {
+                                otherDocCount[0] += removed.getDocCount();
+                                removed.nonCompetitive(reduceContext);
+                            }
+                        }
+                    });
+                    // size is an integer as it should be <= getRequiredSize()
+                    final int size = (int) top.size();
+                    result = new ArrayList<>(size);
+                    for (int i = 0; i < size; i++) {
+                        result.add(top.pop().reduced(AbstractInternalTerms.this::reduceBucket, reduceContext));
                     }
-                });
-                result = top.build();
+                    Collections.reverse(result);
+                }
             } else {
                 result = new ArrayList<>();
                 thisReduceOrder = reduceBuckets(