WIP

Author: elasticsearchmachine

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

@@ -493,7 +493,7 @@ public DenseVectorBlock.Builder newDenseVectorBlockBuilder(int estimatedSize) {
     }
 
     public DenseVectorBlock.Builder newDenseVectorBlockBuilder(int estimatedSize, int dimensions) {
-        return new DenseVectorBlockBuilder(estimatedSize, this);
+        return new DenseVectorBlockBuilder(estimatedSize, this, dimensions);
     }
 
     public DenseVectorBlock newDenseVectorArrayBlock(

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

@@ -28,15 +28,15 @@ final class DenseVectorBlockBuilder extends AbstractBlockBuilder implements Dens
 
 
     DenseVectorBlockBuilder(int estimatedSize, BlockFactory blockFactory, int dimensions) {
-        this(estimatedSize, blockFactory);
+        super(blockFactory);
+        int initialSize = Math.max(estimatedSize, 2);
+        adjustBreaker((RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + elementSize()) * initialSize);
+        values = new float[initialSize][];
         this.dimensions = dimensions;
     }
 
     DenseVectorBlockBuilder(int estimatedSize, BlockFactory blockFactory) {
-        super(blockFactory);
-        int initialSize = Math.max(estimatedSize, 2);
-        adjustBreaker(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER + initialSize * elementSize());
-        values = new float[initialSize][];
+        this(estimatedSize, blockFactory, 0);
     }
 
     @Override
@@ -82,12 +82,14 @@ public DenseVectorBlockBuilder appendNull() {
 
     @Override
     public DenseVectorBlockBuilder beginPositionEntry() {
-        throw new UnsupportedOperationException();
+        super.beginPositionEntry();
+        return this;
     }
 
     @Override
     public DenseVectorBlockBuilder endPositionEntry() {
-        throw new UnsupportedOperationException();
+        super.endPositionEntry();
+        return this;
     }
 
     @Override

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

@@ -21,7 +21,7 @@ final class DenseVectorVectorBuilder extends AbstractVectorBuilder implements De
     DenseVectorVectorBuilder(int estimatedSize, int dimensions, BlockFactory blockFactory) {
         super(blockFactory);
         int initialSize = Math.max(estimatedSize, 2) * dimensions;
-        adjustBreaker(initialSize);
+        adjustBreaker(initialSize * elementSize());
         values = new float[Math.max(estimatedSize, 2)][dimensions];
         this.dimensions = dimensions;
     }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorBuilderTests.java

@@ -29,16 +29,17 @@ public class VectorBuilderTests extends ESTestCase {
     @ParametersFactory
     public static List<Object[]> params() {
         List<Object[]> params = new ArrayList<>();
-        for (ElementType e : ElementType.values()) {
-            if (e == ElementType.UNKNOWN
-                || e == ElementType.NULL
-                || e == ElementType.DOC
-                || e == ElementType.COMPOSITE
-                || e == ElementType.AGGREGATE_METRIC_DOUBLE) {
-                continue;
-            }
-            params.add(new Object[] { e });
-        }
+//        for (ElementType e : ElementType.values()) {
+//            if (e == ElementType.UNKNOWN
+//                || e == ElementType.NULL
+//                || e == ElementType.DOC
+//                || e == ElementType.COMPOSITE
+//                || e == ElementType.AGGREGATE_METRIC_DOUBLE) {
+//                continue;
+//            }
+//            params.add(new Object[] { e });
+//        }
+        params.add(new Object[] { ElementType.DENSE_VECTOR });
         return params;
     }
 
@@ -59,6 +60,7 @@ public void testBuildSmall() {
     }
 
     public void testBuildHuge() {
+        assumeTrue("DENSE_VECTORS will create BigArrays", elementType != ElementType.DENSE_VECTOR);
         testBuild(between(1_000, 50_000));
     }
 
@@ -125,7 +127,7 @@ private Vector.Builder vectorBuilder(int estimatedSize, BlockFactory blockFactor
             case DOUBLE -> blockFactory.newDoubleVectorBuilder(estimatedSize);
             case INT -> blockFactory.newIntVectorBuilder(estimatedSize);
             case LONG -> blockFactory.newLongVectorBuilder(estimatedSize);
-            case DENSE_VECTOR -> blockFactory.newDenseVectorVectorBuilder(estimatedSize, 10);
+            case DENSE_VECTOR -> blockFactory.newDenseVectorVectorBuilder(estimatedSize, 5);
         };
     }
 
@@ -162,6 +164,11 @@ private void fill(Vector.Builder builder, Vector from) {
                     ((LongVector.Builder) builder).appendLong(((LongVector) from).getLong(p));
                 }
             }
+            case DENSE_VECTOR -> {
+                for (int p = 0; p < from.getPositionCount(); p++) {
+                    ((DenseVectorVector.Builder) builder).appendDenseVector(((DenseVectorVector) from).getDenseVector(p));
+                }
+            }
         }
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/VectorFixedBuilderTests.java

@@ -124,7 +124,7 @@ private Vector.Builder vectorBuilder(int size, BlockFactory blockFactory) {
             case FLOAT -> blockFactory.newFloatVectorFixedBuilder(size);
             case INT -> blockFactory.newIntVectorFixedBuilder(size);
             case LONG -> blockFactory.newLongVectorFixedBuilder(size);
-            case DENSE_VECTOR -> blockFactory.newDenseVectorVectorFixedBuilder(10, size);
+            case DENSE_VECTOR -> blockFactory.newDenseVectorVectorFixedBuilder(5, size);
         };
     }
 

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java

@@ -72,7 +72,7 @@ public static Object randomValue(ElementType e) {
             );
             case NULL -> null;
             case COMPOSITE -> throw new IllegalArgumentException("can't make random values for composite");
-            case DENSE_VECTOR -> randomArray(10, 10, Float[]::new, () -> randomFloat());
+            case DENSE_VECTOR -> randomArray(5, 5, Float[]::new, () -> randomFloat());
             case UNKNOWN -> throw new IllegalArgumentException("can't make random values for [" + e + "]");
         };
     }

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/RandomBlock.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.DenseVectorBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.FloatBlock;
@@ -28,6 +29,9 @@
 import java.util.List;
 import java.util.function.Supplier;
 
+import static org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase.randomVector;
+import static org.elasticsearch.test.ESTestCase.randomIntBetween;
+
 /**
  * A block of random values.
  * @param values the values as java object
@@ -155,6 +159,10 @@ public static RandomBlock randomBlock(
                             b.count().appendInt(count);
                             valuesAtPosition.add(new AggregateMetricDoubleBlockBuilder.AggregateMetricDoubleLiteral(min, max, sum, count));
                         }
+                        case DENSE_VECTOR -> {
+                            float[] denseVector = randomVector(5);
+                            ((DenseVectorBlock.Builder) builder).appendDenseVector(denseVector);
+                        }
                         default -> throw new IllegalArgumentException("unsupported element type [" + elementType + "]");
                     }
                 }