Flip impl

Author: nik9000

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

@@ -285,12 +285,14 @@ private void grow(int bucket) {
         long oldMax = values.size();
         assert oldMax % common.bucketSize == 0;
 
-        long newSizeInBuckets = BigArrays.overSize(
-            bucket + 1,
+        long newSize = BigArrays.overSize(
+            ((long) bucket + 1) * common.bucketSize,
             PageCacheRecycler.OBJECT_PAGE_SIZE,
-            RamUsageEstimator.NUM_BYTES_OBJECT_REF * common.bucketSize
+            RamUsageEstimator.NUM_BYTES_OBJECT_REF
         );
-        values = common.bigArrays.resize(values, newSizeInBuckets * common.bucketSize);
+        // Round up to the next full bucket.
+        newSize = (newSize + common.bucketSize - 1) / common.bucketSize;
+        values = common.bigArrays.resize(values, newSize * common.bucketSize);
         // Set the next gather offsets for all newly allocated buckets.
         fillGatherOffsets(oldMax);
     }

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.ByteArray;
 import org.elasticsearch.common.util.ByteUtils;
+import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.IntVector;
@@ -106,7 +107,7 @@ public void collect(BytesRef value, int bucket) {
         // Gathering mode
         long requiredSize = common.endIndex(rootIndex) * IP_LENGTH;
         if (values.size() < requiredSize) {
-            grow(requiredSize);
+            grow(bucket);
         }
         int next = getNextGatherOffset(rootIndex);
         common.assertValidNextOffset(next);
@@ -273,17 +274,23 @@ private void swap(long lhs, long rhs) {
      * Allocate storage for more buckets and store the "next gather offset"
      * for those new buckets.
      */
-    private void grow(long minSize) {
+    private void grow(int bucket) {
         long oldMax = values.size() / IP_LENGTH;
-        values = common.bigArrays.grow(values, minSize);
+        assert oldMax % common.bucketSize == 0;
+
+        int bucketBytes = common.bucketSize * IP_LENGTH;
+        long newSize = BigArrays.overSize(((long) bucket + 1) * bucketBytes, PageCacheRecycler.BYTE_PAGE_SIZE, 1);
+        // Round up to the next full bucket.
+        newSize = (newSize + bucketBytes - 1) / bucketBytes;
+        values = common.bigArrays.resize(values, newSize * bucketBytes);
         // Set the next gather offsets for all newly allocated buckets.
-        setNextGatherOffsets(oldMax - (oldMax % common.bucketSize));
+        fillGatherOffsets(oldMax);
     }
 
     /**
      * Maintain the "next gather offsets" for newly allocated buckets.
      */
-    private void setNextGatherOffsets(long startingAt) {
+    private void fillGatherOffsets(long startingAt) {
         int nextOffset = common.bucketSize - 1;
         for (long bucketRoot = startingAt; bucketRoot < values.size() / IP_LENGTH; bucketRoot += common.bucketSize) {
             setNextGatherOffset(bucketRoot, nextOffset);

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

@@ -269,12 +269,10 @@ $endif$
         long oldMax = values.size();
         assert oldMax % bucketSize == 0;
 
-        long newSizeInBuckets = BigArrays.overSize(
-            bucket + 1,
-            PageCacheRecycler.$TYPE$_PAGE_SIZE,
-            $BYTES$ * bucketSize
-        );
-        values = bigArrays.resize(values, newSizeInBuckets * bucketSize);
+        long newSize = BigArrays.overSize(((long) bucket + 1) * bucketSize, PageCacheRecycler.$TYPE$_PAGE_SIZE, $BYTES$);
+        // Round up to the next full bucket.
+        newSize = (newSize + bucketSize - 1) / bucketSize;
+        values = bigArrays.resize(values, newSize * bucketSize);
         // Set the next gather offsets for all newly allocated buckets.
         fillGatherOffsets(oldMax);
     }