Use bfloat16

Author: thecoop

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

@@ -17,7 +17,6 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
@@ -29,6 +28,7 @@
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.hnsw.OrdinalTranslatedKnnCollector;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.elasticsearch.index.codec.vectors.es91.ES91BFloat16FlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.reflect.OffHeapByteSizeUtils;
 import org.elasticsearch.index.codec.vectors.reflect.OffHeapStats;
 
@@ -41,7 +41,7 @@ public class ES813FlatVectorFormat extends KnnVectorsFormat {
 
     static final String NAME = "ES813FlatVectorFormat";
 
-    private static final FlatVectorsFormat format = new Lucene99FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE);
+    private static final FlatVectorsFormat format = new ES91BFloat16FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE);
 
     /**
      * Sole constructor

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

@@ -16,7 +16,6 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.codecs.hnsw.ScalarQuantizedVectorScorer;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99ScalarQuantizedVectorsWriter;
 import org.apache.lucene.index.ByteVectorValues;
@@ -34,6 +33,7 @@
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
 import org.apache.lucene.util.quantization.QuantizedVectorsReader;
 import org.apache.lucene.util.quantization.ScalarQuantizer;
+import org.elasticsearch.index.codec.vectors.es91.ES91BFloat16FlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.reflect.OffHeapByteSizeUtils;
 import org.elasticsearch.index.codec.vectors.reflect.OffHeapStats;
 import org.elasticsearch.simdvec.VectorScorerFactory;
@@ -50,7 +50,7 @@ public class ES814ScalarQuantizedVectorsFormat extends FlatVectorsFormat {
     static final String NAME = "ES814ScalarQuantizedVectorsFormat";
     private static final int ALLOWED_BITS = (1 << 8) | (1 << 7) | (1 << 4);
 
-    private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE);
+    private static final FlatVectorsFormat rawVectorFormat = new ES91BFloat16FlatVectorsFormat(DefaultFlatVectorScorer.INSTANCE);
 
     static final FlatVectorsScorer flatVectorScorer = new ESFlatVectorsScorer(
         new ScalarQuantizedVectorScorer(DefaultFlatVectorScorer.INSTANCE)

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

@@ -13,7 +13,6 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.ByteVectorValues;
 import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.SegmentReadState;
@@ -24,14 +23,15 @@
 import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
 import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
 import org.apache.lucene.util.quantization.QuantizedByteVectorValues;
+import org.elasticsearch.index.codec.vectors.es91.ES91BFloat16FlatVectorsFormat;
 
 import java.io.IOException;
 
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT;
 
 class ES815BitFlatVectorsFormat extends FlatVectorsFormat {
 
-    private static final FlatVectorsFormat delegate = new Lucene99FlatVectorsFormat(FlatBitVectorScorer.INSTANCE);
+    private static final FlatVectorsFormat delegate = new ES91BFloat16FlatVectorsFormat(FlatBitVectorScorer.INSTANCE);
 
     protected ES815BitFlatVectorsFormat() {
         super("ES815BitFlatVectorsFormat");

server/src/main/java/org/elasticsearch/index/codec/vectors/es816/ES816BinaryQuantizedVectorsFormat.java

@@ -23,9 +23,9 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.elasticsearch.index.codec.vectors.es91.ES91BFloat16FlatVectorsFormat;
 
 import java.io.IOException;
 
@@ -47,7 +47,7 @@ public class ES816BinaryQuantizedVectorsFormat extends FlatVectorsFormat {
     static final String VECTOR_DATA_EXTENSION = "veb";
     static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
 
-    private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
+    private static final FlatVectorsFormat rawVectorFormat = new ES91BFloat16FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/ES818BinaryQuantizedVectorsFormat.java

@@ -23,11 +23,11 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.elasticsearch.core.SuppressForbidden;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
+import org.elasticsearch.index.codec.vectors.es91.ES91BFloat16FlatVectorsFormat;
 
 import java.io.IOException;
 
@@ -110,7 +110,7 @@ private static boolean getUseDirectIO() {
 
     private static final FlatVectorsFormat rawVectorFormat = USE_DIRECT_IO
         ? new DirectIOLucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer())
-        : new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
+        : new ES91BFloat16FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer());
 
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()

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

@@ -9,13 +9,20 @@
 
 package org.elasticsearch.index.codec.vectors.es91;
 
+import org.apache.lucene.util.BitUtil;
+
 class BFloat16 {
 
     public static final int BYTES = Short.BYTES;
 
     public static short floatToBFloat16(float f) {
-        // TODO: maintain NaN if all NaN set bits are in removed section
-        return (short)(Float.floatToIntBits(f) >>> 16);
+        // this does round towards 0
+        // zero - zero exp, zero fraction
+        // denormal - zero exp, non-zero fraction
+        // infinity - all-1 exp, zero fraction
+        // NaN - all-1 exp, non-zero fraction
+        // the Float.NaN constant is 0x7fc0_0000, so this won't turn the most common NaN values into infinities
+        return (short) (Float.floatToIntBits(f) >>> 16);
     }
 
     public static float bFloat16ToFloat(short bf) {
@@ -24,17 +31,24 @@ public static float bFloat16ToFloat(short bf) {
 
     public static short[] floatToBFloat16(float[] f) {
         short[] bf = new short[f.length];
-        for (int i=0; i<f.length; i++) {
+        for (int i = 0; i < f.length; i++) {
             bf[i] = floatToBFloat16(f[i]);
         }
         return bf;
     }
 
     public static float[] bFloat16ToFloat(short[] bf) {
         float[] f = new float[bf.length];
-        for (int i=0; i<bf.length; i++) {
+        for (int i = 0; i < bf.length; i++) {
             f[i] = bFloat16ToFloat(bf[i]);
         }
         return f;
     }
+
+    public static void bFloat16ToFloat(byte[] bfBytes, float[] floats) {
+        assert floats.length * 2 == bfBytes.length;
+        for (int i = 0; i < floats.length; i++) {
+            floats[i] = bFloat16ToFloat((short) BitUtil.VH_LE_SHORT.get(bfBytes, i * 2));
+        }
+    }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es91/ES91BFloat16FlatVectorsFormat.java

@@ -24,8 +24,6 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.codecs.lucene99.ES91BFloat16FlatVectorsReader;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsReader;
-import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;