Fix unit tests and integration tests.

Signed-off-by: Dooyong Kim <[email protected]>
opensearch-project · Oct 18, 2024 · 88509d4 · 88509d4
1 parent c380d84
commit 88509d4
Show file tree

Hide file tree

Showing 10 changed files with 51 additions and 43 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.0](https://github.com/opensearch-project/k-NN/compare/2.x...HEAD)
 ### Features
 ### Enhancements
+* Remove FSDirectory dependency and deprecated FileWatcher [#2182](https://github.com/opensearch-project/k-NN/pull/2182)
 ### Bug Fixes
 ### Infrastructure
 * Removed JDK 11 and 17 version from CI runs [#1921](https://github.com/opensearch-project/k-NN/pull/1921)

diff --git a/src/main/java/org/opensearch/knn/index/KNNIndexShard.java b/src/main/java/org/opensearch/knn/index/KNNIndexShard.java
@@ -213,9 +213,9 @@ List<EngineFileContext> getEngineFileContexts(
         String modelId,
         VectorDataType vectorDataType
     ) throws IOException {
-        // Ex: _0
+        // Ex: 0_
         final String prefix = buildEngineFilePrefix(segmentCommitInfo.info.name);
-        // Ex: my_field.faiss
+        // Ex: _my_field.faiss
         final String suffix = buildEngineFileSuffix(fieldName, fileExtension);
         return segmentCommitInfo.files()
             .stream()

diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java
@@ -562,8 +562,10 @@ public static boolean isFaissAVX2Disabled() {
 
     public static boolean isFaissAVX512Disabled() {
         return Booleans.parseBoolean(
-            KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_DISABLED).toString(),
-            KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE
+            Objects.requireNonNullElse(
+                KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX512_DISABLED),
+                KNN_DEFAULT_FAISS_AVX512_DISABLED_VALUE
+            ).toString()
         );
     }
 

diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryEntryContext.java
@@ -15,14 +15,13 @@
 import org.apache.lucene.store.Directory;
 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.util.IndexUtil;
 import org.opensearch.knn.index.VectorDataType;
 
 import java.io.IOException;
 import java.util.Map;
 import java.util.UUID;
-import java.util.function.Function;
 
 /**
  * Encapsulates all information needed to load a component into native memory.
@@ -123,25 +122,19 @@ public IndexEntryContext(
 
         @Override
         public Integer calculateSizeInKB() {
-            return IndexSizeCalculator.INSTANCE.apply(this);
+            final String indexFileName = NativeMemoryCacheKeyHelper.extractVectorIndexFileName(key);
+            try {
+                final long fileLength = directory.fileLength(indexFileName);
+                return (int) (fileLength / 1024L);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
         }
 
         @Override
         public NativeMemoryAllocation.IndexAllocation load() throws IOException {
             return indexLoadStrategy.load(this);
         }
-
-        private static class IndexSizeCalculator implements Function<IndexEntryContext, Integer> {
-
-            static IndexSizeCalculator INSTANCE = new IndexSizeCalculator();
-
-            IndexSizeCalculator() {}
-
-            @Override
-            public Integer apply(IndexEntryContext indexEntryContext) {
-                return IndexUtil.getFileSizeInKB(indexEntryContext.getKey());
-            }
-        }
     }
 
     public static class TrainingDataEntryContext extends NativeMemoryEntryContext<NativeMemoryAllocation.TrainingDataAllocation> {

diff --git a/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java b/src/main/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategy.java
@@ -78,7 +78,7 @@ public NativeMemoryAllocation.IndexAllocation load(NativeMemoryEntryContext.Inde
             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."
+                    "Invalid cache key was given. The key [" + cacheKey + "] does not contain the corresponding vector file name."
                 );
             }
 

diff --git a/src/main/java/org/opensearch/knn/jni/JNIService.java b/src/main/java/org/opensearch/knn/jni/JNIService.java
@@ -135,7 +135,6 @@ public static void createIndex(
         Map<String, Object> parameters,
         KNNEngine knnEngine
     ) {
-
         if (KNNEngine.NMSLIB == knnEngine) {
             NmslibService.createIndex(ids, vectorsAddress, dim, indexPath, parameters);
             return;

diff --git a/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java b/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java
@@ -11,6 +11,7 @@
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.StringHelper;
 import org.apache.lucene.util.Version;
 import org.mockito.Mockito;
 import org.opensearch.knn.KNNSingleNodeTestCase;
@@ -20,7 +21,9 @@
 import org.opensearch.knn.index.memory.NativeMemoryCacheManager;
 
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
@@ -161,11 +164,17 @@ public void testGetEngineFileContexts() {
             false,
             null,
             Collections.emptyMap(),
-            null,
+            new byte[StringHelper.ID_LENGTH],
             Collections.emptyMap(),
             null
         );
-        segmentInfo.setFiles(files);
+        // Inject 'files' into the segment info instance.
+        // Since SegmentInfo class does trim out its given file list, for example removing segment name from a file name etc,
+        // we can't just use 'setFiles' api to assign the file list. Which will lead this unit test to be fail.
+        final Field setFilesPrivateField = SegmentInfo.class.getDeclaredField("setFiles");
+        setFilesPrivateField.setAccessible(true);
+        setFilesPrivateField.set(segmentInfo, new HashSet<>(files));
+
         final SegmentCommitInfo segmentCommitInfo = new SegmentCommitInfo(segmentInfo, 0, 0, -1, 0, 0, null);
         List<KNNIndexShard.EngineFileContext> included = knnIndexShard.getEngineFileContexts(
             segmentCommitInfo,

diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryEntryContextTests.java
@@ -13,23 +13,21 @@
 
 import com.google.common.collect.ImmutableMap;
 import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.store.MMapDirectory;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.knn.KNNTestCase;
+import org.opensearch.knn.TestUtils;
 import org.opensearch.knn.index.engine.qframe.QuantizationConfig;
-import org.opensearch.knn.index.util.IndexUtil;
 import org.opensearch.knn.index.VectorDataType;
 import org.opensearch.knn.index.engine.KNNEngine;
 
-import java.io.BufferedOutputStream;
 import java.io.IOException;
-import java.io.OutputStream;
-import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Map;
 
-import static java.nio.file.StandardOpenOption.APPEND;
-import static java.nio.file.StandardOpenOption.CREATE;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -46,7 +44,7 @@ public void testIndexEntryContext_load() throws IOException {
         NativeMemoryLoadStrategy.IndexLoadStrategy indexLoadStrategy = mock(NativeMemoryLoadStrategy.IndexLoadStrategy.class);
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
             (Directory) null,
-            "test",
+            TestUtils.createFakeNativeMamoryCacheKey("test"),
             indexLoadStrategy,
             null,
             "test"
@@ -68,36 +66,37 @@ public void testIndexEntryContext_load() throws IOException {
 
     public void testIndexEntryContext_calculateSize() throws IOException {
         // Create a file and write random bytes to it
-        Path tmpFile = createTempFile();
+        final Path tmpDirectory = createTempDir();
+        final Directory directory = new MMapDirectory(tmpDirectory);
+        final String indexFileName = "test.faiss";
         byte[] data = new byte[1024 * 3];
         Arrays.fill(data, (byte) 'c');
 
-        try (OutputStream out = new BufferedOutputStream(Files.newOutputStream(tmpFile, CREATE, APPEND))) {
-            out.write(data, 0, data.length);
-        } catch (IOException x) {
-            fail("Failed to write to file");
+        try (IndexOutput output = directory.createOutput(indexFileName, IOContext.DEFAULT)) {
+            output.writeBytes(data, data.length);
         }
 
         // Get the expected size of this function
-        int expectedSize = IndexUtil.getFileSizeInKB(tmpFile.toAbsolutePath().toString());
+        final long expectedSizeBytes = directory.fileLength(indexFileName);
+        final long expectedSizeKb = expectedSizeBytes / 1024L;
 
         // Check that the indexEntryContext will return the same thing
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
-            (Directory) null,
-            tmpFile.toAbsolutePath().toString(),
+            directory,
+            TestUtils.createFakeNativeMamoryCacheKey(indexFileName),
             null,
             null,
             "test"
         );
 
-        assertEquals(expectedSize, indexEntryContext.calculateSizeInKB().longValue());
+        assertEquals(expectedSizeKb, indexEntryContext.calculateSizeInKB().longValue());
     }
 
     public void testIndexEntryContext_getOpenSearchIndexName() {
         String openSearchIndexName = "test-index";
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
             (Directory) null,
-            "test",
+            TestUtils.createFakeNativeMamoryCacheKey("test"),
             null,
             null,
             openSearchIndexName
@@ -110,7 +109,7 @@ public void testIndexEntryContext_getParameters() {
         Map<String, Object> parameters = ImmutableMap.of("test-1", 10);
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
             (Directory) null,
-            "test",
+            TestUtils.createFakeNativeMamoryCacheKey("test"),
             null,
             parameters,
             "test"

diff --git a/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java b/src/test/java/org/opensearch/knn/index/memory/NativeMemoryLoadStrategyTests.java
@@ -64,7 +64,7 @@ public void testIndexLoadStrategy_load() throws IOException {
         // Setup mock resource manager
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
             luceneDirectory,
-            path,
+            TestUtils.createFakeNativeMamoryCacheKey(indexName),
             NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
             parameters,
             "test"
@@ -110,7 +110,7 @@ public void testLoad_whenFaissBinary_thenSuccess() throws IOException {
         // Setup mock resource manager
         NativeMemoryEntryContext.IndexEntryContext indexEntryContext = new NativeMemoryEntryContext.IndexEntryContext(
             luceneDirectory,
-            path,
+            TestUtils.createFakeNativeMamoryCacheKey(indexName),
             NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
             parameters,
             "test"

diff --git a/src/testFixtures/java/org/opensearch/knn/TestUtils.java b/src/testFixtures/java/org/opensearch/knn/TestUtils.java
@@ -24,6 +24,7 @@
 import org.opensearch.knn.jni.JNIService;
 import org.opensearch.knn.plugin.script.KNNScoringUtil;
 
+import java.util.Base64;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Random;
@@ -186,6 +187,10 @@ public static float[][] getIndexVectors(int docCount, int dimensions, boolean is
         }
     }
 
+    public static String createFakeNativeMamoryCacheKey(final String fileName) {
+        return fileName + "@" + Base64.getEncoder().encodeToString(fileName.getBytes());
+    }
+
     /*
      * Recall is the number of relevant documents retrieved by a search divided by the total number of existing relevant documents.
      * We are similarly calculating recall by verifying number of relevant documents obtained in the search results by comparing with