Support rescoring with expand_nested_docs on the Lucene engine (#3125)

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)
 
 ### Refactoring
 * Wire ResolvedIndexSpec consumers through spec-driven resolution flow [#3421](https://github.com/opensearch-project/k-NN/pull/3421)
@@ -55,3 +56,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)

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) {

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) {

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)
+        );
+    }
 }