diff --git a/test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java b/test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java index 02e8cfd7f16fc..0a0592b5a01f2 100644 --- a/test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java +++ b/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 diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java index f11677d2067f8..a60b26667eb79 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java +++ b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java @@ -193,8 +193,18 @@ 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()) { @@ -202,17 +212,17 @@ public BytesRefBlock build() { } 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; } diff --git a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefVectorBuilder.java b/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefVectorBuilder.java index c5c904715566a..ebff50383df26 100644 --- a/x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefVectorBuilder.java +++ b/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; diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st index 78bebb0fd4ef3..0ccfc45f18664 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st +++ b/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(); diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st index 62859d809f0d3..fc5ebb08483f6 100644 --- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st +++ b/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) { diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderTests.java index f85192aa40a90..2179e68c47832 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockBuilderTests.java +++ b/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())); diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorBuilderTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorBuilderTests.java index a6228119ee524..656d79070f217 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorBuilderTests.java +++ b/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 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()));