ES|QL dense vector field type support (#126456)

Author: carlosdelest

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

@@ -11,13 +11,16 @@
 
 import org.apache.lucene.index.BinaryDocValues;
 import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.ByteArrayStreamInput;
+import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.mapper.BlockLoader.BlockFactory;
 import org.elasticsearch.index.mapper.BlockLoader.BooleanBuilder;
 import org.elasticsearch.index.mapper.BlockLoader.Builder;
@@ -26,6 +29,7 @@
 import org.elasticsearch.index.mapper.BlockLoader.DoubleBuilder;
 import org.elasticsearch.index.mapper.BlockLoader.IntBuilder;
 import org.elasticsearch.index.mapper.BlockLoader.LongBuilder;
+import org.elasticsearch.index.mapper.vectors.VectorEncoderDecoder;
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
 
 import java.io.IOException;
@@ -504,6 +508,87 @@ public String toString() {
         }
     }
 
+    public static class DenseVectorBlockLoader extends DocValuesBlockLoader {
+        private final String fieldName;
+        private final int dimensions;
+
+        public DenseVectorBlockLoader(String fieldName, int dimensions) {
+            this.fieldName = fieldName;
+            this.dimensions = dimensions;
+        }
+
+        @Override
+        public Builder builder(BlockFactory factory, int expectedCount) {
+            return factory.denseVectors(expectedCount, dimensions);
+        }
+
+        @Override
+        public AllReader reader(LeafReaderContext context) throws IOException {
+            FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
+            if (floatVectorValues != null) {
+                return new DenseVectorValuesBlockReader(floatVectorValues, dimensions);
+            }
+            return new ConstantNullsReader();
+        }
+    }
+
+    private static class DenseVectorValuesBlockReader extends BlockDocValuesReader {
+        private final FloatVectorValues floatVectorValues;
+        private final KnnVectorValues.DocIndexIterator iterator;
+        private final int dimensions;
+
+        DenseVectorValuesBlockReader(FloatVectorValues floatVectorValues, int dimensions) {
+            this.floatVectorValues = floatVectorValues;
+            iterator = floatVectorValues.iterator();
+            this.dimensions = dimensions;
+        }
+
+        @Override
+        public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
+            // Doubles from doc values ensures that the values are in order
+            try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count(), dimensions)) {
+                for (int i = 0; i < docs.count(); i++) {
+                    int doc = docs.get(i);
+                    if (doc < iterator.docID()) {
+                        throw new IllegalStateException("docs within same block must be in order");
+                    }
+                    read(doc, builder);
+                }
+                return builder.build();
+            }
+        }
+
+        @Override
+        public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
+            read(docId, (BlockLoader.FloatBuilder) builder);
+        }
+
+        private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
+            if (iterator.advance(doc) == doc) {
+                builder.beginPositionEntry();
+                float[] floats = floatVectorValues.vectorValue(iterator.index());
+                assert floats.length == dimensions
+                    : "unexpected dimensions for vector value; expected " + dimensions + " but got " + floats.length;
+                for (float aFloat : floats) {
+                    builder.appendFloat(aFloat);
+                }
+                builder.endPositionEntry();
+            } else {
+                builder.appendNull();
+            }
+        }
+
+        @Override
+        public int docId() {
+            return iterator.docID();
+        }
+
+        @Override
+        public String toString() {
+            return "BlockDocValuesReader.FloatVectorValuesBlockReader";
+        }
+    }
+
     public static class BytesRefsFromOrdsBlockLoader extends DocValuesBlockLoader {
         private final String fieldName;
 
@@ -752,6 +837,94 @@ public String toString() {
         }
     }
 
