[LuceneOnFaiss Part-8] Added index.knn.memory_optimized_search index setting.

src/main/java/org/opensearch/knn/index/KNNSettings.java

@@ -5,12 +5,14 @@
 
 package org.opensearch.knn.index;
 
+import lombok.NonNull;
 import lombok.extern.log4j.Log4j2;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.OpenSearchParseException;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest;
 import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsResponse;
+import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec;
 import org.opensearch.transport.client.Client;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.service.ClusterService;
@@ -96,6 +98,11 @@ public class KNNSettings {
     public static final String KNN_DERIVED_SOURCE_ENABLED = "index.knn.derived_source.enabled";
     public static final String KNN_INDEX_REMOTE_VECTOR_BUILD = "index.knn.remote_index_build.enabled";
     public static final String KNN_REMOTE_VECTOR_REPO = "knn.remote_index_build.vector_repo";
+    /**
+     * For more details on supported engines, refer to {@link MemoryOptimizedSearchSupportSpec}
+     */
+    public static final String MEMORY_OPTIMIZED_KNN_SEARCH_MODE = "index.knn.memory_optimized_search";
+    public static final boolean DEFAULT_MEMORY_OPTIMIZED_KNN_SEARCH_MODE = false;
 
     /**
      * Default setting values
@@ -121,9 +128,9 @@ public class KNNSettings {
 
     public static final Integer ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_DEFAULT_VALUE = -1;
     public static final Integer KNN_DEFAULT_QUANTIZATION_STATE_CACHE_SIZE_LIMIT_PERCENTAGE = 5; // By default, set aside 5% of the JVM for
-                                                                                                // the limit
+    // the limit
     public static final Integer KNN_MAX_QUANTIZATION_STATE_CACHE_SIZE_LIMIT_PERCENTAGE = 10; // Quantization state cache limit cannot exceed
-                                                                                             // 10% of the JVM heap
+    // 10% of the JVM heap
     public static final Integer KNN_DEFAULT_QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES = 60;
     public static final boolean KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED_VALUE = false;
 
@@ -285,6 +292,12 @@ public class KNNSettings {
      */
     public static final Setting<Boolean> IS_KNN_INDEX_SETTING = Setting.boolSetting(KNN_INDEX, false, IndexScope, Final);
 
+    public static final Setting<Boolean> MEMORY_OPTIMIZED_KNN_SEARCH_MODE_SETTING = Setting.boolSetting(
+        MEMORY_OPTIMIZED_KNN_SEARCH_MODE,
+        false,
+        IndexScope
+    );
+
     /**
      * index_thread_quantity - the parameter specifies how many threads the nms library should use to create the graph.
      * By default, the nms library sets this value to NUM_CORES. However, because ES can spawn NUM_CORES threads for
@@ -577,7 +590,8 @@ public List<Setting<?>> getSettings() {
             KNN_DISK_VECTOR_SHARD_LEVEL_RESCORING_DISABLED_SETTING,
             KNN_DERIVED_SOURCE_ENABLED_SETTING,
             KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING,
-            KNN_REMOTE_VECTOR_REPO_SETTING
+            KNN_REMOTE_VECTOR_REPO_SETTING,
+            MEMORY_OPTIMIZED_KNN_SEARCH_MODE_SETTING
         );
         return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
             .collect(Collectors.toList());
@@ -734,12 +748,25 @@ public static int getEfSearchParam(String index) {
             );
     }
 
+    /**
+     * Return whether memory optimized search enabled for the given index.
+     *
+     * @param indexName The name of target index to test whether if it is on.
+     * @return True if memory optimized search is enabled, otherwise False.
+     */
+    public static boolean isMemoryOptimizedKnnSearchModeEnabled(@NonNull final String indexName) {
+        return KNNSettings.state().clusterService.state()
+            .getMetadata()
+            .index(indexName)
+            .getSettings()
+            .getAsBoolean(MEMORY_OPTIMIZED_KNN_SEARCH_MODE, DEFAULT_MEMORY_OPTIMIZED_KNN_SEARCH_MODE);
+    }
+
     public void setClusterService(ClusterService clusterService) {
         this.clusterService = clusterService;
     }
 
     static class SpaceTypeValidator implements Setting.Validator<String> {
-
         @Override
         public void validate(String value) {
             try {

src/main/java/org/opensearch/knn/index/codec/BasePerFieldKnnVectorsFormat.java

@@ -155,11 +155,12 @@ public KnnVectorsFormat getKnnVectorsFormatForField(final String field) {
     private NativeEngines990KnnVectorsFormat nativeEngineVectorsFormat() {
         // mapperService is already checked for null or valid instance type at caller, hence we don't need
         // addition isPresent check here.
-        int approximateThreshold = getApproximateThresholdValue();
+        final int approximateThreshold = getApproximateThresholdValue();
         return new NativeEngines990KnnVectorsFormat(
             new Lucene99FlatVectorsFormat(FlatVectorScorerUtil.getLucene99FlatVectorsScorer()),
             approximateThreshold,
-            nativeIndexBuildStrategyFactory
+            nativeIndexBuildStrategyFactory,
+            mapperService
         );
     }
 

src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java

@@ -11,6 +11,7 @@
 
 package org.opensearch.knn.index.codec.KNN990Codec;
 
+import lombok.extern.log4j.Log4j2;
 import org.apache.lucene.codecs.KnnVectorsFormat;
 import org.apache.lucene.codecs.KnnVectorsReader;
 import org.apache.lucene.codecs.KnnVectorsWriter;
@@ -19,22 +20,31 @@
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsFormat;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.index.IndexSettings;
+import org.opensearch.index.mapper.MapperService;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategyFactory;
 import org.opensearch.knn.index.engine.KNNEngine;
 
 import java.io.IOException;
+import java.util.Optional;
+
+import static org.opensearch.knn.index.KNNSettings.DEFAULT_MEMORY_OPTIMIZED_KNN_SEARCH_MODE;
+import static org.opensearch.knn.index.KNNSettings.MEMORY_OPTIMIZED_KNN_SEARCH_MODE;
 
 /**
  * This is a Vector format that will be used for Native engines like Faiss and Nmslib for reading and writing vector
  * related data structures.
  */
+@Log4j2
 public class NativeEngines990KnnVectorsFormat extends KnnVectorsFormat {
     /** The format for storing, reading, merging vectors on disk */
     private static FlatVectorsFormat flatVectorsFormat;
     private static final String FORMAT_NAME = "NativeEngines990KnnVectorsFormat";
     private static int approximateThreshold;
     private final NativeIndexBuildStrategyFactory nativeIndexBuildStrategyFactory;
+    private final boolean memoryOptimizedSearchEnabled;
 
     public NativeEngines990KnnVectorsFormat() {
         this(new Lucene99FlatVectorsFormat(new DefaultFlatVectorScorer()));
@@ -49,18 +59,20 @@ public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsForma
     }
 
     public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsFormat, int approximateThreshold) {
-        this(flatVectorsFormat, approximateThreshold, new NativeIndexBuildStrategyFactory());
+        this(flatVectorsFormat, approximateThreshold, new NativeIndexBuildStrategyFactory(), Optional.empty());
     }
 
     public NativeEngines990KnnVectorsFormat(
         final FlatVectorsFormat flatVectorsFormat,
         int approximateThreshold,
-        final NativeIndexBuildStrategyFactory nativeIndexBuildStrategyFactory
+        final NativeIndexBuildStrategyFactory nativeIndexBuildStrategyFactory,
+        final Optional<MapperService> mapperService
     ) {
         super(FORMAT_NAME);
         NativeEngines990KnnVectorsFormat.flatVectorsFormat = flatVectorsFormat;
         NativeEngines990KnnVectorsFormat.approximateThreshold = approximateThreshold;
         this.nativeIndexBuildStrategyFactory = nativeIndexBuildStrategyFactory;
+        this.memoryOptimizedSearchEnabled = isMemoryOptimizedSearchSupported(mapperService);
     }
 
     /**
@@ -85,7 +97,7 @@ public KnnVectorsWriter fieldsWriter(final SegmentWriteState state) throws IOExc
      */
     @Override
     public KnnVectorsReader fieldsReader(final SegmentReadState state) throws IOException {
-        return new NativeEngines990KnnVectorsReader(state, flatVectorsFormat.fieldsReader(state));
+        return new NativeEngines990KnnVectorsReader(state, flatVectorsFormat.fieldsReader(state), memoryOptimizedSearchEnabled);
     }
 
     /**
@@ -107,4 +119,21 @@ public String toString() {
             + approximateThreshold
             + ")";
     }
+
+    private static boolean isMemoryOptimizedSearchSupported(final Optional<MapperService> mapperService) {
+        if (mapperService.isPresent()) {
+            final IndexSettings indexSettings = mapperService.get().getIndexSettings();
+            if (indexSettings != null) {
+                final Settings settings = indexSettings.getSettings();
+                if (settings != null) {
+                    try {
+                        return settings.getAsBoolean(MEMORY_OPTIMIZED_KNN_SEARCH_MODE, DEFAULT_MEMORY_OPTIMIZED_KNN_SEARCH_MODE);
+                    } catch (Throwable th) {
+                        log.error("Failed to get a bool flag of [{}] from settings.", MEMORY_OPTIMIZED_KNN_SEARCH_MODE);
+                    }
+                }
+            }
+        }
+        return false;
+    }
 }