Rescore expanded nested docs on Lucene engine by naykudev · Pull Request #3483 · opensearch-project/k-NN

naykudev · 2026-08-04T23:51:28Z

Description

Fixes a bug where enabling rescoring together with a nested knn_vector field using expand_nested_docs on the Lucene engine (auto-enabled by on_disk / 4x mode) fails to return all child documents.

Previously, KNNQueryFactory silently skipped rescoring for this combination and logged a warning, because RescoreKNNVectorQuery reduces results to k via TopDocs.merge(k, ...) — which would truncate the fully expanded child set produced by ExpandNestedDocsQuery.

This change moves rescoring inside ExpandNestedDocsQuery, mirroring the already-correct native-engine ordering in NativeEngineKnnVectorQuery:

Approximate search over the oversampled parent candidates.
Diversified, full-precision exact rescore that selects the best child per parent, reduced across segments to the top k parents.
Expand all child documents of the surviving top k parents.

The oversampled parent candidates are now retained during segment merging (OSDiversifyingChildrenFloatKnnVectorQuery.mergeLeafResults) whenever rescoring is enabled, so the rescore step has candidates to work with. Only the float Lucene path rescores; the byte path is unchanged.

The filter weight is built once in ExpandNestedDocsQuery and shared across the rescore and expansion passes instead of being rewritten twice. Per-leaf rescore correctness relies on the invariant luceneK >= rescoreK (luceneK is defined as max(rescoreK, efSearch) in KNNQueryFactory); this is now documented and guarded with an assertion so a broken invariant surfaces as a test failure rather than a silent per-leaf truncation.

Why this approach (mirroring the native engine)

The native engines (NativeEngineKnnVectorQuery) already handle nested + rescore + expand_nested_docs correctly, and the fix deliberately makes the Lucene path follow the same ordering rather than teaching the shared RescoreKNNVectorQuery about parents/nesting (which would widen the blast radius to every engine). The key property in the native flow is that the reduce-to-k happens on parents, before expansion — expansion always runs last, on the already-reduced set.

