more fixes and rework

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

@@ -194,15 +194,6 @@ public BooleanBlock build() {
                 block = blockFactory.newBooleanArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, estimatedBytes);
             }
         }
-        /*
-         * 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);
         built();
         return block;
     }

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

@@ -59,15 +59,6 @@ public BooleanVector build() {
             }
             vector = blockFactory.newBooleanArrayVector(values, valueCount, estimatedBytes);
         }
-        /*
-         * 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() - estimatedBytes, false);
         built();
         return vector;
     }

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

@@ -194,15 +194,6 @@ public DoubleBlock build() {
                 block = blockFactory.newDoubleArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, estimatedBytes);
             }
         }
-        /*
-         * 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);
         built();
         return block;
     }

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

@@ -59,15 +59,6 @@ public DoubleVector build() {
             }
             vector = blockFactory.newDoubleArrayVector(values, valueCount, estimatedBytes);
         }
-        /*
-         * 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() - estimatedBytes, false);
         built();
         return vector;
     }

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

@@ -194,15 +194,6 @@ public IntBlock build() {
                 block = blockFactory.newIntArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, estimatedBytes);
             }
         }
-        /*
-         * 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);
         built();
         return block;
     }

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

@@ -59,15 +59,6 @@ public IntVector build() {
             }
             vector = blockFactory.newIntArrayVector(values, valueCount, estimatedBytes);
         }
-        /*
-         * 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() - estimatedBytes, false);
         built();
         return vector;
     }

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

@@ -194,15 +194,6 @@ public LongBlock build() {
                 block = blockFactory.newLongArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, estimatedBytes);
             }
         }
-        /*
-         * 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);
         built();
         return block;
     }

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

@@ -59,15 +59,6 @@ public LongVector build() {
             }
             vector = blockFactory.newLongArrayVector(values, valueCount, estimatedBytes);
         }
-        /*
-         * 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() - estimatedBytes, false);
         built();
         return vector;
     }

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

@@ -138,23 +138,21 @@ BooleanVector.FixedBuilder newBooleanVectorFixedBuilder(int size) {
         return new BooleanVectorFixedBuilder(size, this);
     }
 
-    public BooleanBlock newBooleanArrayBlock(
-        boolean[] values,
-        int positionCount,
-        int[] firstValueIndexes,
-        BitSet nulls,
-        MvOrdering mvOrdering
-    ) {
-        var b = new BooleanArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+    public final BooleanBlock newBooleanArrayBlock(boolean[] values, int pc, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
+        return newBooleanArrayBlock(values, pc, firstValueIndexes, nulls, mvOrdering, 0L);
+    }
+
+    public BooleanBlock newBooleanArrayBlock(boolean[] values, int pc, int[] fvi, BitSet nulls, MvOrdering mvOrder, long preAdjustedBytes) {
+        var b = new BooleanArrayBlock(values, pc, fvi, nulls, mvOrder, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
     public BooleanVector.Builder newBooleanVectorBuilder(int estimatedSize) {
         return new BooleanVectorBuilder(estimatedSize, this);
     }
 
-    public BooleanVector newBooleanArrayVector(boolean[] values, int positionCount) {
+    public final BooleanVector newBooleanArrayVector(boolean[] values, int positionCount) {
         return newBooleanArrayVector(values, positionCount, 0L);
     }
 
@@ -164,19 +162,27 @@ public BooleanVector newBooleanArrayVector(boolean[] values, int positionCount,
         return b;
     }
 
-    public BooleanBlock newConstantBooleanBlockWith(boolean value, int positions) {
+    public final BooleanBlock newConstantBooleanBlockWith(boolean value, int positions) {
+        return newConstantBooleanBlockWith(value, positions, 0L);
+    }
+
+    public BooleanBlock newConstantBooleanBlockWith(boolean value, int positions, long preAdjustedBytes) {
         var b = new ConstantBooleanVector(value, positions, this).asBlock();
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
     public IntBlock.Builder newIntBlockBuilder(int estimatedSize) {
         return new IntBlockBuilder(estimatedSize, this);
     }
 
-    public IntBlock newIntArrayBlock(int[] values, int positionCount, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
-        var b = new IntArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+    public final IntBlock newIntArrayBlock(int[] values, int positionCount, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
+        return newIntArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, 0L);
+    }
+
+    public IntBlock newIntArrayBlock(int[] values, int pc, int[] fvi, BitSet nulls, MvOrdering mvOrdering, long preAdjustedBytes) {
+        var b = new IntArrayBlock(values, pc, fvi, nulls, mvOrdering, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -192,7 +198,7 @@ IntVector.FixedBuilder newIntVectorFixedBuilder(int size) {
      * Creates a new Vector with the given values and positionCount. Equivalent to:
      *   newIntArrayVector(values, positionCount, 0L); // with zero pre-adjusted bytes
      */
