ESQL - Add byte element support for dense_vector data type (elastic#131863)

Author: carlosdelest

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

@@ -10,6 +10,7 @@
 package org.elasticsearch.index.mapper;
 
 import org.apache.lucene.index.BinaryDocValues;
+import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.KnnVectorValues;
@@ -30,11 +31,15 @@
 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.DenseVectorFieldMapper;
+import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType;
 import org.elasticsearch.index.mapper.vectors.VectorEncoderDecoder;
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
 
 import java.io.IOException;
 
+import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.BYTE;
+
 /**
  * A reader that supports reading doc-values from a Lucene segment in Block fashion.
  */
@@ -516,10 +521,12 @@ public String toString() {
     public static class DenseVectorBlockLoader extends DocValuesBlockLoader {
         private final String fieldName;
         private final int dimensions;
+        private final DenseVectorFieldMapper.DenseVectorFieldType fieldType;
 
-        public DenseVectorBlockLoader(String fieldName, int dimensions) {
+        public DenseVectorBlockLoader(String fieldName, int dimensions, DenseVectorFieldMapper.DenseVectorFieldType fieldType) {
             this.fieldName = fieldName;
             this.dimensions = dimensions;
+            this.fieldType = fieldType;
         }
 
         @Override
@@ -529,22 +536,34 @@ public Builder builder(BlockFactory factory, int expectedCount) {
 
         @Override
         public AllReader reader(LeafReaderContext context) throws IOException {
-            FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
-            if (floatVectorValues != null) {
-                return new DenseVectorValuesBlockReader(floatVectorValues, dimensions);
+            switch (fieldType.getElementType()) {
+                case FLOAT -> {
+                    FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
+                    if (floatVectorValues != null) {
+                        return new FloatDenseVectorValuesBlockReader(floatVectorValues, dimensions);
+                    }
+                }
+                case BYTE -> {
+                    ByteVectorValues byteVectorValues = context.reader().getByteVectorValues(fieldName);
+                    if (byteVectorValues != null) {
+                        return new ByteDenseVectorValuesBlockReader(byteVectorValues, dimensions);
+                    }
+                }
             }
+
             return new ConstantNullsReader();
         }
     }
 
-    private static class DenseVectorValuesBlockReader extends BlockDocValuesReader {
-        private final FloatVectorValues floatVectorValues;
-        private final KnnVectorValues.DocIndexIterator iterator;
-        private final int dimensions;
+    private abstract static class DenseVectorValuesBlockReader<T extends KnnVectorValues> extends BlockDocValuesReader {
+
+        protected final T vectorValues;
+        protected final KnnVectorValues.DocIndexIterator iterator;
+        protected final int dimensions;
 
-        DenseVectorValuesBlockReader(FloatVectorValues floatVectorValues, int dimensions) {
-            this.floatVectorValues = floatVectorValues;
-            iterator = floatVectorValues.iterator();
+        DenseVectorValuesBlockReader(T vectorValues, int dimensions) {
+            this.vectorValues = vectorValues;
+            iterator = vectorValues.iterator();
             this.dimensions = dimensions;
         }
 
@@ -569,26 +588,59 @@ private void read(int doc, BlockLoader.FloatBuilder builder) throws IOException
                 builder.appendNull();
             } else if (iterator.docID() == doc || 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);
-                }
+                appendDoc(builder);
                 builder.endPositionEntry();
             } else {
                 builder.appendNull();
             }
         }
 
+        protected abstract void appendDoc(BlockLoader.FloatBuilder builder) throws IOException;
+
         @Override
         public int docId() {
             return iterator.docID();
         }
+    }
+
+    private static class FloatDenseVectorValuesBlockReader extends DenseVectorValuesBlockReader<FloatVectorValues> {
+
+        FloatDenseVectorValuesBlockReader(FloatVectorValues floatVectorValues, int dimensions) {
+            super(floatVectorValues, dimensions);
+        }
+
+        protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
+            float[] floats = vectorValues.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);
+            }
+        }
 
         @Override
         public String toString() {
-            return "BlockDocValuesReader.FloatVectorValuesBlockReader";
+            return "BlockDocValuesReader.FloatDenseVectorValuesBlockReader";
+        }
+    }
+
+    private static class ByteDenseVectorValuesBlockReader extends DenseVectorValuesBlockReader<ByteVectorValues> {
+        ByteDenseVectorValuesBlockReader(ByteVectorValues floatVectorValues, int dimensions) {
+            super(floatVectorValues, dimensions);
+        }
+
+        protected void appendDoc(BlockLoader.FloatBuilder builder) throws IOException {
+            byte[] bytes = vectorValues.vectorValue(iterator.index());
+            assert bytes.length == dimensions
+                : "unexpected dimensions for vector value; expected " + dimensions + " but got " + bytes.length;
+            for (byte aFloat : bytes) {
+                builder.appendFloat(aFloat);
+            }
+        }
+
+        @Override
+        public String toString() {
+            return "BlockDocValuesReader.ByteDenseVectorValuesBlockReader";
         }
     }
 
