[BUG FIX] Add space type default and ef search parameter in warmup

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

@@ -11,17 +11,23 @@
 
 package org.opensearch.knn.index;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.MappingMetadata;
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.common.KNNConstants;
+import org.opensearch.knn.index.util.KNNEngine;
 import org.opensearch.knn.indices.ModelDao;
 import org.opensearch.knn.indices.ModelMetadata;
 
 import java.io.File;
+import java.util.Collections;
 import java.util.Map;
 
 import static org.opensearch.knn.common.KNNConstants.BYTES_PER_KILOBYTES;
+import static org.opensearch.knn.common.KNNConstants.HNSW_ALGO_EF_SEARCH;
+import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
 
 public class IndexUtil {
 
@@ -156,4 +162,26 @@ public static ValidationException validateKnnField(IndexMetadata indexMetadata,
 
         return null;
     }
+
+    /**
+     * Gets the load time parameters for a given engine.
+     *
+     * @param spaceType Space for this particular segment
+     * @param knnEngine Engine used for the native library indices being loaded in
+     * @param indexName Name of OpenSearch index that the segment files belong to
+     * @return load parameters that will be passed to the JNI.
+     */
+    public static Map<String, Object> getParametersAtLoading(SpaceType spaceType, KNNEngine knnEngine, String indexName) {
+        Map<String, Object> loadParameters = Maps.newHashMap(ImmutableMap.of(
+                SPACE_TYPE, spaceType.getValue()
+        ));
+
+        // For nmslib, we need to add the dynamic ef_search parameter that needs to be passed in when the
+        // hnsw graphs are loaded into memory
+        if (KNNEngine.NMSLIB.equals(knnEngine)) {
+            loadParameters.put(HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(indexName));
+        }
+
+        return Collections.unmodifiableMap(loadParameters);
+    }
 }

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

@@ -5,7 +5,6 @@
 
 package org.opensearch.knn.index;
 
-import com.google.common.collect.ImmutableMap;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -31,6 +30,7 @@
 import java.util.stream.Collectors;
 
 import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
