From f38116d1429bca828e62cfea5bd41edbb6e6a626 Mon Sep 17 00:00:00 2001
From: ChrisHegarty <chegar999@gmail.com>
Date: Thu, 21 Sep 2023 10:03:53 +0100
Subject: [PATCH] restore constant null block

---
 .../java/org/elasticsearch/compute/data/Block.java   |  4 ++--
 .../org/elasticsearch/compute/data/BlockFactory.java |  6 ++++++
 .../compute/data/ConstantNullBlock.java              |  7 +++++--
 .../elasticsearch/compute/data/BasicBlockTests.java  | 12 ++++++++++++
 4 files changed, 25 insertions(+), 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
index 0863940b35f59..1982c937f2a17 100644
--- 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
index 6d92826038c03..ba7c95b619cc5 100644
--- 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/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java
index 48a327f303e8f..7ad60d89ed72d 100644
--- a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java
+++ b/x-pack/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
index 6e409bddfb00e..454c0d9111141 100644
--- 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;