Store the raw format name for DiskBBQ in field metadata

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

@@ -53,6 +53,8 @@ public class ES920DiskBBQVectorsFormat extends KnnVectorsFormat {
     public static final String CLUSTER_EXTENSION = "clivf";
     static final String IVF_META_EXTENSION = "mivf";
 
+    static final String RAW_VECTOR_FORMAT = "raw_vector_format";
+
     public static final int VERSION_START = 0;
     public static final int VERSION_CURRENT = VERSION_START;
 
@@ -106,12 +108,18 @@ 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));
+        return new ES920DiskBBQVectorsReader(state);
     }
 
     @Override

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

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
+import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
@@ -25,13 +26,15 @@
 import org.elasticsearch.simdvec.ESVectorUtil;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.Map;
 
 import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.QUERY_BITS;
 import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
 import static org.elasticsearch.index.codec.vectors.BQSpaceUtils.transposeHalfByte;
 import static org.elasticsearch.index.codec.vectors.BQVectorUtils.discretize;
 import static org.elasticsearch.index.codec.vectors.OptimizedScalarQuantizer.DEFAULT_LAMBDA;
+import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.RAW_VECTOR_FORMAT;
 import static org.elasticsearch.simdvec.ES91OSQVectorsScorer.BULK_SIZE;
 
 /**
@@ -40,8 +43,23 @@
  */
 public class ES920DiskBBQVectorsReader extends IVFVectorsReader implements OffHeapStats {
 
-    public ES920DiskBBQVectorsReader(SegmentReadState state, FlatVectorsReader rawVectorsReader) throws IOException {
-        super(state, rawVectorsReader);
+    public ES920DiskBBQVectorsReader(SegmentReadState state) throws IOException {
+        super(state, loadReaders(state));
+    }
+
+    private static Map<String, FlatVectorsReader> loadReaders(SegmentReadState state) throws IOException {
+        Map<String, FlatVectorsReader> readers = new HashMap<>();
+        for (FieldInfo fi : state.fieldInfos) {
+            if (fi.hasVectorValues()) {
+                String formatName = fi.getAttribute(RAW_VECTOR_FORMAT);
+                if (formatName == null) {
+                    throw new IllegalArgumentException("Field does not have " + RAW_VECTOR_FORMAT);
+                }
+                readers.put(fi.name, (FlatVectorsReader) KnnVectorsFormat.forName(formatName).fieldsReader(state));
+            }
+        }
+
+        return Map.copyOf(readers);
     }
 
     CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice) throws IOException {

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

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
+import org.apache.lucene.codecs.KnnFieldVectorsWriter;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
@@ -39,6 +40,8 @@
 import java.util.AbstractList;
 import java.util.Arrays;
 
+import static org.elasticsearch.index.codec.vectors.diskbbq.ES920DiskBBQVectorsFormat.RAW_VECTOR_FORMAT;
+
 /**
  * Default implementation of {@link IVFVectorsWriter}. It uses {@link HierarchicalKMeans} algorithm to
  * partition the vector space, and then stores the centroids and posting list in a sequential
@@ -47,20 +50,29 @@
 public class ES920DiskBBQVectorsWriter extends IVFVectorsWriter {
     private static final Logger logger = LogManager.getLogger(ES920DiskBBQVectorsWriter.class);
 
+    private final String rawVectorFormatName;
     private final int vectorPerCluster;
     private final int centroidsPerParentCluster;
 
     public ES920DiskBBQVectorsWriter(
+        String rawVectorFormatName,
         SegmentWriteState state,
         FlatVectorsWriter rawVectorDelegate,
         int vectorPerCluster,
         int centroidsPerParentCluster
     ) throws IOException {
         super(state, rawVectorDelegate);
+        this.rawVectorFormatName = rawVectorFormatName;
         this.vectorPerCluster = vectorPerCluster;
         this.centroidsPerParentCluster = centroidsPerParentCluster;
     }
 
+    @Override
+    public KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException {
+        fieldInfo.putAttribute(RAW_VECTOR_FORMAT, rawVectorFormatName);
+        return super.addField(fieldInfo);
+    }
+
     @Override
     CentroidOffsetAndLength buildAndWritePostingsLists(
         FieldInfo fieldInfo,

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,
@@ -212,26 +217,34 @@ 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) {
+        FlatVectorsReader reader = rawVectorReaders.get(field);
+        if (reader == null) throw new IllegalArgumentException("No recorded raw vector reader 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 +256,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 +322,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,7 +334,9 @@ 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(

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

@@ -107,7 +107,7 @@ protected IVFVectorsWriter(SegmentWriteState state, FlatVectorsWriter rawVectorD
     }
 
     @Override
-    public final KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException {
+    public KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOException {
         if (fieldInfo.getVectorSimilarityFunction() == VectorSimilarityFunction.COSINE) {
             throw new IllegalArgumentException("IVF does not support cosine similarity");
         }