Browse Source

Use LongArray instead of long[] for owning ordinals when building Internal aggregations (#116874) (#117015)

This commit changes the signature of InternalAggregation#buildAggregations(long[]) to
InternalAggregation#buildAggregations(LongArray) to avoid allocations of humongous arrays.
Ignacio Vera 10 months ago
parent
commit
8a1231c248
49 changed files with 877 additions and 741 deletions
  1. 43 41
      modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/adjacency/AdjacencyMatrixAggregator.java
  2. 3 3
      modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/histogram/AutoDateHistogramAggregator.java
  3. 36 32
      modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/TimeSeriesAggregator.java
  4. 2 1
      modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ChildrenToParentAggregator.java
  5. 6 4
      modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ParentJoinAggregator.java
  6. 2 1
      modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ParentToChildrenAggregator.java
  7. 2 1
      server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java
  8. 4 3
      server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java
  9. 5 3
      server/src/main/java/org/elasticsearch/search/aggregations/Aggregator.java
  10. 4 3
      server/src/main/java/org/elasticsearch/search/aggregations/NonCollectingAggregator.java
  11. 14 12
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java
  12. 102 89
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java
  13. 4 3
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferableBucketAggregator.java
  14. 5 3
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java
  15. 39 36
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregator.java
  16. 69 57
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/countedterms/CountedTermsAggregator.java
  17. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregator.java
  18. 31 28
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregator.java
  19. 3 2
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/global/GlobalAggregator.java
  20. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/AbstractHistogramAggregator.java
  21. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java
  22. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateRangeHistogramAggregator.java
  23. 22 20
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/VariableWidthHistogramAggregator.java
  24. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/missing/MissingAggregator.java
  25. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java
  26. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/ReverseNestedAggregator.java
  27. 53 45
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/prefix/IpPrefixAggregator.java
  28. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java
  29. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregator.java
  30. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java
  31. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregator.java
  32. 2 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplerAggregator.java
  33. 58 52
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java
  34. 1 1
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalSignificantTerms.java
  35. 61 54
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java
  36. 47 42
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/MapStringTermsAggregator.java
  37. 54 50
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/NumericTermsAggregator.java
  38. 68 56
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/StringRareTermsAggregator.java
  39. 5 4
      server/src/main/java/org/elasticsearch/search/aggregations/metrics/MetricsAggregator.java
  40. 2 1
      server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java
  41. 2 1
      server/src/test/java/org/elasticsearch/search/aggregations/AdaptingAggregatorTests.java
  42. 2 1
      server/src/test/java/org/elasticsearch/search/aggregations/AggregatorBaseTests.java
  43. 15 5
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java
  44. 2 1
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregatorTests.java
  45. 2 1
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java
  46. 52 45
      x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/multiterms/MultiTermsAggregator.java
  47. 25 23
      x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java
  48. 4 2
      x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/frequentitemsets/mr/DelegatingCircuitBreakerService.java
  49. 4 3
      x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/frequentitemsets/mr/ItemSetMapReduceAggregator.java

+ 43 - 41
modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/adjacency/AdjacencyMatrixAggregator.java

@@ -15,6 +15,7 @@ import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -177,65 +178,66 @@ public class AdjacencyMatrixAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         // Buckets are ordered into groups - [keyed filters] [key1&key2 intersects]
-        int maxOrd = owningBucketOrds.length * totalNumKeys;
-        int totalBucketsToBuild = 0;
-        for (int ord = 0; ord < maxOrd; ord++) {
+        long maxOrd = owningBucketOrds.size() * totalNumKeys;
+        long totalBucketsToBuild = 0;
+        for (long ord = 0; ord < maxOrd; ord++) {
             if (bucketDocCount(ord) > 0) {
                 totalBucketsToBuild++;
             }
         }
-        long[] bucketOrdsToBuild = new long[totalBucketsToBuild];
-        int builtBucketIndex = 0;
-        for (int ord = 0; ord < maxOrd; ord++) {
-            if (bucketDocCount(ord) > 0) {
-                bucketOrdsToBuild[builtBucketIndex++] = ord;
-            }
-        }
-        assert builtBucketIndex == totalBucketsToBuild;
-        builtBucketIndex = 0;
-        var bucketSubAggs = buildSubAggsForBuckets(bucketOrdsToBuild);
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int owningBucketOrdIdx = 0; owningBucketOrdIdx < owningBucketOrds.length; owningBucketOrdIdx++) {
-            List<InternalAdjacencyMatrix.InternalBucket> buckets = new ArrayList<>(filters.length);
-            for (int i = 0; i < keys.length; i++) {
-                long bucketOrd = bucketOrd(owningBucketOrds[owningBucketOrdIdx], i);
-                long docCount = bucketDocCount(bucketOrd);
-                // Empty buckets are not returned because this aggregation will commonly be used under a
-                // a date-histogram where we will look for transactions over time and can expect many
-                // empty buckets.
-                if (docCount > 0) {
-                    InternalAdjacencyMatrix.InternalBucket bucket = new InternalAdjacencyMatrix.InternalBucket(
-                        keys[i],
-                        docCount,
-                        bucketSubAggs.apply(builtBucketIndex++)
-                    );
-                    buckets.add(bucket);
+        try (LongArray bucketOrdsToBuild = bigArrays().newLongArray(totalBucketsToBuild)) {
+            int builtBucketIndex = 0;
+            for (int ord = 0; ord < maxOrd; ord++) {
+                if (bucketDocCount(ord) > 0) {
+                    bucketOrdsToBuild.set(builtBucketIndex++, ord);
                 }
             }
-            int pos = keys.length;
-            for (int i = 0; i < keys.length; i++) {
-                for (int j = i + 1; j < keys.length; j++) {
-                    long bucketOrd = bucketOrd(owningBucketOrds[owningBucketOrdIdx], pos);
+            assert builtBucketIndex == totalBucketsToBuild;
+            builtBucketIndex = 0;
+            var bucketSubAggs = buildSubAggsForBuckets(bucketOrdsToBuild);
+            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+            for (int owningBucketOrdIdx = 0; owningBucketOrdIdx < results.length; owningBucketOrdIdx++) {
+                List<InternalAdjacencyMatrix.InternalBucket> buckets = new ArrayList<>(filters.length);
+                for (int i = 0; i < keys.length; i++) {
+                    long bucketOrd = bucketOrd(owningBucketOrds.get(owningBucketOrdIdx), i);
                     long docCount = bucketDocCount(bucketOrd);
-                    // Empty buckets are not returned due to potential for very sparse matrices
+                    // Empty buckets are not returned because this aggregation will commonly be used under a
+                    // a date-histogram where we will look for transactions over time and can expect many
+                    // empty buckets.
                     if (docCount > 0) {
-                        String intersectKey = keys[i] + separator + keys[j];
                         InternalAdjacencyMatrix.InternalBucket bucket = new InternalAdjacencyMatrix.InternalBucket(
-                            intersectKey,
+                            keys[i],
                             docCount,
                             bucketSubAggs.apply(builtBucketIndex++)
                         );
                         buckets.add(bucket);
                     }
-                    pos++;
                 }
+                int pos = keys.length;
+                for (int i = 0; i < keys.length; i++) {
+                    for (int j = i + 1; j < keys.length; j++) {
+                        long bucketOrd = bucketOrd(owningBucketOrds.get(owningBucketOrdIdx), pos);
+                        long docCount = bucketDocCount(bucketOrd);
+                        // Empty buckets are not returned due to potential for very sparse matrices
+                        if (docCount > 0) {
+                            String intersectKey = keys[i] + separator + keys[j];
+                            InternalAdjacencyMatrix.InternalBucket bucket = new InternalAdjacencyMatrix.InternalBucket(
+                                intersectKey,
+                                docCount,
+                                bucketSubAggs.apply(builtBucketIndex++)
+                            );
+                            buckets.add(bucket);
+                        }
+                        pos++;
+                    }
+                }
+                results[owningBucketOrdIdx] = new InternalAdjacencyMatrix(name, buckets, metadata());
             }
-            results[owningBucketOrdIdx] = new InternalAdjacencyMatrix(name, buckets, metadata());
+            assert builtBucketIndex == totalBucketsToBuild;
+            return results;
         }
-        assert builtBucketIndex == totalBucketsToBuild;
-        return results;
     }
 
     @Override

+ 3 - 3
modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/histogram/AutoDateHistogramAggregator.java

@@ -141,7 +141,7 @@ abstract class AutoDateHistogramAggregator extends DeferableBucketAggregator {
     protected final InternalAggregation[] buildAggregations(
         LongKeyedBucketOrds bucketOrds,
         LongToIntFunction roundingIndexFor,
-        long[] owningBucketOrds
+        LongArray owningBucketOrds
     ) throws IOException {
         return buildAggregationsForVariableBuckets(
             owningBucketOrds,
@@ -324,7 +324,7 @@ abstract class AutoDateHistogramAggregator extends DeferableBucketAggregator {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
             return buildAggregations(bucketOrds, l -> roundingIdx, owningBucketOrds);
         }
 
@@ -594,7 +594,7 @@ abstract class AutoDateHistogramAggregator extends DeferableBucketAggregator {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
             /*
              * Rebucket before building the aggregation to build as small as result
              * as possible.

+ 36 - 32
modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/timeseries/TimeSeriesAggregator.java

@@ -11,6 +11,8 @@ package org.elasticsearch.aggregations.bucket.timeseries;
 
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.mapper.RoutingPathFields;
@@ -30,6 +32,7 @@ import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
@@ -67,42 +70,43 @@ public class TimeSeriesAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         BytesRef spare = new BytesRef();
-        InternalTimeSeries.InternalBucket[][] allBucketsPerOrd = new InternalTimeSeries.InternalBucket[owningBucketOrds.length][];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-            List<InternalTimeSeries.InternalBucket> buckets = new ArrayList<>();
-            while (ordsEnum.next()) {
-                long docCount = bucketDocCount(ordsEnum.ord());
-                ordsEnum.readValue(spare);
-                InternalTimeSeries.InternalBucket bucket = new InternalTimeSeries.InternalBucket(
-                    BytesRef.deepCopyOf(spare), // Closing bucketOrds will corrupt the bytes ref, so need to make a deep copy here.
-                    docCount,
-                    null,
-                    keyed
-                );
-                bucket.bucketOrd = ordsEnum.ord();
-                buckets.add(bucket);
-                if (buckets.size() >= size) {
-                    break;
+        try (ObjectArray<InternalTimeSeries.InternalBucket[]> allBucketsPerOrd = bigArrays().newObjectArray(owningBucketOrds.size())) {
+            for (long ordIdx = 0; ordIdx < allBucketsPerOrd.size(); ordIdx++) {
+                BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
+                List<InternalTimeSeries.InternalBucket> buckets = new ArrayList<>();
+                while (ordsEnum.next()) {
+                    long docCount = bucketDocCount(ordsEnum.ord());
+                    ordsEnum.readValue(spare);
+                    InternalTimeSeries.InternalBucket bucket = new InternalTimeSeries.InternalBucket(
+                        BytesRef.deepCopyOf(spare), // Closing bucketOrds will corrupt the bytes ref, so need to make a deep copy here.
+                        docCount,
+                        null,
+                        keyed
+                    );
+                    bucket.bucketOrd = ordsEnum.ord();
+                    buckets.add(bucket);
+                    if (buckets.size() >= size) {
+                        break;
+                    }
                 }
+                // NOTE: after introducing _tsid hashing time series are sorted by (_tsid hash, @timestamp) instead of (_tsid, timestamp).
+                // _tsid hash and _tsid might sort differently, and out of order data might result in incorrect buckets due to _tsid value
+                // changes not matching _tsid hash changes. Changes in _tsid hash are handled creating a new bucket as a result of making
+                // the assumption that sorting data results in new buckets whenever there is a change in _tsid hash. This is no true anymore
+                // because we collect data sorted on (_tsid hash, timestamp) but build aggregation results sorted by (_tsid, timestamp).
+                buckets.sort(Comparator.comparing(bucket -> bucket.key));
+                allBucketsPerOrd.set(ordIdx, buckets.toArray(new InternalTimeSeries.InternalBucket[0]));
             }
-            // NOTE: after introducing _tsid hashing time series are sorted by (_tsid hash, @timestamp) instead of (_tsid, timestamp).
-            // _tsid hash and _tsid might sort differently, and out of order data might result in incorrect buckets due to _tsid value
-            // changes not matching _tsid hash changes. Changes in _tsid hash are handled creating a new bucket as a result of making
-            // the assumption that sorting data results in new buckets whenever there is a change in _tsid hash. This is no true anymore
-            // because we collect data sorted on (_tsid hash, timestamp) but build aggregation results sorted by (_tsid, timestamp).
-            buckets.sort(Comparator.comparing(bucket -> bucket.key));
-            allBucketsPerOrd[ordIdx] = buckets.toArray(new InternalTimeSeries.InternalBucket[0]);
-        }
-        buildSubAggsForAllBuckets(allBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
+            buildSubAggsForAllBuckets(allBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
 
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            result[ordIdx] = buildResult(allBucketsPerOrd[ordIdx]);
+            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(allBucketsPerOrd.size())];
+            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                result[ordIdx] = buildResult(allBucketsPerOrd.get(ordIdx));
+            }
+            return result;
         }
-        return result;
     }
 
     @Override
@@ -185,7 +189,7 @@ public class TimeSeriesAggregator extends BucketsAggregator {
     }
 
     InternalTimeSeries buildResult(InternalTimeSeries.InternalBucket[] topBuckets) {
-        return new InternalTimeSeries(name, List.of(topBuckets), keyed, metadata());
+        return new InternalTimeSeries(name, Arrays.asList(topBuckets), keyed, metadata());
     }
 
     @FunctionalInterface

+ 2 - 1
modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ChildrenToParentAggregator.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.join.aggregations;
 
 import org.apache.lucene.search.Query;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
@@ -44,7 +45,7 @@ public class ChildrenToParentAggregator extends ParentJoinAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalParent(

+ 6 - 4
modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ParentJoinAggregator.java

@@ -21,6 +21,7 @@ import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BitArray;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
@@ -115,7 +116,7 @@ public abstract class ParentJoinAggregator extends BucketsAggregator implements
     }
 
     @Override
-    protected void prepareSubAggs(long[] ordsToCollect) throws IOException {
+    protected void prepareSubAggs(LongArray ordsToCollect) throws IOException {
         IndexReader indexReader = searcher().getIndexReader();
         for (LeafReaderContext ctx : indexReader.leaves()) {
             Scorer childDocsScorer = outFilter.scorer(ctx);
@@ -158,9 +159,10 @@ public abstract class ParentJoinAggregator extends BucketsAggregator implements
                  * structure that maps a primitive long to a list of primitive
                  * longs.
                  */
-                for (long owningBucketOrd : ordsToCollect) {
-                    if (collectionStrategy.exists(owningBucketOrd, globalOrdinal)) {
-                        collectBucket(sub, docId, owningBucketOrd);
+                for (long ord = 0; ord < ordsToCollect.size(); ord++) {
+                    long ordToCollect = ordsToCollect.get(ord);
+                    if (collectionStrategy.exists(ordToCollect, globalOrdinal)) {
+                        collectBucket(sub, docId, ordToCollect);
                     }
                 }
             }

+ 2 - 1
modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ParentToChildrenAggregator.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.join.aggregations;
 
 import org.apache.lucene.search.Query;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
@@ -40,7 +41,7 @@ public class ParentToChildrenAggregator extends ParentJoinAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalChildren(

+ 2 - 1
server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java

@@ -27,6 +27,7 @@ import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.concurrent.AtomicArray;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.IndexSettings;
@@ -669,7 +670,7 @@ public class TransportSearchIT extends ESIntegTestCase {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) {
             return new InternalAggregation[] { buildEmptyAggregation() };
         }
 

+ 4 - 3
server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.search.aggregations;
 
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.search.profile.aggregation.InternalAggregationProfileTree;
 
@@ -98,10 +99,10 @@ public abstract class AdaptingAggregator extends Aggregator {
     }
 
     @Override
-    public final InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public final InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         InternalAggregation[] delegateResults = delegate.buildAggregations(owningBucketOrds);
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
+        InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+        for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
             result[ordIdx] = adapt(delegateResults[ordIdx]);
         }
         return result;

+ 5 - 3
server/src/main/java/org/elasticsearch/search/aggregations/Aggregator.java

@@ -13,6 +13,8 @@ import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.search.aggregations.support.AggregationPath;
 import org.elasticsearch.search.sort.SortOrder;
@@ -142,7 +144,7 @@ public abstract class Aggregator extends BucketCollector implements Releasable {
      * @return the results for each ordinal, in the same order as the array
      *         of ordinals
      */
-    public abstract InternalAggregation[] buildAggregations(long[] ordsToCollect) throws IOException;
+    public abstract InternalAggregation[] buildAggregations(LongArray ordsToCollect) throws IOException;
 
     /**
      * Release this aggregation and its sub-aggregations.
@@ -153,11 +155,11 @@ public abstract class Aggregator extends BucketCollector implements Releasable {
      * Build the result of this aggregation if it is at the "top level"
      * of the aggregation tree. If, instead, it is a sub-aggregation of
      * another aggregation then the aggregation that contains it will call
-     * {@link #buildAggregations(long[])}.
+     * {@link #buildAggregations(LongArray)}.
      */
     public final InternalAggregation buildTopLevel() throws IOException {
         assert parent() == null;
-        return buildAggregations(new long[] { 0 })[0];
+        return buildAggregations(BigArrays.NON_RECYCLING_INSTANCE.newLongArray(1, true))[0];
     }
 
     /**

+ 4 - 3
server/src/main/java/org/elasticsearch/search/aggregations/NonCollectingAggregator.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.search.aggregations;
 
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 
 import java.io.IOException;
@@ -39,9 +40,9 @@ public abstract class NonCollectingAggregator extends AggregatorBase {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
             results[ordIdx] = buildEmptyAggregation();
         }
         return results;

+ 14 - 12
server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java

@@ -20,6 +20,7 @@ import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.packed.PackedInts;
 import org.apache.lucene.util.packed.PackedLongValues;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.LongHash;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -146,7 +147,7 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
      * Replay the wrapped collector, but only on a selection of buckets.
      */
     @Override
-    public void prepareSelectedBuckets(long... selectedBuckets) throws IOException {
+    public void prepareSelectedBuckets(LongArray selectedBuckets) throws IOException {
         if (finished == false) {
             throw new IllegalStateException("Cannot replay yet, collection is not finished: postCollect() has not been called");
         }
@@ -154,9 +155,9 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
             throw new IllegalStateException("Already been replayed");
         }
 
-        this.selectedBuckets = new LongHash(selectedBuckets.length, BigArrays.NON_RECYCLING_INSTANCE);
-        for (long ord : selectedBuckets) {
-            this.selectedBuckets.add(ord);
+        this.selectedBuckets = new LongHash(selectedBuckets.size(), BigArrays.NON_RECYCLING_INSTANCE);
+        for (long i = 0; i < selectedBuckets.size(); i++) {
+            this.selectedBuckets.add(selectedBuckets.get(i));
         }
 
         boolean needsScores = scoreMode().needsScores();
@@ -232,21 +233,22 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
      * been collected directly.
      */
     @Override
-    public Aggregator wrap(final Aggregator in) {
+    public Aggregator wrap(final Aggregator in, BigArrays bigArrays) {
         return new WrappedAggregator(in) {
             @Override
-            public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+            public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
                 if (selectedBuckets == null) {
                     throw new IllegalStateException("Collection has not been replayed yet.");
                 }
-                long[] rebasedOrds = new long[owningBucketOrds.length];
-                for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                    rebasedOrds[ordIdx] = selectedBuckets.find(owningBucketOrds[ordIdx]);
-                    if (rebasedOrds[ordIdx] == -1) {
-                        throw new IllegalStateException("Cannot build for a bucket which has not been collected");
+                try (LongArray rebasedOrds = bigArrays.newLongArray(owningBucketOrds.size())) {
+                    for (long ordIdx = 0; ordIdx < owningBucketOrds.size(); ordIdx++) {
+                        rebasedOrds.set(ordIdx, selectedBuckets.find(owningBucketOrds.get(ordIdx)));
+                        if (rebasedOrds.get(ordIdx) == -1) {
+                            throw new IllegalStateException("Cannot build for a bucket which has not been collected");
+                        }
                     }
+                    return in.buildAggregations(rebasedOrds);
                 }
-                return in.buildAggregations(rebasedOrds);
             }
         };
     }

+ 102 - 89
server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java

@@ -10,7 +10,9 @@ package org.elasticsearch.search.aggregations.bucket;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.util.IntArray;
 import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.search.aggregations.AggregationErrors;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -155,22 +157,22 @@ public abstract class BucketsAggregator extends AggregatorBase {
     /**
      * Hook to allow taking an action before building the sub agg results.
      */
-    protected void prepareSubAggs(long[] ordsToCollect) throws IOException {}
+    protected void prepareSubAggs(LongArray ordsToCollect) throws IOException {}
 
     /**
      * Build the results of the sub-aggregations of the buckets at each of
      * the provided ordinals.
      * <p>
      * Most aggregations should probably use something like
-     * {@link #buildSubAggsForAllBuckets(Object[][], ToLongFunction, BiConsumer)}
-     * or {@link #buildAggregationsForVariableBuckets(long[], LongKeyedBucketOrds, BucketBuilderForVariable, ResultBuilderForVariable)}
-     * or {@link #buildAggregationsForFixedBucketCount(long[], int, BucketBuilderForFixedCount, Function)}
-     * or {@link #buildAggregationsForSingleBucket(long[], SingleBucketResultBuilder)}
+     * {@link #buildSubAggsForAllBuckets(ObjectArray, ToLongFunction, BiConsumer)}
+     * or {@link #buildAggregationsForVariableBuckets(LongArray, LongKeyedBucketOrds, BucketBuilderForVariable, ResultBuilderForVariable)}
+     * or {@link #buildAggregationsForFixedBucketCount(LongArray, int, BucketBuilderForFixedCount, Function)}
+     * or {@link #buildAggregationsForSingleBucket(LongArray, SingleBucketResultBuilder)}
      * instead of calling this directly.
      * @return the sub-aggregation results in the same order as the provided
      *         array of ordinals
      */
-    protected final IntFunction<InternalAggregations> buildSubAggsForBuckets(long[] bucketOrdsToCollect) throws IOException {
+    protected final IntFunction<InternalAggregations> buildSubAggsForBuckets(LongArray bucketOrdsToCollect) throws IOException {
         prepareSubAggs(bucketOrdsToCollect);
         InternalAggregation[][] aggregations = new InternalAggregation[subAggregators.length][];
         for (int i = 0; i < subAggregators.length; i++) {
@@ -204,26 +206,28 @@ public abstract class BucketsAggregator extends AggregatorBase {
      * @param setAggs how to set the sub-aggregation results on a bucket
      */
     protected final <B> void buildSubAggsForAllBuckets(
-        B[][] buckets,
+        ObjectArray<B[]> buckets,
         ToLongFunction<B> bucketToOrd,
         BiConsumer<B, InternalAggregations> setAggs
     ) throws IOException {
-        int totalBucketOrdsToCollect = 0;
-        for (B[] bucketsForOneResult : buckets) {
-            totalBucketOrdsToCollect += bucketsForOneResult.length;
+        long totalBucketOrdsToCollect = 0;
+        for (long b = 0; b < buckets.size(); b++) {
+            totalBucketOrdsToCollect += buckets.get(b).length;
         }
-        long[] bucketOrdsToCollect = new long[totalBucketOrdsToCollect];
-        int s = 0;
-        for (B[] bucketsForOneResult : buckets) {
-            for (B bucket : bucketsForOneResult) {
-                bucketOrdsToCollect[s++] = bucketToOrd.applyAsLong(bucket);
+
+        try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(totalBucketOrdsToCollect)) {
+            int s = 0;
+            for (long ord = 0; ord < buckets.size(); ord++) {
+                for (B bucket : buckets.get(ord)) {
+                    bucketOrdsToCollect.set(s++, bucketToOrd.applyAsLong(bucket));
+                }
             }
-        }
-        var results = buildSubAggsForBuckets(bucketOrdsToCollect);
-        s = 0;
-        for (B[] bucket : buckets) {
-            for (int b = 0; b < bucket.length; b++) {
-                setAggs.accept(bucket[b], results.apply(s++));
+            var results = buildSubAggsForBuckets(bucketOrdsToCollect);
+            s = 0;
+            for (long ord = 0; ord < buckets.size(); ord++) {
+                for (B value : buckets.get(ord)) {
+                    setAggs.accept(value, results.apply(s++));
+                }
             }
         }
     }
@@ -237,37 +241,38 @@ public abstract class BucketsAggregator extends AggregatorBase {
      * @param resultBuilder how to build a result from buckets
      */
     protected final <B> InternalAggregation[] buildAggregationsForFixedBucketCount(
-        long[] owningBucketOrds,
+        LongArray owningBucketOrds,
         int bucketsPerOwningBucketOrd,
         BucketBuilderForFixedCount<B> bucketBuilder,
         Function<List<B>, InternalAggregation> resultBuilder
     ) throws IOException {
-        int totalBuckets = owningBucketOrds.length * bucketsPerOwningBucketOrd;
-        long[] bucketOrdsToCollect = new long[totalBuckets];
-        int bucketOrdIdx = 0;
-        for (long owningBucketOrd : owningBucketOrds) {
-            long ord = owningBucketOrd * bucketsPerOwningBucketOrd;
-            for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
-                bucketOrdsToCollect[bucketOrdIdx++] = ord++;
+        try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(owningBucketOrds.size() * bucketsPerOwningBucketOrd)) {
+            int bucketOrdIdx = 0;
+            for (long i = 0; i < owningBucketOrds.size(); i++) {
+                long ord = owningBucketOrds.get(i) * bucketsPerOwningBucketOrd;
+                for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
+                    bucketOrdsToCollect.set(bucketOrdIdx++, ord++);
+                }
             }
-        }
-        bucketOrdIdx = 0;
-        var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int owningOrdIdx = 0; owningOrdIdx < owningBucketOrds.length; owningOrdIdx++) {
-            List<B> buckets = new ArrayList<>(bucketsPerOwningBucketOrd);
-            for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
-                buckets.add(
-                    bucketBuilder.build(
-                        offsetInOwningOrd,
-                        bucketDocCount(bucketOrdsToCollect[bucketOrdIdx]),
-                        subAggregationResults.apply(bucketOrdIdx++)
-                    )
-                );
+            bucketOrdIdx = 0;
+            var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
+
+            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+            for (int owningOrdIdx = 0; owningOrdIdx < results.length; owningOrdIdx++) {
+                List<B> buckets = new ArrayList<>(bucketsPerOwningBucketOrd);
+                for (int offsetInOwningOrd = 0; offsetInOwningOrd < bucketsPerOwningBucketOrd; offsetInOwningOrd++) {
+                    buckets.add(
+                        bucketBuilder.build(
+                            offsetInOwningOrd,
+                            bucketDocCount(bucketOrdsToCollect.get(bucketOrdIdx)),
+                            subAggregationResults.apply(bucketOrdIdx++)
+                        )
+                    );
+                }
+                results[owningOrdIdx] = resultBuilder.apply(buckets);
             }
-            results[owningOrdIdx] = resultBuilder.apply(buckets);
+            return results;
         }
-        return results;
     }
 
     @FunctionalInterface
@@ -280,17 +285,19 @@ public abstract class BucketsAggregator extends AggregatorBase {
      * @param owningBucketOrds owning bucket ordinals for which to build the results
      * @param resultBuilder how to build a result from the sub aggregation results
      */
-    protected final InternalAggregation[] buildAggregationsForSingleBucket(long[] owningBucketOrds, SingleBucketResultBuilder resultBuilder)
-        throws IOException {
+    protected final InternalAggregation[] buildAggregationsForSingleBucket(
+        LongArray owningBucketOrds,
+        SingleBucketResultBuilder resultBuilder
+    ) throws IOException {
         /*
          * It'd be entirely reasonable to call
          * `consumeBucketsAndMaybeBreak(owningBucketOrds.length)`
          * here but we don't because single bucket aggs never have.
          */
         var subAggregationResults = buildSubAggsForBuckets(owningBucketOrds);
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            results[ordIdx] = resultBuilder.build(owningBucketOrds[ordIdx], subAggregationResults.apply(ordIdx));
+        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+            results[ordIdx] = resultBuilder.build(owningBucketOrds.get(ordIdx), subAggregationResults.apply(ordIdx));
         }
         return results;
     }
@@ -307,54 +314,60 @@ public abstract class BucketsAggregator extends AggregatorBase {
      * @param bucketOrds hash of values to the bucket ordinal
      */
     protected final <B> InternalAggregation[] buildAggregationsForVariableBuckets(
-        long[] owningBucketOrds,
+        LongArray owningBucketOrds,
         LongKeyedBucketOrds bucketOrds,
         BucketBuilderForVariable<B> bucketBuilder,
         ResultBuilderForVariable<B> resultBuilder
     ) throws IOException {
         long totalOrdsToCollect = 0;
-        final int[] bucketsInOrd = new int[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            final long bucketCount = bucketOrds.bucketsInOrd(owningBucketOrds[ordIdx]);
-            bucketsInOrd[ordIdx] = (int) bucketCount;
-            totalOrdsToCollect += bucketCount;
-        }
-        if (totalOrdsToCollect > Integer.MAX_VALUE) {
-            // TODO: We should instrument this error. While it is correct for it to be a 400 class IllegalArgumentException, there is not
-            // much the user can do about that. If this occurs with any frequency, we should do something about it.
-            throw new IllegalArgumentException(
-                "Can't collect more than [" + Integer.MAX_VALUE + "] buckets but attempted [" + totalOrdsToCollect + "]"
-            );
-        }
-        long[] bucketOrdsToCollect = new long[(int) totalOrdsToCollect];
-        int b = 0;
-        for (long owningBucketOrd : owningBucketOrds) {
-            LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
-            while (ordsEnum.next()) {
-                bucketOrdsToCollect[b++] = ordsEnum.ord();
+        try (IntArray bucketsInOrd = bigArrays().newIntArray(owningBucketOrds.size())) {
+            for (long ordIdx = 0; ordIdx < owningBucketOrds.size(); ordIdx++) {
+                final long bucketCount = bucketOrds.bucketsInOrd(owningBucketOrds.get(ordIdx));
+                bucketsInOrd.set(ordIdx, (int) bucketCount);
+                totalOrdsToCollect += bucketCount;
             }
-        }
-        var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
-
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        b = 0;
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            List<B> buckets = new ArrayList<>(bucketsInOrd[ordIdx]);
-            LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-            while (ordsEnum.next()) {
-                if (bucketOrdsToCollect[b] != ordsEnum.ord()) {
-                    // If we hit this, something has gone horribly wrong and we need to investigate
-                    throw AggregationErrors.iterationOrderChangedWithoutMutating(
-                        bucketOrds.toString(),
-                        ordsEnum.ord(),
-                        bucketOrdsToCollect[b]
-                    );
+            if (totalOrdsToCollect > Integer.MAX_VALUE) {
+                // TODO: We should instrument this error. While it is correct for it to be a 400 class IllegalArgumentException, there is
+                // not
+                // much the user can do about that. If this occurs with any frequency, we should do something about it.
+                throw new IllegalArgumentException(
+                    "Can't collect more than [" + Integer.MAX_VALUE + "] buckets but attempted [" + totalOrdsToCollect + "]"
+                );
+            }
+            try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(totalOrdsToCollect)) {
+                int b = 0;
+                for (long i = 0; i < owningBucketOrds.size(); i++) {
+                    LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(i));
+                    while (ordsEnum.next()) {
+                        bucketOrdsToCollect.set(b++, ordsEnum.ord());
+                    }
                 }
-                buckets.add(bucketBuilder.build(ordsEnum.value(), bucketDocCount(ordsEnum.ord()), subAggregationResults.apply(b++)));
+                var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
+
+                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+                b = 0;
+                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                    final long owningBucketOrd = owningBucketOrds.get(ordIdx);
+                    List<B> buckets = new ArrayList<>(bucketsInOrd.get(ordIdx));
+                    LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
+                    while (ordsEnum.next()) {
+                        if (bucketOrdsToCollect.get(b) != ordsEnum.ord()) {
+                            // If we hit this, something has gone horribly wrong and we need to investigate
+                            throw AggregationErrors.iterationOrderChangedWithoutMutating(
+                                bucketOrds.toString(),
+                                ordsEnum.ord(),
+                                bucketOrdsToCollect.get(b)
+                            );
+                        }
+                        buckets.add(
+                            bucketBuilder.build(ordsEnum.value(), bucketDocCount(ordsEnum.ord()), subAggregationResults.apply(b++))
+                        );
+                    }
+                    results[ordIdx] = resultBuilder.build(owningBucketOrd, buckets);
+                }
+                return results;
             }
-            results[ordIdx] = resultBuilder.build(owningBucketOrds[ordIdx], buckets);
         }
-        return results;
     }
 
     @FunctionalInterface

+ 4 - 3
server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferableBucketAggregator.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.search.aggregations.bucket;
 
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.BucketCollector;
@@ -65,7 +66,7 @@ public abstract class DeferableBucketAggregator extends BucketsAggregator {
                 }
                 deferredAggregations.add(subAggregators[i]);
                 deferredAggregationNames.add(subAggregators[i].name());
-                subAggregators[i] = deferringCollector.wrap(subAggregators[i]);
+                subAggregators[i] = deferringCollector.wrap(subAggregators[i], bigArrays());
             } else {
                 collectors.add(subAggregators[i]);
             }
@@ -87,7 +88,7 @@ public abstract class DeferableBucketAggregator extends BucketsAggregator {
     /**
      * Build the {@link DeferringBucketCollector}. The default implementation
      * replays all hits against the buckets selected by
-     * {#link {@link DeferringBucketCollector#prepareSelectedBuckets(long...)}.
+     * {#link {@link DeferringBucketCollector#prepareSelectedBuckets(LongArray)}.
      */
     protected DeferringBucketCollector buildDeferringCollector() {
         return new BestBucketsDeferringCollector(topLevelQuery(), searcher(), descendsFromGlobalAggregator(parent()));
@@ -107,7 +108,7 @@ public abstract class DeferableBucketAggregator extends BucketsAggregator {
     }
 
     @Override
-    protected final void prepareSubAggs(long[] bucketOrdsToCollect) throws IOException {
+    protected final void prepareSubAggs(LongArray bucketOrdsToCollect) throws IOException {
         if (deferringCollector != null) {
             deferringCollector.prepareSelectedBuckets(bucketOrdsToCollect);
         }

+ 5 - 3
server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java

@@ -10,6 +10,8 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.BucketCollector;
@@ -37,13 +39,13 @@ public abstract class DeferringBucketCollector extends BucketCollector {
     /**
      * Replay the deferred hits on the selected buckets.
      */
-    public abstract void prepareSelectedBuckets(long... selectedBuckets) throws IOException;
+    public abstract void prepareSelectedBuckets(LongArray selectedBuckets) throws IOException;
 
     /**
      * Wrap the provided aggregator so that it behaves (almost) as if it had
      * been collected directly.
      */
-    public Aggregator wrap(final Aggregator in) {
+    public Aggregator wrap(final Aggregator in, BigArrays bigArrays) {
         return new WrappedAggregator(in);
     }
 
@@ -80,7 +82,7 @@ public abstract class DeferringBucketCollector extends BucketCollector {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
             return in.buildAggregations(owningBucketOrds);
         }
 

+ 39 - 36
server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -35,6 +35,7 @@ import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.RoaringDocIdSet;
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.Rounding;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.index.IndexSortConfig;
@@ -184,50 +185,51 @@ public final class CompositeAggregator extends BucketsAggregator implements Size
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         // Composite aggregator must be at the top of the aggregation tree
-        assert owningBucketOrds.length == 1 && owningBucketOrds[0] == 0L;
+        assert owningBucketOrds.size() == 1 && owningBucketOrds.get(0) == 0L;
         if (deferredCollectors != NO_OP_BUCKET_COLLECTOR) {
             // Replay all documents that contain at least one top bucket (collected during the first pass).
             runDeferredCollections();
         }
 
-        int num = Math.min(size, (int) queue.size());
+        final int num = Math.min(size, (int) queue.size());
         final InternalComposite.InternalBucket[] buckets = new InternalComposite.InternalBucket[num];
-        long[] bucketOrdsToCollect = new long[(int) queue.size()];
-        for (int i = 0; i < queue.size(); i++) {
-            bucketOrdsToCollect[i] = i;
-        }
-        var subAggsForBuckets = buildSubAggsForBuckets(bucketOrdsToCollect);
-        while (queue.size() > 0) {
-            int slot = queue.pop();
-            CompositeKey key = queue.toCompositeKey(slot);
-            InternalAggregations aggs = subAggsForBuckets.apply(slot);
-            long docCount = queue.getDocCount(slot);
-            buckets[(int) queue.size()] = new InternalComposite.InternalBucket(
-                sourceNames,
-                formats,
-                key,
-                reverseMuls,
-                missingOrders,
-                docCount,
-                aggs
-            );
+        try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(queue.size())) {
+            for (int i = 0; i < queue.size(); i++) {
+                bucketOrdsToCollect.set(i, i);
+            }
+            var subAggsForBuckets = buildSubAggsForBuckets(bucketOrdsToCollect);
+            while (queue.size() > 0) {
+                int slot = queue.pop();
+                CompositeKey key = queue.toCompositeKey(slot);
+                InternalAggregations aggs = subAggsForBuckets.apply(slot);
+                long docCount = queue.getDocCount(slot);
+                buckets[(int) queue.size()] = new InternalComposite.InternalBucket(
+                    sourceNames,
+                    formats,
+                    key,
+                    reverseMuls,
+                    missingOrders,
+                    docCount,
+                    aggs
+                );
+            }
+            CompositeKey lastBucket = num > 0 ? buckets[num - 1].getRawKey() : null;
+            return new InternalAggregation[] {
+                new InternalComposite(
+                    name,
+                    size,
+                    sourceNames,
+                    formats,
+                    Arrays.asList(buckets),
+                    lastBucket,
+                    reverseMuls,
+                    missingOrders,
+                    earlyTerminated,
+                    metadata()
+                ) };
         }
-        CompositeKey lastBucket = num > 0 ? buckets[num - 1].getRawKey() : null;
-        return new InternalAggregation[] {
-            new InternalComposite(
-                name,
-                size,
-                sourceNames,
-                formats,
-                Arrays.asList(buckets),
-                lastBucket,
-                reverseMuls,
-                missingOrders,
-                earlyTerminated,
-                metadata()
-            ) };
     }
 
     @Override
@@ -244,6 +246,7 @@ public final class CompositeAggregator extends BucketsAggregator implements Size
             false,
             metadata()
         );
+
     }
 
     private void finishLeaf() {

+ 69 - 57
server/src/main/java/org/elasticsearch/search/aggregations/bucket/countedterms/CountedTermsAggregator.java

@@ -13,6 +13,8 @@ import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
@@ -108,70 +110,80 @@ class CountedTermsAggregator extends TermsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        StringTerms.Bucket[][] topBucketsPerOrd = new StringTerms.Bucket[owningBucketOrds.length][];
-        long[] otherDocCounts = new long[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
-
-            // as users can't control sort order, in practice we'll always sort by doc count descending
-            try (
-                BucketPriorityQueue<StringTerms.Bucket> ordered = new BucketPriorityQueue<>(
-                    size,
-                    bigArrays(),
-                    partiallyBuiltBucketComparator
-                )
-            ) {
-                StringTerms.Bucket spare = null;
-                BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-                Supplier<StringTerms.Bucket> emptyBucketBuilder = () -> new StringTerms.Bucket(new BytesRef(), 0, null, false, 0, format);
-                while (ordsEnum.next()) {
-                    long docCount = bucketDocCount(ordsEnum.ord());
-                    otherDocCounts[ordIdx] += docCount;
-                    if (spare == null) {
-                        spare = emptyBucketBuilder.get();
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        try (
+            LongArray otherDocCounts = bigArrays().newLongArray(owningBucketOrds.size());
+            ObjectArray<StringTerms.Bucket[]> topBucketsPerOrd = bigArrays().newObjectArray(owningBucketOrds.size())
+        ) {
+            for (long ordIdx = 0; ordIdx < topBucketsPerOrd.size(); ordIdx++) {
+                int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
+
+                // as users can't control sort order, in practice we'll always sort by doc count descending
+                try (
+                    BucketPriorityQueue<StringTerms.Bucket> ordered = new BucketPriorityQueue<>(
+                        size,
+                        bigArrays(),
+                        partiallyBuiltBucketComparator
+                    )
+                ) {
+                    StringTerms.Bucket spare = null;
+                    BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
+                    Supplier<StringTerms.Bucket> emptyBucketBuilder = () -> new StringTerms.Bucket(
+                        new BytesRef(),
+                        0,
+                        null,
+                        false,
+                        0,
+                        format
+                    );
+                    while (ordsEnum.next()) {
+                        long docCount = bucketDocCount(ordsEnum.ord());
+                        otherDocCounts.increment(ordIdx, docCount);
+                        if (spare == null) {
+                            spare = emptyBucketBuilder.get();
+                        }
+                        ordsEnum.readValue(spare.getTermBytes());
+                        spare.setDocCount(docCount);
+                        spare.setBucketOrd(ordsEnum.ord());
+                        spare = ordered.insertWithOverflow(spare);
                     }
-                    ordsEnum.readValue(spare.getTermBytes());
-                    spare.setDocCount(docCount);
-                    spare.setBucketOrd(ordsEnum.ord());
-                    spare = ordered.insertWithOverflow(spare);
-                }
 
-                topBucketsPerOrd[ordIdx] = new StringTerms.Bucket[(int) ordered.size()];
-                for (int i = (int) ordered.size() - 1; i >= 0; --i) {
-                    topBucketsPerOrd[ordIdx][i] = ordered.pop();
-                    otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][i].getDocCount();
-                    topBucketsPerOrd[ordIdx][i].setTermBytes(BytesRef.deepCopyOf(topBucketsPerOrd[ordIdx][i].getTermBytes()));
+                    topBucketsPerOrd.set(ordIdx, new StringTerms.Bucket[(int) ordered.size()]);
+                    for (int i = (int) ordered.size() - 1; i >= 0; --i) {
+                        topBucketsPerOrd.get(ordIdx)[i] = ordered.pop();
+                        otherDocCounts.increment(ordIdx, -topBucketsPerOrd.get(ordIdx)[i].getDocCount());
+                        topBucketsPerOrd.get(ordIdx)[i].setTermBytes(BytesRef.deepCopyOf(topBucketsPerOrd.get(ordIdx)[i].getTermBytes()));
+                    }
                 }
             }
-        }
 
-        buildSubAggsForAllBuckets(topBucketsPerOrd, InternalTerms.Bucket::getBucketOrd, InternalTerms.Bucket::setAggregations);
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            final BucketOrder reduceOrder;
-            if (isKeyOrder(order) == false) {
-                reduceOrder = InternalOrder.key(true);
-                Arrays.sort(topBucketsPerOrd[ordIdx], reduceOrder.comparator());
-            } else {
-                reduceOrder = order;
+            buildSubAggsForAllBuckets(topBucketsPerOrd, InternalTerms.Bucket::getBucketOrd, InternalTerms.Bucket::setAggregations);
+            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
+            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                final BucketOrder reduceOrder;
+                if (isKeyOrder(order) == false) {
+                    reduceOrder = InternalOrder.key(true);
+                    Arrays.sort(topBucketsPerOrd.get(ordIdx), reduceOrder.comparator());
+                } else {
+                    reduceOrder = order;
+                }
+                result[ordIdx] = new StringTerms(
+                    name,
+                    reduceOrder,
+                    order,
+                    bucketCountThresholds.getRequiredSize(),
+                    bucketCountThresholds.getMinDocCount(),
+                    metadata(),
+                    format,
+                    bucketCountThresholds.getShardSize(),
+                    false,
+                    otherDocCounts.get(ordIdx),
+                    Arrays.asList(topBucketsPerOrd.get(ordIdx)),
+                    null
+                );
             }
-            result[ordIdx] = new StringTerms(
-                name,
-                reduceOrder,
-                order,
-                bucketCountThresholds.getRequiredSize(),
-                bucketCountThresholds.getMinDocCount(),
-                metadata(),
-                format,
-                bucketCountThresholds.getShardSize(),
-                false,
-                otherDocCounts[ordIdx],
-                Arrays.asList(topBucketsPerOrd[ordIdx]),
-                null
-            );
+            return result;
         }
-        return result;
     }
 
     @Override

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregator.java

@@ -20,6 +20,7 @@ import org.apache.lucene.search.Scorer;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
@@ -208,7 +209,7 @@ public abstract class FiltersAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForFixedBucketCount(
             owningBucketOrds,
             filters.size() + (otherBucketKey == null ? 0 : 1),

+ 31 - 28
server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregator.java

@@ -12,6 +12,8 @@ import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -132,39 +134,40 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid<?>> extends Bu
     protected abstract InternalGeoGridBucket newEmptyBucket();
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        InternalGeoGridBucket[][] topBucketsPerOrd = new InternalGeoGridBucket[owningBucketOrds.length][];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            int size = (int) Math.min(bucketOrds.bucketsInOrd(owningBucketOrds[ordIdx]), shardSize);
-
-            try (BucketPriorityQueue<InternalGeoGridBucket> ordered = new BucketPriorityQueue<>(size, bigArrays())) {
-                InternalGeoGridBucket spare = null;
-                LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-                while (ordsEnum.next()) {
-                    if (spare == null) {
-                        spare = newEmptyBucket();
-                    }
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        try (ObjectArray<InternalGeoGridBucket[]> topBucketsPerOrd = bigArrays().newObjectArray(owningBucketOrds.size())) {
+            for (long ordIdx = 0; ordIdx < topBucketsPerOrd.size(); ordIdx++) {
+                int size = (int) Math.min(bucketOrds.bucketsInOrd(owningBucketOrds.get(ordIdx)), shardSize);
+
+                try (BucketPriorityQueue<InternalGeoGridBucket> ordered = new BucketPriorityQueue<>(size, bigArrays())) {
+                    InternalGeoGridBucket spare = null;
+                    LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
+                    while (ordsEnum.next()) {
+                        if (spare == null) {
+                            spare = newEmptyBucket();
+                        }
 
-                    // need a special function to keep the source bucket
-                    // up-to-date so it can get the appropriate key
-                    spare.hashAsLong = ordsEnum.value();
-                    spare.docCount = bucketDocCount(ordsEnum.ord());
-                    spare.bucketOrd = ordsEnum.ord();
-                    spare = ordered.insertWithOverflow(spare);
-                }
+                        // need a special function to keep the source bucket
+                        // up-to-date so it can get the appropriate key
+                        spare.hashAsLong = ordsEnum.value();
+                        spare.docCount = bucketDocCount(ordsEnum.ord());
+                        spare.bucketOrd = ordsEnum.ord();
+                        spare = ordered.insertWithOverflow(spare);
+                    }
 
-                topBucketsPerOrd[ordIdx] = new InternalGeoGridBucket[(int) ordered.size()];
-                for (int i = (int) ordered.size() - 1; i >= 0; --i) {
-                    topBucketsPerOrd[ordIdx][i] = ordered.pop();
+                    topBucketsPerOrd.set(ordIdx, new InternalGeoGridBucket[(int) ordered.size()]);
+                    for (int i = (int) ordered.size() - 1; i >= 0; --i) {
+                        topBucketsPerOrd.get(ordIdx)[i] = ordered.pop();
+                    }
                 }
             }
+            buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
+            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
+            for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                results[ordIdx] = buildAggregation(name, requiredSize, Arrays.asList(topBucketsPerOrd.get(ordIdx)), metadata());
+            }
+            return results;
         }
-        buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            results[ordIdx] = buildAggregation(name, requiredSize, Arrays.asList(topBucketsPerOrd[ordIdx]), metadata());
-        }
-        return results;
     }
 
     @Override

+ 3 - 2
server/src/main/java/org/elasticsearch/search/aggregations/bucket/global/GlobalAggregator.java

@@ -13,6 +13,7 @@ import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.Weight;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
@@ -60,8 +61,8 @@ public final class GlobalAggregator extends BucketsAggregator implements SingleB
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        assert owningBucketOrds.length == 1 && owningBucketOrds[0] == 0 : "global aggregator can only be a top level aggregator";
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        assert owningBucketOrds.size() == 1 && owningBucketOrds.get(0) == 0 : "global aggregator can only be a top level aggregator";
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalGlobal(

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/AbstractHistogramAggregator.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.search.aggregations.bucket.histogram;
 
 import org.apache.lucene.util.CollectionUtil;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -79,7 +80,7 @@ public abstract class AbstractHistogramAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForVariableBuckets(owningBucketOrds, bucketOrds, (bucketValue, docCount, subAggregationResults) -> {
             double roundKey = Double.longBitsToDouble(bucketValue);
             double key = roundKey * interval + offset;

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -17,6 +17,7 @@ import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.CollectionUtil;
 import org.elasticsearch.common.Rounding;
 import org.elasticsearch.common.Rounding.DateTimeUnit;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasables;
@@ -337,7 +338,7 @@ class DateHistogramAggregator extends BucketsAggregator implements SizedBucketAg
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForVariableBuckets(owningBucketOrds, bucketOrds, (bucketValue, docCount, subAggregationResults) -> {
             return new InternalDateHistogram.Bucket(bucketValue, docCount, keyed, formatter, subAggregationResults);
         }, (owningBucketOrd, buckets) -> {

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateRangeHistogramAggregator.java

@@ -12,6 +12,7 @@ import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.CollectionUtil;
 import org.elasticsearch.common.Rounding;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.fielddata.FieldData;
@@ -163,7 +164,7 @@ class DateRangeHistogramAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForVariableBuckets(
             owningBucketOrds,
             bucketOrds,

+ 22 - 20
server/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/VariableWidthHistogramAggregator.java

@@ -14,6 +14,7 @@ import org.apache.lucene.util.CollectionUtil;
 import org.apache.lucene.util.InPlaceMergeSorter;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.DoubleArray;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -565,34 +566,35 @@ public class VariableWidthHistogramAggregator extends DeferableBucketAggregator
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         int numClusters = collector.finalNumBuckets();
 
-        long[] bucketOrdsToCollect = new long[numClusters];
-        for (int i = 0; i < numClusters; i++) {
-            bucketOrdsToCollect[i] = i;
-        }
+        try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(numClusters)) {
+            for (int i = 0; i < numClusters; i++) {
+                bucketOrdsToCollect.set(i, i);
+            }
 
-        var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
+            var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
 
-        List<InternalVariableWidthHistogram.Bucket> buckets = new ArrayList<>(numClusters);
-        for (int bucketOrd = 0; bucketOrd < numClusters; bucketOrd++) {
-            buckets.add(collector.buildBucket(bucketOrd, subAggregationResults.apply(bucketOrd)));
-        }
+            List<InternalVariableWidthHistogram.Bucket> buckets = new ArrayList<>(numClusters);
+            for (int bucketOrd = 0; bucketOrd < numClusters; bucketOrd++) {
+                buckets.add(collector.buildBucket(bucketOrd, subAggregationResults.apply(bucketOrd)));
+            }
 
-        Function<List<InternalVariableWidthHistogram.Bucket>, InternalAggregation> resultBuilder = bucketsToFormat -> {
-            // The contract of the histogram aggregation is that shards must return
-            // buckets ordered by centroid in ascending order
-            CollectionUtil.introSort(bucketsToFormat, BucketOrder.key(true).comparator());
+            Function<List<InternalVariableWidthHistogram.Bucket>, InternalAggregation> resultBuilder = bucketsToFormat -> {
+                // The contract of the histogram aggregation is that shards must return
+                // buckets ordered by centroid in ascending order
+                CollectionUtil.introSort(bucketsToFormat, BucketOrder.key(true).comparator());
 
-            InternalVariableWidthHistogram.EmptyBucketInfo emptyBucketInfo = new InternalVariableWidthHistogram.EmptyBucketInfo(
-                buildEmptySubAggregations()
-            );
+                InternalVariableWidthHistogram.EmptyBucketInfo emptyBucketInfo = new InternalVariableWidthHistogram.EmptyBucketInfo(
+                    buildEmptySubAggregations()
+                );
 
-            return new InternalVariableWidthHistogram(name, bucketsToFormat, emptyBucketInfo, numBuckets, formatter, metadata());
-        };
+                return new InternalVariableWidthHistogram(name, bucketsToFormat, emptyBucketInfo, numBuckets, formatter, metadata());
+            };
 
-        return new InternalAggregation[] { resultBuilder.apply(buckets) };
+            return new InternalAggregation[] { resultBuilder.apply(buckets) };
+        }
 
     }
 

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/missing/MissingAggregator.java

@@ -8,6 +8,7 @@
  */
 package org.elasticsearch.search.aggregations.bucket.missing;
 
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.fielddata.DocValueBits;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -67,7 +68,7 @@ public class MissingAggregator extends BucketsAggregator implements SingleBucket
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalMissing(

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java

@@ -21,6 +21,7 @@ import org.apache.lucene.search.Weight;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.util.BitSet;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.mapper.NestedObjectMapper;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -124,7 +125,7 @@ public class NestedAggregator extends BucketsAggregator implements SingleBucketA
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalNested(

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/ReverseNestedAggregator.java

@@ -13,6 +13,7 @@ import org.apache.lucene.search.Query;
 import org.apache.lucene.search.join.BitSetProducer;
 import org.apache.lucene.util.BitSet;
 import org.elasticsearch.common.lucene.search.Queries;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.mapper.NestedObjectMapper;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
@@ -86,7 +87,7 @@ public class ReverseNestedAggregator extends BucketsAggregator implements Single
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalReverseNested(

+ 53 - 45
server/src/main/java/org/elasticsearch/search/aggregations/bucket/prefix/IpPrefixAggregator.java

@@ -12,6 +12,8 @@ package org.elasticsearch.search.aggregations.bucket.prefix;
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.CollectionUtil;
+import org.elasticsearch.common.util.IntArray;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.fielddata.FieldData;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
@@ -160,57 +162,63 @@ public final class IpPrefixAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         long totalOrdsToCollect = 0;
-        final int[] bucketsInOrd = new int[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            final long bucketCount = bucketOrds.bucketsInOrd(owningBucketOrds[ordIdx]);
-            bucketsInOrd[ordIdx] = (int) bucketCount;
-            totalOrdsToCollect += bucketCount;
-        }
-
-        long[] bucketOrdsToCollect = new long[(int) totalOrdsToCollect];
-        int b = 0;
-        for (long owningBucketOrd : owningBucketOrds) {
-            BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
-            while (ordsEnum.next()) {
-                bucketOrdsToCollect[b++] = ordsEnum.ord();
+        try (IntArray bucketsInOrd = bigArrays().newIntArray(owningBucketOrds.size())) {
+            for (long ordIdx = 0; ordIdx < owningBucketOrds.size(); ordIdx++) {
+                final long bucketCount = bucketOrds.bucketsInOrd(owningBucketOrds.get(ordIdx));
+                bucketsInOrd.set(ordIdx, (int) bucketCount);
+                totalOrdsToCollect += bucketCount;
             }
-        }
 
-        var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        b = 0;
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            List<InternalIpPrefix.Bucket> buckets = new ArrayList<>(bucketsInOrd[ordIdx]);
-            BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-            while (ordsEnum.next()) {
-                long ordinal = ordsEnum.ord();
-                if (bucketOrdsToCollect[b] != ordinal) {
-                    throw AggregationErrors.iterationOrderChangedWithoutMutating(bucketOrds.toString(), ordinal, bucketOrdsToCollect[b]);
+            try (LongArray bucketOrdsToCollect = bigArrays().newLongArray(totalOrdsToCollect)) {
+                int b = 0;
+                for (long i = 0; i < owningBucketOrds.size(); i++) {
+                    BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(i));
+                    while (ordsEnum.next()) {
+                        bucketOrdsToCollect.set(b++, ordsEnum.ord());
+                    }
+                }
+
+                var subAggregationResults = buildSubAggsForBuckets(bucketOrdsToCollect);
+                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+                b = 0;
+                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                    List<InternalIpPrefix.Bucket> buckets = new ArrayList<>(bucketsInOrd.get(ordIdx));
+                    BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds.get(ordIdx));
+                    while (ordsEnum.next()) {
+                        long ordinal = ordsEnum.ord();
+                        if (bucketOrdsToCollect.get(b) != ordinal) {
+                            throw AggregationErrors.iterationOrderChangedWithoutMutating(
+                                bucketOrds.toString(),
+                                ordinal,
+                                bucketOrdsToCollect.get(b)
+                            );
+                        }
+                        BytesRef ipAddress = new BytesRef();
+                        ordsEnum.readValue(ipAddress);
+                        long docCount = bucketDocCount(ordinal);
+                        buckets.add(
+                            new InternalIpPrefix.Bucket(
+                                config.format(),
+                                BytesRef.deepCopyOf(ipAddress),
+                                keyed,
+                                ipPrefix.isIpv6,
+                                ipPrefix.prefixLength,
+                                ipPrefix.appendPrefixLength,
+                                docCount,
+                                subAggregationResults.apply(b++)
+                            )
+                        );
+
+                        // NOTE: the aggregator is expected to return sorted results
+                        CollectionUtil.introSort(buckets, BucketOrder.key(true).comparator());
+                    }
+                    results[ordIdx] = new InternalIpPrefix(name, config.format(), keyed, minDocCount, buckets, metadata());
                 }
-                BytesRef ipAddress = new BytesRef();
-                ordsEnum.readValue(ipAddress);
-                long docCount = bucketDocCount(ordinal);
-                buckets.add(
-                    new InternalIpPrefix.Bucket(
-                        config.format(),
-                        BytesRef.deepCopyOf(ipAddress),
-                        keyed,
-                        ipPrefix.isIpv6,
-                        ipPrefix.prefixLength,
-                        ipPrefix.appendPrefixLength,
-                        docCount,
-                        subAggregationResults.apply(b++)
-                    )
-                );
-
-                // NOTE: the aggregator is expected to return sorted results
-                CollectionUtil.introSort(buckets, BucketOrder.key(true).comparator());
+                return results;
             }
-            results[ordIdx] = new InternalIpPrefix(name, config.format(), keyed, minDocCount, buckets, metadata());
         }
-        return results;
     }
 
     @Override

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/BinaryRangeAggregator.java

@@ -14,6 +14,7 @@ import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.fielddata.FieldData;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.search.DocValueFormat;
@@ -358,7 +359,7 @@ public final class BinaryRangeAggregator extends BucketsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForFixedBucketCount(
             owningBucketOrds,
             ranges.length,

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -14,6 +14,7 @@ import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.index.fielddata.FieldData;
 import org.elasticsearch.index.fielddata.NumericDoubleValues;
@@ -531,7 +532,7 @@ public abstract class RangeAggregator extends BucketsAggregator {
 
     @Override
     @SuppressWarnings("unchecked")
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForFixedBucketCount(
             owningBucketOrds,
             ranges.length,

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java

@@ -19,6 +19,7 @@ import org.apache.lucene.search.TopScoreDocCollector;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -120,7 +121,7 @@ public class BestDocsDeferringCollector extends DeferringBucketCollector impleme
     }
 
     @Override
-    public void prepareSelectedBuckets(long... selectedBuckets) throws IOException {
+    public void prepareSelectedBuckets(LongArray selectedBuckets) {
         // no-op - deferred aggs processed in postCollection call
     }
 

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/SamplerAggregator.java

@@ -11,6 +11,7 @@ package org.elasticsearch.search.aggregations.bucket.sampler;
 import org.apache.lucene.misc.search.DiversifiedTopDocsCollector;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
@@ -212,7 +213,7 @@ public class SamplerAggregator extends DeferableBucketAggregator implements Sing
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalSampler(

+ 2 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplerAggregator.java

@@ -15,6 +15,7 @@ import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.CheckedSupplier;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -60,7 +61,7 @@ public class RandomSamplerAggregator extends BucketsAggregator implements Single
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return buildAggregationsForSingleBucket(
             owningBucketOrds,
             (owningBucketOrd, subAggregationResults) -> new InternalRandomSampler(

+ 58 - 52
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -23,6 +23,7 @@ import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.LongHash;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.ObjectArrayPriorityQueue;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
@@ -190,7 +191,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return resultStrategy.buildAggregations(owningBucketOrds);
     }
 
@@ -692,61 +693,66 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         B extends InternalMultiBucketAggregation.InternalBucket,
         TB extends InternalMultiBucketAggregation.InternalBucket> implements Releasable {
 
-        private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+
             if (valueCount == 0) { // no context in this reader
-                InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-                for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                    results[ordIdx] = buildNoValuesResult(owningBucketOrds[ordIdx]);
+                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                    results[ordIdx] = buildNoValuesResult(owningBucketOrds.get(ordIdx));
                 }
                 return results;
             }
-
-            B[][] topBucketsPreOrd = buildTopBucketsPerOrd(owningBucketOrds.length);
-            long[] otherDocCount = new long[owningBucketOrds.length];
-            GlobalOrdLookupFunction lookupGlobalOrd = valuesSupplier.get()::lookupOrd;
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                final int size;
-                if (bucketCountThresholds.getMinDocCount() == 0) {
-                    // if minDocCount == 0 then we can end up with more buckets then maxBucketOrd() returns
-                    size = (int) Math.min(valueCount, bucketCountThresholds.getShardSize());
-                } else {
-                    size = (int) Math.min(maxBucketOrd(), bucketCountThresholds.getShardSize());
-                }
-                try (ObjectArrayPriorityQueue<TB> ordered = buildPriorityQueue(size)) {
-                    final int finalOrdIdx = ordIdx;
-                    BucketUpdater<TB> updater = bucketUpdater(owningBucketOrds[ordIdx], lookupGlobalOrd);
-                    collectionStrategy.forEach(owningBucketOrds[ordIdx], new BucketInfoConsumer() {
-                        TB spare = null;
-
-                        @Override
-                        public void accept(long globalOrd, long bucketOrd, long docCount) throws IOException {
-                            otherDocCount[finalOrdIdx] += docCount;
-                            if (docCount >= bucketCountThresholds.getShardMinDocCount()) {
-                                if (spare == null) {
-                                    spare = buildEmptyTemporaryBucket();
+            try (
+                LongArray otherDocCount = bigArrays().newLongArray(owningBucketOrds.size(), true);
+                ObjectArray<B[]> topBucketsPreOrd = buildTopBucketsPerOrd(owningBucketOrds.size())
+            ) {
+                GlobalOrdLookupFunction lookupGlobalOrd = valuesSupplier.get()::lookupOrd;
+                for (long ordIdx = 0; ordIdx < topBucketsPreOrd.size(); ordIdx++) {
+                    final int size;
+                    if (bucketCountThresholds.getMinDocCount() == 0) {
+                        // if minDocCount == 0 then we can end up with more buckets then maxBucketOrd() returns
+                        size = (int) Math.min(valueCount, bucketCountThresholds.getShardSize());
+                    } else {
+                        size = (int) Math.min(maxBucketOrd(), bucketCountThresholds.getShardSize());
+                    }
+                    try (ObjectArrayPriorityQueue<TB> ordered = buildPriorityQueue(size)) {
+                        final long finalOrdIdx = ordIdx;
+                        final long owningBucketOrd = owningBucketOrds.get(ordIdx);
+                        BucketUpdater<TB> updater = bucketUpdater(owningBucketOrd, lookupGlobalOrd);
+                        collectionStrategy.forEach(owningBucketOrd, new BucketInfoConsumer() {
+                            TB spare = null;
+
+                            @Override
+                            public void accept(long globalOrd, long bucketOrd, long docCount) throws IOException {
+                                otherDocCount.increment(finalOrdIdx, docCount);
+                                if (docCount >= bucketCountThresholds.getShardMinDocCount()) {
+                                    if (spare == null) {
+                                        spare = buildEmptyTemporaryBucket();
+                                    }
+                                    updater.updateBucket(spare, globalOrd, bucketOrd, docCount);
+                                    spare = ordered.insertWithOverflow(spare);
                                 }
-                                updater.updateBucket(spare, globalOrd, bucketOrd, docCount);
-                                spare = ordered.insertWithOverflow(spare);
                             }
+                        });
+
+                        // Get the top buckets
+                        topBucketsPreOrd.set(ordIdx, buildBuckets((int) ordered.size()));
+                        for (int i = (int) ordered.size() - 1; i >= 0; --i) {
+                            B bucket = convertTempBucketToRealBucket(ordered.pop(), lookupGlobalOrd);
+                            topBucketsPreOrd.get(ordIdx)[i] = bucket;
+                            otherDocCount.increment(ordIdx, -bucket.getDocCount());
                         }
-                    });
-
-                    // Get the top buckets
-                    topBucketsPreOrd[ordIdx] = buildBuckets((int) ordered.size());
-                    for (int i = (int) ordered.size() - 1; i >= 0; --i) {
-                        topBucketsPreOrd[ordIdx][i] = convertTempBucketToRealBucket(ordered.pop(), lookupGlobalOrd);
-                        otherDocCount[ordIdx] -= topBucketsPreOrd[ordIdx][i].getDocCount();
                     }
                 }
-            }
 
-            buildSubAggs(topBucketsPreOrd);
+                buildSubAggs(topBucketsPreOrd);
 
-            InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                results[ordIdx] = buildResult(owningBucketOrds[ordIdx], otherDocCount[ordIdx], topBucketsPreOrd[ordIdx]);
+                InternalAggregation[] results = new InternalAggregation[Math.toIntExact(topBucketsPreOrd.size())];
+                for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                    results[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCount.get(ordIdx), topBucketsPreOrd.get(ordIdx));
+                }
+                return results;
             }
-            return results;
         }
 
         /**
@@ -781,7 +787,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         /**
          * Build an array to hold the "top" buckets for each ordinal.
          */
-        abstract B[][] buildTopBucketsPerOrd(int size);
+        abstract ObjectArray<B[]> buildTopBucketsPerOrd(long size);
 
         /**
          * Build an array of buckets for a particular ordinal to collect the
@@ -798,7 +804,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
          * Build the sub-aggregations into the buckets. This will usually
          * delegate to {@link #buildSubAggsForAllBuckets}.
          */
-        abstract void buildSubAggs(B[][] topBucketsPreOrd) throws IOException;
+        abstract void buildSubAggs(ObjectArray<B[]> topBucketsPreOrd) throws IOException;
 
         /**
          * Turn the buckets into an aggregation result.
@@ -837,8 +843,8 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         }
 
         @Override
-        StringTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new StringTerms.Bucket[size][];
+        ObjectArray<StringTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -875,7 +881,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         }
 
         @Override
-        void buildSubAggs(StringTerms.Bucket[][] topBucketsPreOrd) throws IOException {
+        void buildSubAggs(ObjectArray<StringTerms.Bucket[]> topBucketsPreOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPreOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
         }
 
@@ -969,8 +975,8 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         }
 
         @Override
-        SignificantStringTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new SignificantStringTerms.Bucket[size][];
+        ObjectArray<SignificantStringTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -1022,7 +1028,7 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         }
 
         @Override
-        void buildSubAggs(SignificantStringTerms.Bucket[][] topBucketsPreOrd) throws IOException {
+        void buildSubAggs(ObjectArray<SignificantStringTerms.Bucket[]> topBucketsPreOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPreOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
         }
 

+ 1 - 1
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/InternalSignificantTerms.java

@@ -62,7 +62,7 @@ public abstract class InternalSignificantTerms<A extends InternalSignificantTerm
         long supersetSize;
         /**
          * Ordinal of the bucket while it is being built. Not used after it is
-         * returned from {@link Aggregator#buildAggregations(long[])} and not
+         * returned from {@link Aggregator#buildAggregations(org.elasticsearch.common.util.LongArray)} and not
          * serialized.
          */
         transient long bucketOrd;

+ 61 - 54
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/LongRareTermsAggregator.java

@@ -12,7 +12,9 @@ import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.LongHash;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.SetBackedScalingCuckooFilter;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.DocValueFormat;
@@ -118,70 +120,75 @@ public class LongRareTermsAggregator extends AbstractRareTermsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         /*
          * Collect the list of buckets, populate the filter with terms
          * that are too frequent, and figure out how to merge sub-buckets.
          */
-        LongRareTerms.Bucket[][] rarestPerOrd = new LongRareTerms.Bucket[owningBucketOrds.length][];
-        SetBackedScalingCuckooFilter[] filters = new SetBackedScalingCuckooFilter[owningBucketOrds.length];
-        long keepCount = 0;
-        long[] mergeMap = new long[(int) bucketOrds.size()];
-        Arrays.fill(mergeMap, -1);
-        long offset = 0;
-        for (int owningOrdIdx = 0; owningOrdIdx < owningBucketOrds.length; owningOrdIdx++) {
-            try (LongHash bucketsInThisOwningBucketToCollect = new LongHash(1, bigArrays())) {
-                filters[owningOrdIdx] = newFilter();
-                List<LongRareTerms.Bucket> builtBuckets = new ArrayList<>();
-                LongKeyedBucketOrds.BucketOrdsEnum collectedBuckets = bucketOrds.ordsEnum(owningBucketOrds[owningOrdIdx]);
-                while (collectedBuckets.next()) {
-                    long docCount = bucketDocCount(collectedBuckets.ord());
-                    // if the key is below threshold, reinsert into the new ords
-                    if (docCount <= maxDocCount) {
-                        LongRareTerms.Bucket bucket = new LongRareTerms.Bucket(collectedBuckets.value(), docCount, null, format);
-                        bucket.bucketOrd = offset + bucketsInThisOwningBucketToCollect.add(collectedBuckets.value());
-                        mergeMap[(int) collectedBuckets.ord()] = bucket.bucketOrd;
-                        builtBuckets.add(bucket);
-                        keepCount++;
-                    } else {
-                        filters[owningOrdIdx].add(collectedBuckets.value());
+        try (
+            ObjectArray<LongRareTerms.Bucket[]> rarestPerOrd = bigArrays().newObjectArray(owningBucketOrds.size());
+            ObjectArray<SetBackedScalingCuckooFilter> filters = bigArrays().newObjectArray(owningBucketOrds.size())
+        ) {
+            try (LongArray mergeMap = bigArrays().newLongArray(bucketOrds.size())) {
+                mergeMap.fill(0, mergeMap.size(), -1);
+                long keepCount = 0;
+                long offset = 0;
+                for (long owningOrdIdx = 0; owningOrdIdx < owningBucketOrds.size(); owningOrdIdx++) {
+                    try (LongHash bucketsInThisOwningBucketToCollect = new LongHash(1, bigArrays())) {
+                        filters.set(owningOrdIdx, newFilter());
+                        List<LongRareTerms.Bucket> builtBuckets = new ArrayList<>();
+                        LongKeyedBucketOrds.BucketOrdsEnum collectedBuckets = bucketOrds.ordsEnum(owningBucketOrds.get(owningOrdIdx));
+                        while (collectedBuckets.next()) {
+                            long docCount = bucketDocCount(collectedBuckets.ord());
+                            // if the key is below threshold, reinsert into the new ords
+                            if (docCount <= maxDocCount) {
+                                LongRareTerms.Bucket bucket = new LongRareTerms.Bucket(collectedBuckets.value(), docCount, null, format);
+                                bucket.bucketOrd = offset + bucketsInThisOwningBucketToCollect.add(collectedBuckets.value());
+                                mergeMap.set(collectedBuckets.ord(), bucket.bucketOrd);
+                                builtBuckets.add(bucket);
+                                keepCount++;
+                            } else {
+                                filters.get(owningOrdIdx).add(collectedBuckets.value());
+                            }
+                        }
+                        rarestPerOrd.set(owningOrdIdx, builtBuckets.toArray(LongRareTerms.Bucket[]::new));
+                        offset += bucketsInThisOwningBucketToCollect.size();
                     }
                 }
-                rarestPerOrd[owningOrdIdx] = builtBuckets.toArray(LongRareTerms.Bucket[]::new);
-                offset += bucketsInThisOwningBucketToCollect.size();
-            }
-        }
 
-        /*
-         * Only merge/delete the ordinals if we have actually deleted one,
-         * to save on some redundant work.
-         */
-        if (keepCount != mergeMap.length) {
-            LongUnaryOperator howToMerge = b -> mergeMap[(int) b];
-            rewriteBuckets(offset, howToMerge);
-            if (deferringCollector() != null) {
-                ((BestBucketsDeferringCollector) deferringCollector()).rewriteBuckets(howToMerge);
+                /*
+                 * Only merge/delete the ordinals if we have actually deleted one,
+                 * to save on some redundant work.
+                 */
+                if (keepCount != mergeMap.size()) {
+                    LongUnaryOperator howToMerge = mergeMap::get;
+                    rewriteBuckets(offset, howToMerge);
+                    if (deferringCollector() != null) {
+                        ((BestBucketsDeferringCollector) deferringCollector()).rewriteBuckets(howToMerge);
+                    }
+                }
             }
-        }
 
-        /*
-         * Now build the results!
-         */
-        buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            Arrays.sort(rarestPerOrd[ordIdx], ORDER.comparator());
-            result[ordIdx] = new LongRareTerms(
-                name,
-                ORDER,
-                metadata(),
-                format,
-                Arrays.asList(rarestPerOrd[ordIdx]),
-                maxDocCount,
-                filters[ordIdx]
-            );
+            /*
+             * Now build the results!
+             */
+            buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
+            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                LongRareTerms.Bucket[] buckets = rarestPerOrd.get(ordIdx);
+                Arrays.sort(buckets, ORDER.comparator());
+                result[ordIdx] = new LongRareTerms(
+                    name,
+                    ORDER,
+                    metadata(),
+                    format,
+                    Arrays.asList(buckets),
+                    maxDocCount,
+                    filters.get(ordIdx)
+                );
+            }
+            return result;
         }
-        return result;
     }
 
     @Override

+ 47 - 42
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/MapStringTermsAggregator.java

@@ -18,6 +18,7 @@ import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.apache.lucene.util.PriorityQueue;
 import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.ObjectArrayPriorityQueue;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -117,7 +118,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return resultStrategy.buildAggregations(owningBucketOrds);
     }
 
@@ -282,45 +283,49 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         implements
             Releasable {
 
-        private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-            B[][] topBucketsPerOrd = buildTopBucketsPerOrd(owningBucketOrds.length);
-            long[] otherDocCounts = new long[owningBucketOrds.length];
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                collectZeroDocEntriesIfNeeded(owningBucketOrds[ordIdx], excludeDeletedDocs);
-                int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
-
-                try (ObjectArrayPriorityQueue<B> ordered = buildPriorityQueue(size)) {
-                    B spare = null;
-                    BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-                    Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningBucketOrds[ordIdx]);
-                    while (ordsEnum.next()) {
-                        long docCount = bucketDocCount(ordsEnum.ord());
-                        otherDocCounts[ordIdx] += docCount;
-                        if (docCount < bucketCountThresholds.getShardMinDocCount()) {
-                            continue;
-                        }
-                        if (spare == null) {
-                            spare = emptyBucketBuilder.get();
+        private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+            try (
+                LongArray otherDocCounts = bigArrays().newLongArray(owningBucketOrds.size(), true);
+                ObjectArray<B[]> topBucketsPerOrd = buildTopBucketsPerOrd(Math.toIntExact(owningBucketOrds.size()))
+            ) {
+                for (long ordIdx = 0; ordIdx < topBucketsPerOrd.size(); ordIdx++) {
+                    long owningOrd = owningBucketOrds.get(ordIdx);
+                    collectZeroDocEntriesIfNeeded(owningOrd, excludeDeletedDocs);
+                    int size = (int) Math.min(bucketOrds.size(), bucketCountThresholds.getShardSize());
+
+                    try (ObjectArrayPriorityQueue<B> ordered = buildPriorityQueue(size)) {
+                        B spare = null;
+                        BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningOrd);
+                        Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningOrd);
+                        while (ordsEnum.next()) {
+                            long docCount = bucketDocCount(ordsEnum.ord());
+                            otherDocCounts.increment(ordIdx, docCount);
+                            if (docCount < bucketCountThresholds.getShardMinDocCount()) {
+                                continue;
+                            }
+                            if (spare == null) {
+                                spare = emptyBucketBuilder.get();
+                            }
+                            updateBucket(spare, ordsEnum, docCount);
+                            spare = ordered.insertWithOverflow(spare);
                         }
-                        updateBucket(spare, ordsEnum, docCount);
-                        spare = ordered.insertWithOverflow(spare);
-                    }
 
-                    topBucketsPerOrd[ordIdx] = buildBuckets((int) ordered.size());
-                    for (int i = (int) ordered.size() - 1; i >= 0; --i) {
-                        topBucketsPerOrd[ordIdx][i] = ordered.pop();
-                        otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][i].getDocCount();
-                        finalizeBucket(topBucketsPerOrd[ordIdx][i]);
+                        topBucketsPerOrd.set(ordIdx, buildBuckets((int) ordered.size()));
+                        for (int i = (int) ordered.size() - 1; i >= 0; --i) {
+                            topBucketsPerOrd.get(ordIdx)[i] = ordered.pop();
+                            otherDocCounts.increment(ordIdx, -topBucketsPerOrd.get(ordIdx)[i].getDocCount());
+                            finalizeBucket(topBucketsPerOrd.get(ordIdx)[i]);
+                        }
                     }
                 }
-            }
 
-            buildSubAggs(topBucketsPerOrd);
-            InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                result[ordIdx] = buildResult(owningBucketOrds[ordIdx], otherDocCounts[ordIdx], topBucketsPerOrd[ordIdx]);
+                buildSubAggs(topBucketsPerOrd);
+                InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
+                for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                    result[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
+                }
+                return result;
             }
-            return result;
         }
 
         /**
@@ -361,7 +366,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         /**
          * Build an array to hold the "top" buckets for each ordinal.
          */
-        abstract B[][] buildTopBucketsPerOrd(int size);
+        abstract ObjectArray<B[]> buildTopBucketsPerOrd(long size);
 
         /**
          * Build an array of buckets for a particular ordinal to collect the
@@ -379,7 +384,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
          * Build the sub-aggregations into the buckets. This will usually
          * delegate to {@link #buildSubAggsForAllBuckets}.
          */
-        abstract void buildSubAggs(B[][] topBucketsPerOrd) throws IOException;
+        abstract void buildSubAggs(ObjectArray<B[]> topBucketsPerOrd) throws IOException;
 
         /**
          * Turn the buckets into an aggregation result.
@@ -501,8 +506,8 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         }
 
         @Override
-        StringTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new StringTerms.Bucket[size][];
+        ObjectArray<StringTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -521,7 +526,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         }
 
         @Override
-        void buildSubAggs(StringTerms.Bucket[][] topBucketsPerOrd) throws IOException {
+        void buildSubAggs(ObjectArray<StringTerms.Bucket[]> topBucketsPerOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
         }
 
@@ -637,8 +642,8 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         }
 
         @Override
-        SignificantStringTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new SignificantStringTerms.Bucket[size][];
+        ObjectArray<SignificantStringTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -657,7 +662,7 @@ public final class MapStringTermsAggregator extends AbstractStringTermsAggregato
         }
 
         @Override
-        void buildSubAggs(SignificantStringTerms.Bucket[][] topBucketsPerOrd) throws IOException {
+        void buildSubAggs(ObjectArray<SignificantStringTerms.Bucket[]> topBucketsPerOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
         }
 

+ 54 - 50
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/NumericTermsAggregator.java

@@ -15,6 +15,7 @@ import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.util.NumericUtils;
 import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.ObjectArrayPriorityQueue;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.Releasables;
@@ -39,7 +40,6 @@ import org.elasticsearch.search.aggregations.support.ValuesSource;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.List;
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
@@ -136,7 +136,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         return resultStrategy.buildAggregations(owningBucketOrds);
     }
 
@@ -163,48 +163,52 @@ public final class NumericTermsAggregator extends TermsAggregator {
     abstract class ResultStrategy<R extends InternalAggregation, B extends InternalMultiBucketAggregation.InternalBucket>
         implements
             Releasable {
-        private InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-            B[][] topBucketsPerOrd = buildTopBucketsPerOrd(owningBucketOrds.length);
-            long[] otherDocCounts = new long[owningBucketOrds.length];
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                collectZeroDocEntriesIfNeeded(owningBucketOrds[ordIdx], excludeDeletedDocs);
-                long bucketsInOrd = bucketOrds.bucketsInOrd(owningBucketOrds[ordIdx]);
-
-                int size = (int) Math.min(bucketsInOrd, bucketCountThresholds.getShardSize());
-                try (ObjectArrayPriorityQueue<B> ordered = buildPriorityQueue(size)) {
-                    B spare = null;
-                    BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-                    Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningBucketOrds[ordIdx]);
-                    while (ordsEnum.next()) {
-                        long docCount = bucketDocCount(ordsEnum.ord());
-                        otherDocCounts[ordIdx] += docCount;
-                        if (docCount < bucketCountThresholds.getShardMinDocCount()) {
-                            continue;
-                        }
-                        if (spare == null) {
-                            spare = emptyBucketBuilder.get();
+        private InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+            try (
+                LongArray otherDocCounts = bigArrays().newLongArray(owningBucketOrds.size(), true);
+                ObjectArray<B[]> topBucketsPerOrd = buildTopBucketsPerOrd(owningBucketOrds.size())
+            ) {
+                for (long ordIdx = 0; ordIdx < topBucketsPerOrd.size(); ordIdx++) {
+                    final long owningBucketOrd = owningBucketOrds.get(ordIdx);
+                    collectZeroDocEntriesIfNeeded(owningBucketOrd, excludeDeletedDocs);
+                    long bucketsInOrd = bucketOrds.bucketsInOrd(owningBucketOrd);
+
+                    int size = (int) Math.min(bucketsInOrd, bucketCountThresholds.getShardSize());
+                    try (ObjectArrayPriorityQueue<B> ordered = buildPriorityQueue(size)) {
+                        B spare = null;
+                        BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
+                        Supplier<B> emptyBucketBuilder = emptyBucketBuilder(owningBucketOrd);
+                        while (ordsEnum.next()) {
+                            long docCount = bucketDocCount(ordsEnum.ord());
+                            otherDocCounts.increment(ordIdx, docCount);
+                            if (docCount < bucketCountThresholds.getShardMinDocCount()) {
+                                continue;
+                            }
+                            if (spare == null) {
+                                spare = emptyBucketBuilder.get();
+                            }
+                            updateBucket(spare, ordsEnum, docCount);
+                            spare = ordered.insertWithOverflow(spare);
                         }
-                        updateBucket(spare, ordsEnum, docCount);
-                        spare = ordered.insertWithOverflow(spare);
-                    }
 
-                    // Get the top buckets
-                    B[] bucketsForOrd = buildBuckets((int) ordered.size());
-                    topBucketsPerOrd[ordIdx] = bucketsForOrd;
-                    for (int b = (int) ordered.size() - 1; b >= 0; --b) {
-                        topBucketsPerOrd[ordIdx][b] = ordered.pop();
-                        otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][b].getDocCount();
+                        // Get the top buckets
+                        B[] bucketsForOrd = buildBuckets((int) ordered.size());
+                        topBucketsPerOrd.set(ordIdx, bucketsForOrd);
+                        for (int b = (int) ordered.size() - 1; b >= 0; --b) {
+                            topBucketsPerOrd.get(ordIdx)[b] = ordered.pop();
+                            otherDocCounts.increment(ordIdx, -topBucketsPerOrd.get(ordIdx)[b].getDocCount());
+                        }
                     }
                 }
-            }
 
-            buildSubAggs(topBucketsPerOrd);
+                buildSubAggs(topBucketsPerOrd);
 
-            InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-            for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-                result[ordIdx] = buildResult(owningBucketOrds[ordIdx], otherDocCounts[ordIdx], topBucketsPerOrd[ordIdx]);
+                InternalAggregation[] result = new InternalAggregation[Math.toIntExact(topBucketsPerOrd.size())];
+                for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                    result[ordIdx] = buildResult(owningBucketOrds.get(ordIdx), otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
+                }
+                return result;
             }
-            return result;
         }
 
         /**
@@ -227,7 +231,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
         /**
          * Build an array to hold the "top" buckets for each ordinal.
          */
-        abstract B[][] buildTopBucketsPerOrd(int size);
+        abstract ObjectArray<B[]> buildTopBucketsPerOrd(long size);
 
         /**
          * Build an array of buckets for a particular ordinal. These arrays
@@ -258,7 +262,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
          * Build the sub-aggregations into the buckets. This will usually
          * delegate to {@link #buildSubAggsForAllBuckets}.
          */
-        abstract void buildSubAggs(B[][] topBucketsPerOrd) throws IOException;
+        abstract void buildSubAggs(ObjectArray<B[]> topBucketsPerOrd) throws IOException;
 
         /**
          * Collect extra entries for "zero" hit documents if they were requested
@@ -297,7 +301,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
         }
 
         @Override
-        final void buildSubAggs(B[][] topBucketsPerOrd) throws IOException {
+        final void buildSubAggs(ObjectArray<B[]> topBucketsPerOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
         }
 
@@ -356,8 +360,8 @@ public final class NumericTermsAggregator extends TermsAggregator {
         }
 
         @Override
-        LongTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new LongTerms.Bucket[size][];
+        ObjectArray<LongTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -397,7 +401,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
                 bucketCountThresholds.getShardSize(),
                 showTermDocCountError,
                 otherDocCount,
-                List.of(topBuckets),
+                Arrays.asList(topBuckets),
                 null
             );
         }
@@ -438,8 +442,8 @@ public final class NumericTermsAggregator extends TermsAggregator {
         }
 
         @Override
-        DoubleTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new DoubleTerms.Bucket[size][];
+        ObjectArray<DoubleTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -479,7 +483,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
                 bucketCountThresholds.getShardSize(),
                 showTermDocCountError,
                 otherDocCount,
-                List.of(topBuckets),
+                Arrays.asList(topBuckets),
                 null
             );
         }
@@ -551,8 +555,8 @@ public final class NumericTermsAggregator extends TermsAggregator {
         }
 
         @Override
-        SignificantLongTerms.Bucket[][] buildTopBucketsPerOrd(int size) {
-            return new SignificantLongTerms.Bucket[size][];
+        ObjectArray<SignificantLongTerms.Bucket[]> buildTopBucketsPerOrd(long size) {
+            return bigArrays().newObjectArray(size);
         }
 
         @Override
@@ -583,7 +587,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
         }
 
         @Override
-        void buildSubAggs(SignificantLongTerms.Bucket[][] topBucketsPerOrd) throws IOException {
+        void buildSubAggs(ObjectArray<SignificantLongTerms.Bucket[]> topBucketsPerOrd) throws IOException {
             buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
         }
 
@@ -601,7 +605,7 @@ public final class NumericTermsAggregator extends TermsAggregator {
                 subsetSizes.get(owningBucketOrd),
                 supersetSize,
                 significanceHeuristic,
-                List.of(topBuckets)
+                Arrays.asList(topBuckets)
             );
         }
 

+ 68 - 56
server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/StringRareTermsAggregator.java

@@ -12,6 +12,8 @@ import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefBuilder;
 import org.elasticsearch.common.util.BytesRefHash;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.SetBackedScalingCuckooFilter;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.fielddata.FieldData;
@@ -119,72 +121,82 @@ public class StringRareTermsAggregator extends AbstractRareTermsAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         /*
          * Collect the list of buckets, populate the filter with terms
          * that are too frequent, and figure out how to merge sub-buckets.
          */
-        StringRareTerms.Bucket[][] rarestPerOrd = new StringRareTerms.Bucket[owningBucketOrds.length][];
-        SetBackedScalingCuckooFilter[] filters = new SetBackedScalingCuckooFilter[owningBucketOrds.length];
-        long keepCount = 0;
-        long[] mergeMap = new long[(int) bucketOrds.size()];
-        Arrays.fill(mergeMap, -1);
-        long offset = 0;
-        for (int owningOrdIdx = 0; owningOrdIdx < owningBucketOrds.length; owningOrdIdx++) {
-            try (BytesRefHash bucketsInThisOwningBucketToCollect = new BytesRefHash(1, bigArrays())) {
-                filters[owningOrdIdx] = newFilter();
-                List<StringRareTerms.Bucket> builtBuckets = new ArrayList<>();
-                BytesKeyedBucketOrds.BucketOrdsEnum collectedBuckets = bucketOrds.ordsEnum(owningBucketOrds[owningOrdIdx]);
-                BytesRef scratch = new BytesRef();
-                while (collectedBuckets.next()) {
-                    collectedBuckets.readValue(scratch);
-                    long docCount = bucketDocCount(collectedBuckets.ord());
-                    // if the key is below threshold, reinsert into the new ords
-                    if (docCount <= maxDocCount) {
-                        StringRareTerms.Bucket bucket = new StringRareTerms.Bucket(BytesRef.deepCopyOf(scratch), docCount, null, format);
-                        bucket.bucketOrd = offset + bucketsInThisOwningBucketToCollect.add(scratch);
-                        mergeMap[(int) collectedBuckets.ord()] = bucket.bucketOrd;
-                        builtBuckets.add(bucket);
-                        keepCount++;
-                    } else {
-                        filters[owningOrdIdx].add(scratch);
+        try (
+            ObjectArray<StringRareTerms.Bucket[]> rarestPerOrd = bigArrays().newObjectArray(owningBucketOrds.size());
+            ObjectArray<SetBackedScalingCuckooFilter> filters = bigArrays().newObjectArray(owningBucketOrds.size())
+        ) {
+            try (LongArray mergeMap = bigArrays().newLongArray(bucketOrds.size())) {
+                mergeMap.fill(0, mergeMap.size(), -1);
+                long keepCount = 0;
+                long offset = 0;
+                for (long owningOrdIdx = 0; owningOrdIdx < owningBucketOrds.size(); owningOrdIdx++) {
+                    try (BytesRefHash bucketsInThisOwningBucketToCollect = new BytesRefHash(1, bigArrays())) {
+                        filters.set(owningOrdIdx, newFilter());
+                        List<StringRareTerms.Bucket> builtBuckets = new ArrayList<>();
+                        BytesKeyedBucketOrds.BucketOrdsEnum collectedBuckets = bucketOrds.ordsEnum(owningBucketOrds.get(owningOrdIdx));
+                        BytesRef scratch = new BytesRef();
+                        while (collectedBuckets.next()) {
+                            collectedBuckets.readValue(scratch);
+                            long docCount = bucketDocCount(collectedBuckets.ord());
+                            // if the key is below threshold, reinsert into the new ords
+                            if (docCount <= maxDocCount) {
+                                StringRareTerms.Bucket bucket = new StringRareTerms.Bucket(
+                                    BytesRef.deepCopyOf(scratch),
+                                    docCount,
+                                    null,
+                                    format
+                                );
+                                bucket.bucketOrd = offset + bucketsInThisOwningBucketToCollect.add(scratch);
+                                mergeMap.set(collectedBuckets.ord(), bucket.bucketOrd);
+                                builtBuckets.add(bucket);
+                                keepCount++;
+                            } else {
+                                filters.get(owningOrdIdx).add(scratch);
+                            }
+                        }
+                        rarestPerOrd.set(owningOrdIdx, builtBuckets.toArray(StringRareTerms.Bucket[]::new));
+                        offset += bucketsInThisOwningBucketToCollect.size();
                     }
                 }
-                rarestPerOrd[owningOrdIdx] = builtBuckets.toArray(StringRareTerms.Bucket[]::new);
-                offset += bucketsInThisOwningBucketToCollect.size();
-            }
-        }
 
-        /*
-         * Only merge/delete the ordinals if we have actually deleted one,
-         * to save on some redundant work.
-         */
-        if (keepCount != mergeMap.length) {
-            LongUnaryOperator howToMerge = b -> mergeMap[(int) b];
-            rewriteBuckets(offset, howToMerge);
-            if (deferringCollector() != null) {
-                ((BestBucketsDeferringCollector) deferringCollector()).rewriteBuckets(howToMerge);
+                /*
+                 * Only merge/delete the ordinals if we have actually deleted one,
+                 * to save on some redundant work.
+                 */
+                if (keepCount != mergeMap.size()) {
+                    LongUnaryOperator howToMerge = mergeMap::get;
+                    rewriteBuckets(offset, howToMerge);
+                    if (deferringCollector() != null) {
+                        ((BestBucketsDeferringCollector) deferringCollector()).rewriteBuckets(howToMerge);
+                    }
+                }
             }
-        }
 
-        /*
-         * Now build the results!
-         */
-        buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            Arrays.sort(rarestPerOrd[ordIdx], ORDER.comparator());
-            result[ordIdx] = new StringRareTerms(
-                name,
-                ORDER,
-                metadata(),
-                format,
-                Arrays.asList(rarestPerOrd[ordIdx]),
-                maxDocCount,
-                filters[ordIdx]
-            );
+            /*
+             * Now build the results!
+             */
+            buildSubAggsForAllBuckets(rarestPerOrd, b -> b.bucketOrd, (b, aggs) -> b.aggregations = aggs);
+            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                StringRareTerms.Bucket[] buckets = rarestPerOrd.get(ordIdx);
+                Arrays.sort(buckets, ORDER.comparator());
+                result[ordIdx] = new StringRareTerms(
+                    name,
+                    ORDER,
+                    metadata(),
+                    format,
+                    Arrays.asList(buckets),
+                    maxDocCount,
+                    filters.get(ordIdx)
+                );
+            }
+            return result;
         }
-        return result;
     }
 
     @Override

+ 5 - 4
server/src/main/java/org/elasticsearch/search/aggregations/metrics/MetricsAggregator.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.search.aggregations.metrics;
 
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorBase;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
@@ -36,10 +37,10 @@ public abstract class MetricsAggregator extends AggregatorBase {
     public abstract InternalAggregation buildAggregation(long owningBucketOrd) throws IOException;
 
     @Override
-    public final InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            results[ordIdx] = buildAggregation(owningBucketOrds[ordIdx]);
+    public final InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+            results[ordIdx] = buildAggregation(owningBucketOrds.get(ordIdx));
         }
         return results;
     }

+ 2 - 1
server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.search.profile.aggregation;
 
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -68,7 +69,7 @@ public class ProfilingAggregator extends Aggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
         Timer timer = profileBreakdown.getNewTimer(AggregationTimingType.BUILD_AGGREGATION);
         InternalAggregation[] result;
         timer.start();

+ 2 - 1
server/src/test/java/org/elasticsearch/search/aggregations/AdaptingAggregatorTests.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.search.aggregations;
 
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.CheckedFunction;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.MapperServiceTestCase;
@@ -113,7 +114,7 @@ public class AdaptingAggregatorTests extends MapperServiceTestCase {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) {
             return new InternalAggregation[] { null };
         }
 

+ 2 - 1
server/src/test/java/org/elasticsearch/search/aggregations/AggregatorBaseTests.java

@@ -15,6 +15,7 @@ import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.index.mapper.DateFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
@@ -47,7 +48,7 @@ public class AggregatorBaseTests extends MapperServiceTestCase {
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) {
             throw new UnsupportedOperationException();
         }
 

+ 15 - 5
server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java

@@ -28,6 +28,8 @@ import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.elasticsearch.common.CheckedBiConsumer;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
 import org.elasticsearch.search.aggregations.BucketCollector;
@@ -77,7 +79,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         collector.preCollection();
         indexSearcher.search(termQuery, collector.asCollector());
         collector.postCollection();
-        collector.prepareSelectedBuckets(0);
+        collector.prepareSelectedBuckets(BigArrays.NON_RECYCLING_INSTANCE.newLongArray(1, true));
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
         for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
@@ -91,7 +93,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         collector.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), collector.asCollector());
         collector.postCollection();
-        collector.prepareSelectedBuckets(0);
+        collector.prepareSelectedBuckets(BigArrays.NON_RECYCLING_INSTANCE.newLongArray(1, true));
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
         for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
@@ -141,7 +143,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
                 }
             }
         }, (deferringCollector, finalCollector) -> {
-            deferringCollector.prepareSelectedBuckets(0, 8, 9);
+            deferringCollector.prepareSelectedBuckets(toLongArray(0, 8, 9));
 
             equalTo(Map.of(0L, List.of(0, 1, 2, 3, 4, 5, 6, 7), 1L, List.of(8), 2L, List.of(9)));
         });
@@ -158,7 +160,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
                 }
             }
         }, (deferringCollector, finalCollector) -> {
-            deferringCollector.prepareSelectedBuckets(0, 8, 9);
+            deferringCollector.prepareSelectedBuckets(toLongArray(0, 8, 9));
 
             assertThat(finalCollector.collection, equalTo(Map.of(0L, List.of(4, 5, 6, 7), 1L, List.of(8), 2L, List.of(9))));
         });
@@ -176,12 +178,20 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
                 }
             }
         }, (deferringCollector, finalCollector) -> {
-            deferringCollector.prepareSelectedBuckets(0, 8, 9);
+            deferringCollector.prepareSelectedBuckets(toLongArray(0, 8, 9));
 
             assertThat(finalCollector.collection, equalTo(Map.of(0L, List.of(0, 1, 2, 3), 1L, List.of(8), 2L, List.of(9))));
         });
     }
 
+    private LongArray toLongArray(long... lons) {
+        LongArray longArray = BigArrays.NON_RECYCLING_INSTANCE.newLongArray(lons.length);
+        for (int i = 0; i < lons.length; i++) {
+            longArray.set(i, lons[i]);
+        }
+        return longArray;
+    }
+
     private void testCase(
         BiFunction<BestBucketsDeferringCollector, LeafBucketCollector, LeafBucketCollector> leafCollector,
         CheckedBiConsumer<BestBucketsDeferringCollector, CollectingBucketCollector, IOException> verify

+ 2 - 1
server/src/test/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregatorTests.java

@@ -16,6 +16,7 @@ import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
@@ -72,7 +73,7 @@ public class BucketsAggregatorTests extends AggregatorTestCase {
                     }
 
                     @Override
-                    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+                    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) {
                         return new InternalAggregation[0];
                     }
 

+ 2 - 1
server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java

@@ -22,6 +22,7 @@ import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.MockPageCacheRecycler;
 import org.elasticsearch.indices.breaker.NoneCircuitBreakerService;
@@ -68,7 +69,7 @@ public class BestDocsDeferringCollectorTests extends AggregatorTestCase {
         collector.preCollection();
         indexSearcher.search(termQuery, collector.asCollector());
         collector.postCollection();
-        collector.prepareSelectedBuckets(0);
+        collector.prepareSelectedBuckets(BigArrays.NON_RECYCLING_INSTANCE.newLongArray(1, true));
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
         for (ScoreDoc scoreDoc : topDocs.scoreDocs) {

+ 52 - 45
x-pack/plugin/analytics/src/main/java/org/elasticsearch/xpack/analytics/multiterms/MultiTermsAggregator.java

@@ -20,6 +20,8 @@ import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.util.LongArray;
+import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.util.ObjectArrayPriorityQueue;
 import org.elasticsearch.core.CheckedConsumer;
 import org.elasticsearch.core.Releasables;
@@ -235,57 +237,62 @@ class MultiTermsAggregator extends DeferableBucketAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        InternalMultiTerms.Bucket[][] topBucketsPerOrd = new InternalMultiTerms.Bucket[owningBucketOrds.length][];
-        long[] otherDocCounts = new long[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            long bucketsInOrd = bucketOrds.bucketsInOrd(owningBucketOrds[ordIdx]);
-
-            int size = (int) Math.min(bucketsInOrd, bucketCountThresholds.getShardSize());
-            try (
-                ObjectArrayPriorityQueue<InternalMultiTerms.Bucket> ordered = new BucketPriorityQueue<>(
-                    size,
-                    bigArrays(),
-                    partiallyBuiltBucketComparator
-                )
-            ) {
-                InternalMultiTerms.Bucket spare = null;
-                BytesRef spareKey = null;
-                BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrds[ordIdx]);
-                while (ordsEnum.next()) {
-                    long docCount = bucketDocCount(ordsEnum.ord());
-                    otherDocCounts[ordIdx] += docCount;
-                    if (docCount < bucketCountThresholds.getShardMinDocCount()) {
-                        continue;
-                    }
-                    if (spare == null) {
-                        spare = new InternalMultiTerms.Bucket(null, 0, null, showTermDocCountError, 0, formats, keyConverters);
-                        spareKey = new BytesRef();
+    public InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        try (
+            LongArray otherDocCounts = bigArrays().newLongArray(owningBucketOrds.size(), true);
+            ObjectArray<InternalMultiTerms.Bucket[]> topBucketsPerOrd = bigArrays().newObjectArray(owningBucketOrds.size())
+        ) {
+            for (long ordIdx = 0; ordIdx < owningBucketOrds.size(); ordIdx++) {
+                final long owningBucketOrd = owningBucketOrds.get(ordIdx);
+                long bucketsInOrd = bucketOrds.bucketsInOrd(owningBucketOrd);
+
+                int size = (int) Math.min(bucketsInOrd, bucketCountThresholds.getShardSize());
+                try (
+                    ObjectArrayPriorityQueue<InternalMultiTerms.Bucket> ordered = new BucketPriorityQueue<>(
+                        size,
+                        bigArrays(),
+                        partiallyBuiltBucketComparator
+                    )
+                ) {
+                    InternalMultiTerms.Bucket spare = null;
+                    BytesRef spareKey = null;
+                    BytesKeyedBucketOrds.BucketOrdsEnum ordsEnum = bucketOrds.ordsEnum(owningBucketOrd);
+                    while (ordsEnum.next()) {
+                        long docCount = bucketDocCount(ordsEnum.ord());
+                        otherDocCounts.increment(ordIdx, docCount);
+                        if (docCount < bucketCountThresholds.getShardMinDocCount()) {
+                            continue;
+                        }
+                        if (spare == null) {
+                            spare = new InternalMultiTerms.Bucket(null, 0, null, showTermDocCountError, 0, formats, keyConverters);
+                            spareKey = new BytesRef();
+                        }
+                        ordsEnum.readValue(spareKey);
+                        spare.terms = unpackTerms(spareKey);
+                        spare.docCount = docCount;
+                        spare.bucketOrd = ordsEnum.ord();
+                        spare = ordered.insertWithOverflow(spare);
                     }
-                    ordsEnum.readValue(spareKey);
-                    spare.terms = unpackTerms(spareKey);
-                    spare.docCount = docCount;
-                    spare.bucketOrd = ordsEnum.ord();
-                    spare = ordered.insertWithOverflow(spare);
-                }
 
-                // Get the top buckets
-                InternalMultiTerms.Bucket[] bucketsForOrd = new InternalMultiTerms.Bucket[(int) ordered.size()];
-                topBucketsPerOrd[ordIdx] = bucketsForOrd;
-                for (int b = (int) ordered.size() - 1; b >= 0; --b) {
-                    topBucketsPerOrd[ordIdx][b] = ordered.pop();
-                    otherDocCounts[ordIdx] -= topBucketsPerOrd[ordIdx][b].getDocCount();
+                    // Get the top buckets
+                    InternalMultiTerms.Bucket[] bucketsForOrd = new InternalMultiTerms.Bucket[(int) ordered.size()];
+                    topBucketsPerOrd.set(ordIdx, bucketsForOrd);
+                    for (int b = (int) ordered.size() - 1; b >= 0; --b) {
+                        InternalMultiTerms.Bucket[] buckets = topBucketsPerOrd.get(ordIdx);
+                        buckets[b] = ordered.pop();
+                        otherDocCounts.increment(ordIdx, -buckets[b].getDocCount());
+                    }
                 }
             }
-        }
 
-        buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
+            buildSubAggsForAllBuckets(topBucketsPerOrd, b -> b.bucketOrd, (b, a) -> b.aggregations = a);
 
-        InternalAggregation[] result = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
-            result[ordIdx] = buildResult(otherDocCounts[ordIdx], topBucketsPerOrd[ordIdx]);
+            InternalAggregation[] result = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+            for (int ordIdx = 0; ordIdx < result.length; ordIdx++) {
+                result[ordIdx] = buildResult(otherDocCounts.get(ordIdx), topBucketsPerOrd.get(ordIdx));
+            }
+            return result;
         }
-        return result;
     }
 
     InternalMultiTerms buildResult(long otherDocCount, InternalMultiTerms.Bucket[] topBuckets) {
@@ -305,7 +312,7 @@ class MultiTermsAggregator extends DeferableBucketAggregator {
             bucketCountThresholds.getShardSize(),
             showTermDocCountError,
             otherDocCount,
-            List.of(topBuckets),
+            Arrays.asList(topBuckets),
             0,
             formats,
             keyConverters,

+ 25 - 23
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java

@@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.util.BytesRefHash;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.core.Releasables;
@@ -110,31 +111,32 @@ public class CategorizeTextAggregator extends DeferableBucketAggregator {
     }
 
     @Override
-    public InternalAggregation[] buildAggregations(long[] ordsToCollect) throws IOException {
-        Bucket[][] topBucketsPerOrd = new Bucket[ordsToCollect.length][];
-        for (int ordIdx = 0; ordIdx < ordsToCollect.length; ordIdx++) {
-            final long ord = ordsToCollect[ordIdx];
-            final TokenListCategorizer categorizer = (ord < categorizers.size()) ? categorizers.get(ord) : null;
-            if (categorizer == null) {
-                topBucketsPerOrd[ordIdx] = new Bucket[0];
-                continue;
+    public InternalAggregation[] buildAggregations(LongArray ordsToCollect) throws IOException {
+        try (ObjectArray<Bucket[]> topBucketsPerOrd = bigArrays().newObjectArray(ordsToCollect.size())) {
+            for (long ordIdx = 0; ordIdx < ordsToCollect.size(); ordIdx++) {
+                final long ord = ordsToCollect.get(ordIdx);
+                final TokenListCategorizer categorizer = (ord < categorizers.size()) ? categorizers.get(ord) : null;
+                if (categorizer == null) {
+                    topBucketsPerOrd.set(ordIdx, new Bucket[0]);
+                    continue;
+                }
+                int size = (int) Math.min(bucketOrds.bucketsInOrd(ordIdx), bucketCountThresholds.getShardSize());
+                topBucketsPerOrd.set(ordIdx, categorizer.toOrderedBuckets(size));
             }
-            int size = (int) Math.min(bucketOrds.bucketsInOrd(ordIdx), bucketCountThresholds.getShardSize());
-            topBucketsPerOrd[ordIdx] = categorizer.toOrderedBuckets(size);
-        }
-        buildSubAggsForAllBuckets(topBucketsPerOrd, Bucket::getBucketOrd, Bucket::setAggregations);
-        InternalAggregation[] results = new InternalAggregation[ordsToCollect.length];
-        for (int ordIdx = 0; ordIdx < ordsToCollect.length; ordIdx++) {
-            results[ordIdx] = new InternalCategorizationAggregation(
-                name,
-                bucketCountThresholds.getRequiredSize(),
-                bucketCountThresholds.getMinDocCount(),
-                similarityThreshold,
-                metadata(),
-                Arrays.asList(topBucketsPerOrd[ordIdx])
-            );
+            buildSubAggsForAllBuckets(topBucketsPerOrd, Bucket::getBucketOrd, Bucket::setAggregations);
+            InternalAggregation[] results = new InternalAggregation[Math.toIntExact(ordsToCollect.size())];
+            for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
+                results[ordIdx] = new InternalCategorizationAggregation(
+                    name,
+                    bucketCountThresholds.getRequiredSize(),
+                    bucketCountThresholds.getMinDocCount(),
+                    similarityThreshold,
+                    metadata(),
+                    Arrays.asList(topBucketsPerOrd.get(ordIdx))
+                );
+            }
+            return results;
         }
-        return results;
     }
 
     @Override

+ 4 - 2
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/frequentitemsets/mr/DelegatingCircuitBreakerService.java

@@ -40,10 +40,12 @@ import java.util.function.Consumer;
  * At the time of writing circuit breakers are a global gauge.)
  *
  * After the map phase and before reduce, the {@link ItemSetMapReduceAggregator} creates instances of
- * {@link InternalItemSetMapReduceAggregation}, see {@link ItemSetMapReduceAggregator#buildAggregations(long[])}.
+ * {@link InternalItemSetMapReduceAggregation}, see
+ * {@link ItemSetMapReduceAggregator#buildAggregations(org.elasticsearch.common.util.LongArray)}.
  *
  * (Note 1: Instead of keeping the existing instance, it would have been possible to deep-copy the object like
- * {@link CardinalityAggregator#buildAggregations(long[])}. I decided against this approach mainly because the deep-copy isn't
+ * {@link CardinalityAggregator#buildAggregations(org.elasticsearch.common.util.LongArray)}.
+ * I decided against this approach mainly because the deep-copy isn't
  * secured by circuit breakers, meaning the node could run out of memory during the deep-copy.)
  * (Note 2: Between {@link ItemSetMapReduceAggregator#doClose()} and serializing {@link InternalItemSetMapReduceAggregation}
  * memory accounting is broken, meaning the agg context gets closed and bytes get returned to the circuit breaker before memory is

+ 4 - 3
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/frequentitemsets/mr/ItemSetMapReduceAggregator.java

@@ -17,6 +17,7 @@ import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.LongObjectPagedHashMap;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.Tuple;
@@ -117,9 +118,9 @@ public abstract class ItemSetMapReduceAggregator<
     }
 
     @Override
-    public final InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
-        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
+    public final InternalAggregation[] buildAggregations(LongArray owningBucketOrds) throws IOException {
+        InternalAggregation[] results = new InternalAggregation[Math.toIntExact(owningBucketOrds.size())];
+        for (int ordIdx = 0; ordIdx < results.length; ordIdx++) {
             results[ordIdx] = buildAggregation(ordIdx);
         }