Native flow (NativeEngineKnnVectorQuery#createWeight):

Approximate search, oversampled — searches with searchK = max(firstPassK, effectiveK), then trims to firstPassK (the oversampled candidate pool).
Full-precision rescore — doRescore(...) runs exact search with useQuantizedVectorsForSearch(false); for nested it gathers the sibling set via getAllSiblings(parentsFilter).
Reduce to top-k parents — reduceToTopK(..., finalK), before the expand block.
Expand — retrieveAll gathers all siblings of the surviving parents and returns every child.

One-to-one mapping to this change:

Native (`NativeEngineKnnVectorQuery`)	This fix (`ExpandNestedDocsQuery`)
`searchK = max(firstPassK, effectiveK)`	`luceneK = max(rescoreK, efSearch)`
`reduceToTopK(..., firstPassK)`	`mergeLeafResults` → `TopDocs.merge(rescoreK, ...)`
`doRescore(...)` w/ `useQuantizedVectorsForSearch(false)`	`rescore(...)` → `diversifyingExactSearch` (full precision)
`reduceToTopK(..., finalK)` then expand	`TopDocs.merge(k, rescored)` then `retrieveAll`
`retrieveAll` / `getAllSiblings`	`retrieveAll` / `getAllSiblings`

It is the same sequence, not shared code: the native path uses KNNWeight.exactSearch + ExactSearcherContext + ResultUtil.reduceToTopK, while the Lucene path uses Lucene-native DiversifyingChildrenFloatKnnVectorQuery.exactSearch + TopDocs.merge.

Related Issues

Resolves #3125

Check List

Added Bug Fixes changelog entry.
Added/updated unit tests (KNNQueryFactoryTests, ExpandNestedEDocsQueryTests, OSDiversifyingChildrenFloatKnnVectorQueryTests) and an integration test (ExpandNestedDocsIT) reproducing [BUG] ExpandNestedDocs with Rescoring enabled for Lucene engine does not return all the docs #3125.
All new and existing tests passed, including the full ExpandNestedDocsIT suite across every engine/mode.
Commits are signed per the DCO using --signoff.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

When rescoring is enabled together with a nested knn_vector field using expand_nested_docs on the Lucene engine (auto-enabled by on_disk / 4x mode), the query dropped child documents. Rescoring was silently skipped because RescoreKNNVectorQuery reduces results to k, which would truncate the fully expanded child set. Move rescoring inside ExpandNestedDocsQuery so it mirrors the native engine ordering: run approximate search over the oversampled candidates, rescore them at full precision and reduce to the top k parents, then expand all child documents of the surviving parents. The oversampled parent candidates are now retained during segment merging so the rescore step has candidates to work with. Fixes opensearch-project#3125 Signed-off-by: ved naykude <vnaykude@amazon.com>

github-actions · 2026-08-04T23:52:59Z

PR Reviewer Guide 🔍

(Review updated until commit `d6e14af`)

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 Duplicate/Incorrect Changelog Entry A new bullet was added referencing "BQ file is not present in the segment" with an empty PR link `[]()`, and a nearly identical bullet already exists just below with a proper link (#3511). This appears to be an accidental duplicate unrelated to this PR's fix and should be removed. * Fix when BQ file is not present in the segment as there is no vectors in the segment []() Constructor Signature Mismatch The class now stores `luceneK` as a separate field, but the existing constructor called from `NestedKnnVectorQueryFactory.createNestedKnnVectorQuery` for the non-expandNested case (`new OSDiversifyingChildrenFloatKnnVectorQuery(fieldName, vector, filterQuery, luceneK, parentFilter, k, rescoreK)`) does not appear in the diff. Verify all constructors initialize the new `luceneK` field; otherwise `mergeLeafResults` will read an uninitialized/zero `luceneK`, causing the `Math.min(rescoreK, luceneK)` safety net to return 0 candidates when the invariant is violated. 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.luceneK = luceneK; this.rescoreK = rescoreK; this.expandNestedDocs = expandNestedDocs; this.parentFilter = parentFilter;

github-actions · 2026-08-04T23:53:17Z

PR Code Suggestions ✨

Latest suggestions up to d6e14af

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Remove duplicate changelog entry with empty link A duplicate changelog entry for the BQ file fix was added with an empty link, alongside the already existing entry with PR #3511. Remove the duplicated line with the empty link to avoid a broken reference and redundant entry. CHANGELOG.md [17] * Return all nested child documents when rescoring is enabled with expand_nested_docs on the Lucene engine [#3483](https://github.com/opensearch-project/k-NN/pull/3483) -* Fix when BQ file is not present in the segment as there is no vectors in the segment []() * Add prefetch for Lucene engine's fp32 and binary vector data type [#3504](https://github.com/opensearch-project/k-NN/pull/3504) * Fix when BQ file is not present in the segment as there is no vectors in the segment [#3511](https://github.com/opensearch-project/k-NN/pull/3511) Suggestion importance[1-10]: 7 __ Why: The duplicate changelog entry with an empty link `[]()` is clearly redundant given the same fix is already documented on line 19 with PR #3511. Removing it improves documentation quality.	Medium
Possible issue	Guard against empty leaves in rescore merge `TopDocs.merge` throws `IllegalArgumentException` when passed an empty array, which can occur if `leafReaderContexts` is empty (e.g., an empty index). Guard against this by returning early with an empty per-leaf list when no rescore tasks were run, matching the empty-result handling in `createWeight`. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [147-151] final TopDocs[] rescored = indexSearcher.getTaskExecutor().invokeAll(rescoreTasks).toArray(TopDocs[]::new); // Reduce across all leaves to the top k parents. final int k = internalNestedKnnVectorQuery.getK(); +if (rescored.length == 0) { + return new ArrayList<>(); +} final TopDocs topKParents = TopDocs.merge(k, rescored); Suggestion importance[1-10]: 3 __ Why: The concern about `TopDocs.merge` with empty array is valid, but in practice `rescore()` is only called after `perLeafResults` is populated and the outer code already handles the case where results are empty via the `topK.scoreDocs.length == 0` check. The edge case is unlikely and the fix is minor.	Low

Previous suggestions

Suggestions up to commit 9aa0deb

Category	Suggestion	Impact
General	Remove duplicated changelog entry A duplicate "Fix when BQ file is not present in the segment as there is no vectors in the segment" entry was introduced with an empty link, likely a merge artifact. Remove the duplicated line with the empty link to keep the changelog clean and avoid a broken reference. CHANGELOG.md [17] * Return all nested child documents when rescoring is enabled with expand_nested_docs on the Lucene engine [#3483](https://github.com/opensearch-project/k-NN/pull/3483) -* Fix when BQ file is not present in the segment as there is no vectors in the segment []() * Add prefetch for Lucene engine's fp32 and binary vector data type [#3504](https://github.com/opensearch-project/k-NN/pull/3504) * Fix when BQ file is not present in the segment as there is no vectors in the segment [#3511](https://github.com/opensearch-project/k-NN/pull/3511) Suggestion importance[1-10]: 6 __ Why: Correctly identifies a duplicate changelog entry with an empty link `[]()` that appears to be a merge artifact, and removing it improves changelog cleanliness.	Low
General	Ensure merge budget matches rescore path When `expandNestedDocs` is true, `mergeLeafResults` is invoked during the approximate pass triggered from `ExpandNestedDocsQuery.knnRewrite`, but `rescoreK` here is the same value used to gate the rescore path in `ExpandNestedDocsQuery`. If a caller sets `expandNestedDocs=true` without wiring `rescoreK` into `ExpandNestedDocsQuery` (as done in `NestedKnnVectorQueryFactory`), the approximate pass will merge to `rescoreK` candidates but the rescore step will be skipped, leaving the final result truncated to `rescoreK` instead of expanding all children. Consider also gating this branch on `expandNestedDocs` matching the actual downstream behavior, or asserting that they are consistent. src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java [95-108] if (rescoreK != RescoreContext.NO_RESCORE_NEEDED) { - ... assert luceneK >= rescoreK : "luceneK (" + luceneK + ") must be >= rescoreK (" + rescoreK + ")"; return TopDocs.merge(Math.min(rescoreK, luceneK), perLeafResults); } Suggestion importance[1-10]: 3 __ Why: The suggestion raises a legitimate concern about consistency between `rescoreK` in the two classes but the `improved_code` is essentially identical to the existing code (only removes comments), and the concern is largely theoretical since `NestedKnnVectorQueryFactory` already wires them correctly.	Low

Suggestions up to commit b4235a1

Category	Suggestion	Impact
Possible issue	Guard rescore merge against empty leaves `TopDocs.merge` requires a non-empty array; when `leafReaderContexts` is empty (or all leaves return empty `TopDocs`), `invokeAll` may return an empty array causing `IllegalArgumentException`. Guard against the empty-leaves case by returning the input `perLeafResults` (or an empty per-leaf list) before calling merge. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [147-151] final TopDocs[] rescored = indexSearcher.getTaskExecutor().invokeAll(rescoreTasks).toArray(TopDocs[]::new); - +if (rescored.length == 0) { + return perLeafResults; +} // Reduce across all leaves to the top k parents. final int k = internalNestedKnnVectorQuery.getK(); final TopDocs topKParents = TopDocs.merge(k, rescored); Suggestion importance[1-10]: 3 __ Why: `TopDocs.merge` in Lucene actually handles empty arrays without throwing, and the empty-leaves case is unlikely in practice since callers iterate over `leafReaderContexts`. The suggestion is a minor defensive addition with limited real impact.	Low

Suggestions up to commit 039a4c2

Category	Suggestion	Impact
General	Guard against non-positive merge size `TopDocs.merge` throws `IllegalArgumentException` if the topN argument is 0 or negative. If `rescoreK` or `luceneK` is 0 (e.g., misconfiguration or empty result), this will fail with an unhelpful message. Consider validating that the merge size is positive or falling back to `k`. src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java [106-107] assert luceneK >= rescoreK : "luceneK (" + luceneK + ") must be >= rescoreK (" + rescoreK + ")"; -return TopDocs.merge(Math.min(rescoreK, luceneK), perLeafResults); +int mergeSize = Math.min(rescoreK, luceneK); +if (mergeSize <= 0) { + return TopDocs.merge(k, perLeafResults); +} +return TopDocs.merge(mergeSize, perLeafResults); Suggestion importance[1-10]: 3 __ Why: `rescoreK` being 0 or negative shouldn't occur in practice since the rescore path is gated by `rescoreK != RescoreContext.NO_RESCORE_NEEDED`, but the defensive fallback to `k` is a minor safety improvement.	Low
General	Guard against invalid leaf index mapping `ReaderUtil.subIndex` expects a list of leaves sorted by `docBase`, but it also requires the list to be passed as `List` where each context's `docBase` matches its position. This is fine for `reader.leaves()`, but be aware that `scoreDoc.doc` could theoretically be `NO_MORE_DOCS` or out of range if a downstream implementation misbehaves. Consider adding a defensive check to fail fast rather than silently mis-routing docs to the wrong leaf. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [161-165] for (ScoreDoc scoreDoc : topKParents.scoreDocs) { final int leafIndex = ReaderUtil.subIndex(scoreDoc.doc, leafReaderContexts); + if (leafIndex < 0 \|\| leafIndex >= leafReaderContexts.size()) { + throw new IllegalStateException("Rescored doc id " + scoreDoc.doc + " is outside any leaf"); + } final LeafReaderContext leafReaderContext = leafReaderContexts.get(leafIndex); survivingPerLeaf.get(leafIndex).put(scoreDoc.doc - leafReaderContext.docBase, NO_SCORE); } Suggestion importance[1-10]: 2 __ Why: `ReaderUtil.subIndex` is a well-established Lucene utility that reliably returns a valid leaf index for any valid doc id. The defensive check offers little practical value since the input `scoreDoc.doc` values come from the trusted rescore pass with rebased global doc ids.	Low

Suggestions up to commit d53006c

Category	Suggestion	Impact
General	Preserve prior merge semantics for non-expand path The non-expandNested rescore path previously merged to `rescoreK`, but now this branch applies unconditionally whenever `rescoreK` is set. For the non-expandNested path there is no follow-up rescore that enforces `luceneK` as an upper bound, so clamping via `Math.min(rescoreK, luceneK)` changes prior semantics when `rescoreK > luceneK` (previously it would merge to `rescoreK`). Consider gating the `Math.min` clamp/assert to the `expandNestedDocs` case to preserve existing behavior for the non-expand path. src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java [95-108] if (rescoreK != RescoreContext.NO_RESCORE_NEEDED) { - // When rescoring is enabled, merge to the oversampled k (rescore budget) rather than the full - // luceneK which may have been expanded by ef_search. For the expandNested path this preserves the - // oversampled parent candidates so ExpandNestedDocsQuery can rescore them at full precision and - // reduce to the top k parents before expanding their child documents. - // Invariant: luceneK >= rescoreK (KNNQueryFactory defines luceneK as max(rescoreK, efSearch)), so - // merging to rescoreK here never returns more candidates than the exact rescore pass can collect. - // The assert catches a broken invariant in tests. The Math.min is a production safety net for when - // assertions are disabled: it keeps this merge budget consistent with the exact rescore budget - // (luceneK) so a future caller that violates the invariant cannot make this step retain more - // candidates than the rescore pass can score, which would otherwise be a silent per-leaf truncation. - assert luceneK >= rescoreK : "luceneK (" + luceneK + ") must be >= rescoreK (" + rescoreK + ")"; - return TopDocs.merge(Math.min(rescoreK, luceneK), perLeafResults); + if (expandNestedDocs) { + assert luceneK >= rescoreK : "luceneK (" + luceneK + ") must be >= rescoreK (" + rescoreK + ")"; + return TopDocs.merge(Math.min(rescoreK, luceneK), perLeafResults); + } + return TopDocs.merge(rescoreK, perLeafResults); } Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern that clamping to `Math.min(rescoreK, luceneK)` changes prior semantics for the non-expand path. However, since `luceneK = max(rescoreK, efSearch) >= rescoreK` is guaranteed by `KNNQueryFactory`, the clamp is effectively a no-op in practice, making this a minor defensive concern.	Low
General	Handle empty rescore result across leaves `ReaderUtil.subIndex` expects a list of `LeafReaderContext`, but the overload that accepts a `List` operates on `docBase` ordering; ensure `leafReaderContexts` is the ordered list returned from `reader.leaves()` (it is here). More importantly, if `topKParents.scoreDocs` is empty (all leaves returned no rescore results), `survivingPerLeaf` will contain only empty maps and `retrieveAll` will still be invoked per leaf — verify `retrieveAll` handles empty per-leaf maps gracefully to avoid producing spurious results or NPEs. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [161-165] for (ScoreDoc scoreDoc : topKParents.scoreDocs) { final int leafIndex = ReaderUtil.subIndex(scoreDoc.doc, leafReaderContexts); final LeafReaderContext leafReaderContext = leafReaderContexts.get(leafIndex); survivingPerLeaf.get(leafIndex).put(scoreDoc.doc - leafReaderContext.docBase, NO_SCORE); } +// Note: if topKParents is empty, survivingPerLeaf contains only empty maps; retrieveAll must handle this. Suggestion importance[1-10]: 3 __ Why: The suggestion only asks the author to verify that `retrieveAll` handles empty per-leaf maps and does not propose a concrete code fix beyond adding a comment. The `improved_code` is essentially identical to the `existing_code`.	Low

Suggestions up to commit d54959c

Category	Suggestion	Impact
Possible issue	Use shardIndex for cross-leaf TopDocs merging `TopDocs.merge` here operates on results whose `ScoreDoc.doc` were rebased to global doc ids, but Lucene's merge assumes segment-local doc ids and expects a `shardIndex` per `TopDocs`. Passing already-globalized ids without setting `shardIndex` on each `TopDocs` may produce incorrect merge/tie-breaking behavior. Consider setting `shardIndex` on each per-leaf `TopDocs` (and keeping segment-local doc ids) before merging, then map back using the shardIndex. src/main/java/org/opensearch/knn/index/query/lucenelib/ExpandNestedDocsQuery.java [149-151] -// Reduce across all leaves to the top k parents. +// Assign shardIndex per leaf and keep segment-local doc ids; TopDocs.merge will use shardIndex for tie-breaking. +for (int i = 0; i < rescored.length; i++) { + for (ScoreDoc sd : rescored[i].scoreDocs) { + sd.shardIndex = i; + } +} final int k = internalNestedKnnVectorQuery.getK(); final TopDocs topKParents = TopDocs.merge(k, rescored); Suggestion importance[1-10]: 6 __ Why: Valid concern about `TopDocs.merge` semantics: Lucene's merge typically expects segment-local ids with `shardIndex` set for tie-breaking, and passing globalized doc ids without `shardIndex` may lead to subtle merge/tie-breaking issues. However, since the code later uses `ReaderUtil.subIndex` to map back to leaves, the globalized ids are intentional, so the impact is limited to potential tie-breaking edge cases.	Low
General	Avoid changing non-nested rescore merge budget The prior behavior for the non-expandNested rescore path merged to `rescoreK` unconditionally; capping with `Math.min(rescoreK, luceneK)` now silently reduces the merge budget for the non-expandNested path if the invariant is ever broken, which can regress recall. Consider limiting the `Math.min` safety net to the `expandNestedDocs` branch to preserve existing non-nested rescore semantics. src/main/java/org/opensearch/knn/index/query/lucenelib/OSDiversifyingChildrenFloatKnnVectorQuery.java [106-107] assert luceneK >= rescoreK : "luceneK (" + luceneK + ") must be >= rescoreK (" + rescoreK + ")"; -return TopDocs.merge(Math.min(rescoreK, luceneK), perLeafResults); +int mergeBudget = expandNestedDocs ? Math.min(rescoreK, luceneK) : rescoreK; +return TopDocs.merge(mergeBudget, perLeafResults); Suggestion importance[1-10]: 3 __ Why: Given the invariant `luceneK >= rescoreK` is guaranteed by `KNNQueryFactory`, the `Math.min` is only a safety net and does not change behavior in practice. The suggestion is a minor stylistic preference with limited real-world impact.	Low

codecov · 2026-08-05T01:32:47Z

Codecov Report

❌ Patch coverage is 97.36842% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 83.99%. Comparing base (5108f83) to head (d6e14af).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...n/index/query/lucenelib/ExpandNestedDocsQuery.java	96.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #3483      +/-   ##
============================================
+ Coverage     83.97%   83.99%   +0.01%     
- Complexity     4505     4518      +13     
============================================
  Files           461      462       +1     
  Lines         16138    16169      +31     
  Branches       2104     2111       +7     
============================================
+ Hits          13552    13581      +29     
- Misses         1797     1798       +1     
- Partials        789      790       +1

☔ 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.

Build the filter weight once in ExpandNestedDocsQuery and share it across the rescore and expansion passes instead of rewriting the filter twice. Document and assert the luceneK >= rescoreK invariant that keeps the per-leaf rescore from silently truncating candidate parents. Signed-off-by: ved naykude <vnaykude@amazon.com>

github-actions · 2026-08-05T20:01:19Z

Persistent review updated to latest commit fb96455

Add tests for the three uncovered lines: - InternalNestedKnnVectorQuery.knnRescoreSearch default throws UnsupportedOperationException for non-rescore implementations. - OSDiversifyingChildrenFloatKnnVectorQuery.mergeLeafResults asserts the luceneK >= rescoreK invariant. - ExpandNestedDocsQuery equals/hashCode account for rescoreK. Signed-off-by: ved naykude <vnaykude@amazon.com>

- Document that expansion re-runs full-precision knnExactSearch, so the rescored scores are recomputed (not propagated) and the final child scores remain full precision; consistent with the native engine. - Cap the rescore merge budget with Math.min(rescoreK, luceneK) so a future invariant violation degrades gracefully instead of silently truncating per-leaf when assertions are disabled. Assert retained for dev-time signal. Signed-off-by: ved naykude <vnaykude@amazon.com>

github-actions · 2026-08-07T19:29:15Z

Persistent review updated to latest commit a4445d2

github-actions · 2026-08-10T22:34:23Z

Persistent review updated to latest commit 449826c

github-actions · 2026-08-11T22:28:24Z

Persistent review updated to latest commit d9e7706

Vikasht34 · 2026-08-12T21:03:21Z

+            // before expanding all of their child documents below.
+            perLeafResults = rescore(searcher, leafReaderContexts, perLeafResults, filterWeight);
+        }
+        TopDocs[] topDocs = retrieveAll(searcher, leafReaderContexts, perLeafResults, filterWeight);


Are we running excat search twice , one on rescore and retrieveAll ?

Yes — with rescore enabled there are two exact-search passes, and they serve different purposes:

rescore() runs a diversified exact search (knnRescoreSearch → best child per parent) purely to select the top-k parents at full precision.

retrieveAll() then runs the non-diversified exact search (knnExactSearch) to expand and score every child of those surviving parents.

Both read raw float vectors, so a child scored in both passes gets the identical score — the second scoring is redundant, not inconsistent, and the final child scores are full precision. This mirrors the native engine (NativeEngineKnnVectorQuery), which also re-runs exact search during expansion rather than propagating rescored scores. Without rescore, only retrieveAll() runs. I added an inline comment at the call site (and expanded the rescore() javadoc) to make this explicit.

Vikasht34 · 2026-08-12T21:04:26Z

+            final LeafReaderContext leafReaderContext = leafReaderContexts.get(leafIndex);
+            survivingPerLeaf.get(leafIndex).put(scoreDoc.doc - leafReaderContext.docBase, scoreDoc.score);
+        }
+        return survivingPerLeaf;


survivingPerLeaf stores scoreDoc.score, but retrieveAll only consumes keySet(), do we even deed to compute score

Good catch — addressed in 0c91ea4. rescore() no longer computes or stores a real score: retrieveAll consumes only keySet() (it re-scores every child via exact search), so the map value is now a NO_SCORE placeholder. The Map<Integer, Float> shape is kept only to match the type retrieveAll already expects. See the updated comments at the top of rescore() and around the survivingPerLeaf population.

Vikasht34

Please make all IT Passes for this

Signed-off-by: naykudev <vnaykude@amazon.com>

github-actions · 2026-08-12T22:31:18Z

Persistent review updated to latest commit 6f161ff

Signed-off-by: ved naykude <vnaykude@amazon.com>

github-actions · 2026-08-12T22:35:07Z

Persistent review updated to latest commit 0c91ea4

Signed-off-by: ved naykude <vnaykude@amazon.com>

github-actions · 2026-08-12T23:03:53Z

Persistent review updated to latest commit 1b020c4

github-actions · 2026-08-13T16:45:23Z

Persistent review updated to latest commit 0d7dd39

github-actions · 2026-08-13T18:09:12Z

Persistent review updated to latest commit 38a6c6b

github-actions · 2026-08-14T00:30:54Z

Persistent review updated to latest commit d54959c

github-actions · 2026-08-15T22:41:14Z

Persistent review updated to latest commit d53006c

github-actions · 2026-08-19T02:51:53Z

Persistent review updated to latest commit 039a4c2

Signed-off-by: naykudev <vnaykude@amazon.com>

github-actions · 2026-08-19T20:22:34Z

Persistent review updated to latest commit b4235a1

Signed-off-by: naykudev <vnaykude@amazon.com>

github-actions · 2026-08-20T16:39:27Z

Persistent review updated to latest commit 9aa0deb

github-actions · 2026-08-20T18:22:08Z

Persistent review updated to latest commit d6e14af

naykudev requested review from 0ctopus13prime, VijayanB, Vikasht34, heemin32, kotwanikunal, luyuncheng, martin-gaievski, naveentatikonda, navneet1v, shatejas and vamshin as code owners August 4, 2026 23:51

naykudev marked this pull request as draft August 5, 2026 18:33

naykudev marked this pull request as ready for review August 5, 2026 20:01

naykudev marked this pull request as draft August 6, 2026 04:02

naykudev marked this pull request as ready for review August 6, 2026 18:45

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

449826c

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

d9e7706

Vikasht34 reviewed Aug 12, 2026

View reviewed changes

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

6f161ff

Signed-off-by: naykudev <vnaykude@amazon.com>

Clarify dual exact-search passes and drop unused rescore score

0c91ea4

Signed-off-by: ved naykude <vnaykude@amazon.com>

Restore changelog entry after main merge

1b020c4

Signed-off-by: ved naykude <vnaykude@amazon.com>

naykudev requested a review from Vikasht34 August 13, 2026 02:11

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

0d7dd39

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

38a6c6b

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

d54959c

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

d53006c

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

039a4c2

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

b4235a1

Signed-off-by: naykudev <vnaykude@amazon.com>

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

9aa0deb

Signed-off-by: naykudev <vnaykude@amazon.com>

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

d6e14af

Uh oh!

Conversation

naykudev commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Why this approach (mirroring the native engine)

Related Issues

Check List

Uh oh!

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

Uh oh!

PR Reviewer Guide 🔍

(Review updated until commit d6e14af)

Uh oh!

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

Uh oh!

PR Code Suggestions ✨

Previous suggestions

Uh oh!

codecov Bot commented Aug 5, 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 Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

github-actions Bot commented Aug 10, 2026

Uh oh!

github-actions Bot commented Aug 11, 2026

Uh oh!

Vikasht34 Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

naykudev Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

Vikasht34 Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

naykudev Aug 18, 2026

Choose a reason for hiding this comment

Uh oh!

Vikasht34 left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 12, 2026

Uh oh!

github-actions Bot commented Aug 12, 2026

Uh oh!

github-actions Bot commented Aug 12, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 14, 2026

Uh oh!

github-actions Bot commented Aug 15, 2026

Uh oh!

github-actions Bot commented Aug 19, 2026

Uh oh!

github-actions Bot commented Aug 19, 2026

Uh oh!

github-actions Bot commented Aug 20, 2026

Uh oh!

github-actions Bot commented Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

naykudev commented Aug 4, 2026 •

edited

Loading

github-actions Bot commented Aug 4, 2026 •

edited

Loading

(Review updated until commit `d6e14af`)

github-actions Bot commented Aug 4, 2026 •

edited

Loading

codecov Bot commented Aug 5, 2026 •

edited

Loading