iter

Author: thecoop

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

@@ -14,12 +14,15 @@
 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;
 
 /**
@@ -119,11 +122,12 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException
 
     @Override
     public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new ES920DiskBBQVectorsReader(state, f -> {
-            var format = supportedFormats.get(f);
-            if (format == null) throw new IllegalArgumentException("Unknown format " + f);
-            return format;
-        });
+        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

@@ -9,7 +9,7 @@
 
 package org.elasticsearch.index.codec.vectors.diskbbq;
 
-import org.apache.lucene.codecs.hnsw.FlatVectorsFormat;
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.VectorSimilarityFunction;
@@ -26,7 +26,6 @@
 
 import java.io.IOException;
 import java.util.Map;
-import java.util.function.Function;
 
 import static org.apache.lucene.codecs.lucene102.Lucene102BinaryQuantizedVectorsFormat.QUERY_BITS;
 import static org.apache.lucene.index.VectorSimilarityFunction.COSINE;
@@ -41,8 +40,8 @@
  */
 public class ES920DiskBBQVectorsReader extends IVFVectorsReader implements OffHeapStats {
 
-    public ES920DiskBBQVectorsReader(SegmentReadState state, Function<String, FlatVectorsFormat> formatLookup) throws IOException {
-        super(state, f -> formatLookup.apply(f).fieldsReader(state));
+    public ES920DiskBBQVectorsReader(SegmentReadState state, Map<String, FlatVectorsReader> rawVectorsReader) throws IOException {
+        super(state, rawVectorsReader);
     }
 
     CentroidIterator getPostingListPrefetchIterator(CentroidIterator centroidIterator, IndexInput postingListSlice) throws IOException {

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

@@ -29,15 +29,13 @@
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.IOFunction;
 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.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -53,11 +51,10 @@ public abstract class IVFVectorsReader extends KnnVectorsReader {
     private final SegmentReadState state;
     private final FieldInfos fieldInfos;
     protected final IntObjectHashMap<FieldEntry> fields;
-    private final IOFunction<String, FlatVectorsReader> rawVectorReaders;
-    private final Map<String, FlatVectorsReader> loadedReaders = new HashMap<>();
+    private final Map<String, FlatVectorsReader> rawVectorReaders;
 
     @SuppressWarnings("this-escape")
-    protected IVFVectorsReader(SegmentReadState state, IOFunction<String, FlatVectorsReader> rawVectorReaders) throws IOException {
+    protected IVFVectorsReader(SegmentReadState state, Map<String, FlatVectorsReader> rawVectorReaders) throws IOException {
         this.state = state;
         this.fieldInfos = state.fieldInfos;
         this.fields = new IntObjectHashMap<>();
@@ -222,20 +219,19 @@ private static VectorEncoding readVectorEncoding(DataInput input) throws IOExcep
 
     @Override
     public final void checkIntegrity() throws IOException {
-        for (var reader : loadedReaders.values()) {
+        for (var reader : rawVectorReaders.values()) {
             reader.checkIntegrity();
         }
         CodecUtil.checksumEntireFile(ivfCentroids);
         CodecUtil.checksumEntireFile(ivfClusters);
     }
 
-    private synchronized FlatVectorsReader getReaderForField(String field) throws IOException {
+    private synchronized FlatVectorsReader getReaderForField(String field) {
         String formatName = fields.get(fieldInfos.fieldInfo(field).number).rawVectorFormatName;
-        FlatVectorsReader reader = loadedReaders.get(formatName);
-        if (reader == null) {
-            reader = rawVectorReaders.apply(formatName);
-            loadedReaders.put(formatName, reader);
-        }
+        FlatVectorsReader reader = rawVectorReaders.get(formatName);
+        if (reader == null) throw new IllegalArgumentException(
+            "Could not find raw vector format [" + formatName + "] for field [" + field + "]"
+        );
         return reader;
     }
 
@@ -343,7 +339,7 @@ public final void search(String field, byte[] target, KnnCollector knnCollector,
 
     @Override
     public void close() throws IOException {
-        List<Closeable> closeables = new ArrayList<>(loadedReaders.values());
+        List<Closeable> closeables = new ArrayList<>(rawVectorReaders.values());
         Collections.addAll(closeables, ivfCentroids, ivfClusters);
         IOUtils.close(closeables);
     }