Base64 support

Author: thecoop

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

@@ -11,7 +11,6 @@
 
 import org.apache.lucene.util.BitUtil;
 
-import java.nio.ByteOrder;
 import java.nio.ShortBuffer;
 
 public final class BFloat16 {
@@ -38,7 +37,6 @@ public static float bFloat16ToFloat(short bf) {
     }
 
     public static void floatToBFloat16(float[] floats, ShortBuffer bFloats) {
-        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
         for (float v : floats) {
             bFloats.put(floatToBFloat16(v));
         }
@@ -52,7 +50,6 @@ public static void bFloat16ToFloat(byte[] bfBytes, float[] floats) {
     }
 
     public static void bFloat16ToFloat(ShortBuffer bFloats, float[] floats) {
-        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
         for (int i = 0; i < floats.length; i++) {
             floats[i] = bFloat16ToFloat(bFloats.get());
         }

server/src/main/java/org/elasticsearch/index/mapper/blockloader/docvalues/DenseVectorBlockLoader.java

@@ -53,7 +53,7 @@ public Builder builder(BlockFactory factory, int expectedCount) {
     @Override
     public AllReader reader(LeafReaderContext context) throws IOException {
         switch (fieldType.getElementType()) {
-            case FLOAT -> {
+            case FLOAT, BFLOAT16 -> {
                 FloatVectorValues floatVectorValues = context.reader().getFloatVectorValues(fieldName);
                 if (floatVectorValues != null) {
                     if (fieldType.isNormalized()) {

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

@@ -1150,20 +1150,26 @@ VectorDataAndMagnitude parseBase64EncodedVector(DocumentParserContext context, I
             throws IOException {
             // BIG_ENDIAN is the default, but just being explicit here
             ByteBuffer byteBuffer = ByteBuffer.wrap(Base64.getDecoder().decode(context.parser().text())).order(ByteOrder.BIG_ENDIAN);
-            if (byteBuffer.remaining() != dims * Float.BYTES) {
+            float[] decodedVector = new float[dims];
+            if (byteBuffer.remaining() == dims * Float.BYTES) {
+                byteBuffer.asFloatBuffer().get(decodedVector);
+            } else if (byteBuffer.remaining() == dims * BFloat16.BYTES) {
+                BFloat16.bFloat16ToFloat(byteBuffer.asShortBuffer(), decodedVector);
+            } else {
                 throw new ParsingException(
                     context.parser().getTokenLocation(),
                     "Failed to parse object: Base64 decoded vector byte length ["
                         + byteBuffer.remaining()
                         + "] does not match the expected length of ["
                         + (dims * Float.BYTES)
+                        + "] or ["
+                        + (dims * BFloat16.BYTES)
                         + "] for dimension count ["
                         + dims
                         + "]"
                 );
             }
-            float[] decodedVector = new float[dims];
-            byteBuffer.asFloatBuffer().get(decodedVector);
+
             dimChecker.accept(decodedVector.length, true);
             VectorData vectorData = VectorData.fromFloats(decodedVector);
             float squaredMagnitude = (float) computeSquaredMagnitude(vectorData);

server/src/main/java/org/elasticsearch/script/VectorScoreScriptUtils.java

@@ -252,7 +252,7 @@ public static final class Hamming {
         @SuppressWarnings("unchecked")
         public Hamming(ScoreScript scoreScript, Object queryVector, String fieldName) {
             DenseVectorDocValuesField field = (DenseVectorDocValuesField) scoreScript.field(fieldName);
-            if (field.getElementType() == DenseVectorFieldMapper.ElementType.FLOAT) {
+            if (field.getElementType() == DenseVectorFieldMapper.ElementType.FLOAT || field.getElementType() == ElementType.BFLOAT16) {
                 throw new IllegalArgumentException("hamming distance is only supported for byte or bit vectors");
             }
             if (queryVector instanceof List) {