Skip to content

Commit

Permalink
Add exact search if no engine files are in segments
Browse files Browse the repository at this point in the history 
When graph is not available, plugin will return empty results. With this change,
exact search will be performed when only no engine file is available in segment.
We also don't need version check or feature flag because, option to not build vector
data structure will only be available post 2.17.
If an index is created using pre 2.17 version, segment will always have engine files
and this feature will never be called during search.

Signed-off-by: Vijayan Balasubramanian <[email protected]>
  • Loading branch information
@VijayanB
VijayanB committed Oct 1, 2024
1 parent d61e7d4 commit 53263a7
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 57 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
* Update Default Rescore Context based on Dimension [#2149](https://github.com/opensearch-project/k-NN/pull/2149)
* KNNIterators should support with and without filters [#2155](https://github.com/opensearch-project/k-NN/pull/2155)
* Adding Support to Enable/Disble Share level Rescoring and Update Oversampling Factor[#2172](https://github.com/opensearch-project/k-NN/pull/2172)
* Add exact search if no native engine files are available [#2136] (https://github.com/opensearch-project/k-NN/pull/2136)
### Bug Fixes
* KNN80DocValues should only be considered for BinaryDocValues fields [#2147](https://github.com/opensearch-project/k-NN/pull/2147)
### Infrastructure
Expand Down
165 changes: 112 additions & 53 deletions src/main/java/org/opensearch/knn/index/query/KNNWeight.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,14 @@

import lombok.extern.log4j.Log4j2;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.FloatVectorValues;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.SegmentReader;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.FilteredDocIdSetIterator;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.VectorScorer;
import org.apache.lucene.search.Weight;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.store.FilterDirectory;
Expand All @@ -33,6 +35,7 @@
import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy;
import org.opensearch.knn.index.engine.KNNEngine;
import org.opensearch.knn.index.quantizationservice.QuantizationService;
import org.opensearch.knn.index.query.ExactSearcher.ExactSearcherContext.ExactSearcherContextBuilder;
import org.opensearch.knn.indices.ModelDao;
import org.opensearch.knn.indices.ModelMetadata;
import org.opensearch.knn.indices.ModelUtil;
Expand Down Expand Up @@ -121,50 +124,27 @@ public Scorer scorer(LeafReaderContext context) throws IOException {
* @return A Map of docId to scores for top k results
*/
public Map<Integer, Float> searchLeaf(LeafReaderContext context, int k) throws IOException {
final BitSet filterBitSet = getFilteredDocsBitSet(context);
int cardinality = filterBitSet.cardinality();
final BitSet acceptedDocs = getFilteredDocsBitSet(context);
int cardinality = acceptedDocs.cardinality();
// We don't need to go to JNI layer if no documents are found which satisfy the filters
// We should give this condition a deeper look that where it should be placed. For now I feel this is a good
// place,
if (filterWeight != null && cardinality == 0) {
return Collections.emptyMap();
}

/*
* The idea for this optimization is to get K results, we need to atleast look at K vectors in the HNSW graph
* The idea for this optimization is to get K results, we need to at least look at K vectors in the HNSW graph
* . Hence, if filtered results are less than K and filter query is present we should shift to exact search.
* This improves the recall.
*/
Map<Integer, Float> docIdsToScoreMap;
final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder()
.k(k)
.isParentHits(true)
.matchedDocs(filterBitSet)
// setting to true, so that if quantization details are present we want to do search on the quantized
// vectors as this flow is used in first pass of search.
.useQuantizedVectorsForSearch(true)
.knnQuery(knnQuery)
.build();
if (filterWeight != null && canDoExactSearch(cardinality)) {
docIdsToScoreMap = exactSearch(context, exactSearcherContext);
} else {
docIdsToScoreMap = doANNSearch(context, filterBitSet, cardinality, k);
if (docIdsToScoreMap == null) {
return Collections.emptyMap();
}
if (canDoExactSearchAfterANNSearch(cardinality, docIdsToScoreMap.size())) {
log.debug(
"Doing ExactSearch after doing ANNSearch as the number of documents returned are less than "
+ "K, even when we have more than K filtered Ids. K: {}, ANNResults: {}, filteredIdCount: {}",
k,
docIdsToScoreMap.size(),
cardinality
);
docIdsToScoreMap = exactSearch(context, exactSearcherContext);
}
if (isFilteredExactSearchPreferred(cardinality)) {
return doExactSearch(context, acceptedDocs, k);
}
if (docIdsToScoreMap.isEmpty()) {
return Collections.emptyMap();
Map<Integer, Float> docIdsToScoreMap = doANNSearch(context, acceptedDocs, cardinality, k);
if (isExactSearchRequire(context, cardinality, docIdsToScoreMap.size())) {
// if filter is provided uses filtered docs else use get all docs
final BitSet docs = filterWeight != null ? acceptedDocs : null;
return doExactSearch(context, docs, k);
}
return docIdsToScoreMap;
}
Expand All @@ -185,13 +165,21 @@ private BitSet getFilteredDocsBitSet(final LeafReaderContext ctx) throws IOExcep
return createBitSet(scorer.iterator(), liveDocs, maxDoc);
}

private BitSet createBitSet(final DocIdSetIterator filteredDocIdsIterator, final Bits liveDocs, int maxDoc) throws IOException {
if (liveDocs == null && filteredDocIdsIterator instanceof BitSetIterator) {
private BitSet getAllDocsBitSet(final LeafReaderContext ctx) throws IOException {
final FloatVectorValues floatVectorValues = ctx.reader().getFloatVectorValues(this.knnQuery.getField());
final Bits liveDocs = ctx.reader().getLiveDocs();
final int maxDoc = ctx.reader().maxDoc();
final VectorScorer scorer = floatVectorValues.scorer(this.knnQuery.getQueryVector());
return createBitSet(scorer.iterator(), liveDocs, maxDoc);
}

private BitSet createBitSet(final DocIdSetIterator iterator, final Bits liveDocs, int maxDoc) throws IOException {
if (liveDocs == null && iterator instanceof BitSetIterator) {
// If we already have a BitSet and no deletions, reuse the BitSet
return ((BitSetIterator) filteredDocIdsIterator).getBitSet();
return ((BitSetIterator) iterator).getBitSet();
}
// Create a new BitSet from matching and live docs
FilteredDocIdSetIterator filterIterator = new FilteredDocIdSetIterator(filteredDocIdsIterator) {
FilteredDocIdSetIterator filterIterator = new FilteredDocIdSetIterator(iterator) {
@Override
protected boolean match(int doc) {
return liveDocs == null || liveDocs.get(doc);
Expand Down Expand Up @@ -221,6 +209,20 @@ private int[] bitSetToIntArray(final BitSet bitSet) {
return intArray;
}

private Map<Integer, Float> doExactSearch(final LeafReaderContext context, final BitSet acceptedDocs, int k) throws IOException {
final ExactSearcherContextBuilder exactSearcherContextBuilder = ExactSearcher.ExactSearcherContext.builder()
.k(k)
.isParentHits(true)
// setting to true, so that if quantization details are present we want to do search on the quantized
// vectors as this flow is used in first pass of search.
.useQuantizedVectorsForSearch(true)
.knnQuery(knnQuery);
if (acceptedDocs != null) {
exactSearcherContextBuilder.matchedDocs(acceptedDocs);
}
return exactSearch(context, exactSearcherContextBuilder.build());
}

private Map<Integer, Float> doANNSearch(
final LeafReaderContext context,
final BitSet filterIdsBitSet,
Expand All @@ -234,7 +236,7 @@ private Map<Integer, Float> doANNSearch(

if (fieldInfo == null) {
log.debug("[KNN] Field info not found for {}:{}", knnQuery.getField(), reader.getSegmentName());
return null;
return Collections.emptyMap();
}

KNNEngine knnEngine;
Expand Down Expand Up @@ -273,8 +275,8 @@ private Map<Integer, Float> doANNSearch(

List<String> engineFiles = KNNCodecUtil.getEngineFiles(knnEngine.getExtension(), knnQuery.getField(), reader.getSegmentInfo().info);
if (engineFiles.isEmpty()) {
log.debug("[KNN] No engine index found for field {} for segment {}", knnQuery.getField(), reader.getSegmentName());
return null;
log.info("[KNN] No native engine files found for field {} for segment {}", knnQuery.getField(), reader.getSegmentName());
return Collections.emptyMap();
}

Path indexPath = PathUtils.get(directory, engineFiles.get(0));
Expand Down Expand Up @@ -363,16 +365,9 @@ private Map<Integer, Float> doANNSearch(
indexAllocation.readUnlock();
indexAllocation.decRef();
}

/*
* Scores represent the distance of the documents with respect to given query vector.
* Lesser the score, the closer the document is to the query vector.
* Since by default results are retrieved in the descending order of scores, to get the nearest
* neighbors we are inverting the scores.
*/
if (results.length == 0) {
log.debug("[KNN] Query yielded 0 results");
return null;
return Collections.emptyMap();
}

return Arrays.stream(results)
Expand All @@ -381,15 +376,18 @@ private Map<Integer, Float> doANNSearch(

/**
* Execute exact search for the given matched doc ids and return the results as a map of docId to score.
*
* @return Map of docId to score for the exact search results.
* @throws IOException If an error occurs during the search.
*/
public Map<Integer, Float> exactSearch(
final LeafReaderContext leafReaderContext,
final ExactSearcher.ExactSearcherContext exactSearcherContext
) throws IOException {
return exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
final Map<Integer, Float> docIdsToScoreMap = exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext);
if (docIdsToScoreMap.isEmpty()) {
return Collections.emptyMap();
}
return docIdsToScoreMap;
}

@Override
Expand All @@ -402,7 +400,10 @@ public static float normalizeScore(float score) {
return -score + 1;
}

private boolean canDoExactSearch(final int filterIdsCount) {
private boolean isFilteredExactSearchPreferred(final int filterIdsCount) {
if (filterWeight == null) {
return false;
}
log.debug(
"Info for doing exact search filterIdsLength : {}, Threshold value: {}",
filterIdsCount,
Expand Down Expand Up @@ -441,14 +442,72 @@ private boolean isExactSearchThresholdSettingSet(int filterThresholdValue) {
return filterThresholdValue != KNNSettings.ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD_DEFAULT_VALUE;
}

/**
* This condition mainly checks whether exact search should be performed or not
* @param context LeafReaderContext
* @param filterIdsCount count of filtered Doc ids
* @param annResultCount Count of Nearest Neighbours we got after doing filtered ANN Search.
* @return boolean - true if exactSearch needs to be done after ANNSearch.
*/
private boolean isExactSearchRequire(final LeafReaderContext context, final int filterIdsCount, final int annResultCount) {
if (annResultCount == 0 && isMissingNativeEngineFiles(context)) {
log.info("Perform exact search after approximate search since no native engine files are available");
return true;
}
if (isFilteredExactSearchRequireAfterANNSearch(filterIdsCount, annResultCount)) {
log.debug(
"Doing ExactSearch after doing ANNSearch as the number of documents returned are less than "
+ "K, even when we have more than K filtered Ids. K: {}, ANNResults: {}, filteredIdCount: {}",
this.knnQuery.getK(),
annResultCount,
filterIdsCount
);
return true;
}
return false;
}

/**
* This condition mainly checks during filtered search we have more than K elements in filterIds but the ANN
* doesn't yeild K nearest neighbors.
* doesn't yield K nearest neighbors.
* @param filterIdsCount count of filtered Doc ids
* @param annResultCount Count of Nearest Neighbours we got after doing filtered ANN Search.
* @return boolean - true if exactSearch needs to be done after ANNSearch.
*/
private boolean canDoExactSearchAfterANNSearch(final int filterIdsCount, final int annResultCount) {
private boolean isFilteredExactSearchRequireAfterANNSearch(final int filterIdsCount, final int annResultCount) {
return filterWeight != null && filterIdsCount >= knnQuery.getK() && knnQuery.getK() > annResultCount;
}

/**
* This condition mainly checks whether segments has native engine files or not
* @return boolean - false if exactSearch needs to be done since no native engine files are in segments.
*/
private boolean isMissingNativeEngineFiles(LeafReaderContext context) {
final SegmentReader reader = Lucene.segmentReader(context.reader());
final FieldInfo fieldInfo = reader.getFieldInfos().fieldInfo(knnQuery.getField());
// if segment has no documents with at least 1 vector field, field info will be null
if (fieldInfo == null) {
return false;
}
final KNNEngine knnEngine = extractKNNEngine(fieldInfo);
final List<String> engineFiles = KNNCodecUtil.getEngineFiles(
knnEngine.getExtension(),
knnQuery.getField(),
reader.getSegmentInfo().info
);
return engineFiles.isEmpty();
}

private KNNEngine extractKNNEngine(FieldInfo fieldInfo) {
final String modelId = fieldInfo.getAttribute(MODEL_ID);
if (modelId == null) {
final String engineName = fieldInfo.attributes().getOrDefault(KNN_ENGINE, KNNEngine.NMSLIB.getName());
return KNNEngine.getEngine(engineName);
}
final ModelMetadata modelMetadata = modelDao.getMetadata(modelId);
if (!ModelUtil.isModelCreated(modelMetadata)) {
throw new RuntimeException("Model \"" + modelId + "\" is not created.");
}
return modelMetadata.getKnnEngine();
}
}
1 change: 0 additions & 1 deletion .../java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngineFieldVectorsWriterTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,5 @@ public void testRamByteUsed_whenValidInput_thenSuccess() {
.create(fieldInfo, InfoStream.getDefault());
// testing for value > 0 as we don't have a concrete way to find out expected bytes. This can OS dependent too.
Assert.assertTrue(byteWriter.ramBytesUsed() > 0);

}
}
7 changes: 4 additions & 3 deletions src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -327,8 +327,9 @@ public void testQueryScoreForFaissWithNonExistingModel() throws IOException {
}

@SneakyThrows
public void testShardWithoutFiles() {
public void testScorer_whenNoVectorFieldsInDocument_thenEmptyScorerIsReturned() {
final KNNQuery query = new KNNQuery(FIELD_NAME, QUERY_VECTOR, K, INDEX_NAME, (BitSetProducer) null);
KNNWeight.initialize(null);
final KNNWeight knnWeight = new KNNWeight(query, 0.0f);

final LeafReaderContext leafReaderContext = mock(LeafReaderContext.class);
Expand Down Expand Up @@ -361,8 +362,8 @@ public void testShardWithoutFiles() {
final FieldInfos fieldInfos = mock(FieldInfos.class);
final FieldInfo fieldInfo = mock(FieldInfo.class);
when(reader.getFieldInfos()).thenReturn(fieldInfos);
when(fieldInfos.fieldInfo(any())).thenReturn(fieldInfo);

// When no knn fields are available , field info for vector field will be null
when(fieldInfos.fieldInfo(FIELD_NAME)).thenReturn(null);
final Scorer knnScorer = knnWeight.scorer(leafReaderContext);
assertEquals(KNNScorer.emptyScorer(knnWeight), knnScorer);
}
Expand Down

0 comments on commit 53263a7

Please sign in to comment.