-    public IntVector newIntArrayVector(int[] values, int positionCount) {
+    public final IntVector newIntArrayVector(int[] values, int positionCount) {
         return newIntArrayVector(values, positionCount, 0L);
     }
 
@@ -213,19 +219,27 @@ public IntVector newIntArrayVector(int[] values, int positionCount, long preAdju
         return b;
     }
 
-    public IntBlock newConstantIntBlockWith(int value, int positions) {
+    public final IntBlock newConstantIntBlockWith(int value, int positions) {
+        return newConstantIntBlockWith(value, positions, 0L);
+    }
+
+    public IntBlock newConstantIntBlockWith(int value, int positions, long preAdjustedBytes) {
         var b = new ConstantIntVector(value, positions, this).asBlock();
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
     public LongBlock.Builder newLongBlockBuilder(int estimatedSize) {
         return new LongBlockBuilder(estimatedSize, this);
     }
 
-    public LongBlock newLongArrayBlock(long[] values, int positionCount, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
-        var b = new LongArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+    public final LongBlock newLongArrayBlock(long[] values, int pc, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
+        return newLongArrayBlock(values, pc, firstValueIndexes, nulls, mvOrdering, 0L);
+    }
+
+    public LongBlock newLongArrayBlock(long[] values, int pc, int[] fvi, BitSet nulls, MvOrdering mvOrdering, long preAdjustedBytes) {
+        var b = new LongArrayBlock(values, pc, fvi, nulls, mvOrdering, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -237,7 +251,7 @@ LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
         return new LongVectorFixedBuilder(size, this);
     }
 
-    public LongVector newLongArrayVector(long[] values, int positionCount) {
+    public final LongVector newLongArrayVector(long[] values, int positionCount) {
         return newLongArrayVector(values, positionCount, 0L);
     }
 
@@ -247,25 +261,28 @@ public LongVector newLongArrayVector(long[] values, int positionCount, long preA
         return b;
     }
 
-    public LongBlock newConstantLongBlockWith(long value, int positions) {
+    public final LongBlock newConstantLongBlockWith(long value, int positions) {
+        return newConstantLongBlockWith(value, positions, 0L);
+    }
+
+    public LongBlock newConstantLongBlockWith(long value, int positions, long preAdjustedBytes) {
         var b = new ConstantLongVector(value, positions, this).asBlock();
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
     public DoubleBlock.Builder newDoubleBlockBuilder(int estimatedSize) {
         return new DoubleBlockBuilder(estimatedSize, this);
     }
 
-    public DoubleBlock newDoubleArrayBlock(
-        double[] values,
-        int positionCount,
-        int[] firstValueIndexes,
-        BitSet nulls,
-        MvOrdering mvOrdering
-    ) {
-        var b = new DoubleArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+    public final DoubleBlock newDoubleArrayBlock(double[] values, int pc, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
+        return newDoubleArrayBlock(values, pc, firstValueIndexes, nulls, mvOrdering, 0L);
+
+    }
+
+    public DoubleBlock newDoubleArrayBlock(double[] values, int pc, int[] fvi, BitSet nulls, MvOrdering mvOrdering, long preAdjustedBytes) {
+        var b = new DoubleArrayBlock(values, pc, fvi, nulls, mvOrdering, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -277,7 +294,7 @@ DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
         return new DoubleVectorFixedBuilder(size, this);
     }
 
-    public DoubleVector newDoubleArrayVector(double[] values, int positionCount) {
+    public final DoubleVector newDoubleArrayVector(double[] values, int positionCount) {
         return newDoubleArrayVector(values, positionCount, 0L);
     }
 
@@ -287,24 +304,22 @@ public DoubleVector newDoubleArrayVector(double[] values, int positionCount, lon
         return b;
     }
 
-    public DoubleBlock newConstantDoubleBlockWith(double value, int positions) {
+    public final DoubleBlock newConstantDoubleBlockWith(double value, int positions) {
+        return newConstantDoubleBlockWith(value, positions, 0L);
+    }
+
+    public DoubleBlock newConstantDoubleBlockWith(double value, int positions, long preAdjustedBytes) {
         var b = new ConstantDoubleVector(value, positions, this).asBlock();
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
     public BytesRefBlock.Builder newBytesRefBlockBuilder(int estimatedSize) {
         return new BytesRefBlockBuilder(estimatedSize, bigArrays, this);
     }
 
-    public BytesRefBlock newBytesRefArrayBlock(
-        BytesRefArray values,
-        int positionCount,
-        int[] firstValueIndexes,
-        BitSet nulls,
-        MvOrdering mvOrdering
-    ) {
-        var b = new BytesRefArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
+    public BytesRefBlock newBytesRefArrayBlock(BytesRefArray values, int pc, int[] firstValueIndexes, BitSet nulls, MvOrdering mvOrdering) {
+        var b = new BytesRefArrayBlock(values, pc, firstValueIndexes, nulls, mvOrdering, this);
         adjustBreaker(b.ramBytesUsed() - values.bigArraysRamBytesUsed(), true);
         return b;
     }

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

@@ -46,6 +46,8 @@ public class DocVector extends AbstractVector implements Vector {
      */
     private int[] shardSegmentDocMapBackwards;
 
+    final DocBlock block;
+
     public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean singleSegmentNonDecreasing) {
         super(shards.getPositionCount(), null);
         this.shards = shards;
@@ -62,6 +64,7 @@ public DocVector(IntVector shards, IntVector segments, IntVector docs, Boolean s
                 "invalid position count [" + shards.getPositionCount() + " != " + docs.getPositionCount() + "]"
             );
         }
+        block = new DocBlock(this);
     }
 
     public IntVector shards() {
@@ -168,7 +171,7 @@ protected void swap(int i, int j) {
 
     @Override
     public DocBlock asBlock() {
-        return new DocBlock(this);
+        return block;
     }
 
     @Override
@@ -218,6 +221,6 @@ public long ramBytesUsed() {
 
     @Override
     public void close() {
-        Releasables.closeExpectNoException(shards, segments, docs);
+        Releasables.closeExpectNoException(shards.asBlock(), segments.asBlock(), docs.asBlock()); // Ugh! we always close blocks
     }
 }

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

@@ -114,22 +114,13 @@ $if(BytesRef)$
         values = null;
 $else$
         if (valueCount == 1) {
-            vector = new Constant$Type$Vector(values[0], 1, blockFactory);
+            vector = blockFactory.newConstant$Type$BlockWith(values[0], 1, estimatedBytes).asVector();
         } else {
             if (values.length - valueCount > 1024 || valueCount < (values.length / 2)) {
                 values = Arrays.copyOf(values, valueCount);
             }
-            vector = new $Type$ArrayVector(values, valueCount, blockFactory);
+            vector = blockFactory.new$Type$ArrayVector(values, valueCount, estimatedBytes);
         }
-        /*
-         * 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() - estimatedBytes, false);
 $endif$
         built();
         return vector;

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/HashAggregationOperator.java

@@ -157,7 +157,9 @@ public void finish() {
             success = true;
         } finally {
             // selected should always be closed
-            Releasables.closeExpectNoException(selected.asBlock()); // we always close blocks, not vectors
+            if (selected != null) {
+                Releasables.closeExpectNoException(selected.asBlock()); // we always close blocks, not vectors
+            }
             if (success == false && blocks != null) {
                 Releasables.closeExpectNoException(blocks);
             }