[LuceneOnFaiss - Part1] Added building blocks for memory optimized search.

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java

@@ -11,6 +11,7 @@
 
 package org.opensearch.knn.index.codec.KNN990Codec;
 
+import lombok.extern.slf4j.Slf4j;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.index.ByteVectorValues;
@@ -22,12 +23,16 @@
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOSupplier;
 import org.apache.lucene.util.IOUtils;
 import org.opensearch.common.UUIDs;
 import org.opensearch.knn.index.codec.util.KNNCodecUtil;
 import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper;
+import org.opensearch.knn.index.engine.KNNEngine;
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 import org.opensearch.knn.index.quantizationservice.QuantizationService;
+import org.opensearch.knn.memoryoptsearch.VectorSearcher;
+import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationState;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateCacheManager;
 import org.opensearch.knn.quantization.models.quantizationState.QuantizationStateReadConfig;
@@ -37,23 +42,89 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+
+import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE;
+import static org.opensearch.knn.index.mapper.KNNVectorFieldMapper.KNN_FIELD;
 
 /**
  * Vectors reader class for reading the flat vectors for native engines. The class provides methods for iterating
  * over the vectors and retrieving their values.
  */
+@Slf4j
 public class NativeEngines990KnnVectorsReader extends KnnVectorsReader {
+    private static final int RESERVE_TWICE_SPACE = 2;
+    private static final float SUFFICIENT_LOAD_FACTOR = 0.6f;
 
     private final FlatVectorsReader flatVectorsReader;
     private Map<String, String> quantizationStateCacheKeyPerField;
     private SegmentReadState segmentReadState;
     private final List<String> cacheKeys;
+    private Map<String, VectorSearcher> vectorSearchers;
 
     public NativeEngines990KnnVectorsReader(final SegmentReadState state, final FlatVectorsReader flatVectorsReader) {
         this.flatVectorsReader = flatVectorsReader;
         this.segmentReadState = state;
         this.cacheKeys = getVectorCacheKeysFromSegmentReaderState(state);
+        this.vectorSearchers = new HashMap<>(RESERVE_TWICE_SPACE * segmentReadState.fieldInfos.size(), SUFFICIENT_LOAD_FACTOR);
         loadCacheKeyMap();
+
+        //
+        // TMP(KDY) : Dynamic update will be covered in part-7. Please refer to
+        // https://github.com/opensearch-project/k-NN/issues/2401#issuecomment-2699777824
+        //
+        final boolean isMemoryOptimizedSearchEnabled = false;
+        if (isMemoryOptimizedSearchEnabled) {
+            loadMemoryOptimizedSearcher();
+        }
+    }
+
+    private IOSupplier<VectorSearcher> getIndexFileNameIfMemoryOptimizedSearchSupported(final FieldInfo fieldInfo) {
+        // Skip non-knn fields.
+        final Map<String, String> attributes = fieldInfo.attributes();
+        if (attributes == null || attributes.containsKey(KNN_FIELD) == false) {
+            return null;
+        }
+
+        // Get engine
+        final String engineName = attributes.getOrDefault(KNN_ENGINE, KNNEngine.DEFAULT.getName());
+        final KNNEngine knnEngine = KNNEngine.getEngine(engineName);
+
+        // Get memory optimized searcher from engine
+        final VectorSearcherFactory searcherFactory = knnEngine.getVectorSearcherFactory();
+        if (searcherFactory == null) {
+            // It's not supported
+            return null;
+        }
+
+        // Start creating searcher
+        final String fileName = KNNCodecUtil.getNativeEngineFileFromFieldInfo(fieldInfo, segmentReadState.segmentInfo);
+        if (fileName != null) {
+            return () -> searcherFactory.createVectorSearcher(segmentReadState.directory, fileName);
+        }
+
+        // Not supported
+        return null;
+    }
+
+    private void loadMemoryOptimizedSearcher() {
+        try {
+            for (FieldInfo fieldInfo : segmentReadState.fieldInfos) {
+                final IOSupplier<VectorSearcher> searcherSupplier = getIndexFileNameIfMemoryOptimizedSearchSupported(fieldInfo);
+                if (searcherSupplier != null) {
+                    final VectorSearcher searcher = Objects.requireNonNull(searcherSupplier.get());
+                    vectorSearchers.put(fieldInfo.getName(), searcher);
+                }
+            }
+        } catch (Exception e) {
+            // Close opened searchers first, then suppress
+            try {
+                IOUtils.closeWhileHandlingException(vectorSearchers.values());
+            } catch (Exception closeException) {
+                log.error(closeException.getMessage(), closeException);
+            }
+            throw new RuntimeException(e);
+        }
     }
 
     /**
@@ -135,6 +206,14 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits
             ((QuantizationConfigKNNCollector) knnCollector).setQuantizationState(quantizationState);
             return;
         }
+
+        // Try with memory optimized searcher
+        final VectorSearcher memoryOptimizedSearcher = vectorSearchers.get(field);
+        if (memoryOptimizedSearcher != null) {
+            memoryOptimizedSearcher.search(target, knnCollector, acceptDocs);
+            return;
+        }
+
         throw new UnsupportedOperationException("Search functionality using codec is not supported with Native Engine Reader");
     }
 
@@ -197,6 +276,9 @@ public void close() throws IOException {
                 quantizationStateCacheManager.evict(cacheKey);
             }
         }
+
+        // TODO(KDY)
+        // Close all memory optimized searchers.
     }
 
     private void loadCacheKeyMap() {

src/main/java/org/opensearch/knn/index/engine/KNNEngine.java

@@ -8,6 +8,7 @@
 import com.google.common.collect.ImmutableSet;
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.index.SpaceType;
+import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
 import org.opensearch.knn.index.engine.faiss.Faiss;
 import org.opensearch.knn.index.engine.lucene.Lucene;
 import org.opensearch.knn.index.engine.nmslib.Nmslib;
@@ -216,4 +217,9 @@ public ResolvedMethodContext resolveMethod(
     public boolean supportsRemoteIndexBuild() {
         return knnLibrary.supportsRemoteIndexBuild();
     }
+
+    @Override
+    public VectorSearcherFactory getVectorSearcherFactory() {
+        return knnLibrary.getVectorSearcherFactory();
+    }
 }

src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java

@@ -7,6 +7,7 @@
 
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.index.SpaceType;
+import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
 
 import java.util.Collections;
 import java.util.List;
@@ -140,11 +141,21 @@ default List<String> mmapFileExtensions() {
         return Collections.emptyList();
     }
 
-    /**
+    /*
      * Returns whether or not the engine implementation supports remote index build
      * @return true if remote index build is supported, false otherwise
      */
     default boolean supportsRemoteIndexBuild() {
         return false;
     }
+
+    /**
+     * Create a new vector searcher factory that compatible with on Lucene search API.
+     * @return New searcher factory that returns {@link org.opensearch.knn.memoryoptsearch.VectorSearcher}
+     *         If it is not supported, it should return null.
+     *         But, if it is supported, the factory shall not return null searcher.
+     */
+    default VectorSearcherFactory getVectorSearcherFactory() {
+        return null;
+    }
 }

