Do not apply memory optimized search for old indices.

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.

qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/MemoryOptimizedSearchIT.java

@@ -0,0 +1,122 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+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;
+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;
+    private static final int K = 5;
+
+    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));
+
+            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);
+
+                    // 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 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
+                    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);
+    }
+}

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;
@@ -31,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<String> SUPPORTED_HNSW_ENCODING = Set.of(ENCODER_FLAT, ENCODER_SQ, ENCODER_BINARY);
 
     /**
@@ -45,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;
             }
 
@@ -80,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;

src/main/java/org/opensearch/knn/index/mapper/EngineFieldMapper.java

@@ -64,7 +64,8 @@ public static EngineFieldMapper createFieldMapper(
         Explicit<Boolean> 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(

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
             );
         }
 

src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java

@@ -11,20 +11,21 @@
 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;
 import org.opensearch.index.query.QueryShardException;
 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;
 
@@ -47,24 +48,45 @@ 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.
+     * 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<String, String> metadata, VectorDataType vectorDataType, KNNMappingConfig annConfig) {
-        super(name, false, false, true, TextSearchInfo.NONE, metadata);
-        this.vectorDataType = vectorDataType;
-        this.knnMappingConfig = annConfig;
+    public KNNVectorFieldType(
+        String name,
+        Map<String, String> metadata,
+        VectorDataType vectorDataType,
+        KNNMappingConfig annConfig,
+        Version indexCreatedVersion
+    ) {
+        this(name, metadata, vectorDataType, annConfig);
         this.memoryOptimizedSearchAvailable = MemoryOptimizedSearchSupportSpec.isSupportedFieldType(
             knnMappingConfig.getKnnMethodContext(),
             annConfig.getQuantizationConfig(),
             annConfig.getModelId()
         );
+        this.indexCreatedVersion = indexCreatedVersion;
+    }
+
+    /**
+     * Constructor for KNNVectorFieldType.
+     *
+     * @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
+     */
+    public KNNVectorFieldType(String name, Map<String, String> metadata, VectorDataType vectorDataType, KNNMappingConfig annConfig) {
+        super(name, false, false, true, TextSearchInfo.NONE, metadata);
+        this.vectorDataType = vectorDataType;
+        this.knnMappingConfig = annConfig;
     }
 
     @Override

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)