+    public static class DenseVectorFromBinaryBlockLoader extends DocValuesBlockLoader {
+        private final String fieldName;
+        private final int dims;
+        private final IndexVersion indexVersion;
+
+        public DenseVectorFromBinaryBlockLoader(String fieldName, int dims, IndexVersion indexVersion) {
+            this.fieldName = fieldName;
+            this.dims = dims;
+            this.indexVersion = indexVersion;
+        }
+
+        @Override
+        public Builder builder(BlockFactory factory, int expectedCount) {
+            return factory.denseVectors(expectedCount, dims);
+        }
+
+        @Override
+        public AllReader reader(LeafReaderContext context) throws IOException {
+            BinaryDocValues docValues = context.reader().getBinaryDocValues(fieldName);
+            if (docValues == null) {
+                return new ConstantNullsReader();
+            }
+            return new DenseVectorFromBinary(docValues, dims, indexVersion);
+        }
+    }
+
+    private static class DenseVectorFromBinary extends BlockDocValuesReader {
+        private final BinaryDocValues docValues;
+        private final IndexVersion indexVersion;
+        private final int dimensions;
+        private final float[] scratch;
+
+        private int docID = -1;
+
+        DenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+            this.docValues = docValues;
+            this.scratch = new float[dims];
+            this.indexVersion = indexVersion;
+            this.dimensions = dims;
+        }
+
+        @Override
+        public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
+            try (BlockLoader.FloatBuilder builder = factory.denseVectors(docs.count(), dimensions)) {
+                for (int i = 0; i < docs.count(); i++) {
+                    int doc = docs.get(i);
+                    if (doc < docID) {
+                        throw new IllegalStateException("docs within same block must be in order");
+                    }
+                    read(doc, builder);
+                }
+                return builder.build();
+            }
+        }
+
+        @Override
+        public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
+            read(docId, (BlockLoader.FloatBuilder) builder);
+        }
+
+        private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException {
+            this.docID = doc;
+            if (false == docValues.advanceExact(doc)) {
+                builder.appendNull();
+                return;
+            }
+            BytesRef bytesRef = docValues.binaryValue();
+            assert bytesRef.length > 0;
+            VectorEncoderDecoder.decodeDenseVector(indexVersion, bytesRef, scratch);
+
+            builder.beginPositionEntry();
+            for (float value : scratch) {
+                builder.appendFloat(value);
+            }
+            builder.endPositionEntry();
+        }
+
+        @Override
+        public int docId() {
+            return docID;
+        }
+
+        @Override
+        public String toString() {
+            return "DenseVectorFromBinary.Bytes";
+        }
+    }
+
     public static class BooleansBlockLoader extends DocValuesBlockLoader {
         private final String fieldName;
 

server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java

@@ -373,6 +373,11 @@ interface BlockFactory {
          */
         DoubleBuilder doubles(int expectedCount);
 
+        /**
+         * Build a builder to load dense vectors without any loading constraints.
+         */
+        FloatBuilder denseVectors(int expectedVectorsCount, int dimensions);
+
         /**
          * Build a builder to load ints as loaded from doc values.
          * Doc values load ints in sorted order.

server/src/main/java/org/elasticsearch/index/mapper/BlockSourceReader.java

@@ -303,6 +303,49 @@ public String toString() {
         }
     }
 
+    /**
+     * Load {@code float}s from {@code _source}.
+     */
+    public static class DenseVectorBlockLoader extends SourceBlockLoader {
+        private final int dimensions;
+
+        public DenseVectorBlockLoader(ValueFetcher fetcher, LeafIteratorLookup lookup, int dimensions) {
+            super(fetcher, lookup);
+            this.dimensions = dimensions;
+        }
+
+        @Override
+        public Builder builder(BlockFactory factory, int expectedCount) {
+            return factory.denseVectors(expectedCount, dimensions);
+        }
+
+        @Override
+        public RowStrideReader rowStrideReader(LeafReaderContext context, DocIdSetIterator iter) {
+            return new DenseVectors(fetcher, iter);
+        }
+
+        @Override
+        protected String name() {
+            return "DenseVectors";
+        }
+    }
+
+    private static class DenseVectors extends BlockSourceReader {
+        DenseVectors(ValueFetcher fetcher, DocIdSetIterator iter) {
+            super(fetcher, iter);
+        }
+
+        @Override
+        protected void append(BlockLoader.Builder builder, Object v) {
+            ((BlockLoader.FloatBuilder) builder).appendFloat(((Number) v).floatValue());
+        }
+
+        @Override
+        public String toString() {
+            return "BlockSourceReader.DenseVectors";
+        }
+    }
+
     /**
      * Load {@code int}s from {@code _source}.
      */