Fix compile

Author: thecoop

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/BFloat16.java renamed to server/src/main/java/org/elasticsearch/index/codec/vectors/BFloat16.java

@@ -7,14 +7,14 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.index.codec.vectors.es92;
+package org.elasticsearch.index.codec.vectors;
 
 import org.apache.lucene.util.BitUtil;
 
 import java.nio.ByteOrder;
 import java.nio.ShortBuffer;
 
-class BFloat16 {
+public class BFloat16 {
 
     public static final int BYTES = Short.BYTES;
 
@@ -47,4 +47,14 @@ public static void bFloat16ToFloat(byte[] bfBytes, float[] floats) {
             floats[i] = bFloat16ToFloat((short) BitUtil.VH_LE_SHORT.get(bfBytes, i * 2));
         }
     }
+
+    public static void bFloat16ToFloat(ShortBuffer bFloats, float[] floats) {
+        assert floats.length == bFloats.remaining();
+        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
+        for (int i = 0; i < floats.length; i++) {
+            floats[i] = bFloat16ToFloat(bFloats.get());
+        }
+    }
+
+    private BFloat16() {}
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/ES92BFloat16FlatVectorsReader.java

@@ -42,6 +42,7 @@
 import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 
 import java.io.IOException;
 import java.util.Map;

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/ES92BFloat16FlatVectorsWriter.java

@@ -46,6 +46,7 @@
 import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
 import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 
 import java.io.Closeable;
 import java.io.IOException;

server/src/main/java/org/elasticsearch/index/codec/vectors/es92/OffHeapBFloat16VectorValues.java

@@ -32,6 +32,7 @@
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.apache.lucene.util.packed.DirectMonotonicReader;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 
 import java.io.IOException;
 

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

@@ -22,6 +22,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.ByteArrayStreamInput;
 import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 import org.elasticsearch.index.mapper.BlockLoader.BlockFactory;
 import org.elasticsearch.index.mapper.BlockLoader.BooleanBuilder;
 import org.elasticsearch.index.mapper.BlockLoader.Builder;
@@ -36,6 +37,9 @@
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.ShortBuffer;
 
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.COSINE_MAGNITUDE_FIELD_SUFFIX;
 
@@ -536,7 +540,8 @@ public Builder builder(BlockFactory factory, int expectedCount) {
         @Override
         public AllReader reader(LeafReaderContext context) throws IOException {
             switch (fieldType.getElementType()) {
-                case FLOAT -> {
+                case FLOAT, BFLOAT16 -> {
+                    // BFloat16 is handled by the implementation of FloatVectorValues
                     FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
                     if (floatVectorValues != null) {
                         if (fieldType.isNormalized()) {
@@ -1052,6 +1057,7 @@ public AllReader reader(LeafReaderContext context) throws IOException {
             }
             return switch (elementType) {
                 case FLOAT -> new FloatDenseVectorFromBinary(docValues, dims, indexVersion);
+                case BFLOAT16 -> new BFloat16DenseVectorFromBinary(docValues, dims, indexVersion);
                 case BYTE -> new ByteDenseVectorFromBinary(docValues, dims, indexVersion);
                 case BIT -> new BitDenseVectorFromBinary(docValues, dims, indexVersion);
             };
@@ -1135,6 +1141,26 @@ public String toString() {
         }
     }
 
+    private static class BFloat16DenseVectorFromBinary extends FloatDenseVectorFromBinary {
+        BFloat16DenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
+            super(docValues, dims, indexVersion);
+        }
+
+        @Override
+        protected void decodeDenseVector(BytesRef bytesRef, float[] scratch) {
+            VectorEncoderDecoder.decodeDenseVector(indexVersion, bytesRef, scratch);
+            ShortBuffer sb = ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length)
+                .order(ByteOrder.LITTLE_ENDIAN)
+                .asShortBuffer();
+            BFloat16.bFloat16ToFloat(sb, scratch);
+        }
+
+        @Override
+        public String toString() {
+            return "BFloat16DenseVectorFromBinary.Bytes";
+        }
+    }
+
     private static class ByteDenseVectorFromBinary extends AbstractDenseVectorFromBinary<byte[]> {
         ByteDenseVectorFromBinary(BinaryDocValues docValues, int dims, IndexVersion indexVersion) {
             this(docValues, dims, indexVersion, dims);

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

@@ -49,6 +49,7 @@
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.IndexVersions;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 import org.elasticsearch.index.codec.vectors.ES813FlatVectorFormat;
 import org.elasticsearch.index.codec.vectors.ES813Int8FlatVectorFormat;
 import org.elasticsearch.index.codec.vectors.ES814HnswScalarQuantizedVectorsFormat;
@@ -1065,17 +1066,17 @@ public ElementType elementType() {
 
         @Override
         public void writeValue(ByteBuffer byteBuffer, float value) {
-            byteBuffer.putShort((short) (Float.floatToIntBits(value) >>> 16));
+            byteBuffer.putShort(BFloat16.floatToBFloat16(value));
         }
 
         @Override
         public void readAndWriteValue(ByteBuffer byteBuffer, XContentBuilder b) throws IOException {
-            b.value(Float.intBitsToFloat(byteBuffer.getShort() << 16));
+            b.value(BFloat16.bFloat16ToFloat(byteBuffer.getShort()));
         }
 
         @Override
         public int getNumBytes(int dimensions) {
-            return dimensions * Short.BYTES;
+            return dimensions * BFloat16.BYTES;
         }
     }
 

server/src/test/java/org/elasticsearch/index/codec/vectors/es92/ES92BinaryQuantizedBFloat16VectorsFormatTests.java

@@ -61,6 +61,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.IndexModule;
 import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.shard.ShardPath;

server/src/test/java/org/elasticsearch/index/codec/vectors/es92/ES92HnswBinaryQuantizedBFloat16VectorsFormatTests.java

@@ -52,6 +52,7 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.IndexModule;
 import org.elasticsearch.index.IndexSettings;
+import org.elasticsearch.index.codec.vectors.BFloat16;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat;
 import org.elasticsearch.index.shard.ShardId;