Skip to content

Commit

Permalink
Added the setting to disable the shard level rescoring and fix the bu…
Browse files Browse the repository at this point in the history
…ild tag

Signed-off-by: Navneet Verma <navneev@amazon.com>
  • Loading branch information
navneet1v committed Sep 28, 2024
1 parent b0d82b7 commit 28e141f
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 5 deletions.
2 changes: 1 addition & 1 deletion build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
24 changes: 23 additions & 1 deletion src/main/java/org/opensearch/knn/index/KNNSettings.java
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -123,6 +125,13 @@ 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,
IndexScope,
Dynamic
);

public static final Setting<String> INDEX_KNN_SPACE_TYPE = Setting.simpleString(
KNN_SPACE_TYPE,
INDEX_KNN_DEFAULT_SPACE_TYPE,
Expand Down Expand Up @@ -441,6 +450,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 + "]");
}

Expand All @@ -461,7 +474,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());
Expand All @@ -483,6 +497,14 @@ public static double getCircuitBreakerUnsetPercentage() {
return KNNSettings.state().getSettingValue(KNNSettings.KNN_CIRCUIT_BREAKER_UNSET_PERCENTAGE);
}

public static boolean isShardLevelRescoringDisabled(String indexName) {
return KNNSettings.state().clusterService.state()
.getMetadata()
.index(indexName)
.getSettings()
.getAsBoolean(KNN_SHARD_LEVEL_RESCORING_DISABLED, true);
}

public static boolean isFaissAVX2Disabled() {
try {
return KNNSettings.state().getSettingValue(KNNSettings.KNN_FAISS_AVX2_DISABLED);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(knnQuery.getIndexName()) == 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()];
Expand Down

0 comments on commit 28e141f

Please sign in to comment.