[Backport 2.x] Preloads .vec and .vex files

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Enhancements
 ### Bug Fixes 
 * Add DocValuesProducers for releasing memory when close index [#1946](https://github.com/opensearch-project/k-NN/pull/1946)
+* Prelaods vec and vex files to address regression in force merge latencies [#2186](https://github.com/opensearch-project/k-NN/pull/2186)
 ### Infrastructure
 ### Documentation
 ### Maintenance

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

@@ -8,7 +8,6 @@
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.index.SpaceType;
 
-import java.util.Collections;
 import java.util.List;
 
 /**
@@ -137,6 +136,6 @@ KNNLibraryIndexingContext getKNNLibraryIndexingContext(
      * @return list of file extensions that will be read/write with mmap
      */
     default List<String> mmapFileExtensions() {
-        return Collections.emptyList();
+        return List.of("vec", "vex");
     }
 }

src/main/java/org/opensearch/knn/index/engine/lucene/Lucene.java

@@ -15,7 +15,6 @@
 import org.opensearch.knn.index.engine.MethodResolver;
 import org.opensearch.knn.index.engine.ResolvedMethodContext;
 
-import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
 
@@ -89,11 +88,6 @@ public Float scoreToRadialThreshold(Float score, SpaceType spaceType) {
         return score;
     }
 
-    @Override
-    public List<String> mmapFileExtensions() {
-        return List.of("vec", "vex");
-    }
-
     @Override
     public ResolvedMethodContext resolveMethod(
         KNNMethodContext knnMethodContext,

src/main/java/org/opensearch/knn/plugin/KNNPlugin.java

@@ -11,6 +11,7 @@
 import org.opensearch.core.action.ActionResponse;
 import org.opensearch.index.codec.CodecServiceFactory;
 import org.opensearch.index.engine.EngineFactory;
+import org.opensearch.index.shard.IndexSettingProvider;
 import org.opensearch.indices.SystemIndexDescriptor;
 import org.opensearch.knn.index.KNNCircuitBreaker;
 import org.opensearch.knn.plugin.search.KNNConcurrentSearchRequestDecider;
@@ -368,14 +369,39 @@ public Settings additionalSettings() {
         // that are set here.
         final List<String> engineSettings = Arrays.stream(KNNEngine.values())
             .flatMap(engine -> engine.mmapFileExtensions().stream())
+            .distinct()
             .collect(Collectors.toList());
+
         final List<String> combinedSettings = Stream.concat(
             IndexModule.INDEX_STORE_HYBRID_MMAP_EXTENSIONS.getDefault(Settings.EMPTY).stream(),
             engineSettings.stream()
         ).collect(Collectors.toList());
+
         return Settings.builder().putList(IndexModule.INDEX_STORE_HYBRID_MMAP_EXTENSIONS.getKey(), combinedSettings).build();
     }
 
+    @Override
+    public Collection<IndexSettingProvider> getAdditionalIndexSettingProviders() {
+        IndexSettingProvider preloadMmapFiles = new IndexSettingProvider() {
+            @Override
+            public Settings getAdditionalIndexSettings(String indexName, boolean isDataStreamIndex, Settings templateAndRequestSettings) {
+
+                if (templateAndRequestSettings.getAsBoolean(KNNSettings.KNN_INDEX, Boolean.FALSE)) {
+                    final List<String> mmapFileExtensions = Arrays.stream(KNNEngine.values())
+                        .flatMap(engine -> engine.mmapFileExtensions().stream())
+                        .distinct()
+                        .collect(Collectors.toList());
+
+                    return Settings.builder().putList(IndexModule.INDEX_STORE_PRE_LOAD_SETTING.getKey(), mmapFileExtensions).build();
+                }
+
+                return Settings.EMPTY;
+            }
+        };
+
+        return List.of(preloadMmapFiles);
+    }
+
     @Override
     public Optional<ConcurrentSearchRequestDecider.Factory> getConcurrentSearchRequestDeciderFactory() {
         return Optional.of(new KNNConcurrentSearchRequestDecider.Factory());

src/test/java/org/opensearch/knn/index/engine/faiss/FaissTests.java

@@ -23,6 +23,7 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 
@@ -367,4 +368,12 @@ public void testMethodAsMapBuilder() throws IOException {
         assertEquals(expectedKNNMethodContext.getVectorValidator(), actualKNNLibraryIndexingContext.getVectorValidator());
     }
 
+    public void testMmapFileExtensions() {
+        final List<String> mMapExtensions = Faiss.INSTANCE.mmapFileExtensions();
+        assertNotNull(mMapExtensions);
+        final List<String> expectedSettings = List.of("vex", "vec");
+        assertTrue(expectedSettings.containsAll(mMapExtensions));
+        assertTrue(mMapExtensions.containsAll(expectedSettings));
+    }
+
 }

src/test/java/org/opensearch/knn/index/engine/nmslib/NMSLibTests.java

@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.engine.nmslib;
+
+import org.opensearch.knn.KNNTestCase;
+
+import java.util.List;
+
+public class NMSLibTests extends KNNTestCase {
+
+    public void testMmapFileExtensions() {
+        final List<String> mmapExtensions = Nmslib.INSTANCE.mmapFileExtensions();
+        assertNotNull(mmapExtensions);
+        final List<String> expectedSettings = List.of("vex", "vec");
+        assertTrue(expectedSettings.containsAll(mmapExtensions));
+        assertTrue(mmapExtensions.containsAll(expectedSettings));
+    }
+}

src/test/java/org/opensearch/knn/plugin/KNNPluginTests.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.plugin;
+
+import org.opensearch.common.settings.Settings;
+import org.opensearch.knn.index.KNNSettings;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.io.IOException;
+import java.util.List;
+
+public class KNNPluginTests extends OpenSearchTestCase {
+
+    public void testKNNPlugin_additionalSettings() throws IOException {
+        try (KNNPlugin knnPlugin = new KNNPlugin()) {
+            Settings additionalSettings = knnPlugin.additionalSettings();
+            List<String> expectedHybrid = List.of("nvd", "dvd", "tim", "tip", "dim", "kdd", "kdi", "cfs", "doc", "vec", "vex");
+
+            Settings settings = Settings.builder().putList("index.store.hybrid.mmap.extensions", expectedHybrid).build();
+            assertEquals(settings, additionalSettings);
+        }
+    }
+
+    public void testKNNPlugin_additionalIndexProviderSettings() throws IOException {
+        try (KNNPlugin knnPlugin = new KNNPlugin()) {
+            Settings additionalSettings = knnPlugin.getAdditionalIndexSettingProviders()
+                .iterator()
+                .next()
+                .getAdditionalIndexSettings("index", false, Settings.builder().put(KNNSettings.KNN_INDEX, Boolean.TRUE).build());
+
+            Settings settings = Settings.builder().putList("index.store.preload", List.of("vec", "vex")).build();
+            assertEquals(settings, additionalSettings);
+
+            additionalSettings = knnPlugin.getAdditionalIndexSettingProviders()
+                .iterator()
+                .next()
+                .getAdditionalIndexSettings("index", false, Settings.builder().build());
+
+            assertEquals(Settings.EMPTY, additionalSettings);
+        }
+    }
+}