From 7d67a14294d4d90720cbf6d51233a434dc3ba208 Mon Sep 17 00:00:00 2001 From: ChrisHegarty Date: Thu, 21 Sep 2023 11:26:08 +0100 Subject: [PATCH] add pre-adjusted vector construction --- .../compute/data/BlockFactory.java | 64 +++++++++++++++++-- .../compute/data/BlockFactoryTests.java | 21 ++++++ 2 files changed, 81 insertions(+), 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 index fae358b8ca899..ca2c4e7c453d0 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 @@ -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 < 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/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java b/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java index 87763e82e98a3..87d6f595328e4 100644 --- a/x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/BlockFactoryTests.java +++ b/x-pack/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));