+import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
 import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileName;
 
 /**
@@ -86,7 +86,7 @@ public void warmup() throws IOException {
                             new NativeMemoryEntryContext.IndexEntryContext(
                                     key,
                                     NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
-                                    ImmutableMap.of(SPACE_TYPE, value.getValue()),
+                                    getParametersAtLoading(value, KNNEngine.getEngineNameFromPath(key), getIndexName()),
                                     getIndexName()
                             ), true);
                 } catch (ExecutionException ex) {
@@ -122,7 +122,10 @@ private Map<String, SpaceType> getEnginePaths(IndexReader indexReader, KNNEngine
 
             for (FieldInfo fieldInfo : reader.getFieldInfos()) {
                 if (fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {
-                    SpaceType spaceType = SpaceType.getSpace(fieldInfo.attributes().get(SPACE_TYPE));
+                    // Space Type will not be present on ES versions 7.1 and 7.4 because the only available space type
+                    // was L2. So, if Space Type is not present, just fall back to L2
+                    String spaceTypeName = fieldInfo.attributes().getOrDefault(SPACE_TYPE, SpaceType.L2.getValue());
+                    SpaceType spaceType = SpaceType.getSpace(spaceTypeName);
                     String engineFileName = buildEngineFileName(reader.getSegmentInfo().info.name,
                             knnEngine.getLatestBuildVersion(), fieldInfo.name, fileExtension);
 

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

@@ -46,6 +46,7 @@
 import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE;
 import static org.opensearch.knn.common.KNNConstants.MODEL_ID;
 import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
+import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
 import static org.opensearch.knn.plugin.stats.KNNCounter.GRAPH_QUERY_ERRORS;
 
 /**
@@ -135,24 +136,13 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
             KNNCounter.GRAPH_QUERY_REQUESTS.increment();
 
             // We need to first get index allocation
-            NativeMemoryAllocation indexAllocation = null;
-
-            Map<String, Object> loadParameters = new HashMap<String, Object>() {{
-                put(SPACE_TYPE, spaceType.getValue());
-            }};
-
-            // For nmslib, we set efSearch from an index setting. This has the advantage of being able to dynamically
-            // update this value, which we cannot do at the moment for mapping parameters.
-            if (knnEngine.equals(KNNEngine.NMSLIB)) {
-                loadParameters.put(KNNConstants.HNSW_ALGO_EF_SEARCH, KNNSettings.getEfSearchParam(knnQuery.getIndexName()));
-            }
-
+            NativeMemoryAllocation indexAllocation;
             try {
                 indexAllocation = nativeMemoryCacheManager.get(
                         new NativeMemoryEntryContext.IndexEntryContext(
                                 indexPath.toString(),
                                 NativeMemoryLoadStrategy.IndexLoadStrategy.getInstance(),
-                                loadParameters,
+                                getParametersAtLoading(spaceType, knnEngine, knnQuery.getIndexName()),
                                 knnQuery.getIndexName()
                         ), true);
             } catch (ExecutionException e) {

src/test/java/org/opensearch/knn/index/IndexUtilTests.java

@@ -0,0 +1,72 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.knn.index;
+
+import com.google.common.collect.ImmutableMap;
+import org.opensearch.cluster.ClusterState;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.cluster.metadata.Metadata;
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.knn.KNNTestCase;
+import org.opensearch.knn.index.util.KNNEngine;
+
+import java.util.Map;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.opensearch.knn.common.KNNConstants.HNSW_ALGO_EF_SEARCH;
+import static org.opensearch.knn.common.KNNConstants.SPACE_TYPE;
+import static org.opensearch.knn.index.IndexUtil.getParametersAtLoading;
+import static org.opensearch.knn.index.KNNSettings.KNN_ALGO_PARAM_EF_SEARCH;
+
+public class IndexUtilTests extends KNNTestCase {
+    public void testGetLoadParameters() {
+        // Test faiss to ensure that space type gets set properly
+        SpaceType spaceType1 = SpaceType.COSINESIMIL;
+        KNNEngine knnEngine1 = KNNEngine.FAISS;
+        String indexName = "my-test-index";
+
+        Map<String, Object> loadParameters = getParametersAtLoading(spaceType1, knnEngine1, indexName);
+        assertEquals(1, loadParameters.size());
+        assertEquals(spaceType1.getValue(), loadParameters.get(SPACE_TYPE));
+
+        // Test nmslib to ensure both space type and ef search are properly set
+        SpaceType spaceType2 = SpaceType.L1;
+        KNNEngine knnEngine2 = KNNEngine.NMSLIB;
+        int efSearchValue = 413;
+
+        // We use the constant for the setting here as opposed to the identifier of efSearch in nmslib jni
+        Map<String, Object> indexSettings = ImmutableMap.of(
+                KNN_ALGO_PARAM_EF_SEARCH, efSearchValue
+        );
+
+        // Because ef search comes from an index setting, we need to mock the long line of calls to get those
+        // index settings
+        Settings settings = Settings.builder().loadFromMap(indexSettings).build();
+        IndexMetadata indexMetadata = mock(IndexMetadata.class);
+        when(indexMetadata.getSettings()).thenReturn(settings);
+        Metadata metadata = mock(Metadata.class);
+        when(metadata.index(anyString())).thenReturn(indexMetadata);
+        ClusterState clusterState = mock(ClusterState.class);
+        when(clusterState.getMetadata()).thenReturn(metadata);
+        ClusterService clusterService = mock(ClusterService.class);
+        when(clusterService.state()).thenReturn(clusterState);
+        KNNSettings.state().setClusterService(clusterService);
+
+        loadParameters = getParametersAtLoading(spaceType2, knnEngine2, indexName);
+        assertEquals(2, loadParameters.size());
+        assertEquals(spaceType2.getValue(), loadParameters.get(SPACE_TYPE));
+        assertEquals(efSearchValue, loadParameters.get(HNSW_ALGO_EF_SEARCH));
+    }
+}