From be683c303b76fa03de9a93570f69ecf62369d8c1 Mon Sep 17 00:00:00 2001 From: Vijayan Balasubramanian Date: Fri, 27 Sep 2024 23:59:43 -0700 Subject: [PATCH] Add exact search if no engine files are in segments 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 --- CHANGELOG.md | 1 + .../opensearch/knn/index/query/KNNWeight.java | 155 ++++++++++++------ .../NativeEngineFieldVectorsWriterTests.java | 1 - .../knn/index/query/KNNWeightTests.java | 7 +- 4 files changed, 107 insertions(+), 57 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4030b913f..5e6ef6558 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java index 1c31ed725..1bdfa398b 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNWeight.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNWeight.java @@ -33,6 +33,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; @@ -121,50 +122,27 @@ public Scorer scorer(LeafReaderContext context) throws IOException { * @return A Map of docId to scores for top k results */ public Map 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 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 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; } @@ -185,13 +163,13 @@ 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 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); @@ -221,6 +199,20 @@ private int[] bitSetToIntArray(final BitSet bitSet) { return intArray; } + private Map 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 doANNSearch( final LeafReaderContext context, final BitSet filterIdsBitSet, @@ -234,7 +226,7 @@ private Map doANNSearch( if (fieldInfo == null) { log.debug("[KNN] Field info not found for {}:{}", knnQuery.getField(), reader.getSegmentName()); - return null; + return Collections.emptyMap(); } KNNEngine knnEngine; @@ -273,8 +265,8 @@ private Map doANNSearch( List 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)); @@ -363,16 +355,9 @@ private Map 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) @@ -381,7 +366,6 @@ private Map 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. */ @@ -389,7 +373,11 @@ public Map exactSearch( final LeafReaderContext leafReaderContext, final ExactSearcher.ExactSearcherContext exactSearcherContext ) throws IOException { - return exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext); + final Map docIdsToScoreMap = exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext); + if (docIdsToScoreMap.isEmpty()) { + return Collections.emptyMap(); + } + return docIdsToScoreMap; } @Override @@ -402,7 +390,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, @@ -441,14 +432,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 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(); + } } diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngineFieldVectorsWriterTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngineFieldVectorsWriterTests.java index 29e3531cf..bdeb1d2b6 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngineFieldVectorsWriterTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngineFieldVectorsWriterTests.java @@ -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); - } } diff --git a/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java b/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java index f92f32406..7a71c44be 100644 --- a/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java +++ b/src/test/java/org/opensearch/knn/index/query/KNNWeightTests.java @@ -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); @@ -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); }