From 72e8250be495ea583e61838814e6851116ca70be Mon Sep 17 00:00:00 2001 From: Vijayan Balasubramanian Date: Thu, 6 Feb 2025 14:10:27 -0800 Subject: [PATCH] Update engine for version 2.19 or above (#2501) Signed-off-by: Balasubramanian (cherry picked from commit 8ee82b9528de4767ffa7bf96f1e51d5895761db7) --- .../opensearch-knn.release-notes-2.19.0.0.md | 1 + .../knn/index/engine/EngineResolver.java | 66 +++++++++++++++---- .../index/mapper/KNNVectorFieldMapper.java | 3 +- .../transport/TrainingModelRequest.java | 2 +- .../knn/index/engine/EngineResolverTests.java | 27 +++++++- .../mapper/KNNVectorFieldMapperTests.java | 2 +- 6 files changed, 82 insertions(+), 19 deletions(-) diff --git a/release-notes/opensearch-knn.release-notes-2.19.0.0.md b/release-notes/opensearch-knn.release-notes-2.19.0.0.md index 8b2f4e0fa..7c972c491 100644 --- a/release-notes/opensearch-knn.release-notes-2.19.0.0.md +++ b/release-notes/opensearch-knn.release-notes-2.19.0.0.md @@ -38,6 +38,7 @@ Compatible with OpenSearch 2.19.0 * Fix Faiss byte vector efficient filter bug [#2448](https://github.com/opensearch-project/k-NN/pull/2448) * Fixing bug where mapping accepts both dimension and model-id [#2410](https://github.com/opensearch-project/k-NN/pull/2410) * Add version check for full field name validation [#2477](https://github.com/opensearch-project/k-NN/pull/2477) +* Update engine for version 2.19 or above [#2501](https://github.com/opensearch-project/k-NN/pull/2501) ### Infrastructure * Updated C++ version in JNI from c++11 to c++17 [#2259](https://github.com/opensearch-project/k-NN/pull/2259) * Upgrade bytebuddy and objenesis version to match OpenSearch core and, update github ci runner for macos [#2279](https://github.com/opensearch-project/k-NN/pull/2279) diff --git a/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java b/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java index 983adf28b..662fd7daf 100644 --- a/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java +++ b/src/main/java/org/opensearch/knn/index/engine/EngineResolver.java @@ -5,8 +5,10 @@ package org.opensearch.knn.index.engine; +import com.google.common.annotations.VisibleForTesting; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.Version; import org.opensearch.knn.index.mapper.CompressionLevel; import org.opensearch.knn.index.mapper.Mode; @@ -22,47 +24,83 @@ public final class EngineResolver { private EngineResolver() {} + @VisibleForTesting + KNNEngine resolveEngine(KNNMethodConfigContext knnMethodConfigContext, KNNMethodContext knnMethodContext, boolean requiresTraining) { + return logAndReturnEngine(resolveKNNEngine(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 + boolean requiresTraining, + Version version ) { - // User configuration gets precedence - if (knnMethodContext != null && knnMethodContext.isEngineConfigured()) { - return logAndReturnEngine(knnMethodContext.getKnnEngine()); + return logAndReturnEngine(resolveKNNEngine(knnMethodConfigContext, knnMethodContext, requiresTraining, version)); + } + + /** + * 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} + */ + private KNNEngine resolveKNNEngine( + 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) { - return logAndReturnEngine(KNNEngine.FAISS); + // 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 logAndReturnEngine(KNNEngine.DEFAULT); + return KNNEngine.DEFAULT; } - // For 1x, we need to default to faiss if mode is provided and use nmslib otherwise + if (compressionLevel == CompressionLevel.x4) { + // Lucene is only engine that supports 4x - so we have to default to it here. + return KNNEngine.LUCENE; + } if (CompressionLevel.isConfigured(compressionLevel) == false || compressionLevel == CompressionLevel.x1) { - return logAndReturnEngine(mode == Mode.ON_DISK ? KNNEngine.FAISS : KNNEngine.NMSLIB); + // For 1x or no compression, we need to default to faiss if mode is provided and use nmslib otherwise based on version check + return resolveEngineForX1OrNoCompression(mode, version); } + return KNNEngine.FAISS; + } - // Lucene is only engine that supports 4x - so we have to default to it here. - if (compressionLevel == CompressionLevel.x4) { - return logAndReturnEngine(KNNEngine.LUCENE); - } + private boolean hasUserConfiguredEngine(KNNMethodContext knnMethodContext) { + return knnMethodContext != null && knnMethodContext.isEngineConfigured(); + } - return logAndReturnEngine(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; } private KNNEngine logAndReturnEngine(KNNEngine knnEngine) { 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 338a913d5..8684b54cc 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java +++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java @@ -541,7 +541,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 index bd2c88347..9a29107eb 100644 --- a/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java +++ b/src/main/java/org/opensearch/knn/plugin/transport/TrainingModelRequest.java @@ -138,7 +138,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 index 291f0c671..376908ce8 100644 --- 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 index b0c3d2ffe..0e4478f0e 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperTests.java @@ -1782,7 +1782,7 @@ public void testTypeParser_whenModeAndCompressionAreSet_thenHandle() throws IOEx assertFalse(builder.getOriginalParameters().isLegacyMapping()); validateBuilderAfterParsing( builder, - KNNEngine.NMSLIB, + KNNEngine.FAISS, SpaceType.L2, VectorDataType.FLOAT, CompressionLevel.x1,