ES|QL dense vector field type support

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,81 @@ public String toString() {
         }
     }
 
+    public static class DenseVectorBlockLoader extends DocValuesBlockLoader {
+        private final String fieldName;
+
+        public DenseVectorBlockLoader(String fieldName) {
+            this.fieldName = fieldName;
+        }
+
+        @Override
+        public Builder builder(BlockFactory factory, int expectedCount) {
+            return factory.floats(expectedCount);
+        }
+
+        @Override
+        public AllReader reader(LeafReaderContext context) throws IOException {
+            FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
+            if (floatVectorValues != null) {
+                return new FloatVectorValuesBlockReader(floatVectorValues);
+            }
+            return new ConstantNullsReader();
+        }
+    }
+
+    private static class FloatVectorValuesBlockReader extends BlockDocValuesReader {
+        private final FloatVectorValues floatVectorValues;
+        private final KnnVectorValues.DocIndexIterator iterator;
+
+        FloatVectorValuesBlockReader(FloatVectorValues floatVectorValues) {
+            this.floatVectorValues = floatVectorValues;
+            iterator = floatVectorValues.iterator();
+        }
+
+        @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.floatsFromDocValues(docs.count())) {
+                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());
+                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 +831,92 @@ 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.floats(expectedCount);
+        }
+
+        @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 float[] scratch;
+
+        private int docID = -1;
+
+        DenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+            this.docValues = docValues;
+            this.scratch = new float[dims];
+            this.indexVersion = indexVersion;
+        }
+
+        @Override
+        public BlockLoader.Block read(BlockFactory factory, Docs docs) throws IOException {
+            try (BlockLoader.FloatBuilder builder = factory.floats(docs.count())) {
+                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,17 @@ interface BlockFactory {
          */
         DoubleBuilder doubles(int expectedCount);
 
+        /**
+         * Build a builder to load floats as loaded from doc values.
+         * Doc values load floats in sorted order.
+         */
+        FloatBuilder floatsFromDocValues(int expectedCount);
+
+        /**
+         * Build a builder to load floats without any loading constraints.
+         */
+        FloatBuilder floats(int expectedCount);
+
         /**
          * 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

@@ -301,6 +301,46 @@ public String toString() {
         }
     }
 
+    /**
+     * Load {@code float}s from {@code _source}.
+     */
+    public static class FloatsBlockLoader extends SourceBlockLoader {
+        public FloatsBlockLoader(ValueFetcher fetcher, LeafIteratorLookup lookup) {
+            super(fetcher, lookup);
+        }
+
+        @Override
+        public Builder builder(BlockFactory factory, int expectedCount) {
+            return factory.floats(expectedCount);
+        }
+
+        @Override
+        public RowStrideReader rowStrideReader(LeafReaderContext context, DocIdSetIterator iter) {
+            return new Floats(fetcher, iter);
+        }
+
+        @Override
+        protected String name() {
+            return "Floats";
+        }
+    }
+
+    private static class Floats extends BlockSourceReader {
+        Floats(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.Floats";
+        }
+    }
+
     /**
      * Load {@code int}s from {@code _source}.
      */

server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java

@@ -51,15 +51,20 @@
 import org.elasticsearch.index.fielddata.FieldDataContext;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.mapper.ArraySourceValueFetcher;
+import org.elasticsearch.index.mapper.BlockDocValuesReader;
+import org.elasticsearch.index.mapper.BlockLoader;
+import org.elasticsearch.index.mapper.BlockSourceReader;
 import org.elasticsearch.index.mapper.DocumentParserContext;
 import org.elasticsearch.index.mapper.FieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MapperBuilderContext;
 import org.elasticsearch.index.mapper.MapperParsingException;
 import org.elasticsearch.index.mapper.MappingParser;
+import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.index.mapper.SimpleMappedFieldType;
 import org.elasticsearch.index.mapper.SourceLoader;
+import org.elasticsearch.index.mapper.SourceValueFetcher;
 import org.elasticsearch.index.mapper.TextSearchInfo;
 import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.query.SearchExecutionContext;
@@ -89,6 +94,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -305,7 +311,8 @@ public DenseVectorFieldMapper build(MapperBuilderContext context) {
                     indexed.getValue(),
                     similarity.getValue(),
                     indexOptions.getValue(),
-                    meta.getValue()
+                    meta.getValue(),
+                    context.isSourceSynthetic()
                 ),
                 builderParams(this, context),
                 indexOptions.getValue(),
@@ -2039,6 +2046,7 @@ public static final class DenseVectorFieldType extends SimpleMappedFieldType {
         private final VectorSimilarity similarity;
         private final IndexVersion indexVersionCreated;
         private final IndexOptions indexOptions;
+        private final boolean isSyntheticSource;
 
         public DenseVectorFieldType(
             String name,
@@ -2048,7 +2056,8 @@ public DenseVectorFieldType(
             boolean indexed,
             VectorSimilarity similarity,
             IndexOptions indexOptions,
-            Map<String, String> meta
+            Map<String, String> meta,
+            boolean isSyntheticSource
         ) {
             super(name, indexed, false, indexed == false, TextSearchInfo.NONE, meta);
             this.elementType = elementType;
@@ -2057,6 +2066,7 @@ public DenseVectorFieldType(
             this.similarity = similarity;
             this.indexVersionCreated = indexVersionCreated;
             this.indexOptions = indexOptions;
+            this.isSyntheticSource = isSyntheticSource;
         }
 
         @Override
@@ -2311,6 +2321,37 @@ int getVectorDimensions() {
         ElementType getElementType() {
             return elementType;
         }
+
+        @Override
+        public BlockLoader blockLoader(MappedFieldType.BlockLoaderContext blContext) {
+            if (elementType != ElementType.FLOAT) {
+                // Just float dense vector support for now
+                return null;
+            }
+
+            if (indexed) {
+                return new BlockDocValuesReader.DenseVectorBlockLoader(name());
+            }
+
+            if (hasDocValues() && (blContext.fieldExtractPreference() != FieldExtractPreference.STORED || isSyntheticSource)) {
+                return new BlockDocValuesReader.DenseVectorFromBinaryBlockLoader(name(), dims, indexVersionCreated);
+            }
+
+            BlockSourceReader.LeafIteratorLookup lookup = BlockSourceReader.lookupMatchingAll();
+            return new BlockSourceReader.FloatsBlockLoader(sourceValueFetcher(blContext.sourcePaths(name())), lookup);
+        }
+
+        private SourceValueFetcher sourceValueFetcher(Set<String> sourcePaths) {
+            return new SourceValueFetcher(sourcePaths, null) {
+                @Override
+                protected Object parseSourceValue(Object value) {
+                    if (value.equals("")) {
+                        return null;
+                    }
+                    return NumberFieldMapper.NumberType.FLOAT.parse(value, false);
+                }
+            };
+        }
     }
 
     private final IndexOptions indexOptions;
@@ -2365,7 +2406,8 @@ public void parse(DocumentParserContext context) throws IOException {
                 fieldType().indexed,
                 fieldType().similarity,
                 fieldType().indexOptions,
-                fieldType().meta()
+                fieldType().meta(),
+                fieldType().isSyntheticSource
             );
             Mapper update = new DenseVectorFieldMapper(
                 leafName(),