diff --git a/CHANGELOG.md b/CHANGELOG.md index bc91f19a29..074ae07651 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Fix exact search and rescore scoring innerproduct and cosinesimil fields with L2 on model based and 2.17 to 2.19 indices [#3537](https://github.com/opensearch-project/k-NN/pull/3537) * Fix native memory leak in MemOptimizedNativeIndexBuildStrategy on failed merges [#3593](https://github.com/opensearch-project/k-NN/pull/3593) * Return model not found instead of the internal model index when getting a model before any model is trained [#3598](https://github.com/opensearch-project/k-NN/pull/3598) +* Support rescoring with `expand_nested_docs` on the Lucene engine, so the final reduction to k counts parent documents instead of child documents [#3579](https://github.com/opensearch-project/k-NN/pull/3579) * Fix double free of off-heap vectors when a native index build fails [#3621](https://github.com/opensearch-project/k-NN/pull/3621) * Honor query cancellation and timeout in memory-optimized search [#3620](https://github.com/opensearch-project/k-NN/pull/3620) @@ -57,3 +58,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), * Add native SIMD cosine scoring for FP16 and SQ formats, removing post-hoc score conversion [#3386](https://github.com/opensearch-project/k-NN/pull/3386) * Skip warmup for warm-tier indices to avoid unnecessary graph loading from remote store [#3565](https://github.com/opensearch-project/k-NN/pull/3565) * Avoid FP16 -> FP32 -> FP16 round trip when merging `half_float` segments [#3610](https://github.com/opensearch-project/k-NN/pull/3610) +* Report `exact_search` timings in the Profile API for nested k-NN queries with `expand_nested_docs` on the Lucene engine [#3579](https://github.com/opensearch-project/k-NN/pull/3579) diff --git a/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java b/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java index 1a20eef84a..a125ddb0bb 100644 --- a/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java +++ b/src/main/java/org/opensearch/knn/index/query/KNNQueryFactory.java @@ -153,16 +153,18 @@ public static Query create(CreateQueryRequest createQueryRequest) { ) ); - // Skip wrapping with RescoreKNNVectorQuery for expandNested as it will reduce the results to k after rescoring - // which will not return all the child documents - // TODO: Skip retrieving child docs in ExpandNestedDocsQuery if rescoring is enabled and instead retrieve them after rescoring and - // reducing to top K. - if (needsRescore && expandNested) { - log.warn("Rescoring is not supported when [{}] is set to true", EXPAND_NESTED); + if (needsRescore == false) { return luceneKnnQuery; } - return needsRescore ? new RescoreKNNVectorQuery(luceneKnnQuery, fieldName, k, vector, shardId) : luceneKnnQuery; + // ExpandNestedDocsQuery owns the rescore stage itself, because rescoring has to happen before the child + // documents are expanded. Wrapping it here would instead cut to k after the expansion, which counts child + // documents rather than parent documents and drops parents that belong in the result. + if (expandNested) { + return luceneKnnQuery; + } + + return new RescoreKNNVectorQuery(luceneKnnQuery, fieldName, k, vector, shardId); } private static int getDimension(float[] floatQueryVector, byte[] byteQueryVector) { diff --git a/src/main/java/org/opensearch/knn/index/query/RescoreKNNVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/RescoreKNNVectorQuery.java index 145b0729b7..7dbb709ccd 100644 --- a/src/main/java/org/opensearch/knn/index/query/RescoreKNNVectorQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/RescoreKNNVectorQuery.java @@ -8,7 +8,6 @@ import com.google.common.annotations.VisibleForTesting; import lombok.extern.log4j.Log4j2; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; @@ -25,7 +24,6 @@ import org.opensearch.knn.index.query.exactsearch.ExactSearcher; import org.opensearch.knn.indices.ModelDao; import org.opensearch.knn.profile.KNNProfileUtil; -import org.opensearch.knn.profile.query.KNNQueryTimingType; import org.opensearch.search.profile.ContextualProfileBreakdown; import org.opensearch.search.profile.query.QueryProfiler; @@ -115,21 +113,15 @@ private TopDocs searchLeaf( if (scorer == null) { return TopDocsCollector.EMPTY_TOPDOCS; } - DocIdSetIterator iterator = scorer.iterator(); - final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder() - .matchedDocsIterator(iterator) - .numberOfMatchedDocs(iterator.cost()) - // setting to false because in re-scoring we want to do exact search on full precision vectors - .useQuantizedVectorsForSearch(false) - .k(k) - .field(field) - .floatQueryVector(queryVector) - .build(); - TopDocs results = (TopDocs) KNNProfileUtil.profileBreakdown( + TopDocs results = QueryUtils.rescoreLeafWithFullPrecision( + searcher, profile, leafReaderContext, - KNNQueryTimingType.EXACT_SEARCH, - () -> searcher.searchLeaf(leafReaderContext, exactSearcherContext) + field, + queryVector, + scorer.iterator(), + k, + null ); if (leafReaderContext.docBase > 0) { for (ScoreDoc scoreDoc : results.scoreDocs) { diff --git a/src/main/java/org/opensearch/knn/index/query/common/QueryUtils.java b/src/main/java/org/opensearch/knn/index/query/common/QueryUtils.java index 4a4e942660..821fc61035 100644 --- a/src/main/java/org/opensearch/knn/index/query/common/QueryUtils.java +++ b/src/main/java/org/opensearch/knn/index/query/common/QueryUtils.java @@ -18,8 +18,13 @@ import org.apache.lucene.util.BitSet; import org.apache.lucene.util.BitSetIterator; import org.apache.lucene.util.Bits; +import org.opensearch.common.Nullable; import org.opensearch.knn.index.query.KNNWeight; +import org.opensearch.knn.index.query.exactsearch.ExactSearcher; import org.opensearch.knn.index.query.iterators.GroupedNestedDocIdSetIterator; +import org.opensearch.knn.profile.KNNProfileUtil; +import org.opensearch.knn.profile.query.KNNQueryTimingType; +import org.opensearch.search.profile.ContextualProfileBreakdown; import java.io.IOException; import java.util.ArrayList; @@ -183,4 +188,58 @@ protected boolean match(int doc) { }; return BitSet.of(filterIterator, maxDoc); } + + /** + * Re-scores an already gathered candidate set of a single leaf against full precision vectors. + * + * The approximate search may have walked the graph over quantized vectors, so its scores are inexact. This + * runs an exact search restricted to the candidates, always over full precision vectors, and reports the + * time it took under {@link KNNQueryTimingType#EXACT_SEARCH} when the search is being profiled. + * + * Only the float, non-radial rescore path is covered, which is what the Lucene engine queries need. The + * native engine queries carry extra context (radius, byte vectors, memory optimized search) and go through + * {@link org.opensearch.knn.index.query.KNNWeight#exactSearch} instead. + * + * Deliberately static rather than an instance method: callers that mock this class to stub out the + * surrounding search steps still exercise the real context construction here. + * + * @param exactSearcher the searcher performing the exact search + * @param profile breakdown to report the exact search time to, or null when the search is not being profiled + * @param leafReaderContext the leaf reader context + * @param field the vector field being searched + * @param floatQueryVector the full precision query vector + * @param matchedDocs the candidates to re-score + * @param k the number of results to keep; pass the candidate count to keep all of them + * @param parentsFilter when non-null, collapses each parent group to its best scoring child before the + * top k cut, so k counts parent documents rather than nested field documents + * @return the re-scored documents, sorted by descending score, with leaf local document IDs + * @throws IOException + */ + public static TopDocs rescoreLeafWithFullPrecision( + final ExactSearcher exactSearcher, + @Nullable final ContextualProfileBreakdown profile, + final LeafReaderContext leafReaderContext, + final String field, + final float[] floatQueryVector, + final DocIdSetIterator matchedDocs, + final int k, + @Nullable final BitSetProducer parentsFilter + ) throws IOException { + final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder() + .matchedDocsIterator(matchedDocs) + .numberOfMatchedDocs(matchedDocs.cost()) + // setting to false because in re-scoring we want to do exact search on full precision vectors + .useQuantizedVectorsForSearch(false) + .k(k) + .field(field) + .floatQueryVector(floatQueryVector) + .parentsFilter(parentsFilter) + .build(); + return (TopDocs) KNNProfileUtil.profileBreakdown( + profile, + leafReaderContext, + KNNQueryTimingType.EXACT_SEARCH, + () -> exactSearcher.searchLeaf(leafReaderContext, exactSearcherContext) + ); + } } diff --git a/src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java b/src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java index 863fd39ed6..96075990de 100644 --- a/src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java @@ -22,12 +22,22 @@ import org.apache.lucene.search.Weight; import org.apache.lucene.util.Bits; import org.opensearch.knn.index.query.common.QueryUtils; +import org.opensearch.knn.index.query.exactsearch.ExactSearcher; +import org.opensearch.knn.index.query.rescore.RescoreContext; +import org.opensearch.knn.indices.ModelDao; +import org.opensearch.knn.profile.KNNProfileUtil; +import org.opensearch.knn.profile.query.KNNQueryTimingType; +import org.opensearch.search.profile.ContextualProfileBreakdown; +import org.opensearch.search.profile.query.QueryProfiler; import java.io.IOException; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; /** * This query is for a nested k-NN field to return multiple nested field documents @@ -37,11 +47,37 @@ * they are reduced to the top k results. Then, it constructs filtered document IDs for nested field documents * from these top k parent documents. Using these document IDs, it executes an exact nearest neighbor search * with a k value of Integer.MAX_VALUE, which provides scores for all specified nested field documents. + * + * When rescoring is enabled, two extra stages run between the approximate search and the expansion. The + * approximate search walked the graph over quantized vectors, so every parent's representative child was + * elected on inexact scores. The rescore stage re-scores all siblings of the oversampled candidates against + * full precision vectors, collapses each parent group back to its single best child, and cuts the result to + * k. Expanding only after that cut is what keeps the list at one row per document until the final top-k is + * decided; cutting to k after the expansion would count child documents instead of parent documents and drop + * parents that belong in the result. This mirrors the ordering that + * {@link org.opensearch.knn.index.query.nativelib.NativeEngineKnnVectorQuery} uses for the native engines. */ @Builder public class ExpandNestedDocsQuery extends Query { final private InternalNestedKnnVectorQuery internalNestedKnnVectorQuery; final private QueryUtils queryUtils; + /** + * Number of oversampled parent candidates the approximate search kept for rescoring. + * {@link RescoreContext#NO_RESCORE_NEEDED} means no rescoring, in which case the expansion runs + * directly on the approximate search results. + */ + @Builder.Default + final private int rescoreK = RescoreContext.NO_RESCORE_NEEDED; + /** + * Full precision query vector. Required when {@link #rescoreK} enables rescoring, otherwise unused. + */ + final private float[] floatQueryVector; + /** + * Runs the full precision rescore pass. Optional test seam; when unset one is created on demand, and only + * when {@link #rescoreK} enables rescoring. It is not created up front because building it reaches for the + * {@link ModelDao} singleton, which is only usable once the plugin has been initialized on a node. + */ + final private ExactSearcher exactSearcher; @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { @@ -49,9 +85,17 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo Weight weight = docAndScoreQuery.createWeight(searcher, scoreMode, boost); IndexReader reader = searcher.getIndexReader(); List leafReaderContexts = reader.leaves(); + Weight filterWeight = getFilterWeight(searcher); + // Both the rescore and the expansion stage need the filter bits of every leaf they touch, so they are + // built once per leaf and shared instead of being rebuilt by each stage. + Map filterBitsByLeaf = new ConcurrentHashMap<>(); + ContextualProfileBreakdown profile = getProfileBreakdown(searcher); List> perLeafResults; perLeafResults = queryUtils.doSearch(searcher, leafReaderContexts, weight); - TopDocs[] topDocs = retrieveAll(searcher, leafReaderContexts, perLeafResults); + if (isRescoreEnabled()) { + perLeafResults = rescoreToTopKParents(searcher, leafReaderContexts, perLeafResults, filterWeight, filterBitsByLeaf, profile); + } + TopDocs[] topDocs = retrieveAll(searcher, leafReaderContexts, perLeafResults, filterWeight, filterBitsByLeaf, profile); int sum = 0; for (TopDocs topDoc : topDocs) { sum += topDoc.scoreDocs.length; @@ -63,26 +107,118 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo return queryUtils.createDocAndScoreQuery(reader, topK).createWeight(searcher, scoreMode, boost); } + private boolean isRescoreEnabled() { + return rescoreK != RescoreContext.NO_RESCORE_NEEDED; + } + + private ExactSearcher resolveExactSearcher() { + return exactSearcher != null ? exactSearcher : new ExactSearcher(ModelDao.OpenSearchKNNModelDao.getInstance()); + } + + /** + * Returns the profile breakdown for this query, or null when the search is not being profiled. The node for + * this query is added to the profile tree by + * {@link org.opensearch.knn.index.query.lucene.LuceneEngineKnnVectorQuery#createWeight} before it delegates + * here, so the lookup always resolves. The breakdown only carries the metrics that + * {@link org.opensearch.knn.plugin.KNNPlugin#getQueryProfileMetricsProvider} registers for this query type. + */ + private ContextualProfileBreakdown getProfileBreakdown(final IndexSearcher indexSearcher) { + QueryProfiler profiler = KNNProfileUtil.getProfiler(indexSearcher); + if (profiler == null) { + return null; + } + return (ContextualProfileBreakdown) profiler.getProfileBreakdown(this); + } + + /** + * Re-scores the oversampled candidates against full precision vectors and reduces them to the top k + * parent documents. + * + */ + private List> rescoreToTopKParents( + final IndexSearcher indexSearcher, + final List leafReaderContexts, + final List> perLeafResults, + final Weight filterWeight, + final Map filterBitsByLeaf, + final ContextualProfileBreakdown profile + ) throws IOException { + final int k = internalNestedKnnVectorQuery.getK(); + final ExactSearcher searcher = resolveExactSearcher(); + List> rescoreTasks = new ArrayList<>(leafReaderContexts.size()); + for (int i = 0; i < perLeafResults.size(); i++) { + LeafReaderContext leafReaderContext = leafReaderContexts.get(i); + Map leafResult = perLeafResults.get(i); + int leafOrd = i; + rescoreTasks.add(() -> { + if (leafResult.isEmpty()) { + return NestedKnnUtil.EMPTY_TOP_DOCS; + } + Bits queryFilter = filterBits(filterBitsByLeaf, leafReaderContext, filterWeight); + DocIdSetIterator allSiblings = queryUtils.getAllSiblings( + leafReaderContext, + leafResult.keySet(), + internalNestedKnnVectorQuery.getParentFilter(), + queryFilter + ); + TopDocs leafTopDocs = QueryUtils.rescoreLeafWithFullPrecision( + searcher, + profile, + leafReaderContext, + internalNestedKnnVectorQuery.getField(), + floatQueryVector, + allSiblings, + k, + // passing the parent filter makes the searcher collapse each parent group to its best child + internalNestedKnnVectorQuery.getParentFilter() + ); + // TopDocs#merge below flattens every leaf into one score sorted array, which throws away the per + // leaf grouping that the expansion stage still needs: it pairs perLeafResults.get(i) with + // leafReaderContexts.get(i). Doc ids cannot recover the leaf, because they are segment local and + // restart at 0 in each leaf, so leaf 0 and leaf 1 can both hold a doc 1. shardIndex is unused on + // this path, so the leaf ordinal is parked there; merge only ever reads it, so it is still set + // when the survivors are regrouped per leaf below. + for (ScoreDoc scoreDoc : leafTopDocs.scoreDocs) { + scoreDoc.shardIndex = leafOrd; + } + return leafTopDocs; + }); + } + TopDocs[] perLeafTopDocs = indexSearcher.getTaskExecutor().invokeAll(rescoreTasks).toArray(TopDocs[]::new); + TopDocs topKParents = TopDocs.merge(k, perLeafTopDocs); + + List> reducedResults = new ArrayList<>(leafReaderContexts.size()); + for (int i = 0; i < leafReaderContexts.size(); i++) { + reducedResults.add(new HashMap<>()); + } + for (ScoreDoc scoreDoc : topKParents.scoreDocs) { + reducedResults.get(scoreDoc.shardIndex).put(scoreDoc.doc, scoreDoc.score); + } + return reducedResults; + } + private TopDocs[] retrieveAll( final IndexSearcher indexSearcher, final List leafReaderContexts, - final List> perLeafResults + final List> perLeafResults, + final Weight filterWeight, + final Map filterBitsByLeaf, + final ContextualProfileBreakdown profile ) throws IOException { // Construct query List> nestedQueryTasks = new ArrayList<>(leafReaderContexts.size()); - Weight filterWeight = getFilterWeight(indexSearcher); for (int i = 0; i < perLeafResults.size(); i++) { LeafReaderContext leafReaderContext = leafReaderContexts.get(i); int finalI = i; nestedQueryTasks.add(() -> { - Bits queryFilter = queryUtils.createBits(leafReaderContext, filterWeight); + Bits queryFilter = filterBits(filterBitsByLeaf, leafReaderContext, filterWeight); DocIdSetIterator allSiblings = queryUtils.getAllSiblings( leafReaderContext, perLeafResults.get(finalI).keySet(), internalNestedKnnVectorQuery.getParentFilter(), queryFilter ); - TopDocs topDocs = internalNestedKnnVectorQuery.knnExactSearch(leafReaderContext, allSiblings); + TopDocs topDocs = scoreAllSiblings(leafReaderContext, allSiblings, profile); // Update doc id from segment id to shard id for (ScoreDoc scoreDoc : topDocs.scoreDocs) { scoreDoc.doc = scoreDoc.doc + leafReaderContext.docBase; @@ -93,6 +229,55 @@ private TopDocs[] retrieveAll( return indexSearcher.getTaskExecutor().invokeAll(nestedQueryTasks).toArray(TopDocs[]::new); } + /** + * Scores every sibling of the surviving parents, each keeping its own score, with no collapsing. + * + */ + private TopDocs scoreAllSiblings( + final LeafReaderContext leafReaderContext, + final DocIdSetIterator allSiblings, + final ContextualProfileBreakdown profile + ) throws IOException { + if (isRescoreEnabled() == false) { + return (TopDocs) KNNProfileUtil.profileBreakdown( + profile, + leafReaderContext, + KNNQueryTimingType.EXACT_SEARCH, + () -> internalNestedKnnVectorQuery.knnExactSearch(leafReaderContext, allSiblings) + ); + } + return QueryUtils.rescoreLeafWithFullPrecision( + resolveExactSearcher(), + profile, + leafReaderContext, + internalNestedKnnVectorQuery.getField(), + floatQueryVector, + allSiblings, + // every sibling keeps its own score, so no top-k cut and no parent filter here + (int) allSiblings.cost(), + null + ); + } + + /** + * Returns the filter bits of a leaf, building them on first use. Reused across the rescore and expansion + * stages, which would otherwise each rebuild the same bit set. + */ + private Bits filterBits(final Map filterBitsByLeaf, final LeafReaderContext leafReaderContext, final Weight filterWeight) + throws IOException { + Bits cached = filterBitsByLeaf.get(leafReaderContext.ord); + if (cached != null) { + return cached; + } + Bits bits = queryUtils.createBits(leafReaderContext, filterWeight); + if (bits == null) { + // Nothing worth caching, and the map does not accept null values. Callers treat null as "no filter". + return null; + } + Bits existing = filterBitsByLeaf.putIfAbsent(leafReaderContext.ord, bits); + return existing != null ? existing : bits; + } + /** * This is copied from {@link org.apache.lucene.search.AbstractKnnVectorQuery#rewrite} */ @@ -119,12 +304,14 @@ public boolean equals(final Object o) { return false; } ExpandNestedDocsQuery other = (ExpandNestedDocsQuery) o; - return internalNestedKnnVectorQuery.equals(other.internalNestedKnnVectorQuery); + // rescoreK is not part of internalNestedKnnVectorQuery's equality, so it has to be compared here. + // Otherwise two queries differing only in oversample_factor would be considered equal by the query cache. + return internalNestedKnnVectorQuery.equals(other.internalNestedKnnVectorQuery) && rescoreK == other.rescoreK; } @Override public int hashCode() { - return internalNestedKnnVectorQuery.hashCode(); + return Objects.hash(internalNestedKnnVectorQuery, rescoreK); } @Override diff --git a/src/main/java/org/opensearch/knn/index/query/lucenelib/InternalNestedKnnFloatVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/lucenelib/InternalNestedKnnFloatVectorQuery.java index 1dbdd41e24..3289edac9e 100644 --- a/src/main/java/org/opensearch/knn/index/query/lucenelib/InternalNestedKnnFloatVectorQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/lucenelib/InternalNestedKnnFloatVectorQuery.java @@ -52,8 +52,7 @@ public InternalNestedKnnFloatVectorQuery( luceneK, parentFilter, k, - rescoreK, - true + rescoreK ); } diff --git a/src/main/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactory.java b/src/main/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactory.java index 54981d709a..f7f9048ca3 100644 --- a/src/main/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactory.java +++ b/src/main/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactory.java @@ -74,7 +74,7 @@ public static Query createNestedKnnVectorQuery( if (expandNestedDocs) { return new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery( new InternalNestedKnnFloatVectorQuery(fieldName, vector, filterQuery, luceneK, parentFilter, k, rescoreK) - ).queryUtils(QueryUtils.getInstance()).build(); + ).queryUtils(QueryUtils.getInstance()).rescoreK(rescoreK).floatQueryVector(vector).build(); } return new OSDiversifyingChildrenFloatKnnVectorQuery(fieldName, vector, filterQuery, luceneK, parentFilter, k, rescoreK); } diff --git a/src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java b/src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java index 5cd749b53c..dbb6abac04 100644 --- a/src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java +++ b/src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java @@ -27,7 +27,6 @@ public final class OSDiversifyingChildrenFloatKnnVectorQuery extends Diversifyin private final int k; private final int rescoreK; - private final boolean expandNestedDocs; private final BitSetProducer parentFilter; public OSDiversifyingChildrenFloatKnnVectorQuery( @@ -38,24 +37,10 @@ public OSDiversifyingChildrenFloatKnnVectorQuery( final BitSetProducer parentFilter, final int k, final int rescoreK - ) { - this(fieldName, vector, filterQuery, luceneK, parentFilter, k, rescoreK, false); - } - - public OSDiversifyingChildrenFloatKnnVectorQuery( - final String fieldName, - final float[] vector, - final Query filterQuery, - final int luceneK, - final BitSetProducer parentFilter, - final int k, - final int rescoreK, - final boolean expandNestedDocs ) { super(fieldName, vector, filterQuery, luceneK, parentFilter); this.k = k; this.rescoreK = rescoreK; - this.expandNestedDocs = expandNestedDocs; this.parentFilter = parentFilter; } @@ -74,8 +59,7 @@ protected TopDocs approximateSearch( @Override protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) { - // TODO: Fix this if condition after adding rescoring logic inside ExpandNestedDocsQuery when rescoring is enabled - if (rescoreK != RescoreContext.NO_RESCORE_NEEDED && !expandNestedDocs) { + if (rescoreK != RescoreContext.NO_RESCORE_NEEDED) { // When rescoring is enabled, merge to oversampled k (rescore budget) rather than the // full luceneK which may have been expanded by ef_search. return TopDocs.merge(rescoreK, perLeafResults); diff --git a/src/main/java/org/opensearch/knn/plugin/KNNPlugin.java b/src/main/java/org/opensearch/knn/plugin/KNNPlugin.java index e0060c3531..970c8f2191 100644 --- a/src/main/java/org/opensearch/knn/plugin/KNNPlugin.java +++ b/src/main/java/org/opensearch/knn/plugin/KNNPlugin.java @@ -56,6 +56,7 @@ import org.opensearch.knn.index.query.RescoreKNNVectorQuery; import org.opensearch.knn.index.query.RescoreRadialSearchQuery; import org.opensearch.knn.index.query.exactsearch.ExactSearcher; +import org.opensearch.knn.index.query.lucenelib.ExpandNestedDocsQuery; import org.opensearch.knn.index.query.nativelib.NativeEngineKnnVectorQuery; import org.opensearch.knn.index.query.parser.KNNQueryBuilderParser; import org.opensearch.knn.index.util.KNNClusterUtil; @@ -222,7 +223,7 @@ public class KNNPlugin extends Plugin @Override public Optional getQueryProfileMetricsProvider() { return Optional.of((searchContext, query) -> { - if (query instanceof KNNQuery || query instanceof RescoreKNNVectorQuery) { + if (query instanceof KNNQuery || query instanceof RescoreKNNVectorQuery || query instanceof ExpandNestedDocsQuery) { return KNNMetrics.getKNNQueryMetrics(); } else if (query instanceof NativeEngineKnnVectorQuery) { return KNNMetrics.getNativeMetrics(); diff --git a/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java b/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java index 0a5877114e..5e3472c5cc 100644 --- a/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java +++ b/src/test/java/org/opensearch/knn/index/query/KNNQueryFactoryTests.java @@ -565,7 +565,7 @@ public void testCreate_whenExpandNestedDocsQueryWithLucene_thenCreateExpandNeste testExpandNestedDocsQuery(KNNEngine.LUCENE, OSDiversifyingChildrenFloatKnnVectorQuery.class, VectorDataType.FLOAT, false); } - public void testCreate_whenRescoreWithExpandNested_thenSkipRescoreWrapper() { + public void testCreate_whenRescoreWithExpandNested_thenExpandNestedDocsQueryOwnsRescore() { QueryShardContext queryShardContext = mock(QueryShardContext.class); BitSetProducer parentFilter = mock(BitSetProducer.class); when(queryShardContext.getParentFilter()).thenReturn(parentFilter); @@ -586,9 +586,28 @@ public void testCreate_whenRescoreWithExpandNested_thenSkipRescoreWrapper() { .build(); Query query = KNNQueryFactory.create(createQueryRequest); - // Should return LuceneEngineKnnVectorQuery without RescoreKNNVectorQuery wrapper + // Rescoring has to happen before the child docs are expanded, so ExpandNestedDocsQuery runs it + // internally rather than being wrapped by RescoreKNNVectorQuery. assertEquals(LuceneEngineKnnVectorQuery.class, query.getClass()); assertFalse(query instanceof RescoreKNNVectorQuery); + Query luceneQuery = ((LuceneEngineKnnVectorQuery) query).getLuceneQuery(); + assertEquals(ExpandNestedDocsQuery.class, luceneQuery.getClass()); + // The rescore-enabled query must be distinguishable from the plain expansion for query caching + assertNotEquals( + luceneQuery, + ((LuceneEngineKnnVectorQuery) KNNQueryFactory.create( + KNNQueryFactory.CreateQueryRequest.builder() + .knnEngine(KNNEngine.LUCENE) + .indexName(testIndexName) + .fieldName(testFieldName) + .vector(testQueryVector) + .vectorDataType(VectorDataType.FLOAT) + .k(testK) + .expandNested(true) + .context(queryShardContext) + .build() + )).getLuceneQuery() + ); } public void testCreate_whenRescoreWithoutExpandNested_thenWrapWithRescoreQuery() { diff --git a/src/test/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedEDocsQueryTests.java b/src/test/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedEDocsQueryTests.java index ecda53e1b8..6894f69cd6 100644 --- a/src/test/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedEDocsQueryTests.java +++ b/src/test/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedEDocsQueryTests.java @@ -13,9 +13,11 @@ import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreDoc; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.TotalHits; import org.apache.lucene.search.Weight; import org.apache.lucene.search.join.BitSetProducer; import org.apache.lucene.store.ByteBuffersDirectory; @@ -25,17 +27,28 @@ import org.mockito.ArgumentCaptor; import org.opensearch.knn.index.query.ResultUtil; import org.opensearch.knn.index.query.common.QueryUtils; +import org.opensearch.knn.index.query.exactsearch.ExactSearcher; +import org.opensearch.knn.index.query.rescore.RescoreContext; +import org.opensearch.knn.profile.query.KNNQueryTimingType; +import org.opensearch.search.internal.ContextIndexSearcher; +import org.opensearch.search.profile.ContextualProfileBreakdown; +import org.opensearch.search.profile.Timer; +import org.opensearch.search.profile.query.QueryProfiler; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.Executors; +import static org.junit.Assert.assertNotEquals; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.atLeast; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -126,12 +139,19 @@ public void testCreateWeight_whenCalled_thenSucceed() { when(queryUtils.getAllSiblings(any(), any(), any(), any())).thenReturn(allSiblings); when(queryUtils.createDocAndScoreQuery(eq(reader), any())).thenReturn(finalQuery); + ExactSearcher exactSearcher = mock(ExactSearcher.class); + // Run - ExpandNestedDocsQuery query = new ExpandNestedDocsQuery(internalQuery, queryUtils); + ExpandNestedDocsQuery query = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .exactSearcher(exactSearcher) + .build(); Weight finalWeigh = query.createWeight(indexSearcher, scoreMode, 1.f); // Verify assertEquals(expectedWeight, finalWeigh); + // No rescore context, so the full precision rescore pass must not run at all + verify(exactSearcher, never()).searchLeaf(any(), any()); verify(queryUtils).createBits(leaf1, filterWeight); verify(queryUtils).createBits(leaf2, filterWeight); verify(queryUtils).getAllSiblings(leaf1, perLeafResults.get(0).keySet(), parentFilter, queryFilterBits); @@ -153,4 +173,357 @@ public void testCreateWeight_whenCalled_thenSucceed() { assertEquals(2, iteratorCaptor.getValue().nextDoc()); assertEquals(DocIdSetIterator.NO_MORE_DOCS, iteratorCaptor.getValue().nextDoc()); } + + /** + * Regression test for https://github.com/opensearch-project/k-NN/issues/3125. + * + * With rescoring enabled the query must re-score the oversampled candidates on full precision vectors, + * collapse each parent group to its best child, and cut to k parent documents before expanding the + * children. This asserts the ordering: the expansion only sees the parents that survived the cut, so the + * final result cannot be truncated to k child documents. + */ + @SneakyThrows + public void testCreateWeight_whenRescoreEnabled_thenRescoreAndCutToTopKParentsBeforeExpanding() { + Directory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document doc1 = new Document(); + doc1.add(new FloatPoint("vector", 1.0f, 2.0f, 3.0f)); + writer.addDocument(doc1); + writer.flush(); + Document doc2 = new Document(); + doc2.add(new FloatPoint("vector", 4.0f, 5.0f, 6.0f)); + writer.addDocument(doc2); + writer.commit(); + } + + IndexReader reader = DirectoryReader.open(directory); + List leaves = reader.leaves(); + assertEquals(2, leaves.size()); + LeafReaderContext leaf1 = leaves.get(0); + LeafReaderContext leaf2 = leaves.get(1); + + int k = 2; + int rescoreK = 4; + float[] queryVector = new float[] { 1.0f, 2.0f, 3.0f }; + + Weight filterWeight = mock(Weight.class); + IndexSearcher indexSearcher = mock(IndexSearcher.class); + when(indexSearcher.getIndexReader()).thenReturn(reader); + when(indexSearcher.getTaskExecutor()).thenReturn(taskExecutor); + when(indexSearcher.createWeight(any(), eq(ScoreMode.COMPLETE_NO_SCORES), eq(1.0F))).thenReturn(filterWeight); + + Weight queryWeight = mock(Weight.class); + ScoreMode scoreMode = mock(ScoreMode.class); + Query docAndScoreQuery = mock(Query.class); + when(docAndScoreQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(queryWeight); + + BitSetProducer parentFilter = mock(BitSetProducer.class); + InternalNestedKnnVectorQuery internalQuery = mock(InternalNestedKnnVectorQuery.class); + when(internalQuery.knnRewrite(indexSearcher)).thenReturn(docAndScoreQuery); + when(internalQuery.getK()).thenReturn(k); + when(internalQuery.getFilter()).thenReturn(mock(Query.class)); + when(internalQuery.getField()).thenReturn("field"); + when(internalQuery.getParentFilter()).thenReturn(parentFilter); + + // Oversampled candidate pool: two candidate parents per leaf. + List> perLeafResults = Arrays.asList( + new HashMap<>(Map.of(1, 20f, 5, 19f)), + new HashMap<>(Map.of(0, 21f, 3, 18f)) + ); + + Bits queryFilterBits = mock(Bits.class); + DocIdSetIterator allSiblings = mock(DocIdSetIterator.class); + QueryUtils queryUtils = mock(QueryUtils.class); + when(queryUtils.doSearch(indexSearcher, reader.leaves(), queryWeight)).thenReturn(perLeafResults); + when(queryUtils.createBits(any(), any())).thenReturn(queryFilterBits); + when(queryUtils.getAllSiblings(any(), any(), any(), any())).thenReturn(allSiblings); + + // Rescore collapses each leaf to one row per parent. Doc 5 in leaf1 and doc 3 in leaf2 lose on full + // precision scores and must be dropped by the cut to k=2 parents. + // Stubs are keyed on the leaf rather than on call order, because TaskExecutor#invokeAll does not + // promise to run the tasks in the order they were submitted. + TopDocs rescoredLeaf1 = topDocs(new ScoreDoc(1, 30f), new ScoreDoc(5, 5f)); + TopDocs rescoredLeaf2 = topDocs(new ScoreDoc(0, 10f), new ScoreDoc(3, 4f)); + // Expansion returns every child of the surviving parents; there must be no cut after this point. + TopDocs expandedLeaf1 = topDocs(new ScoreDoc(0, 30f), new ScoreDoc(1, 28f), new ScoreDoc(2, 27f)); + TopDocs expandedLeaf2 = topDocs(new ScoreDoc(0, 10f), new ScoreDoc(1, 9f), new ScoreDoc(2, 8f)); + // Both stages go through the exact searcher when rescoring. The collapsing stage is the one that carries + // a parent filter, so that is what tells the two apart. + ExactSearcher exactSearcher = mock(ExactSearcher.class); + when(exactSearcher.searchLeaf(any(), any())).thenAnswer(invocation -> { + boolean isFirstLeaf = ((LeafReaderContext) invocation.getArgument(0)).ord == 0; + ExactSearcher.ExactSearcherContext context = invocation.getArgument(1); + if (context.getParentsFilter() != null) { + return isFirstLeaf ? rescoredLeaf1 : rescoredLeaf2; + } + return isFirstLeaf ? expandedLeaf1 : expandedLeaf2; + }); + + Weight expectedWeight = mock(Weight.class); + Query finalQuery = mock(Query.class); + when(finalQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(expectedWeight); + when(queryUtils.createDocAndScoreQuery(eq(reader), any())).thenReturn(finalQuery); + + // Run + ExpandNestedDocsQuery query = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .rescoreK(rescoreK) + .floatQueryVector(queryVector) + .exactSearcher(exactSearcher) + .build(); + Weight finalWeight = query.createWeight(indexSearcher, scoreMode, 1.f); + + assertEquals(expectedWeight, finalWeight); + + // Two collapsing searches plus two expansion searches, one of each per leaf. Lucene's own exact search + // must not be used at all here, because it would score the children on quantized vectors. + ArgumentCaptor contextCaptor = ArgumentCaptor.forClass( + ExactSearcher.ExactSearcherContext.class + ); + verify(exactSearcher, times(4)).searchLeaf(any(), contextCaptor.capture()); + verify(internalQuery, never()).knnExactSearch(any(), any()); + + List collapsing = contextCaptor.getAllValues() + .stream() + .filter(context -> context.getParentsFilter() != null) + .toList(); + List expanding = contextCaptor.getAllValues() + .stream() + .filter(context -> context.getParentsFilter() == null) + .toList(); + assertEquals(2, collapsing.size()); + assertEquals(2, expanding.size()); + + // Every search scores full precision vectors, and every one targets the right field and query vector + for (ExactSearcher.ExactSearcherContext context : contextCaptor.getAllValues()) { + assertFalse(context.isUseQuantizedVectorsForSearch()); + assertEquals("field", context.getField()); + assertTrue(Arrays.equals(queryVector, context.getFloatQueryVector())); + } + // The collapsing stage cuts to k parents; the expansion stage keeps every sibling it is given + for (ExactSearcher.ExactSearcherContext context : collapsing) { + assertEquals(parentFilter, context.getParentsFilter()); + assertEquals(k, context.getK()); + } + + // The rescore stage saw the whole oversampled pool ... + verify(queryUtils).getAllSiblings(leaf1, Set.of(1, 5), parentFilter, queryFilterBits); + verify(queryUtils).getAllSiblings(leaf2, Set.of(0, 3), parentFilter, queryFilterBits); + // ... but the expansion only saw the parents that survived the cut to k. + verify(queryUtils).getAllSiblings(leaf1, Set.of(1), parentFilter, queryFilterBits); + verify(queryUtils).getAllSiblings(leaf2, Set.of(0), parentFilter, queryFilterBits); + + // All children of both surviving parents are returned; nothing is truncated to k. + ArgumentCaptor topDocsCaptor = ArgumentCaptor.forClass(TopDocs.class); + verify(queryUtils).createDocAndScoreQuery(eq(reader), topDocsCaptor.capture()); + assertEquals(6, topDocsCaptor.getValue().scoreDocs.length); + } + + @SneakyThrows + public void testCreateWeight_whenRescoreEnabledAndNoCandidates_thenSkipRescoreForThatLeaf() { + Directory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document doc = new Document(); + doc.add(new FloatPoint("vector", 1.0f, 2.0f, 3.0f)); + writer.addDocument(doc); + writer.commit(); + } + + IndexReader reader = DirectoryReader.open(directory); + assertEquals(1, reader.leaves().size()); + + IndexSearcher indexSearcher = mock(IndexSearcher.class); + when(indexSearcher.getIndexReader()).thenReturn(reader); + when(indexSearcher.getTaskExecutor()).thenReturn(taskExecutor); + + Weight queryWeight = mock(Weight.class); + ScoreMode scoreMode = mock(ScoreMode.class); + Query docAndScoreQuery = mock(Query.class); + when(docAndScoreQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(queryWeight); + + InternalNestedKnnVectorQuery internalQuery = mock(InternalNestedKnnVectorQuery.class); + when(internalQuery.knnRewrite(indexSearcher)).thenReturn(docAndScoreQuery); + when(internalQuery.getK()).thenReturn(2); + // Null filter keeps getFilterWeight from touching the searcher + when(internalQuery.getFilter()).thenReturn(null); + when(internalQuery.getField()).thenReturn("field"); + when(internalQuery.getParentFilter()).thenReturn(mock(BitSetProducer.class)); + when(internalQuery.knnExactSearch(any(), any())).thenReturn(topDocs()); + + QueryUtils queryUtils = mock(QueryUtils.class); + when(queryUtils.doSearch(indexSearcher, reader.leaves(), queryWeight)).thenReturn(List.of(Map.of())); + when(queryUtils.getAllSiblings(any(), any(), any(), any())).thenReturn(DocIdSetIterator.empty()); + + ExactSearcher exactSearcher = mock(ExactSearcher.class); + when(exactSearcher.searchLeaf(any(), any())).thenReturn(topDocs()); + + ExpandNestedDocsQuery query = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .rescoreK(4) + .floatQueryVector(new float[] { 1.0f, 2.0f, 3.0f }) + .exactSearcher(exactSearcher) + .build(); + + Weight weight = query.createWeight(indexSearcher, scoreMode, 1.f); + + // With no candidates there is nothing to collapse, so the rescore stage is skipped for that leaf + ArgumentCaptor contextCaptor = ArgumentCaptor.forClass( + ExactSearcher.ExactSearcherContext.class + ); + verify(exactSearcher, atLeast(0)).searchLeaf(any(), contextCaptor.capture()); + assertTrue( + "no collapsing search should run for a leaf with no candidates", + contextCaptor.getAllValues().stream().noneMatch(context -> context.getParentsFilter() != null) + ); + assertNotNull(weight); + } + + public void testEquals_whenRescoreKDiffers_thenNotEqual() { + InternalNestedKnnVectorQuery internalQuery = mock(InternalNestedKnnVectorQuery.class); + QueryUtils queryUtils = mock(QueryUtils.class); + + ExpandNestedDocsQuery noRescore = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery( + internalQuery + ).queryUtils(queryUtils).rescoreK(RescoreContext.NO_RESCORE_NEEDED).build(); + ExpandNestedDocsQuery rescore = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .rescoreK(10) + .build(); + ExpandNestedDocsQuery sameRescore = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery( + internalQuery + ).queryUtils(queryUtils).rescoreK(10).build(); + + assertNotEquals(noRescore, rescore); + assertNotEquals(noRescore.hashCode(), rescore.hashCode()); + assertEquals(rescore, sameRescore); + assertEquals(rescore.hashCode(), sameRescore.hashCode()); + } + + /** + * When the search is profiled, both exact search stages must be timed under this query's node so that the + * time they take shows up in the Profile API's {@code exact_search} breakdown. + */ + @SneakyThrows + public void testCreateWeight_whenProfilerEnabled_thenTimeBothExactSearchStages() { + Directory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document doc = new Document(); + doc.add(new FloatPoint("vector", 1.0f, 2.0f, 3.0f)); + writer.addDocument(doc); + writer.commit(); + } + + IndexReader reader = DirectoryReader.open(directory); + assertEquals(1, reader.leaves().size()); + LeafReaderContext leaf = reader.leaves().get(0); + + // The profiler is only reachable through a ContextIndexSearcher, so a plain IndexSearcher mock would + // leave profiling switched off. + ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class); + when(indexSearcher.getIndexReader()).thenReturn(reader); + when(indexSearcher.getTaskExecutor()).thenReturn(taskExecutor); + + Timer timer = mock(Timer.class); + ContextualProfileBreakdown breakdown = mock(ContextualProfileBreakdown.class); + when(breakdown.context(leaf)).thenReturn(breakdown); + when(breakdown.getTimer(KNNQueryTimingType.EXACT_SEARCH)).thenReturn(timer); + QueryProfiler profiler = mock(QueryProfiler.class); + when(profiler.getProfileBreakdown(any())).thenReturn(breakdown); + when(indexSearcher.getProfiler()).thenReturn(profiler); + + Weight queryWeight = mock(Weight.class); + ScoreMode scoreMode = mock(ScoreMode.class); + Query docAndScoreQuery = mock(Query.class); + when(docAndScoreQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(queryWeight); + + BitSetProducer parentFilter = mock(BitSetProducer.class); + InternalNestedKnnVectorQuery internalQuery = mock(InternalNestedKnnVectorQuery.class); + when(internalQuery.knnRewrite(indexSearcher)).thenReturn(docAndScoreQuery); + when(internalQuery.getK()).thenReturn(1); + when(internalQuery.getFilter()).thenReturn(null); + when(internalQuery.getField()).thenReturn("field"); + when(internalQuery.getParentFilter()).thenReturn(parentFilter); + + QueryUtils queryUtils = mock(QueryUtils.class); + when(queryUtils.doSearch(indexSearcher, reader.leaves(), queryWeight)).thenReturn(List.of(new HashMap<>(Map.of(1, 20f, 5, 19f)))); + when(queryUtils.getAllSiblings(any(), any(), any(), any())).thenReturn(DocIdSetIterator.all(3)); + + ExactSearcher exactSearcher = mock(ExactSearcher.class); + when(exactSearcher.searchLeaf(any(), any())).thenReturn(topDocs(new ScoreDoc(1, 30f))); + + Query finalQuery = mock(Query.class); + when(finalQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(mock(Weight.class)); + when(queryUtils.createDocAndScoreQuery(eq(reader), any())).thenReturn(finalQuery); + + ExpandNestedDocsQuery query = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .rescoreK(4) + .floatQueryVector(new float[] { 1.0f, 2.0f, 3.0f }) + .exactSearcher(exactSearcher) + .build(); + + query.createWeight(indexSearcher, scoreMode, 1.f); + + // One leaf, two exact search stages: the collapsing rescore and the expansion. Both are timed, so the + // timer is started and stopped once per stage. + verify(exactSearcher, times(2)).searchLeaf(any(), any()); + verify(timer, times(2)).start(); + verify(timer, times(2)).stop(); + } + + /** + * Without a profiler attached the query must not reach for a profile breakdown at all, so that an + * unprofiled search stays on the untimed path. + */ + @SneakyThrows + public void testCreateWeight_whenProfilerDisabled_thenNoProfileBreakdownLookup() { + Directory directory = new ByteBuffersDirectory(); + try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document doc = new Document(); + doc.add(new FloatPoint("vector", 1.0f, 2.0f, 3.0f)); + writer.addDocument(doc); + writer.commit(); + } + + IndexReader reader = DirectoryReader.open(directory); + + ContextIndexSearcher indexSearcher = mock(ContextIndexSearcher.class); + when(indexSearcher.getIndexReader()).thenReturn(reader); + when(indexSearcher.getTaskExecutor()).thenReturn(taskExecutor); + when(indexSearcher.getProfiler()).thenReturn(null); + + Weight queryWeight = mock(Weight.class); + ScoreMode scoreMode = mock(ScoreMode.class); + Query docAndScoreQuery = mock(Query.class); + when(docAndScoreQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(queryWeight); + + InternalNestedKnnVectorQuery internalQuery = mock(InternalNestedKnnVectorQuery.class); + when(internalQuery.knnRewrite(indexSearcher)).thenReturn(docAndScoreQuery); + when(internalQuery.getK()).thenReturn(1); + when(internalQuery.getFilter()).thenReturn(null); + when(internalQuery.getField()).thenReturn("field"); + when(internalQuery.getParentFilter()).thenReturn(mock(BitSetProducer.class)); + when(internalQuery.knnExactSearch(any(), any())).thenReturn(topDocs(new ScoreDoc(1, 30f))); + + QueryUtils queryUtils = mock(QueryUtils.class); + when(queryUtils.doSearch(indexSearcher, reader.leaves(), queryWeight)).thenReturn(List.of(new HashMap<>(Map.of(1, 20f)))); + when(queryUtils.getAllSiblings(any(), any(), any(), any())).thenReturn(DocIdSetIterator.all(3)); + + Query finalQuery = mock(Query.class); + when(finalQuery.createWeight(indexSearcher, scoreMode, 1.f)).thenReturn(mock(Weight.class)); + when(queryUtils.createDocAndScoreQuery(eq(reader), any())).thenReturn(finalQuery); + + ExpandNestedDocsQuery query = new ExpandNestedDocsQuery.ExpandNestedDocsQueryBuilder().internalNestedKnnVectorQuery(internalQuery) + .queryUtils(queryUtils) + .build(); + + query.createWeight(indexSearcher, scoreMode, 1.f); + + verify(indexSearcher).getProfiler(); + verify(internalQuery).knnExactSearch(any(), any()); + } + + private static TopDocs topDocs(ScoreDoc... scoreDocs) { + return new TopDocs(new TotalHits(scoreDocs.length, TotalHits.Relation.EQUAL_TO), scoreDocs); + } } diff --git a/src/test/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactoryTests.java b/src/test/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactoryTests.java index b266d99647..55ae12505b 100644 --- a/src/test/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactoryTests.java +++ b/src/test/java/org/opensearch/knn/index/query/lucenelib/NestedKnnVectorQueryFactoryTests.java @@ -12,6 +12,7 @@ import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery; import org.opensearch.knn.index.query.rescore.RescoreContext; +import static org.junit.Assert.assertNotEquals; import static org.mockito.Mockito.mock; public class NestedKnnVectorQueryFactoryTests extends TestCase { @@ -67,6 +68,43 @@ public void testCreate_whenCalled_thenCreateQuery() { ); } + public void testCreate_whenExpandNestedDocsWithRescore_thenExpandNestedDocsQueryOwnsRescore() { + String fieldName = "field"; + float[] floatVectors = new float[3]; + int luceneK = 10; + int k = 3; + int rescoreK = 6; + Query queryFilter = mock(Query.class); + BitSetProducer parentFilter = mock(BitSetProducer.class); + + Query query = NestedKnnVectorQueryFactory.createNestedKnnVectorQuery( + fieldName, + floatVectors, + luceneK, + queryFilter, + parentFilter, + true, + k, + rescoreK + ); + + assertEquals(ExpandNestedDocsQuery.class, query.getClass()); + // rescoreK is part of the query's identity, so it must not compare equal to the no-rescore variant + assertNotEquals( + query, + NestedKnnVectorQueryFactory.createNestedKnnVectorQuery( + fieldName, + floatVectors, + luceneK, + queryFilter, + parentFilter, + true, + k, + RescoreContext.NO_RESCORE_NEEDED + ) + ); + } + public void testCreate_whenNoExpandNestedDocs_thenDiversifyingQuery() { String fieldName = "field"; byte[] byteVectors = new byte[3]; diff --git a/src/test/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQueryTests.java b/src/test/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQueryTests.java index 2b95adcd42..b90069044d 100644 --- a/src/test/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQueryTests.java +++ b/src/test/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQueryTests.java @@ -15,6 +15,7 @@ import org.apache.lucene.search.join.BitSetProducer; import org.apache.lucene.search.join.DiversifyingChildrenFloatKnnVectorQuery; import org.apache.lucene.search.knn.KnnCollectorManager; +import org.opensearch.knn.index.query.rescore.RescoreContext; import java.io.IOException; @@ -23,60 +24,40 @@ public class OSDiversifyingChildrenFloatKnnVectorQueryTests extends TestCase { - public void testConstructor() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 5; - int rescoreK = 0; - boolean expandNestedDocs = false; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); + private static final String FIELD_NAME = "test_field"; + private static final float[] QUERY_VECTOR = { 1.0f, 2.0f, 3.0f }; - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, + private OSDiversifyingChildrenFloatKnnVectorQuery newQuery(int luceneK, int k, int rescoreK) { + return new OSDiversifyingChildrenFloatKnnVectorQuery( + FIELD_NAME, + QUERY_VECTOR, + mock(Query.class), luceneK, - parentFilter, + mock(BitSetProducer.class), k, - rescoreK, - expandNestedDocs + rescoreK ); - - assertTrue(query instanceof DiversifyingChildrenFloatKnnVectorQuery); } - public void testMergeLeafResultsWithRescoreDisabled() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 3; - int rescoreK = 0; - boolean expandNestedDocs = false; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); - + private static TopDocs[] twoLeavesOfTwo() { ScoreDoc[] scoreDocs1 = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f) }; ScoreDoc[] scoreDocs2 = { new ScoreDoc(3, 0.7f), new ScoreDoc(4, 0.6f) }; + return new TopDocs[] { + new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs1), + new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs2) }; + } - TopDocs topDocs1 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs2); + public void testConstructor() { + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(10, 5, RescoreContext.NO_RESCORE_NEEDED); - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; + assertTrue(query instanceof DiversifyingChildrenFloatKnnVectorQuery); + } - TopDocs result = query.mergeLeafResults(perLeafResults); + public void testMergeLeafResultsWithRescoreDisabled() { + int k = 3; + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(10, k, RescoreContext.NO_RESCORE_NEEDED); + + TopDocs result = query.mergeLeafResults(twoLeavesOfTwo()); assertEquals(k, result.scoreDocs.length); assertTrue(result.scoreDocs[0].score >= result.scoreDocs[1].score); @@ -84,29 +65,10 @@ public void testMergeLeafResultsWithRescoreDisabled() { } public void testMergeLeafResults_withFewerResultsThanK() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 5; - int rescoreK = 0; - boolean expandNestedDocs = false; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(10, 5, RescoreContext.NO_RESCORE_NEEDED); ScoreDoc[] scoreDocs = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f) }; - TopDocs topDocs = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs); - TopDocs[] perLeafResults = { topDocs }; + TopDocs[] perLeafResults = { new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs) }; TopDocs result = query.mergeLeafResults(perLeafResults); @@ -114,153 +76,38 @@ public void testMergeLeafResults_withFewerResultsThanK() { } public void testMergeLeafResultsWithRescoreEnabled() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 3; int rescoreK = 6; - boolean expandNestedDocs = false; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(10, 3, rescoreK); ScoreDoc[] scoreDocs1 = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f), new ScoreDoc(5, 0.5f) }; ScoreDoc[] scoreDocs2 = { new ScoreDoc(3, 0.7f), new ScoreDoc(4, 0.6f), new ScoreDoc(6, 0.4f) }; - TopDocs topDocs1 = new TopDocs(new TotalHits(3, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(3, TotalHits.Relation.EQUAL_TO), scoreDocs2); - - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; + TopDocs[] perLeafResults = { + new TopDocs(new TotalHits(3, TotalHits.Relation.EQUAL_TO), scoreDocs1), + new TopDocs(new TotalHits(3, TotalHits.Relation.EQUAL_TO), scoreDocs2) }; TopDocs result = query.mergeLeafResults(perLeafResults); - // When rescoreK > 0 and expandNestedDocs is false, should trim to rescoreK (not luceneK or k) + // When rescoreK > 0, should trim to rescoreK (not luceneK or k) assertEquals(rescoreK, result.scoreDocs.length); } - public void testMergeLeafResultsWithExpandNestedDocs() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 3; - int rescoreK = 0; - boolean expandNestedDocs = true; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); - - ScoreDoc[] scoreDocs1 = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f) }; - ScoreDoc[] scoreDocs2 = { new ScoreDoc(3, 0.7f), new ScoreDoc(4, 0.6f) }; - - TopDocs topDocs1 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs2); - - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; - - TopDocs result = query.mergeLeafResults(perLeafResults); - - // When expandNestedDocs is true, should reduce to k regardless of rescoreK - assertEquals(k, result.scoreDocs.length); - } - - public void testConstructorWithoutExpandNestedDocs() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 3; - int rescoreK = 0; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - // Test the overloaded constructor without expandNestedDocs parameter - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK - ); - - assertTrue(query instanceof DiversifyingChildrenFloatKnnVectorQuery); - - // Verify it behaves as if expandNestedDocs is false - ScoreDoc[] scoreDocs1 = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f) }; - ScoreDoc[] scoreDocs2 = { new ScoreDoc(3, 0.7f), new ScoreDoc(4, 0.6f) }; - - TopDocs topDocs1 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs2); - - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; - TopDocs result = query.mergeLeafResults(perLeafResults); - - // Should reduce to k (default behavior when expandNestedDocs is false and rescoreK is 0) - assertEquals(k, result.scoreDocs.length); - } - - public void testMergeLeafResultsWithRescoreAndExpandNestedDocs() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; + /** + * Regression test for the expand_nested_docs + rescoring interaction. The candidate pool used to be cut + * to k whenever expand_nested_docs was set, which left rescoring nothing to re-rank. The oversampled + * pool now survives regardless, and ExpandNestedDocsQuery does the cut to k itself after rescoring. + */ + public void testMergeLeafResults_whenRescoreEnabled_thenKeepsOversampledPoolForExpansion() { int k = 3; - int rescoreK = 6; - boolean expandNestedDocs = true; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); - - ScoreDoc[] scoreDocs1 = { new ScoreDoc(1, 0.9f), new ScoreDoc(2, 0.8f) }; - ScoreDoc[] scoreDocs2 = { new ScoreDoc(3, 0.7f), new ScoreDoc(4, 0.6f) }; + int rescoreK = 4; + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(10, k, rescoreK); - TopDocs topDocs1 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(2, TotalHits.Relation.EQUAL_TO), scoreDocs2); - - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; - - TopDocs result = query.mergeLeafResults(perLeafResults); + TopDocs result = query.mergeLeafResults(twoLeavesOfTwo()); - // When expandNestedDocs is true, should reduce to k even if rescoreK > 0 - assertEquals(k, result.scoreDocs.length); + assertEquals(rescoreK, result.scoreDocs.length); } public void testApproximateSearch_whenParentBitSetNull_thenReturnEmptyResults() throws IOException { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; - int luceneK = 10; - int k = 5; - int rescoreK = 0; - Query filterQuery = mock(Query.class); BitSetProducer parentFilter = mock(BitSetProducer.class); LeafReaderContext context = mock(LeafReaderContext.class); KnnCollectorManager knnCollectorManager = mock(KnnCollectorManager.class); @@ -269,13 +116,13 @@ public void testApproximateSearch_whenParentBitSetNull_thenReturnEmptyResults() when(parentFilter.getBitSet(context)).thenReturn(null); OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, + FIELD_NAME, + QUERY_VECTOR, + mock(Query.class), + 10, parentFilter, - k, - rescoreK + 5, + RescoreContext.NO_RESCORE_NEEDED ); TopDocs result = query.approximateSearch(context, acceptDocs, Integer.MAX_VALUE, knnCollectorManager); @@ -286,26 +133,9 @@ public void testApproximateSearch_whenParentBitSetNull_thenReturnEmptyResults() } public void testMergeLeafResultsWithRescoreK_trimsToRescoreKNotLuceneK() { - String fieldName = "test_field"; - float[] queryVector = { 1.0f, 2.0f, 3.0f }; // luceneK=256 (ef_search dominated), rescoreK=200 (oversample dominated), k=100 - int luceneK = 256; - int k = 100; int rescoreK = 200; - boolean expandNestedDocs = false; - Query filterQuery = mock(Query.class); - BitSetProducer parentFilter = mock(BitSetProducer.class); - - OSDiversifyingChildrenFloatKnnVectorQuery query = new OSDiversifyingChildrenFloatKnnVectorQuery( - fieldName, - queryVector, - filterQuery, - luceneK, - parentFilter, - k, - rescoreK, - expandNestedDocs - ); + OSDiversifyingChildrenFloatKnnVectorQuery query = newQuery(256, 100, rescoreK); // Create enough results to exceed rescoreK but be within luceneK ScoreDoc[] scoreDocs1 = new ScoreDoc[128]; @@ -315,10 +145,9 @@ public void testMergeLeafResultsWithRescoreK_trimsToRescoreKNotLuceneK() { scoreDocs2[i] = new ScoreDoc(128 + i, 0.5f - (i * 0.001f)); } - TopDocs topDocs1 = new TopDocs(new TotalHits(128, TotalHits.Relation.EQUAL_TO), scoreDocs1); - TopDocs topDocs2 = new TopDocs(new TotalHits(128, TotalHits.Relation.EQUAL_TO), scoreDocs2); - - TopDocs[] perLeafResults = { topDocs1, topDocs2 }; + TopDocs[] perLeafResults = { + new TopDocs(new TotalHits(128, TotalHits.Relation.EQUAL_TO), scoreDocs1), + new TopDocs(new TotalHits(128, TotalHits.Relation.EQUAL_TO), scoreDocs2) }; TopDocs result = query.mergeLeafResults(perLeafResults); diff --git a/src/test/java/org/opensearch/knn/integ/ExpandNestedDocsWithRescoreIT.java b/src/test/java/org/opensearch/knn/integ/ExpandNestedDocsWithRescoreIT.java new file mode 100644 index 0000000000..30267a8a89 --- /dev/null +++ b/src/test/java/org/opensearch/knn/integ/ExpandNestedDocsWithRescoreIT.java @@ -0,0 +1,427 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.integ; + +import com.google.common.collect.Multimap; +import lombok.SneakyThrows; +import org.apache.hc.core5.http.io.entity.EntityUtils; +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.common.settings.Settings; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.knn.KNNRestTestCase; +import org.opensearch.knn.NestedKnnDocBuilder; +import org.opensearch.knn.index.KNNSettings; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +/** + * Covers nested k-NN search with {@code expand_nested_docs} on a field that has rescoring enabled, which is the + * combination reported in issue 3125. + * + * A compression level of 4x resolves to the Lucene engine and turns rescoring on by default, so these tests + * exercise the Lucene engine's two-phase search together with the nested expansion. The bug was that the final + * reduction to k counted child documents rather than parent documents, so a single parent could consume the + * whole budget and the remaining parents were dropped from the response. + * + * {@link ExpandNestedDocsIT} parameterizes over engines but only ever picks 1x or 32x compression, and 32x + * resolves to Faiss, so this combination had no coverage there. + */ +public class ExpandNestedDocsWithRescoreIT extends KNNRestTestCase { + private static final String INDEX_NAME = "test-index-expand-nested-rescore"; + private static final String NESTED_FIELD = "nested_field"; + private static final String VECTOR_FIELD = "my_vector"; + private static final String NESTED_VECTOR_PATH = NESTED_FIELD + "." + VECTOR_FIELD; + private static final String PARENT_FILTER_FIELD = "parking"; + private static final String CHILD_FILTER_FIELD = "storage"; + private static final String NESTED_FILTER_PATH = NESTED_FIELD + "." + CHILD_FILTER_FIELD; + private static final int DIMENSION = 3; + private static final int CHILDREN_PER_PARENT = 3; + + /** + * The scenario straight out of the issue. Three parents of three children each, k=2, and the query vector is + * an exact match for one of doc 1's children. + * + * Ground truth by l2 distance to [1,1,1]: doc 1 (best child [1,1,1]), then doc 3 (best child [9,9,9], 192), + * then doc 2 (all children [10,10,10], 243). So the response must carry two parents, each with all three of + * its children. The reported bug returned a single parent holding only two of its three children. + */ + @SneakyThrows + public void testExpandNestedDocs_whenRescoreEnabled_thenReturnKParentsWithAllChildren() { + createNestedKnnIndex(); + indexDoc("1", new Object[][] { { 2, 2, 2 }, { 1, 1, 1 }, { 3, 3, 3 } }); + indexDoc("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }); + indexDoc("3", new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }); + refreshIndex(INDEX_NAME); + // A single segment keeps per-leaf top-k effects out of the picture + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String responseBody = search(k, new Object[] { 1, 1, 1 }, null); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(k, docIdToOffsets.keySet().size()); + for (String docId : docIdToOffsets.keySet()) { + assertEquals("parent " + docId + " should carry all of its children", CHILDREN_PER_PARENT, docIdToOffsets.get(docId).size()); + } + assertEquals(k, parseTotalSearchHits(responseBody)); + + // Doc 3 beats doc 2 only on full precision distances (192 vs 243), so getting {1, 3} back in this order + // is what proves the parent was selected on rescored vectors rather than on quantized ones. + assertEquals(List.of("1", "3"), parseIds(responseBody)); + + // The children have to be scored on full precision vectors too, otherwise a quantized child score leaks + // into the parent through the nested score mode. These are the exact values a full precision index + // returns: doc 1 averages 1.0, 0.25 and 0.0769, doc 3 averages 0.00518, 8.4e-6 and 3.7e-6. + List scores = parseScores(responseBody); + assertEquals(0.44230768d, scores.get(0), 1e-7); + assertEquals(0.0017311643d, scores.get(1), 1e-9); + } + + /** + * Same shape, but with more parents than the rescore candidate pool can be satisfied by trivially, so the + * reduction to k parents actually has to discard candidates. Every returned parent must still be complete. + */ + @SneakyThrows + public void testExpandNestedDocs_whenRescoreEnabledWithManyParents_thenReturnKParentsWithAllChildren() { + createNestedKnnIndex(); + for (int i = 1; i <= 150; i++) { + int base = i * 2; + indexDoc( + String.valueOf(i), + new Object[][] { { base, base, base }, { base + 1, base + 1, base + 1 }, { base + 2, base + 2, base + 2 } } + ); + } + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 5; + String responseBody = search(k, new Object[] { 0, 0, 0 }, null); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(k, docIdToOffsets.keySet().size()); + for (String docId : docIdToOffsets.keySet()) { + assertEquals(CHILDREN_PER_PARENT, docIdToOffsets.get(docId).size()); + } + assertEquals(k, parseTotalSearchHits(responseBody)); + } + + /** + * Control: turning rescoring off from the query must not change the shape of the response. This isolates the + * expansion from the rescore stage, so a failure here points at the plain expansion path rather than the fix. + */ + @SneakyThrows + public void testExpandNestedDocs_whenRescoreDisabledFromQuery_thenReturnKParentsWithAllChildren() { + createNestedKnnIndex(); + indexDoc("1", new Object[][] { { 2, 2, 2 }, { 1, 1, 1 }, { 3, 3, 3 } }); + indexDoc("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }); + indexDoc("3", new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }); + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String responseBody = search(k, new Object[] { 1, 1, 1 }, "false"); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(k, docIdToOffsets.keySet().size()); + for (String docId : docIdToOffsets.keySet()) { + assertEquals(CHILDREN_PER_PARENT, docIdToOffsets.get(docId).size()); + } + } + + /** + * An explicit oversample factor has to flow through to the candidate pool rather than being discarded + * because the expansion is enabled. + */ + @SneakyThrows + public void testExpandNestedDocs_whenOversampleFactorProvided_thenReturnKParentsWithAllChildren() { + createNestedKnnIndex(); + indexDoc("1", new Object[][] { { 2, 2, 2 }, { 1, 1, 1 }, { 3, 3, 3 } }); + indexDoc("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }); + indexDoc("3", new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }); + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String responseBody = search(k, new Object[] { 1, 1, 1 }, "{ \"oversample_factor\": 3.0 }"); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(k, docIdToOffsets.keySet().size()); + for (String docId : docIdToOffsets.keySet()) { + assertEquals(CHILDREN_PER_PARENT, docIdToOffsets.get(docId).size()); + } + assertEquals("1", parseIds(responseBody).get(0)); + } + + /** + * A filter on the parent document has to survive the rescore stage: the filtered out parent must not come back + * even though it is the closest one, and the parents that do come back must still carry all of their children. + * + * Doc 4 is an exact match for the query vector, so it would win outright without the filter. With the filter + * applied the ground truth falls back to {1, 3} and the scores are the ones the unfiltered case produces, + * which is what shows the filter narrowed the candidate pool without disturbing the full precision rescoring. + */ + @SneakyThrows + public void testExpandNestedDocs_whenRescoreEnabledWithParentFilter_thenReturnKMatchingParentsWithAllChildren() { + createNestedKnnIndex(); + indexDocWithParentFilter("1", new Object[][] { { 2, 2, 2 }, { 1, 1, 1 }, { 3, 3, 3 } }, true); + indexDocWithParentFilter("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }, true); + indexDocWithParentFilter("3", new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }, true); + indexDocWithParentFilter("4", new Object[][] { { 1, 1, 1 }, { 1, 1, 1 }, { 1, 1, 1 } }, false); + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String filter = String.format(Locale.ROOT, "{\"term\":{\"%s\":true}}", PARENT_FILTER_FIELD); + String responseBody = search(k, new Object[] { 1, 1, 1 }, null, false, filter); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(k, docIdToOffsets.keySet().size()); + assertFalse("the filtered out parent must not be returned", docIdToOffsets.containsKey("4")); + for (String docId : docIdToOffsets.keySet()) { + assertEquals("parent " + docId + " should carry all of its children", CHILDREN_PER_PARENT, docIdToOffsets.get(docId).size()); + } + assertEquals(k, parseTotalSearchHits(responseBody)); + + assertEquals(List.of("1", "3"), parseIds(responseBody)); + + List scores = parseScores(responseBody); + assertEquals(0.44230768d, scores.get(0), 1e-7); + assertEquals(0.0017311643d, scores.get(1), 1e-9); + } + + /** + * A filter on the child documents has to narrow the expansion as well as the candidate selection. Only the + * children matching the filter may be expanded into the inner hits, and the parents have to be ranked on their + * matching children alone. + * + * Ground truth over the filtered children only: doc 1 keeps [1,1,1] and [3,3,3], doc 2 keeps all three + * [10,10,10], doc 3 keeps only [200,200,200] and [300,300,300]. So doc 1 wins, doc 2 is second, and doc 3 + * drops out even though its unfiltered best child [9,9,9] would have beaten all of doc 2's. + */ + @SneakyThrows + public void testExpandNestedDocs_whenRescoreEnabledWithChildFilter_thenReturnOnlyMatchingChildren() { + createNestedKnnIndex(); + indexDocWithChildFilter("1", new Object[][] { { 1, 1, 1 }, { 2, 2, 2 }, { 3, 3, 3 } }, new boolean[] { true, false, true }); + indexDocWithChildFilter("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }, new boolean[] { true, true, true }); + indexDocWithChildFilter( + "3", + new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }, + new boolean[] { false, true, true } + ); + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String filter = String.format( + Locale.ROOT, + "{\"nested\":{\"path\":\"%s\",\"query\":{\"term\":{\"%s\":true}}}}", + NESTED_FIELD, + NESTED_FILTER_PATH + ); + String responseBody = search(k, new Object[] { 1, 1, 1 }, null, false, filter); + + Multimap docIdToOffsets = parseInnerHits(responseBody, NESTED_FIELD); + assertEquals(List.of("1", "2"), parseIds(responseBody)); + assertEquals(k, parseTotalSearchHits(responseBody)); + + // Only the children that pass the filter get expanded, so doc 1 comes back without its middle child + assertEquals(2, docIdToOffsets.get("1").size()); + assertTrue(docIdToOffsets.get("1").containsAll(List.of(0, 2))); + assertEquals(CHILDREN_PER_PARENT, docIdToOffsets.get("2").size()); + + // The excluded child must not contribute to its parent's score either. Averaged over the matching children + // on full precision vectors, doc 1 is (1.0 + 1/13) / 2 and doc 2 is 1/244. + List scores = parseScores(responseBody); + assertEquals(0.53846154d, scores.get(0), 1e-7); + assertEquals(0.0040983607d, scores.get(1), 1e-9); + } + + /** + * The two exact search stages must report their time under the Profile API. This is the only test that can + * catch a missing registration in {@link org.opensearch.knn.plugin.KNNPlugin#getQueryProfileMetricsProvider}: + * without it the breakdown has no {@code exact_search} metric at all, and asking for its timer throws. + */ + @SneakyThrows + public void testExpandNestedDocs_whenProfileEnabled_thenReportExactSearchTime() { + createNestedKnnIndex(); + indexDoc("1", new Object[][] { { 2, 2, 2 }, { 1, 1, 1 }, { 3, 3, 3 } }); + indexDoc("2", new Object[][] { { 10, 10, 10 }, { 10, 10, 10 }, { 10, 10, 10 } }); + indexDoc("3", new Object[][] { { 9, 9, 9 }, { 200, 200, 200 }, { 300, 300, 300 } }); + refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME, 1); + + int k = 2; + String responseBody = search(k, new Object[] { 1, 1, 1 }, null, true); + + // Profiling must not disturb the results + assertEquals(List.of("1", "3"), parseIds(responseBody)); + + // The query sits several levels down the profile tree, under the nested query's block join, so scan for + // the breakdown at any depth rather than assuming a fixed nesting. + List timings = collectBreakdownMetric(responseBody, "exact_search"); + assertFalse("expand_nested_docs with rescoring should report exact_search timings", timings.isEmpty()); + assertTrue("at least one exact_search timing should be non-zero, got " + timings, timings.stream().anyMatch(t -> t > 0L)); + + // The collapsing rescore and the child expansion each run once per weight creation, and a single search + // request creates the weight more than once, so assert that the two stages stay paired rather than + // pinning an exact call count. + List counts = collectBreakdownMetric(responseBody, "exact_search_count"); + long maxCount = counts.stream().mapToLong(Long::longValue).max().orElse(0L); + assertTrue("expected at least one exact_search per stage, got " + counts, maxCount >= 2L); + assertEquals("the collapsing rescore and the child expansion should be counted in pairs, got " + counts, 0L, maxCount % 2L); + } + + /** + * Collects one metric out of every {@code breakdown} object in the profile tree, at any depth. The k-NN query + * is not at a fixed level of the tree, so walking is more robust than a hardcoded path. + */ + @SuppressWarnings("unchecked") + private List collectBreakdownMetric(final String responseBody, final String metric) throws Exception { + Map parsed = createParser(MediaTypeRegistry.getDefaultMediaType().xContent(), responseBody).map(); + List values = new ArrayList<>(); + Deque pending = new ArrayDeque<>(); + pending.push(parsed); + while (pending.isEmpty() == false) { + Object current = pending.pop(); + if (current instanceof Map map) { + Object breakdown = map.get("breakdown"); + if (breakdown instanceof Map breakdownMap && breakdownMap.get(metric) instanceof Number value) { + values.add(value.longValue()); + } + ((Map) map).values().forEach(pending::push); + } else if (current instanceof List list) { + list.forEach(pending::push); + } + } + return values; + } + + /** + * Creates the index from the issue: a nested knn_vector with 4x compression on disk, which resolves to the + * Lucene engine and enables rescoring by default. + */ + private void createNestedKnnIndex() throws Exception { + String mapping = String.format( + Locale.ROOT, + "{\"properties\":{\"%s\":{\"type\":\"nested\",\"properties\":{\"%s\":" + + "{\"type\":\"knn_vector\",\"dimension\":%d,\"space_type\":\"l2\"," + + "\"mode\":\"on_disk\",\"compression_level\":\"4x\"}," + + "\"%s\":{\"type\":\"boolean\"}}}," + + "\"%s\":{\"type\":\"boolean\"}}}", + NESTED_FIELD, + VECTOR_FIELD, + DIMENSION, + CHILD_FILTER_FIELD, + PARENT_FILTER_FIELD + ); + Settings settings = Settings.builder() + .put("number_of_shards", 1) + .put("number_of_replicas", 0) + .put("index.knn", true) + // Force the graph to be built so the query goes through approximate search plus rescoring + .put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0) + .build(); + createKnnIndex(INDEX_NAME, settings, mapping); + } + + private void indexDoc(final String docId, final Object[][] vectors) throws Exception { + addKnnDoc(INDEX_NAME, docId, NestedKnnDocBuilder.create(NESTED_FIELD).addVectors(VECTOR_FIELD, vectors).build()); + } + + /** + * Indexes a parent holding {@code vectors} plus a top level field the k-NN filter can select it by. + */ + private void indexDocWithParentFilter(final String docId, final Object[][] vectors, final boolean filterValue) throws Exception { + addKnnDoc( + INDEX_NAME, + docId, + NestedKnnDocBuilder.create(NESTED_FIELD) + .addVectors(VECTOR_FIELD, vectors) + .addTopLevelField(PARENT_FILTER_FIELD, filterValue) + .build() + ); + } + + /** + * Indexes a parent whose children each carry their own filter value, so a nested filter can select a subset of + * the children of a single parent. + */ + private void indexDocWithChildFilter(final String docId, final Object[][] vectors, final boolean[] filterValues) throws Exception { + assert vectors.length == filterValues.length; + NestedKnnDocBuilder builder = NestedKnnDocBuilder.create(NESTED_FIELD); + for (int i = 0; i < vectors.length; i++) { + builder.addVectorWithMetadata(VECTOR_FIELD, vectors[i], CHILD_FILTER_FIELD, filterValues[i]); + } + addKnnDoc(INDEX_NAME, docId, builder.build()); + } + + /** + * @param rescore raw json for the {@code rescore} clause, or null to leave it out and take the field default + */ + private String search(final int k, final Object[] queryVector, final String rescore) throws Exception { + return search(k, queryVector, rescore, false); + } + + /** + * @param rescore raw json for the {@code rescore} clause, or null to leave it out and take the field default + * @param profile whether to ask for the Profile API breakdown + */ + private String search(final int k, final Object[] queryVector, final String rescore, final boolean profile) throws Exception { + return search(k, queryVector, rescore, profile, null); + } + + /** + * @param rescore raw json for the {@code rescore} clause, or null to leave it out and take the field default + * @param profile whether to ask for the Profile API breakdown + * @param filter raw json for the k-NN query's {@code filter} clause, or null for an unfiltered search + */ + private String search(final int k, final Object[] queryVector, final String rescore, final boolean profile, final String filter) + throws Exception { + String rescoreClause = rescore == null ? "" : String.format(Locale.ROOT, ",\"rescore\":%s", rescore); + String filterClause = filter == null ? "" : String.format(Locale.ROOT, ",\"filter\":%s", filter); + String body = String.format( + Locale.ROOT, + "{\"_source\":false,\"profile\":" + + profile + + ",\"size\":%d,\"query\":{\"nested\":{\"path\":\"%s\"," + + "\"query\":{\"knn\":{\"%s\":{\"vector\":%s,\"k\":%d,\"expand_nested_docs\":true%s%s}}}," + + "\"inner_hits\":{\"size\":100,\"_source\":false}}}}", + k, + NESTED_FIELD, + NESTED_VECTOR_PATH, + java.util.Arrays.toString(queryVector), + k, + rescoreClause, + filterClause + ); + + Request request = new Request("POST", String.format(Locale.ROOT, "/%s/_search", INDEX_NAME)); + request.setJsonEntity(body); + Response response = client().performRequest(request); + assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + return EntityUtils.toString(response.getEntity()); + } + + @Override + public void tearDown() throws Exception { + List indices = List.of(INDEX_NAME); + for (String index : indices) { + try { + deleteKNNIndex(index); + } catch (Exception e) { + // index may not exist for a test that failed during setup + } + } + super.tearDown(); + } +}