add pre-adjusted vector construction

ChrisHegarty · Sep 21, 2023 · 7d67a14 · 7d67a14
1 parent f766b1b
commit 7d67a14
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 4 deletions.
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
@@ -105,6 +105,31 @@ void adjustBreaker(final long delta, final boolean isDataAlreadyCreated) {
         }
     }
 
+    /** Pre-adjusts the breaker for the given position count and element type. Returns the pre-adjusted amount. */
+    public long preAdjustBreakerForBoolean(int positionCount) {
+        long bytes = (long) positionCount * Byte.BYTES;
+        adjustBreaker(bytes, false);
+        return bytes;
+    }
+
+    public long preAdjustBreakerForInt(int positionCount) {
+        long bytes = (long) positionCount * Integer.BYTES;
+        adjustBreaker(bytes, false);
+        return bytes;
+    }
+
+    public long preAdjustBreakerForLong(int positionCount) {
+        long bytes = (long) positionCount * Long.BYTES;
+        adjustBreaker(bytes, false);
+        return bytes;
+    }
+
+    public long preAdjustBreakerForDouble(int positionCount) {
+        long bytes = (long) positionCount * Double.BYTES;
+        adjustBreaker(bytes, false);
+        return bytes;
+    }
+
     public BooleanBlock.Builder newBooleanBlockBuilder(int estimatedSize) {
         return new BooleanBlockBuilder(estimatedSize, this);
     }
@@ -126,8 +151,12 @@ public BooleanVector.Builder newBooleanVectorBuilder(int estimatedSize) {
     }
 
     public BooleanVector newBooleanArrayVector(boolean[] values, int positionCount) {
+        return newBooleanArrayVector(values, positionCount, 0L);
+    }
+
+    public BooleanVector newBooleanArrayVector(boolean[] values, int positionCount, long preAdjustedBytes) {
         var b = new BooleanArrayVector(values, positionCount, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -151,9 +180,28 @@ public IntVector.Builder newIntVectorBuilder(int estimatedSize) {
         return new IntVectorBuilder(estimatedSize, this);
     }
 
+    /**
+     * 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) {
+        return newIntArrayVector(values, positionCount, 0L);
+    }
+
+    /**
+     * Creates a new Vector with the given values and positionCount, where the caller has already
+     * pre-adjusted a number of bytes with the factory's breaker.
+     *
+     * long preAdjustedBytes = blockFactory.preAdjustBreakerForInt(positionCount);
+     * int[] values = new int[positionCount];
+     * for (int i = 0; i &lt; positionCount; i++) {
+     *   values[i] = doWhateverStuff
+     * }
+     * var vector = blockFactory.newIntArrayVector(values, positionCount, preAdjustedBytes);
+     */
+    public IntVector newIntArrayVector(int[] values, int positionCount, long preAdjustedBytes) {
         var b = new IntArrayVector(values, positionCount, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -178,8 +226,12 @@ public LongVector.Builder newLongVectorBuilder(int estimatedSize) {
     }
 
     public LongVector newLongArrayVector(long[] values, int positionCount) {
+        return newLongArrayVector(values, positionCount, 0L);
+    }
+
+    public LongVector newLongArrayVector(long[] values, int positionCount, long preAdjustedBytes) {
         var b = new LongArrayVector(values, positionCount, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 
@@ -210,8 +262,12 @@ public DoubleVector.Builder newDoubleVectorBuilder(int estimatedSize) {
     }
 
     public DoubleVector newDoubleArrayVector(double[] values, int positionCount) {
+        return newDoubleArrayVector(values, positionCount, 0L);
+    }
+
+    public DoubleVector newDoubleArrayVector(double[] values, int positionCount, long preAdjustedBytes) {
         var b = new DoubleArrayVector(values, positionCount, this);
-        adjustBreaker(b.ramBytesUsed(), true);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes, true);
         return b;
     }
 

diff --git a/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java b/...k/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java
@@ -71,6 +71,27 @@ public void checkBreaker() {
         assertThat(breaker.getUsed(), is(0L));
     }
 
+    public void testPreAdjusters() {
+        for (int i = 0; i < 1000; i++) {
+            int positions = randomIntBetween(1, 16384);
+            long preAdjustBytes = blockFactory.preAdjustBreakerForBoolean(positions);
+            assertThat(preAdjustBytes, is((long) positions));
+            blockFactory.adjustBreaker(-preAdjustBytes, true);
+
+            preAdjustBytes = blockFactory.preAdjustBreakerForInt(positions);
+            assertThat(preAdjustBytes, is((long) positions * 4));
+            blockFactory.adjustBreaker(-preAdjustBytes, true);
+
+            preAdjustBytes = blockFactory.preAdjustBreakerForLong(positions);
+            assertThat(preAdjustBytes, is((long) positions * 8));
+            blockFactory.adjustBreaker(-preAdjustBytes, true);
+
+            preAdjustBytes = blockFactory.preAdjustBreakerForDouble(positions);
+            assertThat(preAdjustBytes, is((long) positions * 8));
+            blockFactory.adjustBreaker(-preAdjustBytes, true);
+        }
+    }
+
     public void testIntBlockBuilderWithPossiblyLargeEstimateEmpty() {
         var builder = blockFactory.newIntBlockBuilder(randomIntBetween(0, 2048));
         assertThat(breaker.getUsed(), greaterThan(0L));