ESQL: Array blocks backed by vecs

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/BlockBenchmark.java

@@ -76,6 +76,7 @@ public class BlockBenchmark {
     public static final String[] RELEVANT_TYPE_BLOCK_COMBINATIONS = {
         "boolean/array",
         "boolean/array-multivalue-null",
+        "boolean/big-array",
         "boolean/big-array-multivalue-null",
         "boolean/vector",
         "boolean/vector-big-array",
@@ -86,18 +87,21 @@ public class BlockBenchmark {
         "BytesRef/vector-const",
         "double/array",
         "double/array-multivalue-null",
+        "double/big-array",
         "double/big-array-multivalue-null",
         "double/vector",
         "double/vector-big-array",
         "double/vector-const",
         "int/array",
         "int/array-multivalue-null",
+        "int/big-array",
         "int/big-array-multivalue-null",
         "int/vector",
         "int/vector-big-array",
         "int/vector-const",
         "long/array",
         "long/array-multivalue-null",
+        "long/big-array",
         "long/big-array-multivalue-null",
         "long/vector",
         "long/vector-big-array",
@@ -177,6 +181,23 @@ private static BenchmarkBlocks buildBlocks(String dataType, String blockKind, in
                                 Block.MvOrdering.UNORDERED
                             );
                         }
+                        case "big-array" -> {
+                            BitArray valuesBigArray = new BitArray(totalPositions, BigArrays.NON_RECYCLING_INSTANCE);
+                            for (int i = 0; i < values.length; i++) {
+                                if (values[i]) {
+                                    valuesBigArray.set(i);
+                                }
+                            }
+
+                            blocks[blockIndex] = new BooleanBigArrayBlock(
+                                valuesBigArray,
+                                totalPositions,
+                                null,
+                                null,
+                                Block.MvOrdering.DEDUPLICATED_AND_SORTED_ASCENDING,
+                                blockFactory
+                            );
+                        }
                         case "big-array-multivalue-null" -> {
                             int[] firstValueIndexes = randomFirstValueIndexes(totalPositions);
                             int positionCount = firstValueIndexes.length - 1;
@@ -315,6 +336,21 @@ private static BenchmarkBlocks buildBlocks(String dataType, String blockKind, in
                                 Block.MvOrdering.UNORDERED
                             );
                         }
+                        case "big-array" -> {
+                            DoubleArray valuesBigArray = blockFactory.bigArrays().newDoubleArray(totalPositions, false);
+                            for (int i = 0; i < values.length; i++) {
+                                valuesBigArray.set(i, values[i]);
+                            }
+
+                            blocks[blockIndex] = new DoubleBigArrayBlock(
+                                valuesBigArray,
+                                totalPositions,
+                                null,
+                                null,
+                                Block.MvOrdering.DEDUPLICATED_AND_SORTED_ASCENDING,
+                                blockFactory
+                            );
+                        }
                         case "big-array-multivalue-null" -> {
                             int[] firstValueIndexes = randomFirstValueIndexes(totalPositions);
                             int positionCount = firstValueIndexes.length - 1;
@@ -392,6 +428,21 @@ private static BenchmarkBlocks buildBlocks(String dataType, String blockKind, in
                                 Block.MvOrdering.UNORDERED
                             );
                         }
+                        case "big-array" -> {
+                            IntArray valuesBigArray = blockFactory.bigArrays().newIntArray(totalPositions, false);
+                            for (int i = 0; i < values.length; i++) {
+                                valuesBigArray.set(i, values[i]);
+                            }
+
+                            blocks[blockIndex] = new IntBigArrayBlock(
+                                valuesBigArray,
+                                totalPositions,
+                                null,
+                                null,
+                                Block.MvOrdering.DEDUPLICATED_AND_SORTED_ASCENDING,
+                                blockFactory
+                            );
+                        }
                         case "big-array-multivalue-null" -> {
                             int[] firstValueIndexes = randomFirstValueIndexes(totalPositions);
                             int positionCount = firstValueIndexes.length - 1;
@@ -469,6 +520,21 @@ private static BenchmarkBlocks buildBlocks(String dataType, String blockKind, in
                                 Block.MvOrdering.UNORDERED
                             );
                         }
+                        case "big-array" -> {
+                            LongArray valuesBigArray = blockFactory.bigArrays().newLongArray(totalPositions, false);
+                            for (int i = 0; i < values.length; i++) {
+                                valuesBigArray.set(i, values[i]);
+                            }
+
+                            blocks[blockIndex] = new LongBigArrayBlock(
+                                valuesBigArray,
+                                totalPositions,
+                                null,
+                                null,
+                                Block.MvOrdering.DEDUPLICATED_AND_SORTED_ASCENDING,
+                                blockFactory
+                            );
+                        }
                         case "big-array-multivalue-null" -> {
                             int[] firstValueIndexes = randomFirstValueIndexes(totalPositions);
                             int positionCount = firstValueIndexes.length - 1;
@@ -715,6 +781,7 @@ private static boolean isRandom(String accessType) {
         {
             "boolean/array",
             "boolean/array-multivalue-null",
+            "boolean/big-array",
             "boolean/big-array-multivalue-null",
             "boolean/vector",
             "boolean/vector-big-array",
@@ -725,18 +792,21 @@ private static boolean isRandom(String accessType) {
             "BytesRef/vector-const",
             "double/array",
             "double/array-multivalue-null",
+            "double/big-array",
             "double/big-array-multivalue-null",
             "double/vector",
             "double/vector-big-array",
             "double/vector-const",
             "int/array",
             "int/array-multivalue-null",
+            "int/big-array",
             "int/big-array-multivalue-null",
             "int/vector",
             "int/vector-big-array",
             "int/vector-const",
             "long/array",
             "long/array-multivalue-null",
+            "long/big-array",
             "long/big-array-multivalue-null",
             "long/vector",
             "long/vector-big-array",

server/src/main/java/org/elasticsearch/common/util/BitArray.java

@@ -129,6 +129,10 @@ public boolean get(long index) {
         return (bits.get(wordNum) & bitmask) != 0;
     }
 
+    public long size() {
+        return bits.size() * (long) Long.BYTES * Byte.SIZE;
+    }
+
     private static long wordNum(long index) {
         return index >> 6;
     }

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanArrayBlock.java

@@ -8,19 +8,19 @@
 package org.elasticsearch.compute.data;
 
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.core.Releasables;
 
-import java.util.Arrays;
 import java.util.BitSet;
 
 /**
- * Block implementation that stores an array of boolean.
+ * Block implementation that stores values in a {@link BooleanArrayVector}.
  * This class is generated. Do not edit it.
  */
 final class BooleanArrayBlock extends AbstractArrayBlock implements BooleanBlock {
 
     private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(BooleanArrayBlock.class);
 
-    private final boolean[] values;
+    private final BooleanArrayVector vector;
 
     BooleanArrayBlock(
         boolean[] values,
@@ -31,7 +31,7 @@ final class BooleanArrayBlock extends AbstractArrayBlock implements BooleanBlock
         BlockFactory blockFactory
     ) {
         super(positionCount, firstValueIndexes, nulls, mvOrdering, blockFactory);
-        this.values = values;
+        this.vector = new BooleanArrayVector(values, values.length, blockFactory);
     }
 
     @Override
@@ -41,7 +41,7 @@ public BooleanVector asVector() {
 
     @Override
     public boolean getBoolean(int valueIndex) {
-        return values[valueIndex];
+        return vector.getBoolean(valueIndex);
     }
 
     @Override
@@ -79,7 +79,7 @@ public BooleanBlock expand() {
             incRef();
             return this;
         }
-        // TODO use reference counting to share the values
+        // TODO use reference counting to share the vector
         try (var builder = blockFactory().newBooleanBlockBuilder(firstValueIndexes[getPositionCount()])) {
             for (int pos = 0; pos < getPositionCount(); pos++) {
                 if (isNull(pos)) {
@@ -96,14 +96,13 @@ public BooleanBlock expand() {
         }
     }
 
-    public static long ramBytesEstimated(boolean[] values, int[] firstValueIndexes, BitSet nullsMask) {
-        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(values) + BlockRamUsageEstimator.sizeOf(firstValueIndexes)
-            + BlockRamUsageEstimator.sizeOfBitSet(nullsMask);
+    private long ramBytesUsedOnlyBlock() {
+        return BASE_RAM_BYTES_USED + BlockRamUsageEstimator.sizeOf(firstValueIndexes) + BlockRamUsageEstimator.sizeOfBitSet(nullsMask);
     }
 
     @Override
     public long ramBytesUsed() {
-        return ramBytesEstimated(values, firstValueIndexes, nullsMask);
+        return ramBytesUsedOnlyBlock() + vector.ramBytesUsed();
     }
 
     @Override
@@ -126,13 +125,20 @@ public String toString() {
             + getPositionCount()
             + ", mvOrdering="
             + mvOrdering()
-            + ", values="
-            + Arrays.toString(values)
+            + ", vector="
+            + vector
             + ']';
     }
 
+    @Override
+    public void allowPassingToDifferentDriver() {
+        super.allowPassingToDifferentDriver();
+        vector.allowPassingToDifferentDriver();
+    }
+
     @Override
     public void closeInternal() {
-        blockFactory().adjustBreaker(-ramBytesUsed(), true);
+        blockFactory().adjustBreaker(-ramBytesUsedOnlyBlock(), true);
+        Releasables.closeExpectNoException(vector);
     }
 }

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBigArrayBlock.java

@@ -14,13 +14,14 @@
 import java.util.BitSet;
 
 /**
- * Block implementation that stores values in a BooleanArray.
+ * Block implementation that stores values in a {@link BooleanBigArrayVector}. Does not take ownership of the given
+ * {@link BitArray} and does not adjust circuit breakers to account for it.
  * This class is generated. Do not edit it.
  */
 public final class BooleanBigArrayBlock extends AbstractArrayBlock implements BooleanBlock {
 
     private static final long BASE_RAM_BYTES_USED = 0; // TODO: fix this
-    private final BitArray values;
+    private final BooleanBigArrayVector vector;
 
     public BooleanBigArrayBlock(
         BitArray values,
@@ -31,7 +32,7 @@ public BooleanBigArrayBlock(
         BlockFactory blockFactory
     ) {
         super(positionCount, firstValueIndexes, nulls, mvOrdering, blockFactory);
-        this.values = values;
+        this.vector = new BooleanBigArrayVector(values, (int) values.size(), blockFactory);
     }
 
     @Override
@@ -41,7 +42,7 @@ public BooleanVector asVector() {
 
     @Override
     public boolean getBoolean(int valueIndex) {
-        return values.get(valueIndex);
+        return vector.getBoolean(valueIndex);
     }
 
     @Override
@@ -79,7 +80,7 @@ public BooleanBlock expand() {
             incRef();
             return this;
         }
-        // TODO use reference counting to share the values
+        // TODO use reference counting to share the vector
         try (var builder = blockFactory().newBooleanBlockBuilder(firstValueIndexes[getPositionCount()])) {
             for (int pos = 0; pos < getPositionCount(); pos++) {
                 if (isNull(pos)) {
@@ -96,10 +97,13 @@ public BooleanBlock expand() {
         }
     }
 
+    private long ramBytesUsedOnlyBlock() {
+        return BASE_RAM_BYTES_USED + BlockRamUsageEstimator.sizeOf(firstValueIndexes) + BlockRamUsageEstimator.sizeOfBitSet(nullsMask);
+    }
+
     @Override
     public long ramBytesUsed() {
-        return BASE_RAM_BYTES_USED + RamUsageEstimator.sizeOf(values) + BlockRamUsageEstimator.sizeOf(firstValueIndexes)
-            + BlockRamUsageEstimator.sizeOfBitSet(nullsMask);
+        return ramBytesUsedOnlyBlock() + RamUsageEstimator.sizeOf(vector);
     }
 
     @Override
@@ -123,13 +127,19 @@ public String toString() {
             + ", mvOrdering="
             + mvOrdering()
             + ", ramBytesUsed="
-            + values.ramBytesUsed()
+            + vector.ramBytesUsed()
             + ']';
     }
 
+    @Override
+    public void allowPassingToDifferentDriver() {
+        super.allowPassingToDifferentDriver();
+        vector.allowPassingToDifferentDriver();
+    }
+
     @Override
     public void closeInternal() {
-        blockFactory().adjustBreaker(-ramBytesUsed() + RamUsageEstimator.sizeOf(values), true);
-        Releasables.closeExpectNoException(values);
+        blockFactory().adjustBreaker(-ramBytesUsedOnlyBlock(), true);
+        Releasables.closeExpectNoException(vector);
     }
 }

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBigArrayVector.java

@@ -12,7 +12,8 @@
 import org.elasticsearch.core.Releasable;
 
 /**
- * Vector implementation that defers to an enclosed BooleanArray.
+ * Vector implementation that defers to an enclosed {@link BitArray}.
+ * Does not take ownership of the array and does not adjust circuit breakers to account for it.
  * This class is generated. Do not edit it.
  */
 public final class BooleanBigArrayVector extends AbstractVector implements BooleanVector, Releasable {
@@ -65,6 +66,8 @@ public BooleanVector filter(int... positions) {
 
     @Override
     public void closeInternal() {
+        // The circuit breaker that tracks the values {@link BitArray} is adjusted outside
+        // of this class.
         values.close();
     }
 

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanVectorBlock.java

@@ -10,7 +10,7 @@
 import org.elasticsearch.core.Releasables;
 
 /**
- * Block view of a BooleanVector.
+ * Block view of a {@link BooleanVector}. Cannot represent multi-values or nulls.
  * This class is generated. Do not edit it.
  */
 public final class BooleanVectorBlock extends AbstractVectorBlock implements BooleanBlock {