restore constant null block

ChrisHegarty · Sep 21, 2023 · f38116d · f38116d
1 parent 4666513
commit f38116d
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 4 deletions.
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/Block.java
@@ -123,11 +123,11 @@ enum MvOrdering {
      */
     // Eventually, this should use the GLOBAL breaking instance
     static Block constantNullBlock(int positions) {
-        return new ConstantNullBlock(positions, BlockFactory.getNonBreakingInstance());
+        return constantNullBlock(positions, BlockFactory.getNonBreakingInstance());
     }
 
     static Block constantNullBlock(int positions, BlockFactory blockFactory) {
-        return new ConstantNullBlock(positions, blockFactory);
+        return blockFactory.newConstantNullBlock(positions);
     }
 
     interface Builder {

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java
@@ -262,4 +262,10 @@ public BytesRefBlock newConstantBytesRefBlockWith(BytesRef value, int positions)
         adjustBreaker(b.ramBytesUsed(), true);
         return b;
     }
+
+    public Block newConstantNullBlock(int positions) {
+        var b = new ConstantNullBlock(positions, this);
+        adjustBreaker(b.ramBytesUsed(), true);
+        return b;
+    }
 }
diff --git a/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java b/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java
@@ -7,6 +7,7 @@
 
 package org.elasticsearch.compute.data;
 
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -19,6 +20,8 @@
  */
 public final class ConstantNullBlock extends AbstractBlock {
 
+    private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(ConstantNullBlock.class);
+
     // Eventually, this should use the GLOBAL breaking instance
     ConstantNullBlock(int positionCount) {
         this(positionCount, BlockFactory.getNonBreakingInstance());
@@ -100,7 +103,7 @@ public Block expand() {
 
     @Override
     public long ramBytesUsed() {
-        return 0L;
+        return BASE_RAM_BYTES_USED;
     }
 
     @Override
@@ -123,7 +126,7 @@ public String toString() {
 
     @Override
     public void close() {
-        // no-op
+        blockFactory.adjustBreaker(-ramBytesUsed(), true);
     }
 
     static class Builder implements Block.Builder {

diff --git a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicBlockTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BasicBlockTests.java
@@ -552,6 +552,18 @@ public void testConstantBooleanBlock() {
         }
     }
 
+    public void testConstantNullBlock() {
+        for (int i = 0; i < 100; i++) {
+            assertThat(breaker.getUsed(), is(0L));
+            int positionCount = randomIntBetween(1, 16 * 1024);
+            Block block = Block.constantNullBlock(positionCount, blockFactory);
+            assertThat(positionCount, is(block.getPositionCount()));
+            assertThat(block.getPositionCount(), is(positionCount));
+            assertThat(block.isNull(randomPosition(positionCount)), is(true));
+            releaseAndAssertBreaker(block);
+        }
+    }
+
     public void testSingleValueSparseInt() {
         int positionCount = randomIntBetween(2, 16 * 1024);
         final int builderEstimateSize = randomBoolean() ? randomIntBetween(1, positionCount) : positionCount;