Better now

test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java

@@ -678,6 +678,9 @@ public void addEstimateBytesAndMaybeBreak(long bytes, String label) throws Circu
             while (true) {
                 long old = used.get();
                 long total = old + bytes;
+                if (total < 0) {
+                    throw new AssertionError("total must be >= 0 but was [" + total + "]");
+                }
                 if (total > max.getBytes()) {
                     throw new CircuitBreakingException(ERROR_MESSAGE, bytes, max.getBytes(), Durability.TRANSIENT);
                 }
@@ -689,7 +692,10 @@ public void addEstimateBytesAndMaybeBreak(long bytes, String label) throws Circu
 
         @Override
         public void addWithoutBreaking(long bytes) {
-            used.addAndGet(bytes);
+            long total = used.addAndGet(bytes);
+            if (total < 0) {
+                throw new AssertionError("total must be >= 0 but was [" + total + "]");
+            }
         }
 
         @Override

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

@@ -193,26 +193,36 @@ public BytesRefBlockBuilder mvOrdering(Block.MvOrdering mvOrdering) {
     public BytesRefBlock build() {
         finish();
         BytesRefBlock block;
+        assert estimatedBytes == 0 || firstValueIndexes != null;
         if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
             block = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory).asBlock();
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, false);
             Releasables.closeExpectNoException(values);
         } else {
             if (isDense() && singleValued()) {
                 block = new BytesRefArrayVector(values, positionCount, blockFactory).asBlock();
             } else {
                 block = new BytesRefArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, blockFactory);
             }
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes - values.bigArraysRamBytesUsed(), false);
         }
-        /*
-         * Update the breaker with the actual bytes used.
-         * We pass false below even though we've used the bytes. That's weird,
-         * but if we break here we will throw away the used memory, letting
-         * it be deallocated. The exception will bubble up and the builder will
-         * still technically be open, meaning the calling code should close it
-         * which will return all used memory to the breaker.
-         */
-        blockFactory.adjustBreaker(block.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
-        assert estimatedBytes == 0;
+        values = null;
         built();
         return block;
     }

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

@@ -56,22 +56,31 @@ protected void growValuesArray(int newSize) {
     public BytesRefVector build() {
         finish();
         BytesRefVector vector;
+        assert estimatedBytes == 0;
         if (valueCount == 1) {
             vector = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory);
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(vector.ramBytesUsed(), false);
             Releasables.closeExpectNoException(values);
         } else {
             vector = new BytesRefArrayVector(values, valueCount, blockFactory);
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(vector.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
         }
-        assert estimatedBytes == 0;
-        /*
-         * Update the breaker with the actual bytes used.
-         * We pass false below even though we've used the bytes. That's weird,
-         * but if we break here we will throw away the used memory, letting
-         * it be deallocated. The exception will bubble up and the builder will
-         * still technically be open, meaning the calling code should close it
-         * which will return all used memory to the breaker.
-         */
-        blockFactory.adjustBreaker(vector.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
         values = null;
         built();
         return vector;

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st

@@ -248,18 +248,44 @@ $endif$
     public $Type$Block build() {
         finish();
         $Type$Block block;
-        if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
 $if(BytesRef)$
+        assert estimatedBytes == 0 || firstValueIndexes != null;
+        if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
             block = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory).asBlock();
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, false);
             Releasables.closeExpectNoException(values);
         } else {
+            if (isDense() && singleValued()) {
+                block = new $Type$ArrayVector(values, positionCount, blockFactory).asBlock();
+            } else {
+                block = new $Type$ArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, blockFactory);
+            }
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes - values.bigArraysRamBytesUsed(), false);
+        }
+        values = null;
 $else$
