Update engine for version 2.19 or above

Signed-off-by: Balasubramanian <[email protected]>
opensearch-project · Feb 6, 2025 · 421b09f · 421b09f
1 parent c22ba17
commit 421b09f
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 17 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java b/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.knn.index.engine;
 
+import org.opensearch.Version;
 import org.opensearch.knn.index.mapper.CompressionLevel;
 import org.opensearch.knn.index.mapper.Mode;
 
@@ -30,33 +31,61 @@ public KNNEngine resolveEngine(
         KNNMethodContext knnMethodContext,
         boolean requiresTraining
     ) {
-        // User configuration gets precedence
-        if (knnMethodContext != null && knnMethodContext.isEngineConfigured()) {
+        return resolveEngine(knnMethodConfigContext, knnMethodContext, requiresTraining, Version.CURRENT);
+    }
+
+    /**
+     * Based on the provided {@link Mode} and {@link CompressionLevel}, resolve to a {@link KNNEngine}.
+     *
+     * @param knnMethodConfigContext configuration context
+     * @param knnMethodContext KNNMethodContext
+     * @param requiresTraining whether config requires training
+     * @param version opensearch index version
+     * @return {@link KNNEngine}
+     */
+    public KNNEngine resolveEngine(
+        KNNMethodConfigContext knnMethodConfigContext,
+        KNNMethodContext knnMethodContext,
+        boolean requiresTraining,
+        Version version
+    ) {
+        // Check user configuration first
+        if (hasUserConfiguredEngine(knnMethodContext)) {
             return knnMethodContext.getKnnEngine();
         }
 
-        // Faiss is the only engine that supports training, so we default to faiss here for now
+        // Handle training case
         if (requiresTraining) {
+            // Faiss is the only engine that supports training, so we default to faiss here for now
             return KNNEngine.FAISS;
         }
 
         Mode mode = knnMethodConfigContext.getMode();
         CompressionLevel compressionLevel = knnMethodConfigContext.getCompressionLevel();
+
         // If both mode and compression are not specified, we can just default
         if (Mode.isConfigured(mode) == false && CompressionLevel.isConfigured(compressionLevel) == false) {
             return KNNEngine.DEFAULT;
         }
 
-        // For 1x, we need to default to faiss if mode is provided and use nmslib otherwise
-        if (CompressionLevel.isConfigured(compressionLevel) == false || compressionLevel == CompressionLevel.x1) {
-            return mode == Mode.ON_DISK ? KNNEngine.FAISS : KNNEngine.NMSLIB;
-        }
+        return switch (compressionLevel) {
+            // Lucene is only engine that supports 4x - so we have to default to it here.
+            case x4 -> KNNEngine.LUCENE;
+            // For 1x or no compression, we need to default to faiss if mode is provided and use nmslib otherwise based on version check
+            case x1, CompressionLevel.NOT_CONFIGURED -> resolveEngineForX1OrNoCompression(mode, version);
+            default -> KNNEngine.FAISS;
+        };
+    }
 
-        // Lucene is only engine that supports 4x - so we have to default to it here.
-        if (compressionLevel == CompressionLevel.x4) {
-            return KNNEngine.LUCENE;
-        }
+    private boolean hasUserConfiguredEngine(KNNMethodContext knnMethodContext) {
+        return knnMethodContext != null && knnMethodContext.isEngineConfigured();
+    }
 
-        return KNNEngine.FAISS;
+    private KNNEngine resolveEngineForX1OrNoCompression(Mode mode, Version version) {
+        if (version != null && version.onOrAfter(Version.V_2_19_0)) {
+            return KNNEngine.FAISS;
+        }
+        return mode == Mode.ON_DISK ? KNNEngine.FAISS : KNNEngine.NMSLIB;
     }
+
 }
diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java
@@ -530,7 +530,8 @@ private void resolveKNNMethodComponents(
             KNNEngine resolvedKNNEngine = EngineResolver.INSTANCE.resolveEngine(
                 builder.knnMethodConfigContext,
                 builder.originalParameters.getResolvedKnnMethodContext(),
-                false
+                false,
+                builder.indexCreatedVersion
             );
             setEngine(builder.originalParameters.getResolvedKnnMethodContext(), resolvedKNNEngine);
 

diff --git a/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java b/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java
@@ -134,7 +134,7 @@ public TrainingModelRequest(
             .mode(mode)
             .build();
 
-        KNNEngine knnEngine = EngineResolver.INSTANCE.resolveEngine(knnMethodConfigContext, knnMethodContext, true);
+        KNNEngine knnEngine = EngineResolver.INSTANCE.resolveEngine(knnMethodConfigContext, knnMethodContext, true, Version.CURRENT);
         ResolvedMethodContext resolvedMethodContext = knnEngine.resolveMethod(knnMethodContext, knnMethodConfigContext, true, spaceType);
         this.knnMethodContext = resolvedMethodContext.getKnnMethodContext();
         this.compressionLevel = resolvedMethodContext.getCompressionLevel();

diff --git a/src/test/java/org/opensearch/knn/index/engine/EngineResolverTests.java b/src/test/java/org/opensearch/knn/index/engine/EngineResolverTests.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.knn.index.engine;
 
+import org.opensearch.Version;
 import org.opensearch.knn.KNNTestCase;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.mapper.CompressionLevel;
@@ -41,10 +42,23 @@ public void testResolveEngine_whenModeAndCompressionAreFalse_thenDefault() {
         );
     }
 
-    public void testResolveEngine_whenModeSpecifiedAndCompressionIsNotSpecified_thenNMSLIB() {
+    public void testResolveEngine_whenModeSpecifiedAndCompressionIsNotSpecified_whenVersionBefore2_19_thenNMSLIB() {
         assertEquals(KNNEngine.DEFAULT, ENGINE_RESOLVER.resolveEngine(KNNMethodConfigContext.builder().build(), null, false));
         assertEquals(
             KNNEngine.NMSLIB,
+            ENGINE_RESOLVER.resolveEngine(
+                KNNMethodConfigContext.builder().mode(Mode.IN_MEMORY).build(),
+                new KNNMethodContext(KNNEngine.DEFAULT, SpaceType.UNDEFINED, MethodComponentContext.EMPTY, false),
+                false,
+                Version.V_2_18_0
+            )
+        );
+    }
+
+    public void testResolveEngine_whenModeSpecifiedAndCompressionIsNotSpecified_thenFAISS() {
+        assertEquals(KNNEngine.DEFAULT, ENGINE_RESOLVER.resolveEngine(KNNMethodConfigContext.builder().build(), null, false));
+        assertEquals(
+            KNNEngine.FAISS,
             ENGINE_RESOLVER.resolveEngine(
                 KNNMethodConfigContext.builder().mode(Mode.IN_MEMORY).build(),
                 new KNNMethodContext(KNNEngine.DEFAULT, SpaceType.UNDEFINED, MethodComponentContext.EMPTY, false),
@@ -63,9 +77,18 @@ public void testResolveEngine_whenCompressionIs1x_thenEngineBasedOnMode() {
             )
         );
         assertEquals(
-            KNNEngine.NMSLIB,
+            KNNEngine.FAISS,
             ENGINE_RESOLVER.resolveEngine(KNNMethodConfigContext.builder().compressionLevel(CompressionLevel.x1).build(), null, false)
         );
+        assertEquals(
+            KNNEngine.NMSLIB,
+            ENGINE_RESOLVER.resolveEngine(
+                KNNMethodConfigContext.builder().compressionLevel(CompressionLevel.x1).build(),
+                null,
+                false,
+                Version.V_2_18_0
+            )
+        );
     }
 
     public void testResolveEngine_whenCompressionIs4x_thenEngineIsLucene() {

diff --git a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java
@@ -1769,7 +1769,7 @@ public void testTypeParser_whenModeAndCompressionAreSet_thenHandle() throws IOEx
         assertFalse(builder.getOriginalParameters().isLegacyMapping());
         validateBuilderAfterParsing(
             builder,
-            KNNEngine.NMSLIB,
+            KNNEngine.FAISS,
             SpaceType.L2,
             VectorDataType.FLOAT,
             CompressionLevel.x1,