PR comments

Author: thecoop

qa/vector/src/main/java/org/elasticsearch/test/knn/CmdLineArgs.java

@@ -14,7 +14,7 @@
 import java.nio.ByteOrder;
 import java.nio.ShortBuffer;
 
-public class BFloat16 {
+public final class BFloat16 {
 
     public static final int BYTES = Short.BYTES;
 

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

@@ -88,10 +88,10 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     private final DirectIOCapableFlatVectorsFormat rawVectorFormat;
 
     public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster) {
-        this(vectorPerCluster, centroidsPerParentCluster, false, false);
+        this(vectorPerCluster, centroidsPerParentCluster, false);
     }
 
-    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO, boolean useBFloat16) {
+    public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentCluster, boolean useDirectIO) {
         super(NAME);
         if (vectorPerCluster < MIN_VECTORS_PER_CLUSTER || vectorPerCluster > MAX_VECTORS_PER_CLUSTER) {
             throw new IllegalArgumentException(
@@ -116,7 +116,7 @@ public ES920DiskBBQVectorsFormat(int vectorPerCluster, int centroidsPerParentClu
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
         this.useDirectIO = useDirectIO;
-        this.rawVectorFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
+        this.rawVectorFormat = float32VectorFormat;
     }
 
     /** Constructs a format using the given graph construction parameters and scalar quantization. */

server/src/main/java/org/elasticsearch/index/codec/vectors/diskbbq/ES920DiskBBQVectorsFormat.java

@@ -25,14 +25,12 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
-import org.elasticsearch.index.codec.vectors.DirectIOCapableFlatVectorsFormat;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryFlatVectorsScorer;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsReader;
 import org.elasticsearch.index.codec.vectors.es818.ES818BinaryQuantizedVectorsWriter;
 
 import java.io.IOException;
-import java.util.Map;
 
 /**
  * Copied from Lucene, replace with Lucene's implementation sometime after Lucene 10
@@ -91,57 +89,23 @@ public class ES93BinaryQuantizedVectorsFormat extends ES93GenericFlatVectorsForm
 
     public static final String NAME = "ES93BinaryQuantizedVectorsFormat";
 
-    private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
-    private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
-        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
-    );
-
-    private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
-        float32VectorFormat.getName(),
-        float32VectorFormat,
-        bfloat16VectorFormat.getName(),
-        bfloat16VectorFormat
-    );
-
     private static final ES818BinaryFlatVectorsScorer scorer = new ES818BinaryFlatVectorsScorer(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
 
-    private final boolean useDirectIO;
-    private final DirectIOCapableFlatVectorsFormat rawFormat;
-
     public ES93BinaryQuantizedVectorsFormat() {
         this(false, false);
     }
 
-    public ES93BinaryQuantizedVectorsFormat(boolean useDirectIO, boolean useBFloat16) {
-        super(NAME);
-        this.useDirectIO = useDirectIO;
-        this.rawFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
+    public ES93BinaryQuantizedVectorsFormat(boolean useBFloat16, boolean useDirectIO) {
+        super(NAME, useBFloat16, useDirectIO);
     }
 
     @Override
     protected FlatVectorsScorer flatVectorsScorer() {
         return scorer;
     }
 
-    @Override
-    protected boolean useDirectIOReads() {
-        return useDirectIO;
-    }
-
-    @Override
-    protected DirectIOCapableFlatVectorsFormat writeFlatVectorsFormat() {
-        return rawFormat;
-    }
-
-    @Override
-    protected Map<String, DirectIOCapableFlatVectorsFormat> supportedReadFlatVectorsFormats() {
-        return supportedFormats;
-    }
-
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
         return new ES818BinaryQuantizedVectorsWriter(scorer, super.fieldsWriter(state), state);

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93BinaryQuantizedVectorsFormat.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.index.codec.vectors.es93;
 
+import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
@@ -34,42 +35,45 @@ public abstract class ES93GenericFlatVectorsFormat extends AbstractFlatVectorsFo
         VERSION_CURRENT
     );
 
-    public ES93GenericFlatVectorsFormat(String name) {
-        super(name);
-    }
+    private static final DirectIOCapableFlatVectorsFormat float32VectorFormat = new DirectIOCapableLucene99FlatVectorsFormat(
+        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
+    );
+    private static final DirectIOCapableFlatVectorsFormat bfloat16VectorFormat = new ES93BFloat16FlatVectorsFormat(
+        FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
+    );
 
-    protected abstract DirectIOCapableFlatVectorsFormat writeFlatVectorsFormat();
+    private static final Map<String, DirectIOCapableFlatVectorsFormat> supportedFormats = Map.of(
+        float32VectorFormat.getName(),
+        float32VectorFormat,
+        bfloat16VectorFormat.getName(),
+        bfloat16VectorFormat
+    );
 
-    protected abstract boolean useDirectIOReads();
+    private final DirectIOCapableFlatVectorsFormat writeFormat;
+    private final boolean useDirectIO;
 
-    protected abstract Map<String, DirectIOCapableFlatVectorsFormat> supportedReadFlatVectorsFormats();
+    public ES93GenericFlatVectorsFormat(String name, boolean useBFloat16, boolean useDirectIO) {
+        super(name);
+        writeFormat = useBFloat16 ? bfloat16VectorFormat : float32VectorFormat;
+        this.useDirectIO = useDirectIO;
+    }
 
     @Override
     public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        var flatFormat = writeFlatVectorsFormat();
-        boolean directIO = useDirectIOReads();
-        return new ES93GenericFlatVectorsWriter(META, flatFormat.getName(), directIO, state, flatFormat.fieldsWriter(state));
+        return new ES93GenericFlatVectorsWriter(META, writeFormat.getName(), useDirectIO, state, writeFormat.fieldsWriter(state));
     }
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        var readFormats = supportedReadFlatVectorsFormats();
         return new ES93GenericFlatVectorsReader(META, state, (f, dio) -> {
-            var format = readFormats.get(f);
+            var format = supportedFormats.get(f);
             if (format == null) return null;
             return format.fieldsReader(state, dio);
         });
     }
 
     @Override
     public String toString() {
-        return getName()
-            + "(name="
-            + getName()
-            + ", writeFlatVectorFormat="
-            + writeFlatVectorsFormat()
-            + ", readFlatVectorsFormats="
-            + supportedReadFlatVectorsFormats().values()
-            + ")";
+        return getName() + "(name=" + getName() + ", writeFlatVectorFormat=" + writeFormat + ")";
     }
 }

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93GenericFlatVectorsFormat.java

@@ -53,7 +53,7 @@ public ES93HnswBinaryQuantizedVectorsFormat() {
      */
     public ES93HnswBinaryQuantizedVectorsFormat(int maxConn, int beamWidth, boolean useDirectIO, boolean useBFloat16) {
         super(NAME, maxConn, beamWidth);
-        flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO, useBFloat16);
+        flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useBFloat16, useDirectIO);
     }
 
     /**
@@ -76,7 +76,7 @@ public ES93HnswBinaryQuantizedVectorsFormat(
         ExecutorService mergeExec
     ) {
         super(NAME, maxConn, beamWidth, numMergeWorkers, mergeExec);
-        flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useDirectIO, useBFloat16);
+        flatVectorsFormat = new ES93BinaryQuantizedVectorsFormat(useBFloat16, useDirectIO);
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/codec/vectors/es93/ES93HnswBinaryQuantizedVectorsFormat.java

@@ -21,7 +21,6 @@
 
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.lucene90.IndexedDISI;
-import org.apache.lucene.codecs.lucene95.HasIndexSlice;
 import org.apache.lucene.codecs.lucene95.OrdToDocDISIReaderConfiguration;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.VectorEncoding;
@@ -36,7 +35,7 @@
 
 import java.io.IOException;
 
-abstract class OffHeapBFloat16VectorValues extends FloatVectorValues implements HasIndexSlice {
+abstract class OffHeapBFloat16VectorValues extends FloatVectorValues {
 
     protected final int dimension;
     protected final int size;
@@ -62,7 +61,7 @@ abstract class OffHeapBFloat16VectorValues extends FloatVectorValues implements
         this.byteSize = byteSize;
         this.similarityFunction = similarityFunction;
         this.flatVectorsScorer = flatVectorsScorer;
-        bfloatBytes = new byte[dimension * 2];
+        bfloatBytes = new byte[dimension * BFloat16.BYTES];
         value = new float[dimension];
     }
 
@@ -76,11 +75,6 @@ public int size() {
         return size;
     }
 
-    @Override
-    public IndexInput getSlice() {
-        return slice;
-    }
-
     @Override
     public float[] vectorValue(int targetOrd) throws IOException {
         if (lastOrd == targetOrd) {

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

@@ -2166,8 +2166,7 @@ KnnVectorsFormat getVectorsFormat(ElementType elementType) {
             return new ES920DiskBBQVectorsFormat(
                 clusterSize,
                 ES920DiskBBQVectorsFormat.DEFAULT_CENTROIDS_PER_PARENT_CLUSTER,
-                onDiskRescore,
-                false
+                onDiskRescore
             );
         }