add tests and fixes

ChrisHegarty · Sep 20, 2023 · 3bbf9cc · 3bbf9cc
1 parent c198d53
commit 3bbf9cc
Show file tree

Hide file tree

Showing 19 changed files with 627 additions and 68 deletions.
diff --git a/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java b/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlockBuilder.java
@@ -193,7 +193,6 @@ public BooleanBlock build() {
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanVectorBuilder.java b/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanVectorBuilder.java
@@ -59,7 +59,6 @@ public BooleanVector build() {
             vector = new BooleanArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }

diff --git a/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefArrayBlock.java b/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefArrayBlock.java
@@ -10,6 +10,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 import java.util.BitSet;
 import java.util.stream.IntStream;
@@ -110,6 +111,7 @@ public String toString() {
 
     @Override
     public void close() {
-        blockFactory.adjustBreaker(-ramBytesUsed(), true);
+        blockFactory.adjustBreaker(-(ramBytesUsed() - values.ramBytesUsed()), true);
+        Releasables.closeExpectNoException(values);
     }
 }
diff --git a/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefArrayVector.java b/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefArrayVector.java
@@ -85,6 +85,7 @@ public String toString() {
 
     @Override
     public void close() {
+        blockFactory.adjustBreaker(-BASE_RAM_BYTES_USED, true);
         Releasables.closeExpectNoException(values);
     }
 }
diff --git a/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java b/...l/compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefBlockBuilder.java
@@ -10,6 +10,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 /**
  * Block build of BytesRefBlocks.
@@ -195,16 +196,17 @@ public BytesRefBlock build() {
         finish();
         BytesRefBlock block;
         if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
-            block = new ConstantBytesRefVector(values.get(0, new BytesRef()), 1, blockFactory).asBlock();
+            block = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory).asBlock();
+            Releasables.closeExpectNoException(values);
         } else {
+            estimatedBytes += values.ramBytesUsed();
             if (isDense() && singleValued()) {
                 block = new BytesRefArrayVector(values, positionCount, blockFactory).asBlock();
             } else {
                 block = new BytesRefArrayBlock(values, positionCount, firstValueIndexes, nullsMask, mvOrdering, blockFactory);
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/.../compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefVectorBuilder.java b/.../compute/src/main/generated-src/org/elasticsearch/compute/data/BytesRefVectorBuilder.java
@@ -10,6 +10,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 /**
  * Block build of BytesRefBlocks.
@@ -55,12 +56,13 @@ protected void growValuesArray(int newSize) {
     public BytesRefVector build() {
         BytesRefVector vector;
         if (valueCount == 1) {
-            vector = new ConstantBytesRefVector(values.get(0, new BytesRef()), 1, blockFactory);
+            vector = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory);
+            Releasables.closeExpectNoException(values);
         } else {
+            estimatedBytes = values.ramBytesUsed();
             vector = new BytesRefArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }

diff --git a/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java b/...sql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleBlockBuilder.java
@@ -193,7 +193,6 @@ public DoubleBlock build() {
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleVectorBuilder.java b/...ql/compute/src/main/generated-src/org/elasticsearch/compute/data/DoubleVectorBuilder.java
@@ -59,7 +59,6 @@ public DoubleVector build() {
             vector = new DoubleArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }

diff --git a/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java b/...n/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntBlockBuilder.java
@@ -193,7 +193,6 @@ public IntBlock build() {
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntVectorBuilder.java b/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/data/IntVectorBuilder.java
@@ -59,7 +59,6 @@ public IntVector build() {
             vector = new IntArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }

diff --git a/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java b/.../esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongBlockBuilder.java
@@ -193,7 +193,6 @@ public LongBlock build() {
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/...esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongVectorBuilder.java b/...esql/compute/src/main/generated-src/org/elasticsearch/compute/data/LongVectorBuilder.java
@@ -59,7 +59,6 @@ public LongVector build() {
             vector = new LongArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }

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
@@ -76,6 +76,11 @@ public CircuitBreaker breaker() {
         return breaker;
     }
 
+    // For testing
+    public BigArrays bigArrays() {
+        return bigArrays;
+    }
+
     /**
      * Adjust the circuit breaker with the given delta, if the delta is
      * negative, or checkBreaker is false, the breaker will be adjusted
@@ -238,7 +243,7 @@ public BytesRefBlock newBytesRefArrayBlock(
         MvOrdering mvOrdering
     ) {
         var b = new BytesRefArrayBlock(values, positionCount, firstValueIndexes, nulls, mvOrdering, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - values.ramBytesUsed(), true);
         return b;
     }
 
@@ -248,11 +253,13 @@ public BytesRefVector.Builder newBytesRefVectorBuilder(int estimatedSize) {
 
     public BytesRefVector newBytesRefArrayVector(BytesRefArray values, int positionCount) {
         var b = new BytesRefArrayVector(values, positionCount, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - values.ramBytesUsed(), true);
         return b;
     }
 
     public BytesRefBlock newConstantBytesRefBlockWith(BytesRef value, int positions) {
-        return new ConstantBytesRefVector(value, positions, this).asBlock();
+        var b = new ConstantBytesRefVector(value, positions, this).asBlock();
+        adjustBreaker(b.ramBytesUsed(), true);
+        return b;
     }
 }
diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayBlock.java.st b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayBlock.java.st
@@ -11,6 +11,7 @@ $if(BytesRef)$
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 $else$
 import org.apache.lucene.util.RamUsageEstimator;
@@ -131,6 +132,11 @@ $endif$
 
     @Override
     public void close() {
+    $if(BytesRef)$
+        blockFactory.adjustBreaker(-(ramBytesUsed() - values.ramBytesUsed()), true);
+        Releasables.closeExpectNoException(values);
+    $else$
         blockFactory.adjustBreaker(-ramBytesUsed(), true);
+    $endif$
     }
 }
diff --git a/...ck/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayVector.java.st b/...ck/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayVector.java.st
@@ -110,6 +110,7 @@ $endif$
     @Override
     public void close() {
 $if(BytesRef)$
+        blockFactory.adjustBreaker(-BASE_RAM_BYTES_USED, true);
         Releasables.closeExpectNoException(values);
 $else$
         super.close();

diff --git a/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st b/...k/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BlockBuilder.java.st
@@ -11,6 +11,7 @@ $if(BytesRef)$
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 $else$
 import java.util.Arrays;
@@ -249,8 +250,10 @@ $endif$
         $Type$Block block;
         if (hasNonNullValue && positionCount == 1 && valueCount == 1) {
 $if(BytesRef)$
-            block = new ConstantBytesRefVector(values.get(0, new BytesRef()), 1, blockFactory).asBlock();
+            block = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory).asBlock();
+            Releasables.closeExpectNoException(values);
         } else {
+            estimatedBytes += values.ramBytesUsed();
 $else$
             block = new Constant$Type$Vector(values[0], 1, blockFactory).asBlock();
         } else {
@@ -265,7 +268,6 @@ $endif$
             }
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(block.ramBytesUsed() - estimatedBytes, true);
         return block;
     }

diff --git a/.../plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st b/.../plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-VectorBuilder.java.st
@@ -11,6 +11,7 @@ $if(BytesRef)$
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BytesRefArray;
+import org.elasticsearch.core.Releasables;
 
 $else$
 import java.util.Arrays;
@@ -85,8 +86,10 @@ $endif$
         $Type$Vector vector;
         if (valueCount == 1) {
 $if(BytesRef)$
-            vector = new ConstantBytesRefVector(values.get(0, new BytesRef()), 1, blockFactory);
+            vector = new ConstantBytesRefVector(BytesRef.deepCopyOf(values.get(0, new BytesRef())), 1, blockFactory);
+            Releasables.closeExpectNoException(values);
         } else {
+            estimatedBytes = values.ramBytesUsed();
 $else$
             vector = new Constant$Type$Vector(values[0], 1, blockFactory);
         } else {
@@ -97,7 +100,6 @@ $endif$
             vector = new $Type$ArrayVector(values, valueCount, blockFactory);
         }
         // update the breaker with the actual bytes used.
-        // TODO: verify that this can also give back
         blockFactory.adjustBreaker(vector.ramBytesUsed() - estimatedBytes, true);
         return vector;
     }