Store the raw format name in field metadata

Author: thecoop

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,27 @@
  */
 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

@@ -33,6 +33,7 @@
 import org.elasticsearch.search.vectors.IVFKnnSearchStrategy;
 
 import java.io.IOException;
+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 +47,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,

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");
         }