Add ES92 bfloat16 vector format

Author: thecoop

server/src/main/java/module-info.java

@@ -459,6 +459,8 @@
             org.elasticsearch.index.codec.vectors.es816.ES816HnswBinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat,
             org.elasticsearch.index.codec.vectors.es818.ES818HnswBinaryQuantizedVectorsFormat,
+            org.elasticsearch.index.codec.vectors.es92.ES92BinaryQuantizedBFloat16VectorsFormat,
+            org.elasticsearch.index.codec.vectors.es92.ES92HnswBinaryQuantizedBFloat16VectorsFormat,
             org.elasticsearch.index.codec.vectors.IVFVectorsFormat;
 
     provides org.apache.lucene.codecs.Codec

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

@@ -30,7 +30,7 @@
 /**
  * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
  */
-abstract class BinarizedByteVectorValues extends ByteVectorValues {
+public abstract class BinarizedByteVectorValues extends ByteVectorValues {
 
     /**
      * Retrieve the corrective terms for the given vector ordinal. For the dot-product family of

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

@@ -108,7 +108,7 @@ public RandomVectorScorer getRandomVectorScorer(
         return nonQuantizedDelegate.getRandomVectorScorer(similarityFunction, vectorValues, target);
     }
 
-    RandomVectorScorerSupplier getRandomVectorScorerSupplier(
+    public RandomVectorScorerSupplier getRandomVectorScorerSupplier(
         VectorSimilarityFunction similarityFunction,
         ES818BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues scoringVectors,
         BinarizedByteVectorValues targetVectors
@@ -122,7 +122,7 @@ public String toString() {
     }
 
     /** Vector scorer supplier over binarized vector values */
-    static class BinarizedRandomVectorScorerSupplier implements RandomVectorScorerSupplier {
+    public static class BinarizedRandomVectorScorerSupplier implements RandomVectorScorerSupplier {
         private final ES818BinaryQuantizedVectorsWriter.OffHeapBinarizedQueryVectorValues queryVectors;
         private final BinarizedByteVectorValues targetVectors;
         private final VectorSimilarityFunction similarityFunction;

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

@@ -388,11 +388,11 @@ static FieldEntry create(IndexInput input, VectorEncoding vectorEncoding, Vector
     }
 
     /** Binarized vector values holding row and quantized vector values */
-    protected static final class BinarizedVectorValues extends FloatVectorValues {
+    public static final class BinarizedVectorValues extends FloatVectorValues {
         private final FloatVectorValues rawVectorValues;
         private final BinarizedByteVectorValues quantizedVectorValues;
 
-        BinarizedVectorValues(FloatVectorValues rawVectorValues, BinarizedByteVectorValues quantizedVectorValues) {
+        public BinarizedVectorValues(FloatVectorValues rawVectorValues, BinarizedByteVectorValues quantizedVectorValues) {
             this.rawVectorValues = rawVectorValues;
             this.quantizedVectorValues = quantizedVectorValues;
         }
@@ -437,7 +437,7 @@ public VectorScorer scorer(float[] query) throws IOException {
             return quantizedVectorValues.scorer(query);
         }
 
-        BinarizedByteVectorValues getQuantizedVectorValues() throws IOException {
+        public BinarizedByteVectorValues getQuantizedVectorValues() throws IOException {
             return quantizedVectorValues;
         }
     }

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

@@ -723,7 +723,7 @@ public long ramBytesUsed() {
     }
 
     // When accessing vectorValue method, targerOrd here means a row ordinal.
-    static class OffHeapBinarizedQueryVectorValues {
+    public static class OffHeapBinarizedQueryVectorValues {
         private final IndexInput slice;
         private final int dimension;
         private final int size;
@@ -734,7 +734,7 @@ static class OffHeapBinarizedQueryVectorValues {
         private int lastOrd = -1;
         private int quantizedComponentSum;
 
-        OffHeapBinarizedQueryVectorValues(IndexInput data, int dimension, int size) {
+        public OffHeapBinarizedQueryVectorValues(IndexInput data, int dimension, int size) {
             this.slice = data;
             this.dimension = dimension;
             this.size = size;
@@ -798,7 +798,7 @@ public byte[] vectorValue(int targetOrd) throws IOException {
         }
     }
 
-    static class BinarizedFloatVectorValues extends BinarizedByteVectorValues {
+    public static class BinarizedFloatVectorValues extends BinarizedByteVectorValues {
         private OptimizedScalarQuantizer.QuantizationResult corrections;
         private final byte[] binarized;
         private final int[] initQuantized;
@@ -808,7 +808,7 @@ static class BinarizedFloatVectorValues extends BinarizedByteVectorValues {
 
         private int lastOrd = -1;
 
-        BinarizedFloatVectorValues(FloatVectorValues delegate, OptimizedScalarQuantizer quantizer, float[] centroid) {
+        public BinarizedFloatVectorValues(FloatVectorValues delegate, OptimizedScalarQuantizer quantizer, float[] centroid) {
             this.values = delegate;
             this.quantizer = quantizer;
             this.binarized = new byte[BQVectorUtils.discretize(delegate.dimension(), 64) / 8];
@@ -881,12 +881,16 @@ public int ordToDoc(int ord) {
         }
     }
 
-    static class BinarizedCloseableRandomVectorScorerSupplier implements CloseableRandomVectorScorerSupplier {
+    public static class BinarizedCloseableRandomVectorScorerSupplier implements CloseableRandomVectorScorerSupplier {
         private final RandomVectorScorerSupplier supplier;
         private final KnnVectorValues vectorValues;
         private final Closeable onClose;
 
-        BinarizedCloseableRandomVectorScorerSupplier(RandomVectorScorerSupplier supplier, KnnVectorValues vectorValues, Closeable onClose) {
+        public BinarizedCloseableRandomVectorScorerSupplier(
+            RandomVectorScorerSupplier supplier,
+            KnnVectorValues vectorValues,
+            Closeable onClose
+        ) {
             this.supplier = supplier;
             this.onClose = onClose;
             this.vectorValues = vectorValues;
@@ -913,11 +917,11 @@ public int totalVectorCount() {
         }
     }
 
-    static final class NormalizedFloatVectorValues extends FloatVectorValues {
+    public static final class NormalizedFloatVectorValues extends FloatVectorValues {
         private final FloatVectorValues values;
         private final float[] normalizedVector;
 
-        NormalizedFloatVectorValues(FloatVectorValues values) {
+        public NormalizedFloatVectorValues(FloatVectorValues values) {
             this.values = values;
             this.normalizedVector = new float[values.dimension()];
         }

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

@@ -23,12 +23,12 @@
 import java.util.Collection;
 import java.util.Map;
 
-class MergeReaderWrapper extends FlatVectorsReader {
+public class MergeReaderWrapper extends FlatVectorsReader {
 
     private final FlatVectorsReader mainReader;
     private final FlatVectorsReader mergeReader;
 
-    protected MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) {
+    public MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) {
         super(mainReader.getFlatVectorScorer());
         this.mainReader = mainReader;
         this.mergeReader = mergeReader;

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

@@ -36,7 +36,7 @@
 import java.nio.ByteBuffer;
 
 /** Binarized vector values loaded from off-heap */
-abstract class OffHeapBinarizedVectorValues extends BinarizedByteVectorValues {
+public abstract class OffHeapBinarizedVectorValues extends BinarizedByteVectorValues {
 
     final int dimension;
     final int size;
@@ -151,7 +151,7 @@ public int getVectorByteLength() {
         return numBytes;
     }
 
-    static OffHeapBinarizedVectorValues load(
+    public static OffHeapBinarizedVectorValues load(
         OrdToDocDISIReaderConfiguration configuration,
         int dimension,
         int size,
@@ -197,8 +197,8 @@ static OffHeapBinarizedVectorValues load(
     }
 
     /** Dense off-heap binarized vector values */
-    static class DenseOffHeapVectorValues extends OffHeapBinarizedVectorValues {
-        DenseOffHeapVectorValues(
+    public static class DenseOffHeapVectorValues extends OffHeapBinarizedVectorValues {
+        public DenseOffHeapVectorValues(
             int dimension,
             int size,
             float[] centroid,

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

@@ -0,0 +1,50 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.vectors.es92;
+
+import org.apache.lucene.util.BitUtil;
+
+import java.nio.ByteOrder;
+import java.nio.ShortBuffer;
+
+class BFloat16 {
+
+    public static final int BYTES = Short.BYTES;
+
+    public static short floatToBFloat16(float f) {
+        // this rounds 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) {
+        return Float.intBitsToFloat(bf << 16);
+    }
+
+    public static void floatToBFloat16(float[] floats, ShortBuffer bFloats) {
+        assert bFloats.remaining() == floats.length;
+        assert bFloats.order() == ByteOrder.LITTLE_ENDIAN;
+        for (float v : floats) {
+            bFloats.put(floatToBFloat16(v));
+        }
+    }
+
+    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/es92/ES92BFloat16FlatVectorsFormat.java

@@ -0,0 +1,129 @@
+/*
+ * @notice
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Modifications copyright (C) 2024 Elasticsearch B.V.
+ */
+package org.elasticsearch.index.codec.vectors.es92;
+
+import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
+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.index.SegmentReadState;
+import org.apache.lucene.index.SegmentWriteState;
+import org.apache.lucene.store.FlushInfo;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.MergeInfo;
+import org.elasticsearch.common.util.set.Sets;
+import org.elasticsearch.index.codec.vectors.es818.DirectIOHint;
+import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsFormat;
+import org.elasticsearch.index.codec.vectors.es818.MergeReaderWrapper;
+import org.elasticsearch.index.store.FsDirectoryFactory;
+
+import java.io.IOException;
+import java.util.Set;
+
+public final class ES92BFloat16FlatVectorsFormat extends FlatVectorsFormat {
+
+    static final String NAME = "ES92BFloat16FlatVectorsFormat";
+    static final String META_CODEC_NAME = "ES92BFloat16FlatVectorsFormatMeta";
+    static final String VECTOR_DATA_CODEC_NAME = "ES92BFloat16FlatVectorsFormatData";
+    static final String META_EXTENSION = "vemf";
+    static final String VECTOR_DATA_EXTENSION = "vec";
+
+    public static final int VERSION_START = 0;
+    public static final int VERSION_CURRENT = VERSION_START;
+
+    static final int DIRECT_MONOTONIC_BLOCK_SHIFT = 16;
+    private final FlatVectorsScorer vectorsScorer;
+
+    public ES92BFloat16FlatVectorsFormat(FlatVectorsScorer vectorsScorer) {
+        super(NAME);
+        this.vectorsScorer = vectorsScorer;
+    }
+
+    @Override
+    public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
+        return new ES92BFloat16FlatVectorsWriter(state, vectorsScorer);
+    }
+
+    static boolean shouldUseDirectIO(SegmentReadState state) {
+        return ES818BinaryQuantizedVectorsFormat.USE_DIRECT_IO && FsDirectoryFactory.isHybridFs(state.directory);
+    }
+
+    @Override
+    public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
+        if (shouldUseDirectIO(state) && state.context.context() == IOContext.Context.DEFAULT) {
+            // only override the context for the random-access use case
+            SegmentReadState directIOState = new SegmentReadState(
+                state.directory,
+                state.segmentInfo,
+                state.fieldInfos,
+                new DirectIOContext(state.context.hints()),
+                state.segmentSuffix
+            );
+            // Use mmap for merges and direct I/O for searches.
+            // TODO: Open the mmap file with sequential access instead of random (current behavior).
+            return new MergeReaderWrapper(
+                new ES92BFloat16FlatVectorsReader(directIOState, vectorsScorer),
+                new ES92BFloat16FlatVectorsReader(state, vectorsScorer)
+            );
+        } else {
+            return new ES92BFloat16FlatVectorsReader(state, vectorsScorer);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "ES92BFloat16FlatVectorsFormat(" + "vectorsScorer=" + vectorsScorer + ')';
+    }
+
+    static class DirectIOContext implements IOContext {
+
+        final Set<FileOpenHint> hints;
+
+        DirectIOContext(Set<FileOpenHint> hints) {
+            // always add DirectIOHint to the hints given
+            this.hints = Sets.union(hints, Set.of(DirectIOHint.INSTANCE));
+        }
+
+        @Override
+        public Context context() {
+            return Context.DEFAULT;
+        }
+
+        @Override
+        public MergeInfo mergeInfo() {
+            return null;
+        }
+
+        @Override
+        public FlushInfo flushInfo() {
+            return null;
+        }
+
+        @Override
+        public Set<FileOpenHint> hints() {
+            return hints;
+        }
+
+        @Override
+        public IOContext withHints(FileOpenHint... hints) {
+            return new DirectIOContext(Set.of(hints));
+        }
+    }
+}