allocate whole buckets at a time

Author: nik9000

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/sort/BytesRefBucketedSort.java

@@ -8,10 +8,12 @@
 package org.elasticsearch.compute.data.sort;
 
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.ByteUtils;
 import org.elasticsearch.common.util.ObjectArray;
+import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.IntVector;
@@ -123,7 +125,7 @@ public void collect(BytesRef value, int bucket) {
         // Gathering mode
         long requiredSize = common.endIndex(rootIndex);
         if (values.size() < requiredSize) {
-            grow(requiredSize);
+            grow(bucket);
         }
         int next = getNextGatherOffset(rootIndex);
         common.assertValidNextOffset(next);
@@ -143,13 +145,13 @@ public void collect(BytesRef value, int bucket) {
      */
     public void merge(int bucket, BytesRefBucketedSort other, int otherBucket) {
         long otherRootIndex = other.common.rootIndex(otherBucket);
-        long otherEnd = other.common.endIndex(otherRootIndex);
-        if (otherEnd >= other.values.size()) {
+        if (otherRootIndex >= other.values.size()) {
             // The value was never collected.
             return;
         }
         other.checkInvariant(otherBucket);
         long otherStart = other.startIndex(otherBucket, otherRootIndex);
+        long otherEnd = other.common.endIndex(otherRootIndex);
         // TODO: This can be improved for heapified buckets by making use of the heap structures
         for (long i = otherStart; i < otherEnd; i++) {
             collect(other.values.get(i).bytesRefView(), bucket);
@@ -271,13 +273,21 @@ private void swap(long lhs, long rhs) {
 
     /**
      * Allocate storage for more buckets and store the "next gather offset"
-     * for those new buckets.
+     * for those new buckets. We always grow the storage by whole bucket's
+     * worth of slots at a time. We never allocate space for partial buckets.
      */
-    private void grow(long requiredSize) {
+    private void grow(int bucket) {
         long oldMax = values.size();
-        values = common.bigArrays.grow(values, requiredSize);
+        assert oldMax % common.bucketSize == 0;
+
+        long newSizeInBuckets = BigArrays.overSize(
+            bucket + 1,
+            PageCacheRecycler.OBJECT_PAGE_SIZE,
+            RamUsageEstimator.NUM_BYTES_OBJECT_REF * common.bucketSize
+        );
+        values = common.bigArrays.resize(values, newSizeInBuckets * common.bucketSize);
         // Set the next gather offsets for all newly allocated buckets.
-        fillGatherOffsets(oldMax - (oldMax % common.bucketSize));
+        fillGatherOffsets(oldMax);
     }
 
     /**
@@ -296,6 +306,7 @@ private void fillGatherOffsets(long startingAt) {
             bytes.grow(Integer.BYTES);
             bytes.setLength(Integer.BYTES);
             ByteUtils.writeIntLE(nextOffset, bytes.bytes(), 0);
+            checkInvariant(Math.toIntExact(bucketRoot / common.bucketSize));
         }
     }