@@ -880,11 +932,13 @@ public static class DenseVectorFromBinaryBlockLoader extends DocValuesBlockLoade
         private final String fieldName;
         private final int dims;
         private final IndexVersion indexVersion;
+        private final ElementType elementType;
 
-        public DenseVectorFromBinaryBlockLoader(String fieldName, int dims, IndexVersion indexVersion) {
+        public DenseVectorFromBinaryBlockLoader(String fieldName, int dims, IndexVersion indexVersion, ElementType elementType) {
             this.fieldName = fieldName;
             this.dims = dims;
             this.indexVersion = indexVersion;
+            this.elementType = elementType;
         }
 
         @Override
@@ -898,23 +952,40 @@ public AllReader reader(LeafReaderContext context) throws IOException {
             if (docValues == null) {
                 return new ConstantNullsReader();
             }
-            return new DenseVectorFromBinary(docValues, dims, indexVersion);
+            switch (elementType) {
+                case FLOAT:
+                    return new FloatDenseVectorFromBinary(docValues, dims, indexVersion);
+                case BYTE:
+                    return new ByteDenseVectorFromBinary(docValues, dims, indexVersion);
+                default:
+                    throw new IllegalArgumentException("Unknown element type [" + elementType + "]");
+            }
         }
     }
 
-    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;
+    // Abstract base for dense vector readers
+    private abstract static class AbstractDenseVectorFromBinary<T> extends BlockDocValuesReader {
+        protected final BinaryDocValues docValues;
+        protected final IndexVersion indexVersion;
+        protected final int dimensions;
+        protected final T scratch;
+        protected int docID = -1;
 
-        DenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+        AbstractDenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion, T scratch) {
             this.docValues = docValues;
-            this.scratch = new float[dims];
             this.indexVersion = indexVersion;
             this.dimensions = dims;
+            this.scratch = scratch;
+        }
+
+        @Override
+        public int docId() {
+            return docID;
+        }
+
+        @Override
+        public void read(int docId, BlockLoader.StoredFields storedFields, Builder builder) throws IOException {
+            read(docId, (BlockLoader.FloatBuilder) builder);
         }
 
         @Override
@@ -931,36 +1002,67 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset) throw
             }
         }
 
-        @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)) {
+            if (docValues.advanceExact(doc) == false) {
                 builder.appendNull();
                 return;
             }
             BytesRef bytesRef = docValues.binaryValue();
             assert bytesRef.length > 0;
-            VectorEncoderDecoder.decodeDenseVector(indexVersion, bytesRef, scratch);
+            decodeDenseVector(bytesRef, scratch);
 
             builder.beginPositionEntry();
+            writeScratchToBuilder(scratch, builder);
+            builder.endPositionEntry();
+        }
+
+        protected abstract void decodeDenseVector(BytesRef bytesRef, T scratch);
+
+        protected abstract void writeScratchToBuilder(T scratch, BlockLoader.FloatBuilder builder);
+    }
+
+    private static class FloatDenseVectorFromBinary extends AbstractDenseVectorFromBinary<float[]> {
+        FloatDenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+            super(docValues, dims, indexVersion, new float[dims]);
+        }
+
+        @Override
+        protected void writeScratchToBuilder(float[] scratch, BlockLoader.FloatBuilder builder) {
             for (float value : scratch) {
                 builder.appendFloat(value);
             }
-            builder.endPositionEntry();
         }
 
         @Override
-        public int docId() {
-            return docID;
+        protected void decodeDenseVector(BytesRef bytesRef, float[] scratch) {
+            VectorEncoderDecoder.decodeDenseVector(indexVersion, bytesRef, scratch);
         }
 
         @Override
         public String toString() {
-            return "DenseVectorFromBinary.Bytes";
+            return "FloatDenseVectorFromBinary.Bytes";
+        }
+    }
+
+    private static class ByteDenseVectorFromBinary extends AbstractDenseVectorFromBinary<byte[]> {
+        ByteDenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+            super(docValues, dims, indexVersion, new byte[dims]);
+        }
+
+        @Override
+        public String toString() {
+            return "ByteDenseVectorFromBinary.Bytes";
+        }
+
+        protected void writeScratchToBuilder(byte[] scratch, BlockLoader.FloatBuilder builder) {
+            for (byte value : scratch) {
+                builder.appendFloat(value);
+            }
+        }
+
+        protected void decodeDenseVector(BytesRef bytesRef, byte[] scratch) {
+            VectorEncoderDecoder.decodeDenseVector(indexVersion, bytesRef, scratch);
         }
     }