src/main/java/org/opensearch/knn/index/engine/faiss/Faiss.java

@@ -8,12 +8,14 @@
 import com.google.common.collect.ImmutableMap;
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.SpaceType;
+import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
 import org.opensearch.knn.index.engine.KNNMethod;
 import org.opensearch.knn.index.engine.KNNMethodConfigContext;
 import org.opensearch.knn.index.engine.KNNMethodContext;
 import org.opensearch.knn.index.engine.MethodResolver;
 import org.opensearch.knn.index.engine.NativeLibrary;
 import org.opensearch.knn.index.engine.ResolvedMethodContext;
+import org.opensearch.knn.memoryoptsearch.faiss.FaissMemoryOptimizedSearcherFactory;
 
 import java.util.Map;
 import java.util.function.Function;
@@ -123,4 +125,9 @@ public ResolvedMethodContext resolveMethod(
     public boolean supportsRemoteIndexBuild() {
         return true;
     }
+
+    @Override
+    public VectorSearcherFactory getVectorSearcherFactory() {
+        return new FaissMemoryOptimizedSearcherFactory();
+    }
 }

src/main/java/org/opensearch/knn/memoryoptsearch/VectorSearcher.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch;
+
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.util.Bits;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+/**
+ * This searcher performs vector search on non-Lucene index, for example FAISS index.
+ * Two search APIs will be compatible with Lucene, taking {@link KnnCollector} and {@link Bits}.
+ * In its implementation, it must collect top vectors that is similar to the given query. Make sure to transform the result to similarity
+ * value if internally calculates distance between.
+ */
+public interface VectorSearcher extends Closeable {
+    /**
+     * Return the k nearest neighbor documents as determined by comparison of their vector values for
+     * this field, to the given vector, by the field's similarity function. The score of each document
+     * is derived from the vector similarity in a way that ensures scores are positive and that a
+     * larger score corresponds to a higher ranking.
+     *
+     * <p>The search is allowed to be approximate, meaning the results are not guaranteed to be the
+     * true k closest neighbors. For large values of k (for example when k is close to the total
+     * number of documents), the search may also retrieve fewer than k documents.
+     *
+     * @param target the vector-valued float vector query
+     * @param knnCollector a KnnResults collector and relevant settings for gathering vector results
+     * @param acceptDocs {@link Bits} that represents the allowed documents to match, or {@code null}
+     *     if they are all allowed to match.
+     */
+    void search(float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException;
+
+    /**
+     * Return the k nearest neighbor documents as determined by comparison of their vector values for
+     * this field, to the given vector, by the field's similarity function. The score of each document
+     * is derived from the vector similarity in a way that ensures scores are positive and that a
+     * larger score corresponds to a higher ranking.
+     *
+     * <p>The search is allowed to be approximate, meaning the results are not guaranteed to be the
+     * true k closest neighbors. For large values of k (for example when k is close to the total
+     * number of documents), the search may also retrieve fewer than k documents.
+     *
+     * @param target the vector-valued byte vector query
+     * @param knnCollector a KnnResults collector and relevant settings for gathering vector results
+     * @param acceptDocs {@link Bits} that represents the allowed documents to match, or {@code null}
+     *     if they are all allowed to match.
+     */
+    void search(byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException;
+}

src/main/java/org/opensearch/knn/memoryoptsearch/VectorSearcherFactory.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch;
+
+import org.apache.lucene.store.Directory;
+
+import java.io.IOException;
+
+/**
+ * Factory to create {@link VectorSearcher}.
+ * Provided parameters will have {@link Directory} and a file name where implementation can rely on it to open an input stream.
+ */
+public interface VectorSearcherFactory {
+    /**
+     * Create a non-null {@link VectorSearcher} with given Lucene's {@link Directory}.
+     *
+     * @param directory Lucene's Directory.
+     * @param fileName Logical file name to load.
+     * @return It must return a non-null {@link VectorSearcher}
+     * @throws IOException
+     */
+    VectorSearcher createVectorSearcher(Directory directory, String fileName) throws IOException;
+}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissMemoryOptimizedSearcher.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch.faiss;
+
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+import org.opensearch.knn.memoryoptsearch.VectorSearcher;
+
+import java.io.IOException;
+
+/**
+ * This searcher directly reads FAISS index file via the provided {@link IndexInput} then perform vector search on it.
+ */
+public class FaissMemoryOptimizedSearcher implements VectorSearcher {
+    private final IndexInput indexInput;
+
+    public FaissMemoryOptimizedSearcher(IndexInput indexInput) {
+        this.indexInput = indexInput;
+    }
+
+    @Override
+    public void search(float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        // TODO(KDY) : This will be covered in subsequent parts.
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public void search(byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        // TODO(KDY) : This will be covered in subsequent parts.
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    @Override
+    public void close() throws IOException {
+        indexInput.close();
+    }
+}

src/main/java/org/opensearch/knn/memoryoptsearch/faiss/FaissMemoryOptimizedSearcherFactory.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.memoryoptsearch.faiss;
+
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.ReadAdvice;
+import org.opensearch.knn.memoryoptsearch.VectorSearcher;
+import org.opensearch.knn.memoryoptsearch.VectorSearcherFactory;
+
+import java.io.IOException;
+
+/**
+ * This factory returns {@link VectorSearcher} that performs vector search directly on FAISS index.
+ * Note that we pass `RANDOM` as advice to prevent the underlying storage from performing read-ahead. Since vector search naturally accesses
+ * random vector locations, read-ahead does not improve performance. By passing the `RANDOM` context, we explicitly indicate that
+ * this searcher will access vectors randomly.
+ */
+public class FaissMemoryOptimizedSearcherFactory implements VectorSearcherFactory {
+    @Override
+    public VectorSearcher createVectorSearcher(final Directory directory, final String fileName) throws IOException {
+        final IndexInput indexInput = directory.openInput(
+            fileName,
+            new IOContext(IOContext.Context.DEFAULT, null, null, ReadAdvice.RANDOM)
+        );
+        return new FaissMemoryOptimizedSearcher(indexInput);
+    }
+}