diff --git a/src/main/java/org/opensearch/knn/index/KNNIndexShard.java b/src/main/java/org/opensearch/knn/index/KNNIndexShard.java index 3a771cbf6..772a7c38e 100644 --- a/src/main/java/org/opensearch/knn/index/KNNIndexShard.java +++ b/src/main/java/org/opensearch/knn/index/KNNIndexShard.java @@ -12,12 +12,15 @@ import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfo; import org.apache.lucene.index.SegmentReader; import org.apache.lucene.store.Directory; import org.opensearch.common.lucene.Lucene; import org.opensearch.index.engine.Engine; import org.opensearch.index.shard.IndexShard; import org.opensearch.knn.common.FieldInfoExtractor; +import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper; import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; import org.opensearch.knn.index.memory.NativeMemoryAllocation; @@ -28,7 +31,6 @@ import java.io.IOException; import java.util.ArrayList; -import java.util.Collection; import java.util.List; import java.util.Optional; import java.util.concurrent.ExecutionException; @@ -91,14 +93,18 @@ public void warmup() throws IOException { try (Engine.Searcher searcher = indexShard.acquireSearcher("knn-warmup")) { getAllEngineFileContexts(searcher.getIndexReader()).forEach((engineFileContext) -> { try { + final String cacheKey = NativeMemoryCacheKeyHelper.constructCacheKey( + engineFileContext.vectorFileName, + engineFileContext.segmentInfo + ); nativeMemoryCacheManager.get( new NativeMemoryEntryContext.IndexEntryContext( directory, - engineFileContext.getIndexPath(), + cacheKey, NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(), getParametersAtLoading( engineFileContext.getSpaceType(), - KNNEngine.getEngineNameFromPath(engineFileContext.getIndexPath()), + KNNEngine.getEngineNameFromPath(engineFileContext.getVectorFileName()), getIndexName(), engineFileContext.getVectorDataType() ), @@ -130,9 +136,13 @@ public void clearCache() { indexAllocation.writeLock(); log.info("[KNN] Evicting index from cache: [{}]", indexName); try (Engine.Searcher searcher = indexShard.acquireSearcher(INDEX_SHARD_CLEAR_CACHE_SEARCHER)) { - getAllEngineFileContexts(searcher.getIndexReader()).forEach( - (engineFileContext) -> nativeMemoryCacheManager.invalidate(engineFileContext.getIndexPath()) - ); + getAllEngineFileContexts(searcher.getIndexReader()).forEach((engineFileContext) -> { + final String cacheKey = NativeMemoryCacheKeyHelper.constructCacheKey( + engineFileContext.vectorFileName, + engineFileContext.segmentInfo + ); + nativeMemoryCacheManager.invalidate(cacheKey); + }); } catch (IOException ex) { log.error("[KNN] Failed to evict index from cache: [{}]", indexName, ex); throw new RuntimeException(ex); @@ -176,8 +186,7 @@ List getEngineFileContexts(IndexReader indexReader, KNNEngine String modelId = fieldInfo.attributes().getOrDefault(MODEL_ID, null); engineFiles.addAll( getEngineFileContexts( - reader.getSegmentInfo().files(), - reader.getSegmentInfo().info.name, + reader.getSegmentInfo(), fieldInfo.name, fileExtension, spaceType, @@ -197,20 +206,22 @@ List getEngineFileContexts(IndexReader indexReader, KNNEngine @VisibleForTesting List getEngineFileContexts( - Collection files, - String segmentName, + SegmentCommitInfo segmentCommitInfo, String fieldName, String fileExtension, SpaceType spaceType, String modelId, VectorDataType vectorDataType - ) { - String prefix = buildEngineFilePrefix(segmentName); - String suffix = buildEngineFileSuffix(fieldName, fileExtension); - return files.stream() + ) throws IOException { + // Ex: _0 + final String prefix = buildEngineFilePrefix(segmentCommitInfo.info.name); + // Ex: my_field.faiss + final String suffix = buildEngineFileSuffix(fieldName, fileExtension); + return segmentCommitInfo.files() + .stream() .filter(fileName -> fileName.startsWith(prefix)) .filter(fileName -> fileName.endsWith(suffix)) - .map(fileName -> new EngineFileContext(spaceType, modelId, fileName, vectorDataType)) + .map(vectorFileName -> new EngineFileContext(spaceType, modelId, vectorFileName, vectorDataType, segmentCommitInfo.info)) .collect(Collectors.toList()); } @@ -220,7 +231,8 @@ List getEngineFileContexts( static class EngineFileContext { private final SpaceType spaceType; private final String modelId; - private final String indexPath; + private final String vectorFileName; private final VectorDataType vectorDataType; + private final SegmentInfo segmentInfo; } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java index 7dc55feb9..83f22959a 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesProducer.java @@ -16,20 +16,23 @@ import org.apache.lucene.index.BinaryDocValues; import org.apache.lucene.index.DocValuesType; import org.apache.lucene.index.FieldInfo; +import java.io.IOException; import org.apache.lucene.index.NumericDocValues; import org.apache.lucene.index.SegmentReadState; import org.apache.lucene.index.SortedDocValues; import org.apache.lucene.index.SortedNumericDocValues; import org.apache.lucene.index.SortedSetDocValues; import org.opensearch.knn.index.codec.util.KNNCodecUtil; +import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper; import org.opensearch.knn.index.memory.NativeMemoryCacheManager; -import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; +import static org.opensearch.knn.index.mapper.KNNVectorFieldMapper.KNN_FIELD; + @Log4j2 public class KNN80DocValuesProducer extends DocValuesProducer { private final DocValuesProducer delegate; @@ -39,17 +42,23 @@ public KNN80DocValuesProducer(DocValuesProducer delegate, SegmentReadState state this.delegate = delegate; for (FieldInfo field : state.fieldInfos) { + if (!field.attributes().containsKey(KNN_FIELD)) { + continue; + } // Only segments that contains BinaryDocValues and doesn't have vector values should be considered. // By default, we don't create BinaryDocValues for knn field anymore. However, users can set doc_values = true // to create binary doc values explicitly like any other field. Hence, we only want to include fields // where approximate search is possible only by BinaryDocValues. - if (field.hasVectorValues() || field.getDocValuesType() != DocValuesType.BINARY) { + if (field.getDocValuesType() != DocValuesType.BINARY || field.hasVectorValues()) { continue; } + final String vectorIndexFileName = KNNCodecUtil.getEngineFileFromFieldInfo(field, state.segmentInfo); - if (vectorIndexFileName != null) { - fieldNameToVectorFileName.putIfAbsent(field.getName(), vectorIndexFileName); + if (vectorIndexFileName == null) { + continue; } + final String cacheKey = NativeMemoryCacheKeyHelper.constructCacheKey(vectorIndexFileName, state.segmentInfo); + fieldNameToVectorFileName.putIfAbsent(field.getName(), cacheKey); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java index 191416101..87fb60db9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsReader.java @@ -25,6 +25,7 @@ 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.memory.NativeMemoryCacheManager; import org.opensearch.knn.index.quantizationservice.QuantizationService; import org.opensearch.knn.quantization.models.quantizationState.QuantizationState; @@ -35,6 +36,8 @@ import java.util.HashMap; import java.util.Map; +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. @@ -55,10 +58,18 @@ public NativeEngines990KnnVectorsReader(final SegmentReadState state, final Flat private void fillIndexToVectorFileName(SegmentReadState state) { for (FieldInfo field : state.fieldInfos) { + if (!field.attributes().containsKey(KNN_FIELD)) { + continue; + } + if (!field.hasVectorValues()) { + continue; + } final String vectorIndexFileName = KNNCodecUtil.getEngineFileFromFieldInfo(field, state.segmentInfo); - if (vectorIndexFileName != null) { - fieldNameToVectorFileName.putIfAbsent(field.getName(), vectorIndexFileName); + if (vectorIndexFileName == null) { + continue; } + final String cacheKey = NativeMemoryCacheKeyHelper.constructCacheKey(vectorIndexFileName, state.segmentInfo); + fieldNameToVectorFileName.putIfAbsent(field.getName(), cacheKey); } } diff --git a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java index 828840865..2ee87ecb6 100644 --- a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java +++ b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java @@ -19,7 +19,6 @@ import java.util.List; import java.util.stream.Collectors; -import static org.opensearch.knn.common.KNNConstants.MODEL_ID; import static org.opensearch.knn.index.mapper.KNNVectorFieldMapper.KNN_FIELD; public class KNNCodecUtil { @@ -116,9 +115,6 @@ public static String getEngineFileFromFieldInfo(FieldInfo field, SegmentInfo seg * @return if and only if Native Engine we return specific engine, else return null */ private static KNNEngine getNativeKNNEngine(@NonNull FieldInfo field) { - if (field.attributes().containsKey(MODEL_ID)) { - return null; - } final KNNEngine engine = FieldInfoExtractor.extractKNNEngine(field); if (KNNEngine.getEnginesThatCreateCustomSegmentFiles().contains(engine)) { return engine; diff --git a/src/main/java/org/opensearch/knn/index/codec/util/NativeMemoryCacheKeyHelper.java b/src/main/java/org/opensearch/knn/index/codec/util/NativeMemoryCacheKeyHelper.java new file mode 100644 index 000000000..3bff185ea --- /dev/null +++ b/src/main/java/org/opensearch/knn/index/codec/util/NativeMemoryCacheKeyHelper.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.index.codec.util; + +import org.apache.lucene.index.SegmentInfo; + +import java.util.Base64; + +public final class NativeMemoryCacheKeyHelper { + private NativeMemoryCacheKeyHelper() {} + + /** + * Construct a unique cache key for look-up operation in {@link org.opensearch.knn.index.memory.NativeMemoryCacheManager} + * + * @param vectorIndexFileName Vector index file name. Ex: _0_165_test_field.faiss. + * @param segmentInfo Segment info object representing a logical segment unit containing a vector index. + * @return Unique cache key that can be used for look-up and invalidating in + * {@link org.opensearch.knn.index.memory.NativeMemoryCacheManager} + */ + public static String constructCacheKey(final String vectorIndexFileName, final SegmentInfo segmentInfo) { + final String segmentId = Base64.getEncoder().encodeToString(segmentInfo.getId()); + final String cacheKey = vectorIndexFileName + "@" + segmentId; + return cacheKey; + } + + public static String extractVectorIndexFileName(final String cacheKey) { + final int indexOfDelimiter = cacheKey.indexOf('@'); + if (indexOfDelimiter != -1) { + final String vectorFileName = cacheKey.substring(0, indexOfDelimiter); + return vectorFileName; + } + return null; + } +} diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java index 00bf023f9..3e2e6b665 100644 --- a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java +++ b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java @@ -80,26 +80,26 @@ public static class IndexEntryContext extends NativeMemoryEntryContext parameters, String openSearchIndexName ) { - this(directory, indexPath, indexLoadStrategy, parameters, openSearchIndexName, null); + this(directory, vectorIndexCacheKey, indexLoadStrategy, parameters, openSearchIndexName, null); } /** * Constructor * * @param directory Lucene directory to create required IndexInput/IndexOutput to access files. - * @param indexPath path to index file. Also used as key in cache. + * @param vectorIndexCacheKey Cache key for {@link NativeMemoryCacheManager}. It must contain a vector file name. * @param indexLoadStrategy strategy to load index into memory * @param parameters load time parameters * @param openSearchIndexName opensearch index associated with index @@ -107,13 +107,13 @@ public IndexEntryContext( */ public IndexEntryContext( Directory directory, - String indexPath, + String vectorIndexCacheKey, NativeMemoryLoadStrategy.IndexLoadStrategy indexLoadStrategy, Map parameters, String openSearchIndexName, String modelId ) { - super(indexPath); + super(vectorIndexCacheKey); this.directory = directory; this.indexLoadStrategy = indexLoadStrategy; this.openSearchIndexName = openSearchIndexName; diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java index 2c49c4879..aa4b5ae89 100644 --- a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java +++ b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java @@ -16,6 +16,7 @@ import org.apache.lucene.store.IOContext; import org.apache.lucene.store.IndexInput; import org.opensearch.core.action.ActionListener; +import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper; import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import org.opensearch.knn.index.store.IndexInputWithBuffer; import org.opensearch.knn.index.util.IndexUtil; @@ -71,16 +72,25 @@ private IndexLoadStrategy() { @Override public NativeMemoryAllocation.IndexAllocation load(NativeMemoryEntryContext.IndexEntryContext indexEntryContext) throws IOException { - // Ex: _0_NativeEngines990KnnVectorsFormat_0.vec - final String vectorFileName = indexEntryContext.getKey(); - final KNNEngine knnEngine = KNNEngine.getEngineNameFromPath(vectorFileName); + // Extract vector file name from the given cache key. + // Ex: _0_165_my_field.faiss@1vaqiupVUwvkXAG4Qc/RPg== + final String cacheKey = indexEntryContext.getKey(); + final String vectorFileName = NativeMemoryCacheKeyHelper.extractVectorIndexFileName(cacheKey); + if (vectorFileName == null) { + throw new IllegalStateException( + "Invalid cache key was given. The key [" + cacheKey + "] does not contain corresponding vector file name." + ); + } + // Prepare for opening index input from directory. + final KNNEngine knnEngine = KNNEngine.getEngineNameFromPath(vectorFileName); final Directory directory = indexEntryContext.getDirectory(); final int indexSizeKb = Math.toIntExact(directory.fileLength(vectorFileName) / 1024); + // Try to open an index input then pass it down to native engine for loading an index. try (IndexInput readStream = directory.openInput(vectorFileName, IOContext.READONCE)) { - IndexInputWithBuffer indexInputWithBuffer = new IndexInputWithBuffer(readStream); - long indexAddress = JNIService.loadIndex(indexInputWithBuffer, indexEntryContext.getParameters(), knnEngine); + final IndexInputWithBuffer indexInputWithBuffer = new IndexInputWithBuffer(readStream); + final long indexAddress = JNIService.loadIndex(indexInputWithBuffer, indexEntryContext.getParameters(), knnEngine); return createIndexAllocation(indexEntryContext, knnEngine, indexAddress, indexSizeKb, vectorFileName); } diff --git a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java index a90093da3..b5b2a5d22 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java @@ -26,6 +26,7 @@ import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; import org.opensearch.knn.index.codec.util.KNNCodecUtil; +import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper; import org.opensearch.knn.index.memory.NativeMemoryAllocation; import org.opensearch.knn.index.memory.NativeMemoryCacheManager; import org.opensearch.knn.index.memory.NativeMemoryEntryContext; @@ -274,6 +275,7 @@ private Map doANNSearch( } final String vectorIndexFileName = engineFiles.get(0); + final String cacheKey = NativeMemoryCacheKeyHelper.constructCacheKey(vectorIndexFileName, reader.getSegmentInfo().info); final KNNQueryResult[] results; KNNCounter.GRAPH_QUERY_REQUESTS.increment(); @@ -284,7 +286,7 @@ private Map doANNSearch( indexAllocation = nativeMemoryCacheManager.get( new NativeMemoryEntryContext.IndexEntryContext( reader.directory(), - vectorIndexFileName, + cacheKey, NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(), getParametersAtLoading( spaceType, diff --git a/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java b/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java index aaeee31e1..9b26e0537 100644 --- a/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java @@ -8,6 +8,11 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import lombok.SneakyThrows; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfo; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.Version; +import org.mockito.Mockito; import org.opensearch.knn.KNNSingleNodeTestCase; import org.opensearch.index.IndexService; import org.opensearch.index.engine.Engine; @@ -15,8 +20,7 @@ import org.opensearch.knn.index.memory.NativeMemoryCacheManager; import java.io.IOException; -import java.nio.file.Path; -import java.nio.file.Paths; +import java.util.Collections; import java.util.List; import java.util.Set; import java.util.concurrent.ExecutionException; @@ -113,11 +117,14 @@ public void testGetAllEngineFileContexts() throws IOException, ExecutionExceptio searcher = indexShard.acquireSearcher("test-hnsw-paths-2"); engineFileContexts = knnIndexShard.getAllEngineFileContexts(searcher.getIndexReader()); assertEquals(1, engineFileContexts.size()); - List paths = engineFileContexts.stream().map(KNNIndexShard.EngineFileContext::getIndexPath).collect(Collectors.toList()); + List paths = engineFileContexts.stream() + .map(KNNIndexShard.EngineFileContext::getVectorFileName) + .collect(Collectors.toList()); assertTrue(paths.get(0).contains("hnsw") || paths.get(0).contains("hnswc")); searcher.close(); } + @SneakyThrows public void testGetEngineFileContexts() { // Check that the correct engine paths are being returned by the KNNIndexShard String segmentName = "_0"; @@ -143,20 +150,34 @@ public void testGetEngineFileContexts() { KNNIndexShard knnIndexShard = new KNNIndexShard(null); - Path path = Paths.get(""); - List included = knnIndexShard.getEngineFileContexts( - files, + final Directory dummyDirectory = Mockito.mock(Directory.class); + final SegmentInfo segmentInfo = new SegmentInfo( + dummyDirectory, + Version.LATEST, + null, segmentName, + 0, + false, + false, + null, + Collections.emptyMap(), + null, + Collections.emptyMap(), + null + ); + segmentInfo.setFiles(files); + final SegmentCommitInfo segmentCommitInfo = new SegmentCommitInfo(segmentInfo, 0, 0, -1, 0, 0, null); + List included = knnIndexShard.getEngineFileContexts( + segmentCommitInfo, fieldName, fileExt, - path, spaceType, modelId, vectorDataType ); assertEquals(includedFileNames.size(), included.size()); - included.stream().map(KNNIndexShard.EngineFileContext::getIndexPath).forEach(o -> assertTrue(includedFileNames.contains(o))); + included.stream().map(KNNIndexShard.EngineFileContext::getVectorFileName).forEach(o -> assertTrue(includedFileNames.contains(o))); } @SneakyThrows diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java index 3d9969a1e..e6fcb643d 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java @@ -185,8 +185,6 @@ public void testMultiFieldsKnnIndex(Codec codec) throws Exception { writer.flush(); IndexReader reader = writer.getReader(); writer.close(); - ResourceWatcherService resourceWatcherService = createDisabledResourceWatcherService(); - NativeMemoryLoadStrategy.IndexLoadStrategy.initialize(resourceWatcherService); List hnswfiles = Arrays.stream(dir.listAll()).filter(x -> x.contains("hnsw")).collect(Collectors.toList()); // there should be 2 hnsw index files created. one for test_vector and one for my_vector @@ -213,7 +211,6 @@ public void testMultiFieldsKnnIndex(Codec codec) throws Exception { reader.close(); dir.close(); - resourceWatcherService.close(); NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close(); } @@ -298,8 +295,6 @@ public void testBuildFromModelTemplate(Codec codec) throws IOException, Executio // Make sure that search returns the correct results KNNWeight.initialize(modelDao); - ResourceWatcherService resourceWatcherService = createDisabledResourceWatcherService(); - NativeMemoryLoadStrategy.IndexLoadStrategy.initialize(resourceWatcherService); float[] query = { 10.0f, 10.0f, 10.0f }; IndexSearcher searcher = new IndexSearcher(reader); TopDocs topDocs = searcher.search(new KNNQuery(fieldName, query, 4, "dummy", (BitSetProducer) null), 10); @@ -311,7 +306,6 @@ public void testBuildFromModelTemplate(Codec codec) throws IOException, Executio reader.close(); dir.close(); - resourceWatcherService.close(); NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close(); } } @@ -422,8 +416,6 @@ public void testKnnVectorIndex( writer.addDocument(doc1); IndexReader reader1 = writer.getReader(); writer.close(); - ResourceWatcherService resourceWatcherService = createDisabledResourceWatcherService(); - NativeMemoryLoadStrategy.IndexLoadStrategy.initialize(resourceWatcherService); verify(perFieldKnnVectorsFormatSpy, atLeastOnce()).getKnnVectorsFormatForField(eq(FIELD_NAME_TWO)); verify(perFieldKnnVectorsFormatSpy, atLeastOnce()).getMaxDimensions(eq(FIELD_NAME_TWO)); @@ -444,7 +436,6 @@ public void testKnnVectorIndex( reader1.close(); dir.close(); - resourceWatcherService.close(); NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance().close(); } } diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryAllocationTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryAllocationTests.java index d9dfb019c..db6231adf 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryAllocationTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryAllocationTests.java @@ -26,8 +26,6 @@ import org.opensearch.knn.index.util.IndexUtil; import org.opensearch.knn.jni.JNICommons; import org.opensearch.knn.jni.JNIService; -import org.opensearch.watcher.FileWatcher; -import org.opensearch.watcher.WatcherHandle; import java.nio.file.Path; import java.util.Arrays; @@ -40,7 +38,6 @@ import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicReference; -import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.opensearch.knn.common.featureflags.KNNFeatureFlags.KNN_FORCE_EVICT_CACHE_ENABLED_SETTING; @@ -86,10 +83,6 @@ public void testIndexAllocation_close() throws InterruptedException { // Load index into memory long memoryAddress = JNIService.loadIndex(path, parameters, knnEngine); - @SuppressWarnings("unchecked") - WatcherHandle watcherHandle = (WatcherHandle) mock(WatcherHandle.class); - doNothing().when(watcherHandle).stop(); - ExecutorService executorService = Executors.newSingleThreadExecutor(); NativeMemoryAllocation.IndexAllocation indexAllocation = new NativeMemoryAllocation.IndexAllocation( executorService, @@ -97,8 +90,7 @@ public void testIndexAllocation_close() throws InterruptedException { IndexUtil.getFileSizeInKB(path), knnEngine, path, - "test", - watcherHandle + "test" ); indexAllocation.close(); @@ -147,10 +139,6 @@ public void testClose_whenBinaryFiass_thenSuccess() { // Load index into memory long memoryAddress = JNIService.loadIndex(path, parameters, knnEngine); - @SuppressWarnings("unchecked") - WatcherHandle watcherHandle = (WatcherHandle) mock(WatcherHandle.class); - doNothing().when(watcherHandle).stop(); - ExecutorService executorService = Executors.newSingleThreadExecutor(); NativeMemoryAllocation.IndexAllocation indexAllocation = new NativeMemoryAllocation.IndexAllocation( executorService, @@ -159,7 +147,6 @@ public void testClose_whenBinaryFiass_thenSuccess() { knnEngine, path, "test", - watcherHandle, null, true ); @@ -189,8 +176,7 @@ public void testIndexAllocation_getMemoryAddress() { 0, null, "test", - "test", - null + "test" ); assertEquals(memoryAddress, indexAllocation.getMemoryAddress()); @@ -205,8 +191,7 @@ public void testIndexAllocation_readLock() throws InterruptedException { 0, null, "test", - "test", - null + "test" ); int initialValue = 10; @@ -232,7 +217,6 @@ public void testIndexAllocation_readLock() throws InterruptedException { } public void testIndexAllocation_closeDefault() { - WatcherHandle watcherHandle = (WatcherHandle) mock(WatcherHandle.class); ExecutorService executorService = Executors.newFixedThreadPool(2); AtomicReference expectedException = new AtomicReference<>(); @@ -243,8 +227,7 @@ public void testIndexAllocation_closeDefault() { 0, null, "test", - "test", - watcherHandle + "test" ); executorService.submit(nonBlockingIndexAllocation::readLock); @@ -261,7 +244,6 @@ public void testIndexAllocation_closeDefault() { public void testIndexAllocation_closeBlocking() throws InterruptedException, ExecutionException { // Prepare mocking and a thread pool. - WatcherHandle watcherHandle = (WatcherHandle) mock(WatcherHandle.class); ExecutorService executorService = Executors.newSingleThreadExecutor(); // Enable `KNN_FORCE_EVICT_CACHE_ENABLED_SETTING` to force it to block other threads. @@ -273,8 +255,7 @@ public void testIndexAllocation_closeBlocking() throws InterruptedException, Exe 0, null, "test", - "test", - watcherHandle + "test" ); // Acquire a read lock @@ -309,8 +290,7 @@ public void testIndexAllocation_writeLock() throws InterruptedException { 0, null, "test", - "test", - null + "test" ); int initialValue = 10; @@ -342,8 +322,7 @@ public void testIndexAllocation_getSize() { size, null, "test", - "test", - null + "test" ); assertEquals(size, indexAllocation.getSizeInKB()); @@ -357,8 +336,7 @@ public void testIndexAllocation_getKnnEngine() { 0, knnEngine, "test", - "test", - null + "test" ); assertEquals(knnEngine, indexAllocation.getKnnEngine()); @@ -372,8 +350,7 @@ public void testIndexAllocation_getIndexPath() { 0, null, indexPath, - "test", - null + "test" ); assertEquals(indexPath, indexAllocation.getVectorFileName()); @@ -387,8 +364,7 @@ public void testIndexAllocation_getOsIndexName() { 0, null, "test", - osIndexName, - null + osIndexName ); assertEquals(osIndexName, indexAllocation.getOpenSearchIndexName()); diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java index 85eaf3322..4baf66cb4 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java @@ -118,8 +118,7 @@ public void testGetIndexSizeInKilobytes() throws ExecutionException, IOException indexEntryWeight, null, key, - indexName, - null + indexName ); NativeMemoryEntryContext.IndexEntryContext indexEntryContext = mock(NativeMemoryEntryContext.IndexEntryContext.class); @@ -152,8 +151,7 @@ public void testGetIndexSizeAsPercentage() throws ExecutionException, IOExceptio indexEntryWeight, null, key, - indexName, - null + indexName ); NativeMemoryEntryContext.IndexEntryContext indexEntryContext = mock(NativeMemoryEntryContext.IndexEntryContext.class); @@ -231,8 +229,7 @@ public void testGetIndexGraphCount() throws ExecutionException, IOException { indexEntryWeight, null, key1, - indexName1, - null + indexName1 ); NativeMemoryEntryContext.IndexEntryContext indexEntryContext = mock(NativeMemoryEntryContext.IndexEntryContext.class); @@ -247,8 +244,7 @@ public void testGetIndexGraphCount() throws ExecutionException, IOException { indexEntryWeight, null, key2, - indexName1, - null + indexName1 ); indexEntryContext = mock(NativeMemoryEntryContext.IndexEntryContext.class); @@ -263,8 +259,7 @@ public void testGetIndexGraphCount() throws ExecutionException, IOException { indexEntryWeight, null, key3, - indexName2, - null + indexName2 ); indexEntryContext = mock(NativeMemoryEntryContext.IndexEntryContext.class); @@ -408,8 +403,7 @@ public void testGetIndicesCacheStats() throws IOException, ExecutionException { size1, null, testKey1, - indexName1, - null + indexName1 ); NativeMemoryAllocation.IndexAllocation indexAllocation2 = new NativeMemoryAllocation.IndexAllocation( @@ -418,8 +412,7 @@ public void testGetIndicesCacheStats() throws IOException, ExecutionException { size2, null, testKey2, - indexName1, - null + indexName1 ); NativeMemoryAllocation.IndexAllocation indexAllocation3 = new NativeMemoryAllocation.IndexAllocation( @@ -428,8 +421,7 @@ public void testGetIndicesCacheStats() throws IOException, ExecutionException { size1, null, testKey3, - indexName2, - null + indexName2 ); NativeMemoryAllocation.IndexAllocation indexAllocation4 = new NativeMemoryAllocation.IndexAllocation( @@ -438,8 +430,7 @@ public void testGetIndicesCacheStats() throws IOException, ExecutionException { size2, null, testKey4, - indexName2, - null + indexName2 ); NativeMemoryEntryContext.IndexEntryContext indexEntryContext1 = mock(NativeMemoryEntryContext.IndexEntryContext.class); diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java index 72cab9a1b..8bda2f0f1 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java @@ -58,8 +58,7 @@ public void testIndexEntryContext_load() throws IOException { 10, KNNEngine.DEFAULT, "test-path", - "test-name", - null + "test-name" ); when(indexLoadStrategy.load(indexEntryContext)).thenReturn(indexAllocation); diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java index 373afddc7..e78a86166 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java @@ -28,7 +28,6 @@ import org.opensearch.knn.index.engine.KNNEngine; import org.opensearch.knn.training.FloatTrainingDataConsumer; import org.opensearch.knn.training.VectorReader; -import org.opensearch.watcher.ResourceWatcherService; import java.io.IOException; import java.nio.file.Path; @@ -38,10 +37,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.eq; import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; public class NativeMemoryLoadStrategyTests extends KNNTestCase { @@ -66,10 +62,6 @@ public void testIndexLoadStrategy_load() throws IOException { TestUtils.createIndex(ids, memoryAddress, dimension, path, parameters, knnEngine); // Setup mock resource manager - ResourceWatcherService resourceWatcherService = mock(ResourceWatcherService.class); - doReturn(null).when(resourceWatcherService).add(any()); - NativeMemoryLoadStrategy.IndexLoadStrategy.initialize(resourceWatcherService); - NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext( luceneDirectory, path, @@ -116,10 +108,6 @@ public void testLoad_whenFaissBinary_thenSuccess() throws IOException { TestUtils.createIndex(ids, memoryAddress, dimension, path, parameters, knnEngine); // Setup mock resource manager - ResourceWatcherService resourceWatcherService = mock(ResourceWatcherService.class); - doReturn(null).when(resourceWatcherService).add(any()); - NativeMemoryLoadStrategy.IndexLoadStrategy.initialize(resourceWatcherService); - NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext( luceneDirectory, path,