From 12757cc50bbe9f0ccdbb91700b7d13472efedce1 Mon Sep 17 00:00:00 2001 From: Dooyong Kim Date: Wed, 1 Oct 2025 09:43:23 -0700 Subject: [PATCH 1/2] Do not apply memory optimized serach for old indices. Signed-off-by: Dooyong Kim --- CHANGELOG.md | 1 + .../knn/bwc/MemoryOptimizedSearchIT.java | 100 ++++++++++++++++++ .../MemoryOptimizedSearchSupportSpec.java | 9 +- .../knn/index/mapper/EngineFieldMapper.java | 6 +- .../index/mapper/KNNVectorFieldMapper.java | 3 +- .../knn/index/mapper/KNNVectorFieldType.java | 35 ++++-- .../mapper/KNNVectorFieldMapperTests.java | 21 ++-- ...MemoryOptimizedSearchSupportSpecTests.java | 93 +++++++++++++--- 8 files changed, 233 insertions(+), 35 deletions(-) create mode 100644 qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java diff --git a/CHANGELOG.md b/CHANGELOG.md index c493ae43a9..29fd9a5ba3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Fix local ref leak in JNI [#2916](https://github.com/opensearch-project/k-NN/pull/2916) * Fix rescoring logic for nested exact search [#2921](https://github.com/opensearch-project/k-NN/pull/2921) * Update Visitor to delegate for other fields [#2925](https://github.com/opensearch-project/k-NN/pull/2925) +* Fix blocking old indices created before 2.18 to use memory optimized search. [#2918](https://github.com/opensearch-project/k-NN/pull/2918) ### Refactoring * Refactored the KNN Stat files for better readability. diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java new file mode 100644 index 0000000000..e0aee9a865 --- /dev/null +++ b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java @@ -0,0 +1,100 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.bwc; + +import org.opensearch.Version; +import org.opensearch.common.settings.Settings; +import org.opensearch.knn.index.KNNSettings; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.test.rest.OpenSearchRestTestCase; + +import java.io.IOException; +import java.util.Collections; + +import static org.opensearch.knn.TestUtils.NODES_BWC_CLUSTER; +import static org.opensearch.knn.common.KNNConstants.FAISS_NAME; +import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW; +import static org.opensearch.knn.common.KNNConstants.NMSLIB_NAME; + +public class MemoryOptimizedSearchIT extends AbstractRestartUpgradeTestCase { + private static final String TEST_FIELD = "target_field"; + private static final int DIMENSION = 128; + private static final int NUM_DOCS = 200; + + public void testMemoryOptimizedSearchWithFaiss() throws Exception { + doTestMemoryOptimizedSearch(FAISS_NAME); + } + + public void testMemoryOptimizedSearchWithNmslib() throws Exception { + doTestMemoryOptimizedSearch(NMSLIB_NAME); + } + + private void doTestMemoryOptimizedSearch(final String engine) throws Exception { + waitForClusterHealthGreen(NODES_BWC_CLUSTER); + if (isRunningAgainstOldCluster()) { + // Create index and ingest some data + createIndexAndIngestData(engine); + + // Force merge to a single segment. + forceMergeKnnIndex(testIndex); + } else { + String versionString = getBWCVersion().get(); + versionString = versionString.replaceAll("-SNAPSHOT", ""); + final Version oldVersion = Version.fromString(versionString); + + // Turn on memory optimized search + turnOnMemoryOptSearch(); + + // Validate warm-up is done without an issue + knnWarmup(Collections.singletonList(testIndex)); + + // Validate search + validateKNNSearch(testIndex, TEST_FIELD, DIMENSION, NUM_DOCS, 5); + + if (engine.equals(NMSLIB_NAME)) { + // Memory optimized search does not support NMSLIB + assertEquals(1, getTotalGraphsInCache()); + } else { + // Validate memory optimized search applied conditionally + if (oldVersion.compareTo(Version.V_2_17_0) < 0) { + // For indices created version < 2.19 are not supported + assertEquals(1, getTotalGraphsInCache()); + } else { + // Memory optimized search is applied, no off-heap graph is expected + assertEquals(0, getTotalGraphsInCache()); + } + } + + // Delete index + deleteKNNIndex(testIndex); + } + } + + private void turnOnMemoryOptSearch() throws IOException { + // Close index + closeKNNIndex(testIndex); + + // Update settings + OpenSearchRestTestCase.updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.MEMORY_OPTIMIZED_KNN_SEARCH_MODE, true)); + + // Reopen index again + OpenSearchRestTestCase.openIndex(testIndex); + } + + private void createIndexAndIngestData(final String engine) throws IOException { + final Settings indexSettings = getKNNDefaultIndexSettings(); + + // Create an index + createKnnIndex( + testIndex, + indexSettings, + createKnnIndexMapping(TEST_FIELD, DIMENSION, METHOD_HNSW, engine, SpaceType.L2.getValue()) + ); + + // Ingest 200 docs + addKNNDocs(testIndex, TEST_FIELD, DIMENSION, 0, NUM_DOCS); + } +} diff --git a/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java b/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java index e221d171d5..0ad5bf9a7d 100644 --- a/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java +++ b/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java @@ -5,6 +5,7 @@ package org.opensearch.knn.index.engine; +import org.opensearch.Version; import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.engine.qframe.QuantizationConfig; import org.opensearch.knn.index.mapper.CompressionLevel; @@ -69,8 +70,14 @@ public static boolean isSupportedFieldType(final KNNVectorFieldType fieldType, f public static boolean isSupportedFieldType( final Optional methodContextOpt, final QuantizationConfig quantizationConfig, - final Optional modelId + final Optional modelId, + final Version indexCreatedVersion ) { + // Since LuceneOnFaiss logic sits in newer codec, we don't support LuceneOnFaiss for older codec whose version < 2.19 + if (indexCreatedVersion.before(Version.V_2_17_0)) { + return false; + } + // PQ is not supported. if (modelId.isPresent()) { return false; diff --git a/src/main/java/org/opensearch/knn/index/mapper/EngineFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/EngineFieldMapper.java index 47ec08d8c9..99c6a37e4e 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/EngineFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/EngineFieldMapper.java @@ -64,7 +64,8 @@ public static EngineFieldMapper createFieldMapper( Explicit ignoreMalformed, boolean stored, boolean hasDocValues, - OriginalMappingParameters originalMappingParameters + OriginalMappingParameters originalMappingParameters, + Version indexCreatedVersion ) { KNNMethodContext methodContext = originalMappingParameters.getResolvedKnnMethodContext(); KNNLibraryIndexingContext libraryContext = methodContext.getKnnEngine() @@ -112,7 +113,8 @@ public QuantizationConfig getQuantizationConfig() { public KNNLibraryIndexingContext getKnnLibraryIndexingContext() { return libraryContext; } - } + }, + indexCreatedVersion ); return new EngineFieldMapper( diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java index 4048df2b99..ff237d2e70 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java @@ -336,7 +336,8 @@ public KNNVectorFieldMapper build(BuilderContext context) { ignoreMalformed, stored.getValue(), hasDocValues.get(), - originalParameters + originalParameters, + indexCreatedVersion ); } diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java index 8bd5d3ab11..cbb67b424b 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java @@ -11,9 +11,10 @@ import org.apache.lucene.search.FieldExistsQuery; import org.apache.lucene.search.Query; import org.apache.lucene.util.BytesRef; +import org.opensearch.Version; import org.opensearch.index.fielddata.IndexFieldData; -import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.ArraySourceValueFetcher; +import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.TextSearchInfo; import org.opensearch.index.mapper.ValueFetcher; import org.opensearch.index.query.QueryShardContext; @@ -21,10 +22,10 @@ import org.opensearch.knn.index.KNNVectorIndexFieldData; import org.opensearch.knn.index.VectorDataType; import org.opensearch.knn.index.engine.KNNMethodContext; +import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec; import org.opensearch.knn.index.query.rescore.RescoreContext; import org.opensearch.knn.indices.ModelDao; import org.opensearch.knn.indices.ModelMetadata; -import org.opensearch.knn.index.engine.MemoryOptimizedSearchSupportSpec; import org.opensearch.search.aggregations.support.CoreValuesSourceType; import org.opensearch.search.lookup.SearchLookup; @@ -48,6 +49,31 @@ public class KNNVectorFieldType extends MappedFieldType { // Whether this field type can be benefit from memory optimized search? boolean memoryOptimizedSearchAvailable; + /** + * Constructor for KNNVectorFieldType with index created version. + * + * @param name name of the field + * @param metadata metadata of the field + * @param vectorDataType data type of the vector + * @param annConfig configuration context for the ANN index + * @param indexCreatedVersion Index created version. + */ + public KNNVectorFieldType( + String name, + Map metadata, + VectorDataType vectorDataType, + KNNMappingConfig annConfig, + Version indexCreatedVersion + ) { + this(name, metadata, vectorDataType, annConfig); + this.memoryOptimizedSearchAvailable = MemoryOptimizedSearchSupportSpec.isSupportedFieldType( + knnMappingConfig.getKnnMethodContext(), + annConfig.getQuantizationConfig(), + annConfig.getModelId(), + indexCreatedVersion + ); + } + /** * Constructor for KNNVectorFieldType. * @@ -60,11 +86,6 @@ public KNNVectorFieldType(String name, Map metadata, VectorDataT super(name, false, false, true, TextSearchInfo.NONE, metadata); this.vectorDataType = vectorDataType; this.knnMappingConfig = annConfig; - this.memoryOptimizedSearchAvailable = MemoryOptimizedSearchSupportSpec.isSupportedFieldType( - knnMappingConfig.getKnnMethodContext(), - annConfig.getQuantizationConfig(), - annConfig.getModelId() - ); } @Override diff --git a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java index 8da145a321..7ea79f7c45 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java @@ -1435,7 +1435,8 @@ private void doTestMethodFieldMapper_saveBestMatchingVectorSimilarityFunction( new Explicit<>(true, true), false, false, - originalMappingParameters + originalMappingParameters, + CURRENT ); methodFieldMapper.parseCreateField(parseContext, dimension, dataType); final IndexableField field1 = document.getFields().get(0); @@ -1504,7 +1505,8 @@ public void testMethodFieldMapperParseCreateField_validInput_thenDifferentFieldT new Explicit<>(true, true), false, false, - originalMappingParameters + originalMappingParameters, + CURRENT ); methodFieldMapper.parseCreateField(parseContext, dimension, dataType); @@ -1545,7 +1547,8 @@ public void testMethodFieldMapperParseCreateField_validInput_thenDifferentFieldT new Explicit<>(true, true), false, false, - originalMappingParameters + originalMappingParameters, + CURRENT ); methodFieldMapper.parseCreateField(parseContext, dimension, dataType); @@ -1724,7 +1727,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() { new Explicit<>(true, true), false, true, - originalMappingParameters + originalMappingParameters, + CURRENT ); luceneFieldMapper.parseCreateField(parseContext, TEST_DIMENSION, VectorDataType.FLOAT); @@ -1786,7 +1790,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withFloats() { new Explicit<>(true, true), false, false, - originalMappingParameters + originalMappingParameters, + CURRENT ); luceneFieldMapper.parseCreateField(parseContext, TEST_DIMENSION, VectorDataType.FLOAT); @@ -1840,7 +1845,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withBytes() { new Explicit<>(true, true), false, true, - originalMappingParameters + originalMappingParameters, + CURRENT ) ); doReturn(Optional.of(TEST_BYTE_VECTOR)).when(luceneFieldMapper) @@ -1894,7 +1900,8 @@ public void testLuceneFieldMapper_parseCreateField_docValues_withBytes() { new Explicit<>(true, true), false, false, - originalMappingParameters + originalMappingParameters, + CURRENT ) ); doReturn(Optional.of(TEST_BYTE_VECTOR)).when(luceneFieldMapper) diff --git a/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java b/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java index 3db8251ba5..258aff909f 100644 --- a/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java +++ b/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java @@ -6,6 +6,7 @@ package org.opensearch.knn.memoryoptsearch; import org.mockito.MockedStatic; +import org.opensearch.Version; import org.opensearch.knn.KNNTestCase; import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; @@ -21,6 +22,7 @@ import org.opensearch.knn.index.mapper.Mode; import org.opensearch.knn.quantization.enums.ScalarQuantizationType; +import java.util.Arrays; import java.util.Collections; import java.util.Map; import java.util.Optional; @@ -48,7 +50,8 @@ public void testLuceneEngineIsIsSupportedFieldType() { VectorDataType.FLOAT, mock(MethodComponentContext.class), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); mustIsSupportedFieldType( @@ -58,7 +61,8 @@ public void testLuceneEngineIsIsSupportedFieldType() { VectorDataType.BYTE, mock(MethodComponentContext.class), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); mustIsSupportedFieldType( @@ -68,7 +72,8 @@ public void testLuceneEngineIsIsSupportedFieldType() { VectorDataType.FLOAT, mock(MethodComponentContext.class), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); mustIsSupportedFieldType( @@ -78,11 +83,50 @@ public void testLuceneEngineIsIsSupportedFieldType() { VectorDataType.BYTE, mock(MethodComponentContext.class), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); } + public void testNewIndicesSupportsMemoryOptimizedSearch() { + for (final Version version : Arrays.asList(Version.V_2_19_0, Version.V_3_0_0, Version.V_3_1_0, Version.V_3_2_0, Version.CURRENT)) { + mustIsSupportedFieldType( + new TestingSpec( + KNNEngine.FAISS, + SpaceType.L2, + VectorDataType.FLOAT, + new MethodComponentContext( + METHOD_HNSW, + Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) + ), + NO_QUANTIZATION, + NO_MODEL_ID, + version + ) + ); + } + } + + public void testOldIndicesDoesNotSupportsMemoryOptimizedSearch() { + for (final Version version : Arrays.asList(Version.V_2_16_0, Version.V_2_15_0, Version.V_2_2_0)) { + mustNotIsSupportedFieldType( + new TestingSpec( + KNNEngine.FAISS, + SpaceType.L2, + VectorDataType.FLOAT, + new MethodComponentContext( + METHOD_HNSW, + Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) + ), + NO_QUANTIZATION, + NO_MODEL_ID, + version + ) + ); + } + } + public void testFaissIsSupportedFieldTypeCases() { // HNSW,float, L2|IP, Flat // HNSW,float, L2|IP, SQ @@ -99,7 +143,8 @@ public void testFaissIsSupportedFieldTypeCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -113,7 +158,8 @@ public void testFaissIsSupportedFieldTypeCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_SQ, Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -127,7 +173,8 @@ public void testFaissIsSupportedFieldTypeCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -141,7 +188,8 @@ public void testFaissIsSupportedFieldTypeCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_BINARY, Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -155,7 +203,8 @@ public void testFaissIsSupportedFieldTypeCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_SQ, Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); } @@ -171,7 +220,8 @@ public void testFaissQuantizationCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_SQ, Collections.emptyMap())) ), QuantizationConfig.builder().quantizationType(ScalarQuantizationType.ONE_BIT).build(), - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -185,7 +235,8 @@ public void testFaissQuantizationCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_SQ, Collections.emptyMap())) ), QuantizationConfig.builder().quantizationType(ScalarQuantizationType.TWO_BIT).build(), - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -199,7 +250,8 @@ public void testFaissQuantizationCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_SQ, Collections.emptyMap())) ), QuantizationConfig.builder().quantizationType(ScalarQuantizationType.FOUR_BIT).build(), - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); } @@ -216,7 +268,8 @@ public void testFaissUnsupportedCases() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext("DUMMY_KEY", Collections.emptyMap())) ), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); @@ -228,7 +281,8 @@ public void testFaissUnsupportedCases() { VectorDataType.FLOAT, new MethodComponentContext(METHOD_HNSW, Map.of(METHOD_ENCODER_PARAMETER, new Object())), NO_QUANTIZATION, - NO_MODEL_ID + NO_MODEL_ID, + Version.CURRENT ) ); } @@ -245,7 +299,8 @@ public void testPQNotIsSupportedFieldType() { Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) ), mock(QuantizationConfig.class), - Optional.of("model_id") + Optional.of("model_id"), + Version.CURRENT ) ); } @@ -312,7 +367,8 @@ private void doTest(final TestingSpec testingSpec, final boolean expected) { final boolean isSupported = MemoryOptimizedSearchSupportSpec.isSupportedFieldType( testingSpec.methodComponentContext, testingSpec.quantizationConfig, - testingSpec.modelId + testingSpec.modelId, + testingSpec.version ); assertEquals(expected, isSupported); } @@ -324,6 +380,7 @@ private static class TestingSpec { final Optional methodComponentContext; final QuantizationConfig quantizationConfig; final Optional modelId; + final Version version; private TestingSpec( final KNNEngine knnEngine, @@ -331,7 +388,8 @@ private TestingSpec( final VectorDataType vectorDataType, final MethodComponentContext methodComponentContext, final QuantizationConfig quantizationConfig, - final Optional modelId + final Optional modelId, + final Version version ) { this.knnEngine = knnEngine; this.spaceType = spaceType; @@ -340,6 +398,7 @@ private TestingSpec( this.methodComponentContext = Optional.of(methodContext); this.quantizationConfig = quantizationConfig; this.modelId = modelId; + this.version = version; } } } From e788f7ea905c9f1d9ad2ed4a636b301ec6dced1d Mon Sep 17 00:00:00 2001 From: Dooyong Kim Date: Wed, 15 Oct 2025 14:36:58 -0700 Subject: [PATCH 2/2] Fix MemoryOptimizedSearchSupportSpec to throw an exception for old indices. Signed-off-by: Dooyong Kim --- .../knn/bwc/MemoryOptimizedSearchIT.java | 38 ++++-- .../MemoryOptimizedSearchSupportSpec.java | 30 +++-- .../knn/index/mapper/KNNVectorFieldType.java | 5 +- .../knn/index/query/KNNQueryBuilderTests.java | 21 +++- ...MemoryOptimizedSearchSupportSpecTests.java | 110 ++++++++++-------- 5 files changed, 129 insertions(+), 75 deletions(-) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java index e0aee9a865..a0d768ba0c 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java @@ -6,12 +6,15 @@ package org.opensearch.knn.bwc; import org.opensearch.Version; +import org.opensearch.client.ResponseException; import org.opensearch.common.settings.Settings; import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.test.rest.OpenSearchRestTestCase; import java.io.IOException; +import java.util.Arrays; import java.util.Collections; import static org.opensearch.knn.TestUtils.NODES_BWC_CLUSTER; @@ -23,6 +26,7 @@ public class MemoryOptimizedSearchIT extends AbstractRestartUpgradeTestCase { private static final String TEST_FIELD = "target_field"; private static final int DIMENSION = 128; private static final int NUM_DOCS = 200; + private static final int K = 5; public void testMemoryOptimizedSearchWithFaiss() throws Exception { doTestMemoryOptimizedSearch(FAISS_NAME); @@ -51,16 +55,34 @@ private void doTestMemoryOptimizedSearch(final String engine) throws Exception { // Validate warm-up is done without an issue knnWarmup(Collections.singletonList(testIndex)); - // Validate search - validateKNNSearch(testIndex, TEST_FIELD, DIMENSION, NUM_DOCS, 5); + if (oldVersion.compareTo(Version.V_2_17_0) < 0) { + if (engine.equals(NMSLIB_NAME)) { + // Search should work regardless + validateKNNSearch(testIndex, TEST_FIELD, DIMENSION, NUM_DOCS, K); - if (engine.equals(NMSLIB_NAME)) { - // Memory optimized search does not support NMSLIB - assertEquals(1, getTotalGraphsInCache()); + // Memory optimized search is applied, but NMSLIB should load off-heap index. + assertEquals(1, getTotalGraphsInCache()); + } else { + // For FAISS engine, indices created < 2.17 should throw an exception. + final float[] queryVector = new float[DIMENSION]; + Arrays.fill(queryVector, (float) NUM_DOCS); + + // Exception should be thrown + assertThrows( + ResponseException.class, + () -> searchKNNIndex( + testIndex, + KNNQueryBuilder.builder().k(K).methodParameters(null).fieldName(TEST_FIELD).vector(queryVector).build(), + K + ) + ); + } } else { - // Validate memory optimized search applied conditionally - if (oldVersion.compareTo(Version.V_2_17_0) < 0) { - // For indices created version < 2.19 are not supported + // Validate search, should be fine + validateKNNSearch(testIndex, TEST_FIELD, DIMENSION, NUM_DOCS, K); + + if (engine.equals(NMSLIB_NAME)) { + // Memory optimized search is applied, but NMSLIB should load off-heap index. assertEquals(1, getTotalGraphsInCache()); } else { // Memory optimized search is applied, no off-heap graph is expected diff --git a/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java b/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java index 0ad5bf9a7d..34770dca97 100644 --- a/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java +++ b/src/main/java/org/opensearch/knn/index/engine/MemoryOptimizedSearchSupportSpec.java @@ -32,6 +32,7 @@ * The overall logic will be made based on the given method context and quantization configuration. */ public class MemoryOptimizedSearchSupportSpec { + private static final Version MIN_VERSION_SUPPORTS_MEM_OPT_SEARCH = Version.V_2_17_0; private static final Set SUPPORTED_HNSW_ENCODING = Set.of(ENCODER_FLAT, ENCODER_SQ, ENCODER_BINARY); /** @@ -46,6 +47,22 @@ public class MemoryOptimizedSearchSupportSpec { public static boolean isSupportedFieldType(final KNNVectorFieldType fieldType, final String indexName) { if (fieldType.isMemoryOptimizedSearchAvailable()) { if (KNNSettings.isMemoryOptimizedKnnSearchModeEnabled(indexName)) { + final boolean shouldBlockMemoryOptimizedSearch = fieldType.getIndexCreatedVersion() == null + || fieldType.getIndexCreatedVersion().before(MIN_VERSION_SUPPORTS_MEM_OPT_SEARCH); + if (shouldBlockMemoryOptimizedSearch) { + // Memory-optimized search is enabled, but some existing indices were created before + // the minimum version that supports this feature. Throw an exception to clearly + // notify the user of the incompatibility. + throw new IllegalStateException( + "Memory optimized search does not support old indices created before " + + MIN_VERSION_SUPPORTS_MEM_OPT_SEARCH.toString() + + ". Index [" + + indexName + + "] was created in " + + fieldType.getIndexCreatedVersion().toString() + ); + } + return true; } @@ -70,14 +87,8 @@ public static boolean isSupportedFieldType(final KNNVectorFieldType fieldType, f public static boolean isSupportedFieldType( final Optional methodContextOpt, final QuantizationConfig quantizationConfig, - final Optional modelId, - final Version indexCreatedVersion + final Optional modelId ) { - // Since LuceneOnFaiss logic sits in newer codec, we don't support LuceneOnFaiss for older codec whose version < 2.19 - if (indexCreatedVersion.before(Version.V_2_17_0)) { - return false; - } - // PQ is not supported. if (modelId.isPresent()) { return false; @@ -87,11 +98,6 @@ public static boolean isSupportedFieldType( final KNNMethodContext methodContext = methodContextOpt.get(); final KNNEngine engine = methodContext.getKnnEngine(); - // We support Lucene engine - if (engine == KNNEngine.LUCENE) { - return true; - } - // We don't support non-FAISS engine if (engine != KNNEngine.FAISS) { return false; diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java index cbb67b424b..161a7cd011 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java @@ -48,6 +48,7 @@ public class KNNVectorFieldType extends MappedFieldType { VectorDataType vectorDataType; // Whether this field type can be benefit from memory optimized search? boolean memoryOptimizedSearchAvailable; + Version indexCreatedVersion; /** * Constructor for KNNVectorFieldType with index created version. @@ -69,9 +70,9 @@ public KNNVectorFieldType( this.memoryOptimizedSearchAvailable = MemoryOptimizedSearchSupportSpec.isSupportedFieldType( knnMappingConfig.getKnnMethodContext(), annConfig.getQuantizationConfig(), - annConfig.getModelId(), - indexCreatedVersion + annConfig.getModelId() ); + this.indexCreatedVersion = indexCreatedVersion; } /** diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java index b862a06429..9a7cf6b9da 100644 --- a/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java +++ b/src/test/java/org/opensearch/knn/index/query/KNNQueryBuilderTests.java @@ -15,10 +15,13 @@ import org.mockito.MockedStatic; import org.opensearch.Version; import org.opensearch.cluster.ClusterModule; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.io.stream.BytesStreamOutput; -import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; import org.opensearch.core.common.io.stream.NamedWriteableAwareStreamInput; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.StreamInput; @@ -79,16 +82,21 @@ public class KNNQueryBuilderTests extends KNNTestCase { private static final Float MIN_SCORE = 0.5f; private static final TermQueryBuilder TERM_QUERY = QueryBuilders.termQuery("field", "value"); private static final float[] QUERY_VECTOR = new float[] { 1.0f, 2.0f, 3.0f, 4.0f }; - protected static final String TEXT_FIELD_NAME = "some_field"; - protected static final String TEXT_VALUE = "some_value"; @Before @Override public void setUp() throws Exception { super.setUp(); - ClusterSettings clusterSettings = mock(ClusterSettings.class); - when(clusterService.getClusterSettings()).thenReturn(clusterSettings); - KNNSettings.state().setClusterService(clusterService); + + // Mocked index setting as MemoryOptimizedSearchSupportSpec needs this. + final ClusterState clusterState = mock(ClusterState.class); + final Metadata metadata = mock(Metadata.class); + final IndexMetadata indexMetadata = mock(IndexMetadata.class); + final Settings settings = mock(Settings.class); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.getMetadata()).thenReturn(metadata); + when(metadata.index(anyString())).thenReturn(indexMetadata); + when(indexMetadata.getSettings()).thenReturn(settings); } public void testInvalidK() { @@ -712,6 +720,7 @@ private void do_testDoToQuery_whenMemoryOptimizedSearchIsEnabled( // Field type KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + when(mockKNNVectorField.getIndexCreatedVersion()).thenReturn(Version.CURRENT); when(mockQueryShardContext.fieldMapper(anyString())).thenReturn(mockKNNVectorField); when(mockKNNVectorField.isMemoryOptimizedSearchAvailable()).thenReturn(memoryOptimizedSearchEnabledInField); when(mockKNNVectorField.getVectorDataType()).thenReturn(vectorDataType); diff --git a/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java b/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java index 258aff909f..692d2c169d 100644 --- a/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java +++ b/src/test/java/org/opensearch/knn/memoryoptsearch/MemoryOptimizedSearchSupportSpecTests.java @@ -43,7 +43,7 @@ public class MemoryOptimizedSearchSupportSpecTests extends KNNTestCase { public void testLuceneEngineIsIsSupportedFieldType() { // Lucene + any configurations must be supported. - mustIsSupportedFieldType( + mustNotIsSupportedFieldType( new TestingSpec( KNNEngine.LUCENE, SpaceType.INNER_PRODUCT, @@ -54,7 +54,7 @@ public void testLuceneEngineIsIsSupportedFieldType() { Version.CURRENT ) ); - mustIsSupportedFieldType( + mustNotIsSupportedFieldType( new TestingSpec( KNNEngine.LUCENE, SpaceType.INNER_PRODUCT, @@ -65,7 +65,7 @@ public void testLuceneEngineIsIsSupportedFieldType() { Version.CURRENT ) ); - mustIsSupportedFieldType( + mustNotIsSupportedFieldType( new TestingSpec( KNNEngine.LUCENE, SpaceType.L2, @@ -76,7 +76,7 @@ public void testLuceneEngineIsIsSupportedFieldType() { Version.CURRENT ) ); - mustIsSupportedFieldType( + mustNotIsSupportedFieldType( new TestingSpec( KNNEngine.LUCENE, SpaceType.L2, @@ -108,25 +108,6 @@ public void testNewIndicesSupportsMemoryOptimizedSearch() { } } - public void testOldIndicesDoesNotSupportsMemoryOptimizedSearch() { - for (final Version version : Arrays.asList(Version.V_2_16_0, Version.V_2_15_0, Version.V_2_2_0)) { - mustNotIsSupportedFieldType( - new TestingSpec( - KNNEngine.FAISS, - SpaceType.L2, - VectorDataType.FLOAT, - new MethodComponentContext( - METHOD_HNSW, - Map.of(METHOD_ENCODER_PARAMETER, new MethodComponentContext(ENCODER_FLAT, Collections.emptyMap())) - ), - NO_QUANTIZATION, - NO_MODEL_ID, - version - ) - ); - } - } - public void testFaissIsSupportedFieldTypeCases() { // HNSW,float, L2|IP, Flat // HNSW,float, L2|IP, SQ @@ -306,39 +287,75 @@ public void testPQNotIsSupportedFieldType() { } public void testIsSupportedFieldTypeDuringSearch() { - // |----------------------|-------------|---------------||-----------| - // | field type supported | mem_opt_src | on_disk && 1x || supported | - // |----------------------|-------------|---------------||-----------| - // | true | true | true || true | - // | true | true | false || true | - // | true | false | true || true | - // | true | false | false || false | - // | false | true | true || false | - // | false | true | false || false | - // | false | false | true || false | - // | false | false | false || false | - // |----------------------|-------------|---------------||-----------| - - doTestIsSupportedFieldTypeDuringSearch(true, true, true, true); - doTestIsSupportedFieldTypeDuringSearch(true, true, false, true); - doTestIsSupportedFieldTypeDuringSearch(true, false, true, true); - doTestIsSupportedFieldTypeDuringSearch(true, false, false, false); - doTestIsSupportedFieldTypeDuringSearch(false, true, true, false); - doTestIsSupportedFieldTypeDuringSearch(false, true, false, false); - doTestIsSupportedFieldTypeDuringSearch(false, false, true, false); - doTestIsSupportedFieldTypeDuringSearch(false, false, false, false); + // @formatter:off + /* + |----------------------|-------------|---------------||-----------| + | field type supported | mem_opt_src | on_disk && 1x || supported | + |----------------------|-------------|---------------||-----------| + | true | true | true || true | + | true | true | false || true | + | true | false | true || true | + | true | false | false || false | + | false | true | true || false | + | false | true | false || false | + | false | false | true || false | + | false | false | false || false | + |----------------------|-------------|---------------||-----------| + */ + // @formatter:on + + doTestIsSupportedFieldTypeDuringSearch(true, true, true, true, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(true, true, false, true, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(true, false, true, true, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(true, false, false, false, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(false, true, true, false, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(false, true, false, false, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(false, false, true, false, Version.CURRENT); + doTestIsSupportedFieldTypeDuringSearch(false, false, false, false, Version.CURRENT); + } + + public void testShouldThrowExceptionRegardless() { + // @formatter:off + /* + |----------------------|-------------|---------------|----------||-----------| + | field type supported | mem_opt_src | on_disk && 1x | version || supported | + |----------------------|-------------|---------------|----------||-----------| + | true | true | true | old || false | + | true | true | false | old || false | + | true | false | true | old || false | + | true | false | false | old || false | + | false | true | true | old || false | + | false | true | false | old || false | + | false | false | true | old || false | + | false | false | false | old || false | + |----------------------|-------------|---------------|----------||-----------| + */ + // @formatter:on + + assertThrows(IllegalStateException.class, () -> doTestIsSupportedFieldTypeDuringSearch(true, true, true, true, Version.V_2_16_0)); + assertThrows(IllegalStateException.class, () -> doTestIsSupportedFieldTypeDuringSearch(true, true, false, true, Version.V_2_16_0)); + + // It's ok! MemOptSrch is turned off. + doTestIsSupportedFieldTypeDuringSearch(false, true, false, false, Version.V_2_16_0); + doTestIsSupportedFieldTypeDuringSearch(false, true, true, false, Version.V_2_16_0); + doTestIsSupportedFieldTypeDuringSearch(false, false, true, false, Version.V_2_16_0); + doTestIsSupportedFieldTypeDuringSearch(false, false, false, false, Version.V_2_16_0); + doTestIsSupportedFieldTypeDuringSearch(false, false, true, false, Version.V_2_16_0); + doTestIsSupportedFieldTypeDuringSearch(false, false, false, false, Version.V_2_16_0); } public void doTestIsSupportedFieldTypeDuringSearch( final boolean fieldTypeSupported, final boolean memoryOptSrchSupported, final boolean onDiskWith1x, - final boolean expected + final boolean expected, + final Version version ) { try (MockedStatic knnSettingsMockedStatic = mockStatic(KNNSettings.class)) { knnSettingsMockedStatic.when(() -> KNNSettings.isMemoryOptimizedKnnSearchModeEnabled(any())).thenReturn(memoryOptSrchSupported); final KNNVectorFieldType fieldType = mock(KNNVectorFieldType.class); + when(fieldType.getIndexCreatedVersion()).thenReturn(version); when(fieldType.isMemoryOptimizedSearchAvailable()).thenReturn(fieldTypeSupported); final KNNMappingConfig mappingConfig = mock(KNNMappingConfig.class); @@ -367,8 +384,7 @@ private void doTest(final TestingSpec testingSpec, final boolean expected) { final boolean isSupported = MemoryOptimizedSearchSupportSpec.isSupportedFieldType( testingSpec.methodComponentContext, testingSpec.quantizationConfig, - testingSpec.modelId, - testingSpec.version + testingSpec.modelId ); assertEquals(expected, isSupported); }