ESQL: Estimate memory usage on `Block.Builder` (#107923)

This adds a method to `Block.Builder` that estimates the number of bytes
that'll be used by the `Block` that it builds. It's not always accurate,
but it is directional.

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

@@ -20,7 +20,7 @@ import java.util.BitSet;
  */
 final class BooleanArrayBlock extends AbstractArrayBlock implements BooleanBlock {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(BooleanArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(BooleanArrayBlock.class);
 
     private final BooleanArrayVector vector;
 

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

@@ -46,6 +46,11 @@ final class BooleanVectorFixedBuilder implements BooleanVector.FixedBuilder {
             );
     }
 
+    @Override
+    public long estimatedBytes() {
+        return ramBytesUsed(values.length);
+    }
+
     @Override
     public BooleanVector build() {
         if (nextIndex < 0) {

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

@@ -23,7 +23,7 @@ import java.util.BitSet;
  */
 final class BytesRefArrayBlock extends AbstractArrayBlock implements BytesRefBlock {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(BytesRefArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(BytesRefArrayBlock.class);
 
     private final BytesRefArrayVector vector;
 

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

@@ -140,6 +140,11 @@ final class BytesRefBlockBuilder extends AbstractBlockBuilder implements BytesRe
         return this;
     }
 
+    @Override
+    public long estimatedBytes() {
+        return super.estimatedBytes() + BytesRefArrayBlock.BASE_RAM_BYTES_USED + values.ramBytesUsed();
+    }
+
     private BytesRefBlock buildFromBytesArray() {
         assert estimatedBytes == 0 || firstValueIndexes != null;
         final BytesRefBlock theBlock;

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

@@ -20,7 +20,7 @@ import java.util.BitSet;
  */
 final class DoubleArrayBlock extends AbstractArrayBlock implements DoubleBlock {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(DoubleArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(DoubleArrayBlock.class);
 
     private final DoubleArrayVector vector;
 

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

@@ -46,6 +46,11 @@ final class DoubleVectorFixedBuilder implements DoubleVector.FixedBuilder {
             );
     }
 
+    @Override
+    public long estimatedBytes() {
+        return ramBytesUsed(values.length);
+    }
+
     @Override
     public DoubleVector build() {
         if (nextIndex < 0) {

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

@@ -20,7 +20,7 @@ import java.util.BitSet;
  */
 final class IntArrayBlock extends AbstractArrayBlock implements IntBlock {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(IntArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(IntArrayBlock.class);
 
     private final IntArrayVector vector;
 

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

@@ -223,13 +223,6 @@ public sealed interface IntBlock extends Block permits IntArrayBlock, IntVectorB
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-        /**
-         * An estimate of the number of bytes the {@link IntBlock} created by
-         * {@link #build} will use. This may overestimate the size but shouldn't
-         * underestimate it.
-         */
-        long estimatedBytes();
-
         @Override
         IntBlock build();
     }

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

@@ -182,9 +182,4 @@ final class IntBlockBuilder extends AbstractBlockBuilder implements IntBlock.Bui
             throw e;
         }
     }
-
-    @Override
-    public long estimatedBytes() {
-        return estimatedBytes;
-    }
 }

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

@@ -46,6 +46,11 @@ final class IntVectorFixedBuilder implements IntVector.FixedBuilder {
             );
     }
 
+    @Override
+    public long estimatedBytes() {
+        return ramBytesUsed(values.length);
+    }
+
     @Override
     public IntVector build() {
         if (nextIndex < 0) {

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

@@ -20,7 +20,7 @@ import java.util.BitSet;
  */
 final class LongArrayBlock extends AbstractArrayBlock implements LongBlock {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(LongArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(LongArrayBlock.class);
 
     private final LongArrayVector vector;
 

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

@@ -46,6 +46,11 @@ final class LongVectorFixedBuilder implements LongVector.FixedBuilder {
             );
     }
 
+    @Override
+    public long estimatedBytes() {
+        return ramBytesUsed(values.length);
+    }
+
     @Override
     public LongVector build() {
         if (nextIndex < 0) {

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

@@ -120,6 +120,11 @@ abstract class AbstractBlockBuilder implements Block.Builder {
         }
     }
 
+    @Override
+    public long estimatedBytes() {
+        return estimatedBytes;
+    }
+
     /**
      * Called during implementations of {@link Block.Builder#build} as a last step
      * to mark the Builder as closed and make sure that further closes don't double

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

@@ -62,6 +62,11 @@ abstract class AbstractVectorBuilder implements Vector.Builder {
         }
     }
 
+    @Override
+    public long estimatedBytes() {
+        return estimatedBytes;
+    }
+
     /**
      * Called during implementations of {@link Block.Builder#build} as a last step
      * to mark the Builder as closed and make sure that further closes don't double

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

@@ -200,6 +200,13 @@ public interface Block extends Accountable, BlockLoader.Block, NamedWriteable, R
          */
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
+        /**
+         * An estimate of the number of bytes the {@link Block} created by
+         * {@link #build} will use. This may overestimate the size but shouldn't
+         * underestimate it.
+         */
+        long estimatedBytes();
+
         /**
          * Builds the block. This method can be called multiple times.
          */

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

@@ -182,6 +182,11 @@ final class ConstantNullBlock extends AbstractNonThreadSafeRefCounted
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return BASE_RAM_BYTES_USED;
+        }
+
         @Override
         public Block build() {
             if (closed) {

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

@@ -160,6 +160,11 @@ public class DocBlock extends AbstractVectorBlock implements Block {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return DocVector.BASE_RAM_BYTES_USED + shards.estimatedBytes() + segments.estimatedBytes() + docs.estimatedBytes();
+        }
+
         @Override
         public DocBlock build() {
             // Pass null for singleSegmentNonDecreasing so we calculate it when we first need it.

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

@@ -18,7 +18,7 @@ import java.util.Objects;
  */
 public final class DocVector extends AbstractVector implements Vector {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(DocVector.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(DocVector.class);
 
     /**
      * Per position memory cost to build the shard segment doc map required

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

@@ -151,13 +151,23 @@ public class SingletonOrdinalsBuilder implements BlockLoader.SingletonOrdinalsBu
         }
     }
 
+    @Override
+    public long estimatedBytes() {
+        /*
+         * This is a *terrible* estimate because we have no idea how big the
+         * values in the ordinals are.
+         */
+        long overhead = shouldBuildOrdinalsBlock() ? 5 : 20;
+        return ords.length * overhead;
+    }
+
     @Override
     public BytesRefBlock build() {
-        if (ords.length >= 2 * docValues.getValueCount() && ords.length >= 32) {
-            return buildOrdinal();
-        } else {
-            return buildRegularBlock();
-        }
+        return shouldBuildOrdinalsBlock() ? buildOrdinal() : buildRegularBlock();
+    }
+
+    boolean shouldBuildOrdinalsBlock() {
+        return ords.length >= 2 * docValues.getValueCount() && ords.length >= 32;
     }
 
     @Override

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

@@ -62,6 +62,13 @@ public interface Vector extends Accountable, RefCounted, Releasable {
      * This is {@link Releasable} and should be released after building the vector or if building the vector fails.
      */
     interface Builder extends Releasable {
+        /**
+         * An estimate of the number of bytes the {@link Vector} created by
+         * {@link #build} will use. This may overestimate the size but shouldn't
+         * underestimate it.
+         */
+        long estimatedBytes();
+
         /**
          * Builds the block. This method can be called multiple times.
          */

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

@@ -32,7 +32,7 @@ $endif$
  */
 final class $Type$ArrayBlock extends AbstractArrayBlock implements $Type$Block {
 
-    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance($Type$ArrayBlock.class);
+    static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance($Type$ArrayBlock.class);
 
     private final $Type$ArrayVector vector;
 

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

@@ -277,15 +277,6 @@ $endif$
         @Override
         Builder mvOrdering(Block.MvOrdering mvOrdering);
 
-$if(int)$
-        /**
-         * An estimate of the number of bytes the {@link IntBlock} created by
-         * {@link #build} will use. This may overestimate the size but shouldn't
-         * underestimate it.
-         */
-        long estimatedBytes();
-
-$endif$
         @Override
         $Type$Block build();
     }

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

@@ -188,6 +188,11 @@ $endif$
     }
 
 $if(BytesRef)$
+    @Override
+    public long estimatedBytes() {
+        return super.estimatedBytes() + BytesRefArrayBlock.BASE_RAM_BYTES_USED + values.ramBytesUsed();
+    }
+
     private $Type$Block buildFromBytesArray() {
         assert estimatedBytes == 0 || firstValueIndexes != null;
         final $Type$Block theBlock;
@@ -295,11 +300,5 @@ $if(BytesRef)$
     public void extraClose() {
         Releasables.closeExpectNoException(values);
     }
-$elseif(int)$
-
-    @Override
-    public long estimatedBytes() {
-        return estimatedBytes;
-    }
 $endif$
 }

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

@@ -46,6 +46,11 @@ final class $Type$VectorFixedBuilder implements $Type$Vector.FixedBuilder {
             );
     }
 
+    @Override
+    public long estimatedBytes() {
+        return ramBytesUsed(values.length);
+    }
+
     @Override
     public $Type$Vector build() {
         if (nextIndex < 0) {

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

@@ -21,8 +21,11 @@ import org.elasticsearch.test.ESTestCase;
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.hamcrest.Matchers.both;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.lessThan;
 
 public class BlockBuilderTests extends ESTestCase {
     @ParametersFactory
@@ -58,6 +61,10 @@ public class BlockBuilderTests extends ESTestCase {
         for (int i = 0; i < numEntries; i++) {
             builder.appendNull();
         }
+        assertThat(
+            builder.estimatedBytes(),
+            both(greaterThan(blockFactory.breaker().getUsed() - 1024)).and(lessThan(blockFactory.breaker().getUsed() + 1024))
+        );
         try (Block block = builder.build()) {
             assertThat(block.getPositionCount(), is(numEntries));
             for (int p = 0; p < numEntries; p++) {
@@ -113,6 +120,10 @@ public class BlockBuilderTests extends ESTestCase {
         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());
+            assertThat(
+                builder.estimatedBytes(),
+                both(greaterThan(blockFactory.breaker().getUsed() - 1024)).and(lessThan(blockFactory.breaker().getUsed() + 1024))
+            );
             try (Block built = builder.build()) {
                 assertThat(built, equalTo(random.block()));
                 assertThat(blockFactory.breaker().getUsed(), equalTo(built.ramBytesUsed()));

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

@@ -113,6 +113,11 @@ public abstract class TestBlockBuilder implements Block.Builder {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return builder.estimatedBytes();
+        }
+
         @Override
         public IntBlock build() {
             return builder.build();
@@ -168,6 +173,11 @@ public abstract class TestBlockBuilder implements Block.Builder {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return builder.estimatedBytes();
+        }
+
         @Override
         public LongBlock build() {
             return builder.build();
@@ -223,6 +233,11 @@ public abstract class TestBlockBuilder implements Block.Builder {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return builder.estimatedBytes();
+        }
+
         @Override
         public DoubleBlock build() {
             return builder.build();
@@ -278,6 +293,11 @@ public abstract class TestBlockBuilder implements Block.Builder {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return builder.estimatedBytes();
+        }
+
         @Override
         public BytesRefBlock build() {
             return builder.build();
@@ -336,6 +356,11 @@ public abstract class TestBlockBuilder implements Block.Builder {
             return this;
         }
 
+        @Override
+        public long estimatedBytes() {
+            return builder.estimatedBytes();
+        }
+
         @Override
         public BooleanBlock build() {
             return builder.build();