Added the setting to disable the shard level rescoring and fix the bu…

…ild tag Signed-off-by: Navneet Verma <[email protected]>
navneet1v · Sep 28, 2024 · e722556 · e722556
1 parent b0d82b7
commit e722556
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 5 deletions.
diff --git a/build.gradle b/build.gradle
@@ -13,7 +13,7 @@ buildscript {
     ext {
         // build.version_qualifier parameter applies to knn plugin artifacts only. OpenSearch version must be set
         // explicitly as 'opensearch.version' property, for instance opensearch.version=2.0.0-rc1-SNAPSHOT
-        opensearch_version = System.getProperty("opensearch.version", "2.17.1-SNAPSHOT")
+        opensearch_version = System.getProperty("opensearch.version", "2.17.0-SNAPSHOT")
         version_qualifier = System.getProperty("build.version_qualifier", "")
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")

diff --git a/src/main/java/org/opensearch/knn/index/KNNSettings.java b/src/main/java/org/opensearch/knn/index/KNNSettings.java
@@ -86,6 +86,7 @@ public class KNNSettings {
     public static final String KNN_FAISS_AVX2_DISABLED = "knn.faiss.avx2.disabled";
     public static final String QUANTIZATION_STATE_CACHE_SIZE_LIMIT = "knn.quantization.cache.size.limit";
     public static final String QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES = "knn.quantization.cache.expiry.minutes";
+    public static final String KNN_SHARD_LEVEL_RESCORING_DISABLED = "index.knn.shard_level_rescoring_disabled";
 
     /**
      * Default setting values
@@ -109,6 +110,7 @@ public class KNNSettings {
     public static final Integer KNN_MAX_QUANTIZATION_STATE_CACHE_SIZE_LIMIT_PERCENTAGE = 10; // Quantization state cache limit cannot exceed
                                                                                              // 10% of the JVM heap
     public static final Integer KNN_DEFAULT_QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES = 60;
+    public static final boolean KNN_SHARD_LEVEL_RESCORING_DISABLED_DEFAULT_VALUE = true;
 
     /**
      * Settings Definition
@@ -123,6 +125,12 @@ public class KNNSettings {
         Setting.Property.NodeScope
     );
 
+    public static final Setting<Boolean> KNN_SHARD_LEVEL_RESCORING_DISABLED_SETTING = Setting.boolSetting(
+        KNN_SHARD_LEVEL_RESCORING_DISABLED,
+        KNN_SHARD_LEVEL_RESCORING_DISABLED_DEFAULT_VALUE,
+        Setting.Property.NodeScope
+    );
+
     public static final Setting<String> INDEX_KNN_SPACE_TYPE = Setting.simpleString(
         KNN_SPACE_TYPE,
         INDEX_KNN_DEFAULT_SPACE_TYPE,
@@ -441,6 +449,10 @@ private Setting<?> getSetting(String key) {
             return QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING;
         }
 
+        if (KNN_SHARD_LEVEL_RESCORING_DISABLED.equals(key)) {
+            return KNN_SHARD_LEVEL_RESCORING_DISABLED_SETTING;
+        }
+
         throw new IllegalArgumentException("Cannot find setting by key [" + key + "]");
     }
 
@@ -461,7 +473,8 @@ public List<Setting<?>> getSettings() {
             KNN_FAISS_AVX2_DISABLED_SETTING,
             KNN_VECTOR_STREAMING_MEMORY_LIMIT_PCT_SETTING,
             QUANTIZATION_STATE_CACHE_SIZE_LIMIT_SETTING,
-            QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING
+            QUANTIZATION_STATE_CACHE_EXPIRY_TIME_MINUTES_SETTING,
+            KNN_SHARD_LEVEL_RESCORING_DISABLED_SETTING
         );
         return Stream.concat(settings.stream(), Stream.concat(getFeatureFlags().stream(), dynamicCacheSettings.values().stream()))
             .collect(Collectors.toList());
@@ -483,6 +496,10 @@ public static double getCircuitBreakerUnsetPercentage() {
         return KNNSettings.state().getSettingValue(KNNSettings.KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE);
     }
 
+    public static boolean isShardLevelRescoringDisabled() {
+        return KNNSettings.state().getSettingValue(KNNSettings.KNN_SHARD_LEVEL_RESCORING_DISABLED);
+    }
+
     public static boolean isFaissAVX2Disabled() {
         try {
             return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX2_DISABLED);

diff --git a/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java
@@ -20,6 +20,7 @@
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
 import org.opensearch.common.StopWatch;
+import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.query.ExactSearcher;
 import org.opensearch.knn.index.query.KNNQuery;
 import org.opensearch.knn.index.query.KNNWeight;
@@ -63,12 +64,16 @@ public Weight createWeight(IndexSearcher indexSearcher, ScoreMode scoreMode, flo
         } else {
             int firstPassK = rescoreContext.getFirstPassK(finalK);
             perLeafResults = doSearch(indexSearcher, leafReaderContexts, knnWeight, firstPassK);
-            ResultUtil.reduceToTopK(perLeafResults, firstPassK);
-
+            if (KNNSettings.isShardLevelRescoringDisabled() == false) {
+                log.info("Shard Level Rescoring is enabled. Reducing docs to {}", firstPassK);
+                ResultUtil.reduceToTopK(perLeafResults, firstPassK);
+            } else {
+                log.info("Shard Level Rescoring is not enabled.");
+            }
             StopWatch stopWatch = new StopWatch().start();
             perLeafResults = doRescore(indexSearcher, leafReaderContexts, knnWeight, perLeafResults, finalK);
             long rescoreTime = stopWatch.stop().totalTime().millis();
-            log.debug("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size());
+            log.info("Rescoring results took {} ms. oversampled k:{}, segments:{}", rescoreTime, firstPassK, leafReaderContexts.size());
         }
         ResultUtil.reduceToTopK(perLeafResults, finalK);
         TopDocs[] topDocs = new TopDocs[perLeafResults.size()];