diff --git a/src/main/java/org/opensearch/knn/index/IndexUtil.java b/src/main/java/org/opensearch/knn/index/IndexUtil.java index a80464cdb..1782d3384 100644 --- a/src/main/java/org/opensearch/knn/index/IndexUtil.java +++ b/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 getParametersAtLoading(SpaceType spaceType, KNNEngine knnEngine, String indexName) { + Map 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); + } } diff --git a/src/main/java/org/opensearch/knn/index/KNNIndexShard.java b/src/main/java/org/opensearch/knn/index/KNNIndexShard.java index eede05356..36ac9097a 100644 --- a/src/main/java/org/opensearch/knn/index/KNNIndexShard.java +++ b/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.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 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); diff --git a/src/main/java/org/opensearch/knn/index/KNNWeight.java b/src/main/java/org/opensearch/knn/index/KNNWeight.java index 8288ef130..38fcf7ee6 100644 --- a/src/main/java/org/opensearch/knn/index/KNNWeight.java +++ b/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; /** @@ -134,24 +135,13 @@ public Scorer scorer(LeafReaderContext context) throws IOException { KNNCounter.GRAPH_QUERY_REQUESTS.increment(); // We need to first get index allocation - NativeMemoryAllocation indexAllocation = null; - - Map loadParameters = new HashMap() {{ - 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) { diff --git a/src/test/java/org/opensearch/knn/index/IndexUtilTests.java b/src/test/java/org/opensearch/knn/index/IndexUtilTests.java new file mode 100644 index 000000000..e90e300f7 --- /dev/null +++ b/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 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 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)); + } +}