Support rescoring with expand_nested_docs on the Lucene engine (#3125) by benkim1028 · Pull Request #3579 · opensearch-project/k-NN

benkim1028 · 2026-09-18T06:12:24Z

Description

expand_nested_docs combined with rescoring on the Lucene engine did not return all matching documents (#3125). With k=2 the response came back with 1 parent instead of 2 — a single parent's children filled the whole result and the other parents were dropped.

Root cause: ExpandNestedDocsQuery expanded child documents before rescoring, so the final reduction to k counted child rows instead of parent rows. On main this was worked around by silently dropping rescoring for this combination entirely — so parents were also selected on quantized distances, and child scores were quantized.

Fix: make ExpandNestedDocsQuery own the rescore stage and run it before expansion — the same ordering NativeEngineKnnVectorQuery already uses for the faiss/nmslib engines. The one rule: keep the list at one row per parent document until after the final top-k cut.

Pipeline before (Lucene engine, rescoring requested):

Approximate search → best child per parent
Expand every child of those parents
Reduce to k ← counts child rows, so it drops parents
- and on main, rescoring was skipped entirely → parents chosen on quantized distances, child scores quantized

Pipeline after (mirrors NativeEngineKnnVectorQuery):

Approximate search → best child per parent (unchanged)
Expand candidates to all siblings — getAllSiblings, from NativeEngineKnnVectorQuery.doRescore
Rescore on full-precision vectors and collapse each parent to its single best child — ExactSearcher with parentsFilter set, from doRescore
Reduce to top k parents — one row per parent, so the cut counts documents — mirrors native's ResultUtil.reduceToTopK
Expand the k winners to all their children — getAllSiblings, from NativeEngineKnnVectorQuery.retrieveAll
Exact search on full-precision vectors, no collapse → every child keeps its own score — from retrieveAll
Merge without truncation — sum of rows, not k — mirrors native's getMergeTopN

Testing

Verified against the running cluster with the issue's exact repro: lucene_4x now returns 2 parents × 6 children with scores byte-identical to a full-precision faiss baseline.

Related Issues

Resolves #3125

Check List

New functionality includes testing. — unit tests + new ExpandNestedDocsWithRescoreIT
New functionality has been documented. — n/a, bug fix (no API change)
API changes companion pull request created. — n/a
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created. — n/a

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

github-actions · 2026-09-18T06:13:43Z

PR Reviewer Guide 🔍

(Review updated until commit `0d6910f`)

Here are some key observations to aid the review process:

🧪 PR contains tests
🔒 No security concerns identified
✅ No TODO sections
🔀 No multiple PR themes
⚡ Recommended focus areas for review Possible Issue `rescoreToTopKParents` repurposes `ScoreDoc.shardIndex` to carry the leaf ordinal across the `TopDocs.merge` call, but `TopDocs.merge(k, perLeafTopDocs)` is the single-arg-shard-aware overload that in Lucene assigns `shardIndex` based on the array position when `setShardIndex` is enabled (and leaves it alone otherwise). If the merge path overwrites `shardIndex`, the subsequent `reducedResults.get(scoreDoc.shardIndex).put(...)` will route docs to the wrong leaf map, silently corrupting the rescored results. Please confirm the Lucene `TopDocs.merge` overload used here preserves the pre-set `shardIndex` rather than reassigning it, and add a test that exercises more than one leaf so this is caught. 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<Map<Integer, Float>> 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; Equality/Caching `equals`/`hashCode` include `rescoreK` but ignore `floatQueryVector`. Two queries with the same `rescoreK` and identical `internalNestedKnnVectorQuery` but different `floatQueryVector` arrays would be considered equal, which could cause an incorrect query cache hit if the query vector is not fully reflected in `internalNestedKnnVectorQuery.equals`. Verify the internal query's equality already captures the vector; if not, `floatQueryVector` must be part of equality/hash to avoid returning cached results computed for a different query vector. public boolean equals(final Object o) { if (!sameClassAs(o)) { return false; } ExpandNestedDocsQuery other = (ExpandNestedDocsQuery) o; // 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 Objects.hash(internalNestedKnnVectorQuery, rescoreK); }

github-actions · 2026-09-18T06:14:09Z

PR Code Suggestions ✨

Latest suggestions up to 0d6910f

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Cache null filter bits to avoid recomputation The cache never stores null, so a leaf whose filter is null will have `createBits` recomputed on every call (once in the rescore stage and once in the expansion stage for the same leaf). More importantly, concurrent rescore/expansion tasks for the same leaf may both call `createBits` before either stores a value. Consider using a sentinel or `Map<Integer, Optional>` / `computeIfAbsent` to also cache the "no filter" result and avoid duplicate computation. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [268-278] -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; +return filterBitsByLeaf.computeIfAbsent(leafReaderContext.ord, ord -> { + try { + return queryUtils.createBits(leafReaderContext, filterWeight); + } catch (IOException e) { + throw new UncheckedIOException(e); + } +}); Suggestion importance[1-10]: 4 __ Why: Valid observation about not caching null bits causing recomputation, but the impact is minor since `createBits` for a null filter is cheap, and the suggested `computeIfAbsent` with checked-exception wrapping changes the error handling contract.	Low
General	Document shardIndex reuse invariant Overwriting `shardIndex` on the ScoreDocs returned by the exact searcher mutates state that may be reused or observed elsewhere, and relies on an undocumented invariant that `TopDocs.merge` only reads (never writes) `shardIndex`. Since `TopDocs.merge` with a non-null shardIndex array of ScoreDocs is fragile, consider using the `TopDocs.merge(int, TopDocs[])` overload with explicit per-leaf tagging via a dedicated structure, or at minimum assert the invariant with a comment reference to the Lucene version. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [181-184] for (ScoreDoc scoreDoc : leafTopDocs.scoreDocs) { + // shardIndex is used purely as a leaf ordinal carrier; TopDocs.merge preserves it unchanged. scoreDoc.shardIndex = leafOrd; } return leafTopDocs; Suggestion importance[1-10]: 2 __ Why: The existing code already has a detailed comment explaining the `shardIndex` usage; the suggestion only adds a marginal clarifying comment without a functional change.	Low

Previous suggestions

Suggestions up to commit 0d6910f

Category	Suggestion	Impact
General	Guard merge against null per-leaf TopDocs `TopDocs.merge(k, perLeafTopDocs)` may throw when any leaf returned `NestedKnnUtil.EMPTY_TOP_DOCS` because merge requires each input to be sorted or non-empty with consistent structure, and empty leaves also have `shardIndex` unset. Ensure every per-leaf `TopDocs` is non-null and that empty leaves are either filtered out or produce an empty but valid `TopDocs` so the merge is safe. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [187-197] TopDocs[] perLeafTopDocs = indexSearcher.getTaskExecutor().invokeAll(rescoreTasks).toArray(TopDocs[]::new); +for (int i = 0; i < perLeafTopDocs.length; i++) { + if (perLeafTopDocs[i] == null) { + perLeafTopDocs[i] = NestedKnnUtil.EMPTY_TOP_DOCS; + } +} TopDocs topKParents = TopDocs.merge(k, perLeafTopDocs); List<Map<Integer, Float>> 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; Suggestion importance[1-10]: 3 __ Why: The rescore tasks always return a non-null `TopDocs` (either `NestedKnnUtil.EMPTY_TOP_DOCS` or the result from `rescoreLeafWithFullPrecision`), so the null-guard is largely unnecessary. The concern about `shardIndex` not being set on empty leaves is minor since empty TopDocs have no scoreDocs to iterate.	Low
General	Clarify filter bits cache scope and semantics The cache key `leafReaderContext.ord` is only unique within a single `IndexReader`; if `createWeight` is ever called across different readers in the same instance it could collide. More importantly, the first `get` is racy with `putIfAbsent`, which is fine, but the method may still build the bits twice unnecessarily. Consider `computeIfAbsent` with a wrapper to handle the `null` sentinel, or document that the map is scoped to one `createWeight` invocation. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [277-278] +Bits existing = filterBitsByLeaf.putIfAbsent(leafReaderContext.ord, bits); +return existing != null ? existing : bits; - Suggestion importance[1-10]: 2 __ Why: The `improved_code` is identical to the `existing_code`, making this suggestion non-actionable. The cache is scoped to a single `createWeight` call via a local map, so the concern about cross-reader collisions does not apply.	Low

Suggestions up to commit 72111ec

Category	Suggestion	Impact
General	Avoid returning shared empty TopDocs instance When a leaf's `leafResult` is empty the task returns `NestedKnnUtil.EMPTY_TOP_DOCS` whose `scoreDocs` is a shared empty array, but later code sets `scoreDoc.shardIndex =` `leafOrd` on returned docs. While the empty array has no elements and is safe here, `TopDocs.merge` requires each entry to be non-null; verify EMPTY_TOP_DOCS is non-null (and consider returning a fresh empty TopDocs to avoid any shared-instance mutation risk if callers ever iterate differently). src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [153-163] rescoreTasks.add(() -> { if (leafResult.isEmpty()) { - return NestedKnnUtil.EMPTY_TOP_DOCS; + return new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]); } Bits queryFilter = filterBits(filterBitsByLeaf, leafReaderContext, filterWeight); DocIdSetIterator allSiblings = queryUtils.getAllSiblings( leafReaderContext, leafResult.keySet(), internalNestedKnnVectorQuery.getParentFilter(), queryFilter ); Suggestion importance[1-10]: 2 __ Why: The suggestion is speculative; `NestedKnnUtil.EMPTY_TOP_DOCS` with an empty `scoreDocs` array is safe for `TopDocs.merge` and the empty array has no elements to mutate. The suggestion itself admits it is a "verify" request with no concrete bug identified.	Low

Suggestions up to commit 92d780a

Category	Suggestion	Impact
General	Avoid returning shared empty TopDocs singleton `NestedKnnUtil.EMPTY_TOP_DOCS` is likely a shared singleton; returning it from multiple leaf tasks and then passing them to `TopDocs.merge` may be fine, but using a shared mutable `TopDocs` (with a shared `scoreDocs` array) risks subtle bugs if any downstream code mutates `shardIndex` or `doc`. Consider returning a fresh empty `TopDocs` per leaf to match the pattern used elsewhere and avoid potential shared-state issues (similar to the `PerLeafResult.EMPTY_RESULT` NPE fix noted in CHANGELOG). src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [154-156] rescoreTasks.add(() -> { if (leafResult.isEmpty()) { - return NestedKnnUtil.EMPTY_TOP_DOCS; + return new TopDocs(new org.apache.lucene.search.TotalHits(0, org.apache.lucene.search.TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]); } Suggestion importance[1-10]: 3 __ Why: The suggestion raises a plausible concern about shared mutable state in `NestedKnnUtil.EMPTY_TOP_DOCS`, but it is speculative without evidence that the singleton has a mutable `scoreDocs` array that gets mutated downstream. The impact is minor and the concern may not be real.	Low
General	Verify leaf-ord caching correctness under concurrency `ConcurrentHashMap` does not permit null values, and `filterBits` already handles null by early-returning without caching. However, if `queryUtils.createBits` returns a non-null `Bits` whose `ord` collides across concurrent tasks, `putIfAbsent` is correct but the bits computed by the loser are discarded unused. This is fine functionally, but note that leaf ord is unique per leaf, so collisions only happen when the same leaf runs rescore and expansion concurrently — ensure the two phases do not run concurrently (currently they are sequential: rescore completes before `retrieveAll`), which is the case, so this is safe. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [91] +Map<Integer, Bits> filterBitsByLeaf = new ConcurrentHashMap<>(); - Suggestion importance[1-10]: 1 __ Why: The suggestion is purely a verification/analysis note with identical `existing_code` and `improved_code`, offering no actual code change or improvement.	Low

Suggestions up to commit 2591782

Category	Suggestion	Impact
General	Avoid sharing mutable empty TopDocs instance When a leaf has no candidates, returning `NestedKnnUtil.EMPTY_TOP_DOCS` (a shared static instance) is fine for merging, but any downstream code that mutates `scoreDoc.shardIndex` on these results would corrupt the shared empty array. Although `scoreDocs` is empty here so no mutation occurs, confirm that `EMPTY_TOP_DOCS` is never mutated elsewhere; otherwise return a fresh empty `TopDocs` to be safe. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [154-156] rescoreTasks.add(() -> { if (leafResult.isEmpty()) { - return NestedKnnUtil.EMPTY_TOP_DOCS; + return new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]); } Suggestion importance[1-10]: 3 __ Why: The concern is speculative; `EMPTY_TOP_DOCS` has an empty `scoreDocs` array, so the shardIndex mutation loop is a no-op. The suggestion is defensive but adds minor value without evidence of actual mutation elsewhere.	Low
General	Use a more robust cache key for leaves `leafReaderContext.ord` is the ordinal within the parent reader. If the query is executed against sub-readers at different levels (e.g., per-leaf vs. top-level), two different leaves could share the same `ord`, causing bit-set collisions. Consider keying the cache by `leafReaderContext.id()` or the `LeafReaderContext` itself to guarantee uniqueness across all contexts seen by this query. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [277-278] +Bits existing = filterBitsByLeaf.putIfAbsent(leafReaderContext.ord, bits); +return existing != null ? existing : bits; - Suggestion importance[1-10]: 2 __ Why: Within a single `createWeight` call, all leaves come from the same `IndexReader#leaves()` call, so `ord` values are unique. The `existing_code` and `improved_code` are also identical, which further reduces the suggestion's value.	Low

Suggestions up to commit ad0bf3f

Category	Suggestion	Impact
Possible issue	Verify shardIndex tag survives TopDocs.merge The `ScoreDoc` instances returned here are mutated with a `shardIndex` tag and later used to look up the leaf via `scoreDoc.shardIndex`. However, `TopDocs.merge(k,` `perLeafTopDocs)` with any non-null `shardIndex` semantics may rewrite `shardIndex` on the returned ScoreDocs (Lucene's merge sets shardIndex when the input docs have shardIndex set and i is passed). Verify that the shardIndex tag survives `TopDocs.merge` — if Lucene overwrites it with the array index, the subsequent `reducedResults.get(scoreDoc.shardIndex)` will populate the wrong leaf map and silently corrupt results. Consider using a side map keyed by identity or encoding the leaf ord into the doc id. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [180-183] +for (ScoreDoc scoreDoc : leafTopDocs.scoreDocs) { + scoreDoc.shardIndex = leafOrd; +} +return leafTopDocs; - Suggestion importance[1-10]: 7 __ Why: This raises a legitimate concern about `TopDocs.merge` potentially overwriting `shardIndex`. The PR comment claims merge preserves it, but this is worth verifying as it could cause silent data corruption. The improved_code is identical to existing_code so it's only a verification request, limiting the score.	Medium
General	Avoid constructing searcher twice per query `resolveExactSearcher()` is called from both `rescoreToTopKParents` and `scoreAllSiblings` in `createWeight`, and each call constructs a new `ExactSearcher` when `exactSearcher` is unset. Cache the resolved instance (e.g., in a local variable passed through, or memoize once at the start of `createWeight`) to avoid building two `ExactSearcher` instances per query execution, each reaching for the `ModelDao` singleton. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [114-116] private ExactSearcher resolveExactSearcher() { - return exactSearcher != null ? exactSearcher : new ExactSearcher(ModelDao.OpenSearchKNNModelDao.getInstance()); + if (exactSearcher != null) { + return exactSearcher; + } + return new ExactSearcher(ModelDao.OpenSearchKNNModelDao.getInstance()); } Suggestion importance[1-10]: 3 __ Why: Minor optimization to avoid constructing `ExactSearcher` twice. The improved_code is functionally equivalent to the existing code and doesn't actually implement the memoization suggested in the description.	Low

Vikasht34 · 2026-09-18T06:23:04Z

Congratulations @benkim1028 for First PR !!

github-actions · 2026-09-18T17:04:04Z

Persistent review updated to latest commit 3d88dd1

codecov · 2026-09-18T18:47:26Z

Codecov Report

❌ Patch coverage is 97.46835% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.11%. Comparing base (53a3ab0) to head (0d6910f).

Files with missing lines	Patch %	Lines
...n/index/query/lucenelib/ExpandNestedDocsQuery.java	96.61%	0 Missing and 2 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #3579      +/-   ##
============================================
+ Coverage     83.06%   83.11%   +0.05%     
- Complexity     4968     5015      +47     
============================================
  Files           489      489              
  Lines         17740    17793      +53     
  Branches       2389     2399      +10     
============================================
+ Hits          14735    14788      +53     
+ Misses         2106     2100       -6     
- Partials        899      905       +6

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

github-actions · 2026-09-22T16:35:42Z

Persistent review updated to latest commit 84441fc

naveentatikonda · 2026-09-23T06:54:14Z

+                    .build();
+                TopDocs leafTopDocs = searcher.searchLeaf(leafReaderContext, exactSearcherContext);
+                for (ScoreDoc scoreDoc : leafTopDocs.scoreDocs) {
+                    scoreDoc.shardIndex = leafOrd;


Can you pls help me understand why are we setting this shardIndex ?

Won't TopDocs.merge(k, perLeafTopDocs) on the next line already sets scoreDoc.shardIndex, can you pls confirm ?

As far as I understood, there is no downstream method that sets shardIndex. I tested without this line. There was no doc returned even though there were 6 docs and 2 parents should have been returned.

can we add comment on why we are doing this?

Yes. I added the comment above.

naveentatikonda · 2026-09-23T06:56:17Z

+                    internalNestedKnnVectorQuery.getParentFilter(),
+                    queryFilter
+                );
+                final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder()


Looks like you missed the profile breakdown for the exact search per leaf

Got it I will add that.

Added profile breakdown to all exact search. Tested and verified that {"exact_search" : 1143583, "exact_search_count" : 4} is added to the ExpandNestedDocsQuery's breakdown.

naveentatikonda · 2026-09-23T07:00:37Z

+                );
+                final ExactSearcher.ExactSearcherContext exactSearcherContext = ExactSearcher.ExactSearcherContext.builder()
+                    .matchedDocsIterator(allSiblings)
+                    .numberOfMatchedDocs(allSiblings.cost())


Looks like a major chunk of this rescore logic duplicates RescoreKNNVectorQuery.searchLeaf, can you see if we can extract this into a util method ?

Sounds good. Will see if there is simple way to extract them out to a single util function.

I extracted them to a single util function. But I think it did not add too much value. Let me know if you think reverting this change is better.

naveentatikonda · 2026-09-23T07:04:43Z

+    private static final int DIMENSION = 3;
+    private static final int CHILDREN_PER_PARENT = 3;
+
+    /**


can you add an IT to test filter + expand_nested_docs + rescore ?

Sounds good. Will add that.

Added two new tests covering both parent filter and child filter.

testExpandNestedDocs_whenRescoreEnabledWithParentFilter_thenReturnKMatchingParentsWithAllChildren

testExpandNestedDocs_whenRescoreEnabledWithChildFilter_thenReturnOnlyMatchingChildren

github-actions · 2026-09-24T18:42:08Z

Persistent review updated to latest commit a14f4c3

github-actions · 2026-09-24T18:59:21Z

Persistent review updated to latest commit 5006617

github-actions · 2026-09-29T19:37:59Z

Persistent review updated to latest commit 0cc7c13

github-actions · 2026-09-30T18:26:57Z

Persistent review updated to latest commit b75ba15

github-actions · 2026-10-01T18:19:28Z

Persistent review updated to latest commit 393554a

naveentatikonda

LGTM!

But, @benkim1028 as discussed offline please refactor this as a followup task using the wrapper approach by breaking them down into individual components and reuse them for both Lucene and NativeEngine query to avoid code duplication. Also, please create a github issue for tracking it. Thanks!

github-actions · 2026-10-06T18:38:03Z

Persistent review updated to latest commit a1da462

benkim1028 · 2026-10-06T19:21:00Z

LGTM!

But, @benkim1028 as discussed offline please refactor this as a followup task using the wrapper approach by breaking them down into individual components and reuse them for both Lucene and NativeEngine query to avoid code duplication. Also, please create a github issue for tracking it. Thanks!

@naveentatikonda An issue has been created for the refactoring: #3619

github-actions · 2026-10-06T23:13:41Z

Persistent review updated to latest commit ad0bf3f

github-actions · 2026-10-06T23:20:41Z

Persistent review updated to latest commit 2591782

VijayanB

LGTM

github-actions · 2026-10-08T17:02:59Z

PR Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit 2faf128.

⛔ Hard block: Issues at High severity or above will block this PR from merging.

'Diff too large, requires skip by maintainers after manual review'

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

The merge-base changed after approval.

github-actions · 2026-10-08T17:05:51Z

Persistent review updated to latest commit 92d780a

…earch-project#3125) With rescoring enabled, expand_nested_docs on the Lucene engine cut to k after expanding children, so k counted child documents rather than parent documents: a single parent could consume the whole budget and the remaining parents were dropped. KNNQueryFactory also silently skipped the rescore wrapper whenever expand_nested_docs was set. ExpandNestedDocsQuery now owns the rescore stage, in the same order NativeEngineKnnVectorQuery uses: approximate search over the oversampled candidate pool, a full precision rescore that collapses each parent group to its best child and cuts to k parents, and only then the child expansion. The expansion also scores children on full precision vectors when rescoring is enabled, so a quantized child score no longer leaks into the parent through the nested score mode. The collapsing pass tags each hit with its leaf ordinal in shardIndex, because TopDocs#merge flattens the per leaf results and the expansion has to regroup them by leaf. Removes the short circuit in KNNQueryFactory and the dead expandNestedDocs branch in OSDiversifyingChildrenFloatKnnVectorQuery. The full precision rescore of a leaf is extracted into QueryUtils#rescoreLeafWithFullPrecision and shared by RescoreKNNVectorQuery and both ExpandNestedDocsQuery stages. Profiling: all exact searches in ExpandNestedDocsQuery are timed under KNNQueryTimingType.EXACT_SEARCH, and the query is registered in KNNPlugin#getQueryProfileMetricsProvider. The registration is required, because without it a profiled query fails looking up the timer. Testing: unit tests for the rescore ordering, empty leaves, query equality and profiling. ExpandNestedDocsWithRescoreIT covers the issue's scenario on a 4x on_disk index, which resolves to the Lucene engine, asserting parent ids and exact full precision scores. It also covers many parents, rescore disabled, an explicit oversample factor, parent and child filters, and the Profile API breakdown. Signed-off-by: Ben Kim <kimsong@amazon.com>

github-actions · 2026-10-08T17:40:01Z

Persistent review updated to latest commit 72111ec

Signed-off-by: Ben Kim <benkim1028@gmail.com>

github-actions · 2026-10-08T23:26:08Z

Persistent review updated to latest commit 0d6910f

github-actions · 2026-10-09T16:48:04Z

Persistent review updated to latest commit 0d6910f

benkim1028 requested review from 0ctopus13prime, VijayanB, Vikasht34, heemin32, kotwanikunal, luyuncheng, martin-gaievski, naveentatikonda, navneet1v, shatejas and vamshin as code owners September 18, 2026 06:12

naveentatikonda reviewed Sep 23, 2026

View reviewed changes

benkim1028 force-pushed the fix/3125-expand-nested-rescore branch from a14f4c3 to 5006617 Compare September 24, 2026 18:57

naveentatikonda previously approved these changes Oct 2, 2026

View reviewed changes

benkim1028 dismissed naveentatikonda’s stale review via a1da462 October 6, 2026 18:36

VijayanB previously approved these changes Oct 7, 2026

View reviewed changes

benkim1028 force-pushed the fix/3125-expand-nested-rescore branch from 2591782 to 2faf128 Compare October 8, 2026 17:02

benkim1028 force-pushed the fix/3125-expand-nested-rescore branch from 92d780a to 72111ec Compare October 8, 2026 17:38

VijayanB previously approved these changes Oct 8, 2026

View reviewed changes

Merge branch 'main' into fix/3125-expand-nested-rescore

0d6910f

Signed-off-by: Ben Kim <benkim1028@gmail.com>

benkim1028 dismissed VijayanB’s stale review via 0d6910f October 8, 2026 23:24

kotwanikunal approved these changes Oct 9, 2026

View reviewed changes

VijayanB approved these changes Oct 9, 2026

View reviewed changes

benkim1028 closed this Oct 9, 2026

benkim1028 reopened this Oct 9, 2026

VijayanB merged commit 7318951 into opensearch-project:main Oct 9, 2026
161 of 165 checks passed

Conversation

benkim1028 commented Sep 18, 2026

Description

Testing

Related Issues

Check List

Uh oh!

github-actions Bot commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Reviewer Guide 🔍

(Review updated until commit 0d6910f)

Uh oh!

github-actions Bot commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Code Suggestions ✨

Previous suggestions

Uh oh!

Vikasht34 commented Sep 18, 2026

Uh oh!

github-actions Bot commented Sep 18, 2026

Uh oh!

codecov Bot commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Sep 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026

Uh oh!

github-actions Bot commented Oct 1, 2026

Uh oh!

naveentatikonda left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

benkim1028 commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

github-actions Bot commented Oct 6, 2026

Uh oh!

VijayanB left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 8, 2026

PR Code Analyzer ❗

github-actions Bot commented Sep 18, 2026 •

edited

Loading

(Review updated until commit `0d6910f`)

github-actions Bot commented Sep 18, 2026 •

edited

Loading

codecov Bot commented Sep 18, 2026 •

edited

Loading

benkim1028 commented Oct 6, 2026 •

edited

Loading