+        if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
             block = new Constant$Type$Vector(values[0], 1, blockFactory).asBlock();
         } else {
             if (values.length - valueCount > 1024 || valueCount < (values.length / 2)) {
                 values = Arrays.copyOf(values, valueCount);
             }
-$endif$
             if (isDense() && singleValued()) {
                 block = new $Type$ArrayVector(values, positionCount, blockFactory).asBlock();
             } else {
@@ -274,10 +300,6 @@ $endif$
          * still technically be open, meaning the calling code should close it
          * which will return all used memory to the breaker.
          */
-$if(BytesRef)$
-        blockFactory.adjustBreaker(block.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
-        assert estimatedBytes == 0;
-$else$
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, false);
 $endif$
         built();

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st

@@ -86,22 +86,31 @@ $endif$
         finish();
         $Type$Vector vector;
 $if(BytesRef)$
+        assert estimatedBytes == 0;
         if (valueCount == 1) {
             vector = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory);
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(vector.ramBytesUsed(), false);
             Releasables.closeExpectNoException(values);
         } else {
             vector = new $Type$ArrayVector(values, valueCount, blockFactory);
+            /*
+             * Update the breaker with the actual bytes used.
+             * We pass false below even though we've used the bytes. That's weird,
+             * but if we break here we will throw away the used memory, letting
+             * it be deallocated. The exception will bubble up and the builder will
+             * still technically be open, meaning the calling code should close it
+             * which will return all used memory to the breaker.
+             */
+            blockFactory.adjustBreaker(vector.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
         }
-        assert estimatedBytes == 0;
-        /*
-         * Update the breaker with the actual bytes used.
-         * We pass false below even though we've used the bytes. That's weird,
-         * but if we break here we will throw away the used memory, letting
-         * it be deallocated. The exception will bubble up and the builder will
-         * still technically be open, meaning the calling code should close it
-         * which will return all used memory to the breaker.
-         */
-        blockFactory.adjustBreaker(vector.ramBytesUsed() - values.bigArraysRamBytesUsed(), false);
         values = null;
 $else$
         if (valueCount == 1) {

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderTests.java

@@ -9,8 +9,6 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import com.carrotsearch.randomizedtesting.annotations.Repeat;
-
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -75,10 +73,46 @@ public void testCloseWithoutBuilding() {
         assertThat(blockFactory.breaker().getUsed(), equalTo(0L));
     }
 
-    public void testSingleBuild() {
+    public void testBuildSmallSingleValued() {
+        testBuild(between(1, 100), false, 1);
+    }
+
+    public void testBuildHugeSingleValued() {
+        testBuild(between(1_000, 50_000), false, 1);
+    }
+
+    public void testBuildSmallSingleValuedNullable() {
+        testBuild(between(1, 100), true, 1);
+    }
+
+    public void testBuildHugeSingleValuedNullable() {
+        testBuild(between(1_000, 50_000), true, 1);
+    }
+
+    public void testBuildSmallMultiValued() {
+        testBuild(between(1, 100), false, 3);
+    }
+
+    public void testBuildHugeMultiValued() {
+        testBuild(between(1_000, 50_000), false, 3);
+    }
+
+    public void testBuildSmallMultiValuedNullable() {
+        testBuild(between(1, 100), true, 3);
+    }
+
+    public void testBuildHugeMultiValuedNullable() {
+        testBuild(between(1_000, 50_000), true, 3);
+    }
+
+    public void testBuildSingle() {
+        testBuild(1, false, 1);
+    }
+
+    private void testBuild(int size, boolean nullable, int maxValueCount) {
         BlockFactory blockFactory = BlockFactoryTests.blockFactory(ByteSizeValue.ofGb(1));
-        try (Block.Builder builder = elementType.newBlockBuilder(10, blockFactory)) {
-            BasicBlockTests.RandomBlock random = BasicBlockTests.randomBlock(elementType, 10, false, 1, 1, 0, 0);
+        try (Block.Builder builder = elementType.newBlockBuilder(randomBoolean() ? size : 1, blockFactory)) {
+            BasicBlockTests.RandomBlock random = BasicBlockTests.randomBlock(elementType, size, nullable, 1, maxValueCount, 0, 0);
             builder.copyFrom(random.block(), 0, random.block().getPositionCount());
             try (Block built = builder.build()) {
                 assertThat(built, equalTo(random.block()));

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorBuilderTests.java

@@ -9,10 +9,6 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import com.carrotsearch.randomizedtesting.annotations.Repeat;
-
-import com.carrotsearch.randomizedtesting.annotations.Seed;
-
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
@@ -28,7 +24,6 @@
 
 import static org.hamcrest.Matchers.equalTo;
 
-@Seed("7E594D56CF8843C5:4B7F19DE1F496E74")
 public class VectorBuilderTests extends ESTestCase {
     @ParametersFactory
     public static List<Object[]> params() {
@@ -54,10 +49,22 @@ public void testCloseWithoutBuilding() {
         assertThat(blockFactory.breaker().getUsed(), equalTo(0L));
     }
 
-    public void testSingleBuild() {
+    public void testBuildSmall() {
+        testBuild(between(1, 100));
+    }
+
+    public void testBuildHuge() {
+        testBuild(between(1_000, 50_000));
+    }
+
+    public void testBuildSingle() {
+        testBuild(1);
+    }
+
+    private void testBuild(int size) {
         BlockFactory blockFactory = BlockFactoryTests.blockFactory(ByteSizeValue.ofGb(1));
-        try (Vector.Builder builder = vectorBuilder(10, blockFactory)) {
-            BasicBlockTests.RandomBlock random = BasicBlockTests.randomBlock(elementType, 10, false, 1, 1, 0, 0);
+        try (Vector.Builder builder = vectorBuilder(randomBoolean() ? size : 1, blockFactory)) {
+            BasicBlockTests.RandomBlock random = BasicBlockTests.randomBlock(elementType, size, false, 1, 1, 0, 0);
             fill(builder, random.block().asVector());
             try (Vector built = builder.build()) {
                 assertThat(built, equalTo(random.block().asVector()));