Store the raw format name for DiskBBQ in field metadata

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

@@ -14,12 +14,16 @@
 import org.apache.lucene.codecs.KnnVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorScorerUtil;
 import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
 
 /**
  * Codec format for Inverted File Vector indexes. This index expects to break the dimensional space
@@ -59,6 +63,7 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     private static final FlatVectorsFormat rawVectorFormat = new Lucene99FlatVectorsFormat(
         FlatVectorScorerUtil.getLucene99FlatVectorsScorer()
     );
+    private static final Map<String, FlatVectorsFormat> supportedFormats = Map.of(rawVectorFormat.getName(), rawVectorFormat);
 
     // This dynamically sets the cluster probe based on the `k` requested and the number of clusters.
     // useful when searching with 'efSearch' type parameters instead of requiring a specific ratio.
@@ -106,12 +111,23 @@ public ES920DiskBBQVectorsFormat() {
 
     @Override
     public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException {
-        return new ES920DiskBBQVectorsWriter(state, rawVectorFormat.fieldsWriter(state), vectorPerCluster, centroidsPerParentCluster);
+        return new ES920DiskBBQVectorsWriter(
+            rawVectorFormat.getName(),
+            state,
+            rawVectorFormat.fieldsWriter(state),
+            vectorPerCluster,
+            centroidsPerParentCluster
+        );
     }
 
     @Override
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new ES920DiskBBQVectorsReader(state, rawVectorFormat.fieldsReader(state));
+        Map<String, FlatVectorsReader> readers = Maps.newHashMapWithExpectedSize(supportedFormats.size());
+        for (var fe : supportedFormats.entrySet()) {
+            readers.put(fe.getKey(), fe.getValue().fieldsReader(state));
+        }
+
+        return new ES920DiskBBQVectorsReader(state, Collections.unmodifiableMap(readers));
     }
 
     @Override

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

@@ -40,7 +40,7 @@
  */
 public class ES920DiskBBQVectorsReader extends IVFVectorsReader implements OffHeapStats {
 
-    public ES920DiskBBQVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader) throws IOException {
+    public ES920DiskBBQVectorsReader(SegmentReadState state, Map<String, FlatVectorsReader> rawVectorsReader) throws IOException {
         super(state, rawVectorsReader);
     }
 

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

@@ -51,12 +51,13 @@ public class ES920DiskBBQVectorsWriter extends IVFVectorsWriter {
     private final int centroidsPerParentCluster;
 
     public ES920DiskBBQVectorsWriter(
+        String rawVectorFormatName,
         SegmentWriteState state,
         FlatVectorsWriter rawVectorDelegate,
         int vectorPerCluster,
         int centroidsPerParentCluster
     ) throws IOException {
-        super(state, rawVectorDelegate);
+        super(state, rawVectorFormatName, rawVectorDelegate);
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
     }

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

@@ -32,7 +32,12 @@
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.search.vectors.IVFKnnSearchStrategy;
 
+import java.io.Closeable;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
 import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.DYNAMIC_VISIT_RATIO;
@@ -46,14 +51,14 @@ public abstract class IVFVectorsReader extends KnnVectorsReader {
     private final SegmentReadState state;
     private final FieldInfos fieldInfos;
     protected final IntObjectHashMap<FieldEntry> fields;
-    private final FlatVectorsReader rawVectorsReader;
+    private final Map<String, FlatVectorsReader> rawVectorReaders;
 
     @SuppressWarnings("this-escape")
-    protected IVFVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader) throws IOException {
+    protected IVFVectorsReader(SegmentReadState state, Map<String, FlatVectorsReader> rawVectorReaders) throws IOException {
         this.state = state;
         this.fieldInfos = state.fieldInfos;
-        this.rawVectorsReader = rawVectorsReader;
         this.fields = new IntObjectHashMap<>();
+        this.rawVectorReaders = rawVectorReaders;
         String meta = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
             state.segmentSuffix,
@@ -156,6 +161,7 @@ private void readFields(ChecksumIndexInput meta) throws IOException {
     }
 
     private FieldEntry readField(IndexInput input, FieldInfo info) throws IOException {
+        final String rawVectorFormat = input.readString();
         final VectorEncoding vectorEncoding = readVectorEncoding(input);
         final VectorSimilarityFunction similarityFunction = readSimilarityFunction(input);
         if (similarityFunction != info.getVectorSimilarityFunction()) {
@@ -182,6 +188,7 @@ private FieldEntry readField(IndexInput input, FieldInfo info) throws IOExceptio
             globalCentroidDp = Float.intBitsToFloat(input.readInt());
         }
         return new FieldEntry(
+            rawVectorFormat,
             similarityFunction,
             vectorEncoding,
             numCentroids,
@@ -212,26 +219,43 @@ private static VectorEncoding readVectorEncoding(DataInput input) throws IOExcep
 
     @Override
     public final void checkIntegrity() throws IOException {
-        rawVectorsReader.checkIntegrity();
+        for (var reader : rawVectorReaders.values()) {
+            reader.checkIntegrity();
+        }
         CodecUtil.checksumEntireFile(ivfCentroids);
         CodecUtil.checksumEntireFile(ivfClusters);
     }
 
+    private FlatVectorsReader getReaderForField(String field) {
+        FieldInfo fieldInfo = fieldInfos.fieldInfo(field);
+        if (fieldInfo.getVectorEncoding() != VectorEncoding.FLOAT32) {
+            // IVF only works on floats, so pass through any others straight
+            // to the first flat vectors format we can find
+            return rawVectorReaders.values().iterator().next();
+        }
+        var formatName = fields.get(fieldInfo.number).rawVectorFormatName;
+        FlatVectorsReader reader = rawVectorReaders.get(formatName);
+        if (reader == null) throw new IllegalArgumentException(
+            "Could not find raw vector format [" + formatName + "] for field [" + field + "]"
+        );
+        return reader;
+    }
+
     @Override
     public final FloatVectorValues getFloatVectorValues(String field) throws IOException {
-        return rawVectorsReader.getFloatVectorValues(field);
+        return getReaderForField(field).getFloatVectorValues(field);
     }
 
     @Override
     public final ByteVectorValues getByteVectorValues(String field) throws IOException {
-        return rawVectorsReader.getByteVectorValues(field);
+        return getReaderForField(field).getByteVectorValues(field);
     }
 
     @Override
     public final void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
         if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32) == false) {
-            rawVectorsReader.search(field, target, knnCollector, acceptDocs);
+            getReaderForField(field).search(field, target, knnCollector, acceptDocs);
             return;
         }
         if (fieldInfo.getVectorDimension() != target.length) {
@@ -243,7 +267,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
         if (acceptDocs instanceof BitSet bitSet) {
             percentFiltered = Math.max(0f, Math.min(1f, (float) bitSet.approximateCardinality() / bitSet.length()));
         }
-        int numVectors = rawVectorsReader.getFloatVectorValues(field).size();
+        int numVectors = getReaderForField(field).getFloatVectorValues(field).size();
         float visitRatio = DYNAMIC_VISIT_RATIO;
         // Search strategy may be null if this is being called from checkIndex (e.g. from a test)
         if (knnCollector.getSearchStrategy() instanceof IVFKnnSearchStrategy ivfSearchStrategy) {
@@ -309,7 +333,7 @@ public final void search(String field, float[] target, KnnCollector knnCollector
     @Override
     public final void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
         final FieldInfo fieldInfo = state.fieldInfos.fieldInfo(field);
-        final ByteVectorValues values = rawVectorsReader.getByteVectorValues(field);
+        final ByteVectorValues values = getReaderForField(field).getByteVectorValues(field);
         for (int i = 0; i < values.size(); i++) {
             final float score = fieldInfo.getVectorSimilarityFunction().compare(target, values.vectorValue(i));
             knnCollector.collect(values.ordToDoc(i), score);
@@ -321,10 +345,13 @@ public final void search(String field, byte[] target, KnnCollector knnCollector,
 
     @Override
     public void close() throws IOException {
-        IOUtils.close(rawVectorsReader, ivfCentroids, ivfClusters);
+        List<Closeable> closeables = new ArrayList<>(rawVectorReaders.values());
+        Collections.addAll(closeables, ivfCentroids, ivfClusters);
+        IOUtils.close(closeables);
     }
 
     protected record FieldEntry(
+        String rawVectorFormatName,
         VectorSimilarityFunction similarityFunction,
         VectorEncoding vectorEncoding,
         int numCentroids,

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

@@ -51,10 +51,13 @@ public abstract class IVFVectorsWriter extends KnnVectorsWriter {
     private final List<FieldWriter> fieldWriters = new ArrayList<>();
     private final IndexOutput ivfCentroids, ivfClusters;
     private final IndexOutput ivfMeta;
+    private final String rawVectorFormatName;
     private final FlatVectorsWriter rawVectorDelegate;
 
     @SuppressWarnings("this-escape")
-    protected IVFVectorsWriter(SegmentWriteState state, FlatVectorsWriter rawVectorDelegate) throws IOException {
+    protected IVFVectorsWriter(SegmentWriteState state, String rawVectorFormatName, FlatVectorsWriter rawVectorDelegate)
+        throws IOException {
+        this.rawVectorFormatName = rawVectorFormatName;
         this.rawVectorDelegate = rawVectorDelegate;
         final String metaFileName = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
@@ -476,6 +479,7 @@ private void writeMeta(
         float[] globalCentroid
     ) throws IOException {
         ivfMeta.writeInt(field.number);
+        ivfMeta.writeString(rawVectorFormatName);
         ivfMeta.writeInt(field.getVectorEncoding().ordinal());
         ivfMeta.writeInt(distFuncToOrd(field.getVectorSimilarityFunction()));
         ivfMeta.writeInt(numCentroids);