Remove FSDirectory hard dependency in vector search.

0ctopus13prime · Sep 13, 2024 · 6cdde7f · 6cdde7f
1 parent 88a02f9
commit 6cdde7f
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 16 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryAllocation.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryAllocation.java
@@ -184,7 +184,9 @@ private void cleanup() {
 
             this.closed = true;
 
-            watcherHandle.stop();
+            if (watcherHandle != null) {
+                watcherHandle.stop();
+            }
 
             // memoryAddress is sometimes initialized to 0. If this is ever the case, freeing will surely fail.
             if (memoryAddress != 0) {

diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java
@@ -13,7 +13,6 @@
 
 import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.store.Directory;
-import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.opensearch.core.action.ActionListener;
@@ -130,16 +129,15 @@ public NativeMemoryAllocation.IndexAllocation load(NativeMemoryEntryContext.Inde
 
             // Get the vector index path (note that it's logical path)
             final String indexPath = indexEntryContext.getKey();
-            FileWatcher fileWatcher = new FileWatcher(Path.of(indexPath));
-            fileWatcher.addListener(indexFileOnDeleteListener);
-            fileWatcher.init();
 
             KNNEngine knnEngine = KNNEngine.getEngineNameFromPath(indexPath);
 
             if (knnEngine != KNNEngine.FAISS) {
                 return loadFileBasedVectorIndex(indexEntryContext);
             }
 
+            final int indexSize = (int) directory.fileLength(indexPath);
+
             try (IndexInput readStream = directory.openInput(indexPath, IOContext.READONCE)) {
                 IndexInputWithBuffer indexInputWithBuffer = new IndexInputWithBuffer(readStream);
                 long indexAddress = JNIService.loadIndex(indexInputWithBuffer, indexEntryContext.getParameters(), knnEngine);
@@ -152,15 +150,14 @@ public NativeMemoryAllocation.IndexAllocation load(NativeMemoryEntryContext.Inde
                     JNIService.setSharedIndexState(indexAddress, sharedIndexState.getSharedIndexStateAddress(), knnEngine);
                 }
 
-                final WatcherHandle<FileWatcher> watcherHandle = resourceWatcherService.add(fileWatcher);
                 return new NativeMemoryAllocation.IndexAllocation(
                     executor,
                     indexAddress,
-                    indexEntryContext.calculateSizeInKB(),
+                    indexSize,
                     knnEngine,
-                    indexPath.toString(),
+                    indexPath,
                     indexEntryContext.getOpenSearchIndexName(),
-                    watcherHandle,
+                    null,
                     sharedIndexState,
                     IndexUtil.isBinaryIndex(knnEngine, indexEntryContext.getParameters())
                 );

diff --git a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java
@@ -15,13 +15,10 @@
 import org.apache.lucene.search.FilteredDocIdSetIterator;
 import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.Weight;
-import org.apache.lucene.store.FSDirectory;
-import org.apache.lucene.store.FilterDirectory;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.BitSetIterator;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.FixedBitSet;
-import org.opensearch.common.io.PathUtils;
 import org.opensearch.common.lucene.Lucene;
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.KNNSettings;
@@ -43,7 +40,6 @@
 import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
 
 import java.io.IOException;
-import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
@@ -223,7 +219,6 @@ private Map<Integer, Float> doANNSearch(
         final int k
     ) throws IOException {
         final SegmentReader reader = Lucene.segmentReader(context.reader());
-        String directory = ((FSDirectory) FilterDirectory.unwrap(reader.directory())).getDirectory().toString();
 
         FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(knnQuery.getField());
 
@@ -269,7 +264,7 @@ private Map<Integer, Float> doANNSearch(
             return null;
         }
 
-        Path indexPath = PathUtils.get(directory, engineFiles.get(0));
+        final String indexPath = engineFiles.get(0);
         final KNNQueryResult[] results;
         KNNCounter.GRAPH_QUERY_REQUESTS.increment();
 
@@ -279,7 +274,7 @@ private Map<Integer, Float> doANNSearch(
             indexAllocation = nativeMemoryCacheManager.get(
                 new NativeMemoryEntryContext.IndexEntryContext(
                     reader.directory(),
-                    indexPath.toString(),
+                    indexPath,
                     NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
                     getParametersAtLoading(
                         spaceType,