Update Default Rescore Context based on Dimesion (#2149) (#2150)

(cherry picked from commit a1227ea) Signed-off-by: Naveen Tatikonda <[email protected]> Co-authored-by: Vikasht34 <[email protected]>
opensearch-project · Sep 26, 2024 · db37c93 · db37c93
1 parent 4447a43
commit db37c93
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Add AVX512 support to k-NN for FAISS library [#2069](https://github.com/opensearch-project/k-NN/pull/2069)
 ### Enhancements
 * Add short circuit if no live docs are in segments [#2059](https://github.com/opensearch-project/k-NN/pull/2059)
+* Update Default Rescore Context based on Dimension [#2149](https://github.com/opensearch-project/k-NN/pull/2149)
 ### Bug Fixes
 * KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147)
 ### Infrastructure

diff --git a/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java b/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java
@@ -94,9 +94,34 @@ public static boolean isConfigured(CompressionLevel compressionLevel) {
         return compressionLevel != null && compressionLevel != NOT_CONFIGURED;
     }
 
-    public RescoreContext getDefaultRescoreContext(Mode mode) {
+    /**
+     * Returns the appropriate {@link RescoreContext} based on the given {@code mode} and {@code dimension}.
+     *
+     * <p>If the {@code mode} is present in the valid {@code modesForRescore} set, the method checks the value of
+     * {@code dimension}:
+     * <ul>
+     *     <li>If {@code dimension} is less than or equal to 1000, it returns a {@link RescoreContext} with an
+     *         oversample factor of 5.0f.</li>
+     *     <li>If {@code dimension} is greater than 1000, it returns the default {@link RescoreContext} associated with
+     *         the {@link CompressionLevel}. If no default is set, it falls back to {@link RescoreContext#getDefault()}.</li>
+     * </ul>
+     * If the {@code mode} is not valid, the method returns {@code null}.
+     *
+     * @param mode      The {@link Mode} for which to retrieve the {@link RescoreContext}.
+     * @param dimension The dimensional value that determines the {@link RescoreContext} behavior.
+     * @return          A {@link RescoreContext} with an oversample factor of 5.0f if {@code dimension} is less than
+     *                  or equal to 1000, the default {@link RescoreContext} if greater, or {@code null} if the mode
+     *                  is invalid.
+     */
+    public RescoreContext getDefaultRescoreContext(Mode mode, int dimension) {
         if (modesForRescore.contains(mode)) {
-            return defaultRescoreContext;
+            // Adjust RescoreContext based on dimension
+            if (dimension <= RescoreContext.DIMENSION_THRESHOLD) {
+                // For dimensions <= 1000, return a RescoreContext with 5.0f oversample factor
+                return RescoreContext.builder().oversampleFactor(RescoreContext.OVERSAMPLE_FACTOR_BELOW_DIMENSION_THRESHOLD).build();
+            } else {
+                return defaultRescoreContext;
+            }
         }
         return null;
     }

diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldType.java
@@ -93,6 +93,10 @@ public RescoreContext resolveRescoreContext(RescoreContext userProvidedContext)
         if (userProvidedContext != null) {
             return userProvidedContext;
         }
-        return getKnnMappingConfig().getCompressionLevel().getDefaultRescoreContext(getKnnMappingConfig().getMode());
+        KNNMappingConfig knnMappingConfig = getKnnMappingConfig();
+        int dimension = knnMappingConfig.getDimension();
+        CompressionLevel compressionLevel = knnMappingConfig.getCompressionLevel();
+        Mode mode = knnMappingConfig.getMode();
+        return compressionLevel.getDefaultRescoreContext(mode, dimension);
     }
 }
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
@@ -21,6 +21,8 @@ public final class RescoreContext {
     public static final float MIN_OVERSAMPLE_FACTOR = 1.0f;
 
     public static final int MAX_FIRST_PASS_RESULTS = 10000;
+    public static final int DIMENSION_THRESHOLD = 1000;
+    public static final float OVERSAMPLE_FACTOR_BELOW_DIMENSION_THRESHOLD = 5.0f;
 
     // Todo:- We will improve this in upcoming releases
     public static final int MIN_FIRST_PASS_RESULTS = 100;

diff --git a/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java b/src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java
@@ -44,26 +44,65 @@ public void testIsConfigured() {
     public void testGetDefaultRescoreContext() {
         // Test rescore context for ON_DISK mode
         Mode mode = Mode.ON_DISK;
+        int belowThresholdDimension = 500; // A dimension below the threshold
+        int aboveThresholdDimension = 1500; // A dimension above the threshold
 
-        // x32 should have RescoreContext with an oversample factor of 3.0f
-        RescoreContext rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode);
+        // x32 with dimension <= 1000 should have an oversample factor of 5.0f
+        RescoreContext rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x32 with dimension > 1000 should have an oversample factor of 3.0f
+        rescoreContext = CompressionLevel.x32.getDefaultRescoreContext(mode, aboveThresholdDimension);
         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);
+        // x16 with dimension <= 1000 should have an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x16 with dimension > 1000 should have an oversample factor of 3.0f
+        rescoreContext = CompressionLevel.x16.getDefaultRescoreContext(mode, aboveThresholdDimension);
         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);
+        // x8 with dimension <= 1000 should have an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+
+        // x8 with dimension > 1000 should have an oversample factor of 2.0f
+        rescoreContext = CompressionLevel.x8.getDefaultRescoreContext(mode, aboveThresholdDimension);
         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));
+        // x4 with dimension <= 1000 should have an oversample factor of 5.0f (though it doesn't have its own RescoreContext)
+        rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+        // x4 with dimension > 1000 should return null (no RescoreContext is configured for x4)
+        rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        // Other compression levels should behave similarly with respect to dimension
+
+        rescoreContext = CompressionLevel.x2.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
+        // x2 with dimension > 1000 should return null
+        rescoreContext = CompressionLevel.x2.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        rescoreContext = CompressionLevel.x1.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
+        // x1 with dimension > 1000 should return null
+        rescoreContext = CompressionLevel.x1.getDefaultRescoreContext(mode, aboveThresholdDimension);
+        assertNull(rescoreContext);
+
+        // NOT_CONFIGURED with dimension <= 1000 should return a RescoreContext with an oversample factor of 5.0f
+        rescoreContext = CompressionLevel.NOT_CONFIGURED.getDefaultRescoreContext(mode, belowThresholdDimension);
+        assertNull(rescoreContext);
+
     }
 }