diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java index b8aecc5a5..96fd71121 100644 --- a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java +++ b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryCacheManager.java @@ -328,6 +328,7 @@ public NativeMemoryAllocation get(NativeMemoryEntryContext nativeMemoryEntryC // Cache Miss // Evict before put + nativeMemoryEntryContext.preload(); synchronized (this) { if (getCacheSizeInKilobytes() + nativeMemoryEntryContext.calculateSizeInKB() >= maxWeight) { Iterator lruIterator = accessRecencyQueue.iterator(); 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 0af13fb46..25da3f2a7 100644 --- a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java +++ b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java @@ -12,12 +12,16 @@ package org.opensearch.knn.index.memory; import lombok.Getter; +import lombok.Setter; import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.Nullable; import org.opensearch.knn.index.codec.util.NativeMemoryCacheKeyHelper; import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import org.opensearch.knn.index.VectorDataType; +import org.opensearch.knn.index.store.IndexInputWithBuffer; import java.io.IOException; import java.util.Map; @@ -26,7 +30,7 @@ /** * Encapsulates all information needed to load a component into native memory. */ -public abstract class NativeMemoryEntryContext { +public abstract class NativeMemoryEntryContext implements AutoCloseable { protected final String key; @@ -55,6 +59,18 @@ public String getKey() { */ public abstract Integer calculateSizeInKB(); + /** + * Preloads the entry by opening the indexInput + */ + + public abstract void preload(); + + /** + * Provides the capability to close the closable objects in the {@link NativeMemoryEntryContext} + */ + @Override + public void close() {} + /** * Loads entry into memory. * @@ -75,6 +91,18 @@ public static class IndexEntryContext extends NativeMemoryEntryContext { @@ -192,6 +265,11 @@ public Integer calculateSizeInKB() { return size; } + @Override + public void preload() { + return; + } + @Override public NativeMemoryAllocation.TrainingDataAllocation load() { return trainingLoadStrategy.load(this); @@ -278,6 +356,11 @@ public Integer calculateSizeInKB() { return size; } + @Override + public void preload() { + return; + } + @Override public NativeMemoryAllocation.AnonymousAllocation load() throws IOException { return loadStrategy.load(this); 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 8cbdb4fd7..9e3097c84 100644 --- a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java +++ b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java @@ -13,12 +13,9 @@ import lombok.extern.log4j.Log4j2; import org.apache.lucene.store.Directory; -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; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.index.engine.KNNEngine; @@ -88,10 +85,16 @@ public NativeMemoryAllocation.IndexAllocation load(NativeMemoryEntryContext.Inde 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)) { - final IndexInputWithBuffer indexInputWithBuffer = new IndexInputWithBuffer(readStream); - final long indexAddress = JNIService.loadIndex(indexInputWithBuffer, indexEntryContext.getParameters(), knnEngine); - + // preload takes care of opening the indexInput file + if (!indexEntryContext.isPreloaded()) { + throw new IllegalStateException("Index [" + indexEntryContext.getOpenSearchIndexName() + "] is not preloaded"); + } + try (indexEntryContext) { + final long indexAddress = JNIService.loadIndex( + indexEntryContext.indexInputWithBuffer, + indexEntryContext.getParameters(), + knnEngine + ); return createIndexAllocation(indexEntryContext, knnEngine, indexAddress, indexSizeKb, vectorFileName); } } 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 5fe41c88c..9c7570f37 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryCacheManagerTests.java @@ -542,6 +542,9 @@ public Integer calculateSizeInKB() { return size; } + @Override + public void preload() {} + @Override public TestNativeMemoryAllocation load() throws IOException { return new TestNativeMemoryAllocation(size, memoryAddress); 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 5379abc74..19054b795 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java @@ -61,6 +61,8 @@ public void testIndexEntryContext_load() throws IOException { when(indexLoadStrategy.load(indexEntryContext)).thenReturn(indexAllocation); + // since we are returning mock instance, set indexEntryContext.isPreloaded to true. + indexEntryContext.setPreloaded(true); assertEquals(indexAllocation, indexEntryContext.load()); } @@ -292,6 +294,11 @@ public Integer calculateSizeInKB() { return size; } + @Override + public void preload() { + return; + } + @Override public TestNativeMemoryAllocation load() throws IOException { return null; 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 735974bd1..7da68f9ee 100644 --- a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java +++ b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java @@ -69,8 +69,7 @@ public void testIndexLoadStrategy_load() throws IOException { ); // Load - NativeMemoryAllocation.IndexAllocation indexAllocation = NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance() - .load(indexEntryContext); + NativeMemoryAllocation.IndexAllocation indexAllocation = indexEntryContext.load(); // Confirm that the file was loaded by querying float[] query = new float[dimension]; @@ -115,8 +114,7 @@ public void testLoad_whenFaissBinary_thenSuccess() throws IOException { ); // Load - NativeMemoryAllocation.IndexAllocation indexAllocation = NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance() - .load(indexEntryContext); + NativeMemoryAllocation.IndexAllocation indexAllocation = indexEntryContext.load(); // Verify assertTrue(indexAllocation.isBinaryIndex());