diff --git a/src/main/java/org/opensearch/knn/index/engine/MethodComponent.java b/src/main/java/org/opensearch/knn/index/engine/MethodComponent.java index 75e18a243..d456ea89f 100644 --- a/src/main/java/org/opensearch/knn/index/engine/MethodComponent.java +++ b/src/main/java/org/opensearch/knn/index/engine/MethodComponent.java @@ -11,6 +11,8 @@ import org.opensearch.common.ValidationException; import org.opensearch.knn.common.KNNConstants; import org.opensearch.knn.index.VectorDataType; +import org.opensearch.knn.index.mapper.CompressionLevel; +import org.opensearch.knn.index.mapper.Mode; import org.opensearch.knn.index.util.IndexHyperParametersUtil; import java.util.HashMap; @@ -328,6 +330,15 @@ public static Map getParameterMapWithDefaultsAdded( Map parametersWithDefaultsMap = new HashMap<>(); Map userProvidedParametersMap = methodComponentContext.getParameters(); Version indexCreationVersion = knnMethodConfigContext.getVersionCreated(); + Mode mode = knnMethodConfigContext.getMode(); + CompressionLevel compressionLevel = knnMethodConfigContext.getCompressionLevel(); + + // Check if the mode is ON_DISK and the compression level is one of the binary quantization levels (x32, x16, or x8). + // This determines whether to use binary quantization-specific values for parameters like ef_search and ef_construction. + boolean isOnDiskWithBinaryQuantization = (compressionLevel == CompressionLevel.x32 + || compressionLevel == CompressionLevel.x16 + || compressionLevel == CompressionLevel.x8); + for (Parameter parameter : methodComponent.getParameters().values()) { if (methodComponentContext.getParameters().containsKey(parameter.getName())) { parametersWithDefaultsMap.put(parameter.getName(), userProvidedParametersMap.get(parameter.getName())); @@ -335,12 +346,27 @@ public static Map getParameterMapWithDefaultsAdded( // Picking the right values for the parameters whose values are different based on different index // created version. if (parameter.getName().equals(KNNConstants.METHOD_PARAMETER_EF_SEARCH)) { - parametersWithDefaultsMap.put(parameter.getName(), IndexHyperParametersUtil.getHNSWEFSearchValue(indexCreationVersion)); + if (isOnDiskWithBinaryQuantization) { + parametersWithDefaultsMap.put(parameter.getName(), IndexHyperParametersUtil.getBinaryQuantizationEFSearchValue()); + } else { + parametersWithDefaultsMap.put( + parameter.getName(), + IndexHyperParametersUtil.getHNSWEFSearchValue(indexCreationVersion) + ); + } } else if (parameter.getName().equals(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION)) { - parametersWithDefaultsMap.put( - parameter.getName(), - IndexHyperParametersUtil.getHNSWEFConstructionValue(indexCreationVersion) - ); + if (isOnDiskWithBinaryQuantization) { + parametersWithDefaultsMap.put( + parameter.getName(), + IndexHyperParametersUtil.getBinaryQuantizationEFConstructionValue() + ); + } else { + parametersWithDefaultsMap.put( + parameter.getName(), + IndexHyperParametersUtil.getHNSWEFConstructionValue(indexCreationVersion) + ); + } + } else { Object value = parameter.getDefaultValue(); if (value != null) { diff --git a/src/main/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolver.java b/src/main/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolver.java index 90e938eb3..c976a0959 100644 --- a/src/main/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolver.java +++ b/src/main/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolver.java @@ -65,8 +65,6 @@ public ResolvedMethodContext resolveMethod( // Fill in parameters for the encoder and then the method. resolveEncoder(resolvedKNNMethodContext, knnMethodConfigContext, encoderMap); - resolveMethodParams(resolvedKNNMethodContext.getMethodComponentContext(), knnMethodConfigContext, method); - // From the resolved method context, get the compression level and validate it against the passed in // configuration CompressionLevel resolvedCompressionLevel = resolveCompressionLevelFromMethodContext( @@ -77,6 +75,9 @@ public ResolvedMethodContext resolveMethod( // Validate that resolved compression doesnt have any conflicts validateCompressionConflicts(knnMethodConfigContext.getCompressionLevel(), resolvedCompressionLevel); + knnMethodConfigContext.setCompressionLevel(resolvedCompressionLevel); + resolveMethodParams(resolvedKNNMethodContext.getMethodComponentContext(), knnMethodConfigContext, method); + return ResolvedMethodContext.builder() .knnMethodContext(resolvedKNNMethodContext) .compressionLevel(resolvedCompressionLevel) diff --git a/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java b/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java index cc80bb1ed..222e042b6 100644 --- a/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java +++ b/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java @@ -25,9 +25,9 @@ public enum CompressionLevel { x1(1, "1x", null, Collections.emptySet()), x2(2, "2x", null, Collections.emptySet()), x4(4, "4x", null, Collections.emptySet()), - x8(8, "8x", new RescoreContext(1.5f), Set.of(Mode.ON_DISK)), - x16(16, "16x", new RescoreContext(2.0f), Set.of(Mode.ON_DISK)), - x32(32, "32x", new RescoreContext(2.0f), Set.of(Mode.ON_DISK)); + x8(8, "8x", new RescoreContext(2.0f), Set.of(Mode.ON_DISK)), + x16(16, "16x", new RescoreContext(3.0f), Set.of(Mode.ON_DISK)), + x32(32, "32x", new RescoreContext(3.0f), Set.of(Mode.ON_DISK)); // Internally, an empty string is easier to deal with them null. However, from the mapping, // we do not want users to pass in the empty string and instead want null. So we make the conversion herex diff --git a/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java b/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java index 9fe2ddbc5..5e57b4aba 100644 --- a/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java +++ b/src/main/java/org/opensearch/knn/index/query/rescore/RescoreContext.java @@ -22,6 +22,9 @@ public final class RescoreContext { public static final int MAX_FIRST_PASS_RESULTS = 10000; + // Todo:- We will improve this in upcoming releases + public static final int MIN_FIRST_PASS_RESULTS = 100; + @Builder.Default private float oversampleFactor = DEFAULT_OVERSAMPLE_FACTOR; @@ -40,6 +43,6 @@ public static RescoreContext getDefault() { * @return The number of results to return for the first pass of rescoring */ public int getFirstPassK(int finalK) { - return Math.min(MAX_FIRST_PASS_RESULTS, (int) Math.ceil(finalK * oversampleFactor)); + return Math.min(MAX_FIRST_PASS_RESULTS, Math.max(MIN_FIRST_PASS_RESULTS, (int) Math.ceil(finalK * oversampleFactor))); } } diff --git a/src/main/java/org/opensearch/knn/index/util/IndexHyperParametersUtil.java b/src/main/java/org/opensearch/knn/index/util/IndexHyperParametersUtil.java index af842788a..c88a28e5c 100644 --- a/src/main/java/org/opensearch/knn/index/util/IndexHyperParametersUtil.java +++ b/src/main/java/org/opensearch/knn/index/util/IndexHyperParametersUtil.java @@ -28,6 +28,8 @@ public class IndexHyperParametersUtil { private static final int INDEX_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION_OLD_VALUE = 512; private static final int INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH_OLD_VALUE = 512; + private static final int INDEX_BINARY_QUANTIZATION_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION = 256; + private static final int INDEX_BINARY_QUANTIZATION_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH = 256; /** * Returns the default value of EF Construction that should be used for the input index version. After version 2.12.0 @@ -76,4 +78,22 @@ public static int getHNSWEFSearchValue(@NonNull final Version indexVersion) { ); return KNNSettings.INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH; } + + /* + * Returns the default value of EF Construction that should be used with Binary Quantization. + * + * @return default value of EF Construction + */ + public static int getBinaryQuantizationEFConstructionValue() { + return INDEX_BINARY_QUANTIZATION_KNN_DEFAULT_ALGO_PARAM_EF_CONSTRUCTION; + } + + /* + * Returns the default value of EF Search that should be used with Binary Quantization. + * + * @return default value of EF Search + */ + public static int getBinaryQuantizationEFSearchValue() { + return INDEX_BINARY_QUANTIZATION_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH; + } } diff --git a/src/test/java/org/opensearch/knn/index/engine/MethodComponentTests.java b/src/test/java/org/opensearch/knn/index/engine/MethodComponentTests.java index 7730095c7..e247fb70a 100644 --- a/src/test/java/org/opensearch/knn/index/engine/MethodComponentTests.java +++ b/src/test/java/org/opensearch/knn/index/engine/MethodComponentTests.java @@ -11,6 +11,9 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.knn.index.VectorDataType; +import org.opensearch.knn.index.mapper.CompressionLevel; +import org.opensearch.knn.index.mapper.Mode; +import org.opensearch.knn.index.util.IndexHyperParametersUtil; import java.io.IOException; import java.util.Map; @@ -214,4 +217,41 @@ public void testBuilder() { .getLibraryParameters() ); } + + /** + * Test the new flow where EF_SEARCH and EF_CONSTRUCTION are set for ON_DISK mode + * with binary quantization compression levels. + */ + public void testGetParameterMapWithDefaultsAdded_forOnDiskWithBinaryQuantization() { + // Set up MethodComponent and context + String methodName = "test-method"; + String parameterEFSearch = "ef_search"; + String parameterEFConstruction = "ef_construction"; + + MethodComponent methodComponent = MethodComponent.Builder.builder(methodName) + .addParameter(parameterEFSearch, new Parameter.IntegerParameter(parameterEFSearch, 512, (v, context) -> v > 0)) + .addParameter(parameterEFConstruction, new Parameter.IntegerParameter(parameterEFConstruction, 512, (v, context) -> v > 0)) + .build(); + + // Simulate ON_DISK mode and binary quantization compression levels + KNNMethodConfigContext knnMethodConfigContext = KNNMethodConfigContext.builder() + .versionCreated(Version.CURRENT) + .mode(Mode.ON_DISK) // ON_DISK mode + .compressionLevel(CompressionLevel.x32) // Binary quantization compression level + .build(); + + MethodComponentContext methodComponentContext = new MethodComponentContext(methodName, Map.of()); + + // Retrieve parameter map with defaults added + Map resultMap = MethodComponent.getParameterMapWithDefaultsAdded( + methodComponentContext, + methodComponent, + knnMethodConfigContext + ); + + // Check that binary quantization values are used + assertEquals(IndexHyperParametersUtil.getBinaryQuantizationEFSearchValue(), resultMap.get(parameterEFSearch)); + assertEquals(IndexHyperParametersUtil.getBinaryQuantizationEFConstructionValue(), resultMap.get(parameterEFConstruction)); + } + } diff --git a/src/test/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolverTests.java b/src/test/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolverTests.java index ad466d4bb..3a33736fa 100644 --- a/src/test/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolverTests.java +++ b/src/test/java/org/opensearch/knn/index/engine/faiss/FaissMethodResolverTests.java @@ -136,6 +136,33 @@ public void testResolveMethod_whenValid_thenResolve() { SpaceType.L2 ); validateResolveMethodContext(resolvedMethodContext, CompressionLevel.x1, SpaceType.L2, ENCODER_FLAT, false); + + KNNMethodConfigContext knnMethodConfigContext = KNNMethodConfigContext.builder() + .vectorDataType(VectorDataType.FLOAT) + .versionCreated(Version.CURRENT) + .build(); + + resolvedMethodContext = TEST_RESOLVER.resolveMethod( + new KNNMethodContext( + KNNEngine.FAISS, + SpaceType.L2, + new MethodComponentContext( + METHOD_HNSW, + Map.of( + METHOD_ENCODER_PARAMETER, + new MethodComponentContext( + QFrameBitEncoder.NAME, + Map.of(QFrameBitEncoder.BITCOUNT_PARAM, CompressionLevel.x8.numBitsForFloat32()) + ) + ) + ) + ), + knnMethodConfigContext, + false, + SpaceType.L2 + ); + assertEquals(knnMethodConfigContext.getCompressionLevel(), CompressionLevel.x8); + validateResolveMethodContext(resolvedMethodContext, CompressionLevel.x8, SpaceType.L2, QFrameBitEncoder.NAME, true); } private void validateResolveMethodContext( diff --git a/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java b/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java index 07475109a..9eb302b13 100644 --- a/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java +++ b/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java @@ -7,6 +7,7 @@ import org.opensearch.core.common.Strings; import org.opensearch.knn.KNNTestCase; +import org.opensearch.knn.index.query.rescore.RescoreContext; public class CompressionLevelTests extends KNNTestCase { @@ -39,4 +40,30 @@ public void testIsConfigured() { assertFalse(CompressionLevel.isConfigured(null)); assertTrue(CompressionLevel.isConfigured(CompressionLevel.x1)); } + + public void testGetDefaultRescoreContext() { + // Test rescore context for ON_DISK mode + Mode mode = Mode.ON_DISK; + + // x32 should have RescoreContext with an oversample factor of 3.0f + RescoreContext rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode); + assertNotNull(rescoreContext); + assertEquals(3.0f, rescoreContext.getOversampleFactor(), 0.0f); + + // x16 should have RescoreContext with an oversample factor of 3.0f + rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode); + assertNotNull(rescoreContext); + assertEquals(3.0f, rescoreContext.getOversampleFactor(), 0.0f); + + // x8 should have RescoreContext with an oversample factor of 2.0f + rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode); + assertNotNull(rescoreContext); + assertEquals(2.0f, rescoreContext.getOversampleFactor(), 0.0f); + + // Other compression levels should not have a RescoreContext for ON_DISK mode + assertNull(CompressionLevel.x4.getDefaultRescoreContext(mode)); + assertNull(CompressionLevel.x2.getDefaultRescoreContext(mode)); + assertNull(CompressionLevel.x1.getDefaultRescoreContext(mode)); + assertNull(CompressionLevel.NOT_CONFIGURED.getDefaultRescoreContext(mode)); + } } diff --git a/src/test/java/org/opensearch/knn/index/query/rescore/RescoreContextTests.java b/src/test/java/org/opensearch/knn/index/query/rescore/RescoreContextTests.java index 6d872ddb8..fd94667db 100644 --- a/src/test/java/org/opensearch/knn/index/query/rescore/RescoreContextTests.java +++ b/src/test/java/org/opensearch/knn/index/query/rescore/RescoreContextTests.java @@ -8,6 +8,7 @@ import org.opensearch.knn.KNNTestCase; import static org.opensearch.knn.index.query.rescore.RescoreContext.MAX_FIRST_PASS_RESULTS; +import static org.opensearch.knn.index.query.rescore.RescoreContext.MIN_FIRST_PASS_RESULTS; public class RescoreContextTests extends KNNTestCase { @@ -17,10 +18,43 @@ public void testGetFirstPassK() { int finalK = 100; assertEquals(260, rescoreContext.getFirstPassK(finalK)); finalK = 1; - assertEquals(3, rescoreContext.getFirstPassK(finalK)); + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); finalK = 0; - assertEquals(0, rescoreContext.getFirstPassK(finalK)); + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); finalK = MAX_FIRST_PASS_RESULTS; assertEquals(MAX_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); } + + public void testGetFirstPassKWithMinPassK() { + float oversample = 2.6f; + RescoreContext rescoreContext = RescoreContext.builder().oversampleFactor(oversample).build(); + + // Case 1: Test with a finalK that results in a value greater than MIN_FIRST_PASS_RESULTS + int finalK = 100; + assertEquals(260, rescoreContext.getFirstPassK(finalK)); + + // Case 2: Test with a very small finalK that should result in a value less than MIN_FIRST_PASS_RESULTS + finalK = 1; + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); + + // Case 3: Test with finalK = 0, should return 0 + finalK = 0; + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); + + // Case 4: Test with finalK = MAX_FIRST_PASS_RESULTS, should cap at MAX_FIRST_PASS_RESULTS + finalK = MAX_FIRST_PASS_RESULTS; + assertEquals(MAX_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); + + // Case 5: Test where finalK * oversample is smaller than MIN_FIRST_PASS_RESULTS + finalK = 10; + oversample = 0.5f; // This will result in 5, which is less than MIN_FIRST_PASS_RESULTS + rescoreContext = RescoreContext.builder().oversampleFactor(oversample).build(); + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); + + // Case 6: Test where finalK * oversample results in exactly MIN_FIRST_PASS_RESULTS + finalK = 100; + oversample = 1.0f; // This will result in exactly 100 (MIN_FIRST_PASS_RESULTS) + rescoreContext = RescoreContext.builder().oversampleFactor(oversample).build(); + assertEquals(MIN_FIRST_PASS_RESULTS, rescoreContext.getFirstPassK(finalK)); + } } diff --git a/src/test/java/org/opensearch/knn/index/util/IndexHyperParametersUtilTests.java b/src/test/java/org/opensearch/knn/index/util/IndexHyperParametersUtilTests.java index 508b8765c..96eb41476 100644 --- a/src/test/java/org/opensearch/knn/index/util/IndexHyperParametersUtilTests.java +++ b/src/test/java/org/opensearch/knn/index/util/IndexHyperParametersUtilTests.java @@ -41,4 +41,12 @@ public void testGetHNSWEFSearchValue_withDifferentValues_thenSuccess() { IndexHyperParametersUtil.getHNSWEFConstructionValue(Version.CURRENT) ); } + + public void testGetBinaryQuantizationEFValues_thenSuccess() { + // Test for Binary Quantization EF Construction value + Assert.assertEquals(256, IndexHyperParametersUtil.getBinaryQuantizationEFConstructionValue()); + + // Test for Binary Quantization EF Search value + Assert.assertEquals(256, IndexHyperParametersUtil.getBinaryQuantizationEFSearchValue()); + } }