Update rescore context for 4x compression

Since rescore is supported for lucene engine,
update rescore context for 4x compression to
use 1x oversampling factor.

Signed-off-by: Vijayan B <[email protected]>

Author: VijayanB

release-notes/opensearch-knn.release-notes-3.1.0.0.md

@@ -12,6 +12,7 @@ Compatible with OpenSearch 3.1.0.0
 * [Remote Vector Index Build] Add tuned repository upload/download configurations per benchmarking results [#2662](https://github.com/opensearch-project/k-NN/pull/2662)
 * [Remote Vector Index Build] Add segment size upper bound setting and prepare other settings for GA [#2734](https://github.com/opensearch-project/k-NN/pull/2734)
 * [Remote Vector Index Build] Make `index.knn.remote_index_build.enabled` default to true [#2743](https://github.com/opensearch-project/k-NN/pull/2743)
+* Update rescore context for 4X Compression [#2750](https://github.com/opensearch-project/k-NN/pull/2750)
 
 ### Bug Fixes
 * [BUGFIX] Fix KNN Quantization state cache have an invalid weight threshold [#2666](https://github.com/opensearch-project/k-NN/pull/2666)

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

@@ -5,8 +5,10 @@
 
 package org.opensearch.knn.index.mapper;
 
+import com.google.common.annotations.VisibleForTesting;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
+import org.opensearch.Version;
 import org.opensearch.core.common.Strings;
 import org.opensearch.knn.index.query.rescore.RescoreContext;
 
@@ -24,7 +26,7 @@ public enum CompressionLevel {
     NOT_CONFIGURED(-1, "", null, Collections.emptySet()),
     x1(1, "1x", null, Collections.emptySet()),
     x2(2, "2x", null, Collections.emptySet()),
-    x4(4, "4x", null, Collections.emptySet()),
+    x4(4, "4x", new RescoreContext(1.0f, false, true), Set.of(Mode.ON_DISK)),
     x8(8, "8x", new RescoreContext(2.0f, false, true), Set.of(Mode.ON_DISK)),
     x16(16, "16x", new RescoreContext(3.0f, false, true), Set.of(Mode.ON_DISK)),
     x32(32, "32x", new RescoreContext(3.0f, false, true), Set.of(Mode.ON_DISK)),
@@ -105,20 +107,29 @@ public static boolean isConfigured(CompressionLevel compressionLevel) {
      *                  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) {
+    public RescoreContext getDefaultRescoreContext(Mode mode, int dimension, Version version) {
+        //TODO move this to separate class called resolver to resolve rescore context
         if (modesForRescore.contains(mode)) {
+            if( this == x4 && version.before(Version.V_3_1_0)){
+                // For index created before 3.1, context was always null and mode is empty
+                return null;
+            }
             // 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)
                     .userProvided(false)
                     .build();
-            } else {
-                return defaultRescoreContext;
             }
+            return defaultRescoreContext;
         }
         return null;
     }
 
+    @VisibleForTesting
+    RescoreContext getDefaultRescoreContext(Mode mode, int dimension){
+        return getDefaultRescoreContext(mode, dimension, Version.CURRENT);
+    }
+
 }

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

@@ -125,7 +125,7 @@ public RescoreContext resolveRescoreContext(RescoreContext userProvidedContext)
         int dimension = knnMappingConfig.getDimension();
         CompressionLevel compressionLevel = knnMappingConfig.getCompressionLevel();
         Mode mode = knnMappingConfig.getMode();
-        return compressionLevel.getDefaultRescoreContext(mode, dimension);
+        return compressionLevel.getDefaultRescoreContext(mode, dimension, knnMappingConfig.getIndexCreatedVersion());
     }
 
     /**

src/test/java/org/opensearch/knn/index/mapper/CompressionLevelTests.java

@@ -90,12 +90,18 @@ public void testGetDefaultRescoreContext() {
         assertTrue(rescoreContext.isRescoreEnabled());
         assertFalse(rescoreContext.isUserProvided());
 
-        // x4 with dimension <= 1000 should have an oversample factor of 5.0f (though it doesn't have its own RescoreContext)
+        // x4 with dimension <= 1000 should have an oversample factor of 5.0f
         rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, belowThresholdDimension);
-        assertNull(rescoreContext);
-        // x4 with dimension > 1000 should return null (no RescoreContext is configured for x4)
+        assertNotNull(rescoreContext);
+        assertEquals(5.0f, rescoreContext.getOversampleFactor(), 0.0f);
+        assertTrue(rescoreContext.isRescoreEnabled());
+        assertFalse(rescoreContext.isUserProvided());
+        // x4 with dimension > 1000 should have an oversample factor of 1.0f
         rescoreContext = CompressionLevel.x4.getDefaultRescoreContext(mode, aboveThresholdDimension);
-        assertNull(rescoreContext);
+        assertNotNull(rescoreContext);
+        assertEquals(1.0f, rescoreContext.getOversampleFactor(), 0.0f);
+        assertTrue(rescoreContext.isRescoreEnabled());
+        assertFalse(rescoreContext.isUserProvided());
         // Other compression levels should behave similarly with respect to dimension
         rescoreContext = CompressionLevel.x2.getDefaultRescoreContext(mode, belowThresholdDimension);
         assertNull(rescoreContext);

src/test/java/org/opensearch/knn/integ/ModeAndCompressionIT.java

@@ -650,8 +650,7 @@ private void validateSearch(
         String exactSearchResponseBody = EntityUtils.toString(exactSearchResponse.getEntity());
         List<Float> exactSearchKnnResults = parseSearchResponseScore(exactSearchResponseBody, FIELD_NAME);
         assertEquals(NUM_DOCS, exactSearchKnnResults.size());
-
-        if (CompressionLevel.x4.getName().equals(compressionLevelString) == false && Mode.ON_DISK.getName().equals(mode)) {
+        if(Mode.ON_DISK.getName().equals(mode)){
             Assert.assertEquals(exactSearchKnnResults, knnResults);
         }
 
@@ -681,7 +680,7 @@ private void validateSearch(
         responseBody = EntityUtils.toString(response.getEntity());
         knnResults = parseSearchResponseScore(responseBody, FIELD_NAME);
         assertEquals(K, knnResults.size());
-        if (CompressionLevel.x4.getName().equals(compressionLevelString) == false && Mode.ON_DISK.getName().equals(mode)) {
+        if (Mode.ON_DISK.getName().equals(mode)) {
             Assert.assertEquals(exactSearchKnnResults, knnResults);
         }
     }