Add Vector and Blocks for Dense Vector to factories and interfaces

Author: carlosdelest

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

@@ -487,4 +487,78 @@ public final AggregateMetricDoubleBlock newAggregateMetricDoubleBlock(
     public long maxPrimitiveArrayBytes() {
         return maxPrimitiveArrayBytes;
     }
+
+    public DenseVectorBlock.Builder newDenseVectorBlockBuilder(int estimatedSize) {
+        return new DenseVectorBlockBuilder(estimatedSize, this);
+    }
+
+    public final DenseVectorBlock newDenseVectorArrayBlock(
+        float[][] values,
+        int dimensions,
+        int positionCount,
+        int[] firstValueIndexes,
+        BitSet nulls,
+        MvOrdering mvOrdering
+    ) {
+        return newDenseVectorArrayBlock(values, dimensions, positionCount, firstValueIndexes, nulls, mvOrdering, 0L);
+    }
+
+    public DenseVectorBlock newDenseVectorArrayBlock(
+        float[][] values,
+        int dimensions,
+        int pc,
+        int[] fvi,
+        BitSet nulls,
+        MvOrdering mvOrdering,
+        long preAdjustedBytes
+    ) {
+        var b = new DenseVectorArrayBlock(values, pc, fvi, nulls, mvOrdering, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes);
+        return b;
+    }
+
+    public DenseVectorVector.Builder newDenseVectorVectorBuilder(int estimatedSize, int dimensions) {
+        return new DenseVectorVectorBuilder(dimensions, estimatedSize, this);
+    }
+
+    /**
+     * Build a {@link DenseVectorVector.FixedBuilder} that never grows.
+     */
+    public DenseVectorVector.FixedBuilder newDenseVectorVectorFixedBuilder(int dimensions, int size) {
+        return new DenseVectorVectorFixedBuilder(dimensions, size, this);
+    }
+
+    /**
+     * 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.preAdjustBreakerForDenseVector(positionCount);
+     * int[] values = new int[positionCount];
+     * for (int i = 0; i < positionCount; i++) {
+     *   values[i] = doWhateverStuff
+     * }
+     * var vector = blockFactory.newDenseVectorArrayVector(values, positionCount, preAdjustedBytes);
+     */
+    public DenseVectorVector newDenseVectorArrayVector(float[][] values, int positionCount, long preAdjustedBytes) {
+        var b = new DenseVectorArrayVector(values, positionCount, this);
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes);
+        return b;
+    }
+
+    public final DenseVectorBlock newConstantDenseVectorBlockWith(float[] value, int positions) {
+        return newConstantDenseVectorBlockWith(value, positions, 0L);
+    }
+
+    public DenseVectorBlock newConstantDenseVectorBlockWith(float[] value, int positions, long preAdjustedBytes) {
+        var b = new ConstantDenseVectorVector(value, positions, this).asBlock();
+        adjustBreaker(b.ramBytesUsed() - preAdjustedBytes);
+        return b;
+    }
+
+    public DenseVectorVector newConstantDenseVectorVector(float[] value, int positions) {
+        adjustBreaker(ConstantDenseVectorVector.RAM_BYTES_USED);
+        var v = new ConstantDenseVectorVector(value, positions, this);
+        assert v.ramBytesUsed() == ConstantDenseVectorVector.RAM_BYTES_USED;
+        return v;
+    }
 }

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

@@ -299,6 +299,7 @@ yield new AggregateMetricDoubleLiteral(
                     aggBlock.countBlock().getInt(offset)
                 );
             }
+            case DENSE_VECTOR -> ((DenseVectorBlock) block).getDenseVector(offset);
             case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
         };
     }

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

@@ -26,7 +26,8 @@ public final class ConstantNullBlock extends AbstractNonThreadSafeRefCounted
         FloatBlock,
         DoubleBlock,
         BytesRefBlock,
-        AggregateMetricDoubleBlock {
+        AggregateMetricDoubleBlock,
+        DenseVectorBlock {
 
     private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(ConstantNullBlock.class);
     private final int positionCount;
@@ -292,6 +293,12 @@ public long getLong(int valueIndex) {
         throw new UnsupportedOperationException("null block");
     }
 
+    @Override
+    public float[] getDenseVector(int valueIndex) {
+        assert false : "null block";
+        throw new UnsupportedOperationException("null block");
+    }
+
     @Override
     public int getTotalValueCount() {
         return 0;

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

@@ -24,7 +24,8 @@ public final class ConstantNullVector extends AbstractVector
         DoubleVector,
         FloatVector,
         IntVector,
-        LongVector {
+        LongVector,
+        DenseVectorVector {
 
     private ConstantNullVector(int positionCount, BlockFactory blockFactory) {
         super(positionCount, blockFactory);
@@ -66,6 +67,11 @@ public ConstantNullBlock keepMask(BooleanVector mask) {
         throw new UnsupportedOperationException("null vector");
     }
 
+    @Override
+    public int dimensions() {
+        return 0;
+    }
+
     @Override
     public ReleasableIterator<ConstantNullBlock> lookup(IntBlock positions, ByteSizeValue targetBlockSize) {
         assert false : "null vector";
@@ -108,6 +114,12 @@ public long getLong(int position) {
         throw new UnsupportedOperationException("null vector");
     }
 
+    @Override
+    public float[] getDenseVector(int position) {
+        assert false : "null vector";
+        throw new UnsupportedOperationException("null vector");
+    }
+
     @Override
     public int min() {
         assert false : "null vector";