Check to see if Lucene's search budget has exhausted when deciding to exact search

[MrFlap]

Description

Lucene engine has a check such that it falls back to exact search if we hit a visit limit before finding qualified candidates. This change takes the conditional from Lucene and applies the same logic to Memory Optimized Search.

After this change we will see a more aggressive exact search strategy in MOS when filtered docs are sparse in the graph. We should see a recall gain in low filtering percentages (~1-5%).

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 29d324e)

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 Null Pointer Risk exhaustedSearch can be null if knnExplanation.getExhaustedSearch(context.id()) returns null (e.g., when the leaf was not previously recorded). Line 212 checks exhaustedSearch in a boolean context without a null check, which will throw a NullPointerException if the value is null. final Boolean exhaustedSearch = knnExplanation.getExhaustedSearch(context.id()); if (annResult != null && annResult == 0 && isMissingNativeEngineFiles(context)) { sb.append(KNNConstants.EXACT_SEARCH).append(" since no native engine files are available"); } if (annResult != null && isFilteredExactSearchRequireAfterANNSearch(cardinality, exhaustedSearch, annResult)) { boolean isExactSearchDisabled = KNNSettings.isKnnIndexFaissEfficientFilterExactSearchDisabled(knnQuery.getIndexName()); if (isExactSearchDisabled) { sb.append(KNNConstants.ANN_SEARCH) .append(", it is not falling back to exact search after ") .append(KNNConstants.ANN_SEARCH) .append(" search since exact search is disabled,"); } else if (exhaustedSearch) { sb.append(KNNConstants.EXACT_SEARCH).append(" since lucene vector search has exhausted number of steps."); Null Pointer Risk At line 689, isFilteredExactSearchRequireAfterANNSearch is called with annSearchBudgetExhausted which can be null if knnExplanation.getExhaustedSearch(context.id()) returns null. The method uses this boolean parameter in a return statement without null checking, risking a NullPointerException when unboxing. if (isFilteredExactSearchRequireAfterANNSearch(filterIdsCount, annSearchBudgetExhausted, annResultCount)) {

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 29d324e

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Compute budget exhaustion unconditionally The annSearchBudgetExhausted flag is only set when knnQuery.isExplain() is true. When explain is disabled, this variable will always be false, potentially skipping necessary exact search fallback. Store this flag unconditionally to ensure correct search behavior. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [351-361] +int annResultsCount = topDocs.scoreDocs.length; +boolean annSearchBudgetExhausted = topDocs.totalHits.relation() != TotalHits.Relation.EQUAL_TO; + +if (knnQuery.isExplain()) { + knnExplanation.addLeafResult(context.id(), annResultsCount); + knnExplanation.addExhaustedSearch(context.id(), annSearchBudgetExhausted); +} + if (isExactSearchRequire(context, filterCardinality, annSearchBudgetExhausted, annResultsCount)) { - final BitSetIterator docs = filterWeight != null ? new BitSetIterator(filterBitSet, filterCardinality) : null; - final TopDocs result = doExactSearch(context, docs, filterCardinality, k); Suggestion importance[1-10]: 10 __ Why: This is a critical correctness issue. The annSearchBudgetExhausted flag is only computed and stored when knnQuery.isExplain() is true (lines 354-357), but it's used unconditionally in isExactSearchRequire at line 361. When explain is disabled, the variable will always be false, causing the system to incorrectly skip exact search fallback even when the search budget was exhausted. This directly impacts search correctness.	High
	Handle null exhaustedSearch value Add a null check for exhaustedSearch before passing it to isFilteredExactSearchRequireAfterANNSearch. If the value is null (e.g., when explain is disabled), the method will receive null and may cause unexpected behavior or NPE in boolean operations. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [201-205] final Boolean exhaustedSearch = knnExplanation.getExhaustedSearch(context.id()); if (annResult != null && annResult == 0 && isMissingNativeEngineFiles(context)) { sb.append(KNNConstants.EXACT_SEARCH).append(" since no native engine files are available"); } -if (annResult != null && isFilteredExactSearchRequireAfterANNSearch(cardinality, exhaustedSearch, annResult)) { +if (annResult != null && isFilteredExactSearchRequireAfterANNSearch(cardinality, exhaustedSearch != null && exhaustedSearch, annResult)) { Suggestion importance[1-10]: 9 __ Why: The exhaustedSearch variable can be null when explain is disabled, but it's passed directly to isFilteredExactSearchRequireAfterANNSearch which expects a boolean primitive. This will cause an NPE when the method tries to use it in boolean operations. The suggestion correctly identifies this critical bug and provides the proper null-safe fix.	High

---

Previous suggestions

Suggestions up to commit fb4a05c

Category	Suggestion	Impact
Possible issue	Track exhaustedSearch regardless of explain mode The exhaustedSearch flag is only populated when isExplain() is true, but it's used unconditionally in isExactSearchRequire. This creates a critical bug where exact search decisions depend on whether explain mode is enabled, leading to inconsistent search behavior. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [354-361] -if (knnQuery.isExplain()) { - knnExplanation.addLeafResult(context.id(), annResultsCount); - knnExplanation.addExhaustedSearch(context.id(), exhaustedSearch); -} -... +knnExplanation.addLeafResult(context.id(), annResultsCount); +knnExplanation.addExhaustedSearch(context.id(), exhaustedSearch); + if (isExactSearchRequire(context, filterCardinality, exhaustedSearch, annResultsCount)) { Suggestion importance[1-10]: 10 __ Why: This identifies a critical bug where exhaustedSearch is only populated when knnQuery.isExplain() is true (lines 354-357), but is used unconditionally in isExactSearchRequire at line 361. This causes search behavior to differ based on whether explain mode is enabled, which is a serious correctness issue affecting core search functionality.	High

Suggestions up to commit a71fc1f

Category	Suggestion	Impact
Possible issue	Add null safety check Add null check for exhaustedSearch before using it in the condition. If getExhaustedSearch() returns null, this will cause a NullPointerException when unboxing the Boolean. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [212-214] -} else if(exhaustedSearch) { +} else if(Boolean.TRUE.equals(exhaustedSearch)) { sb.append(KNNConstants.EXACT_SEARCH) .append(" since lucene vector search has exhausted number of steps."); Suggestion importance[1-10]: 8 __ Why: The exhaustedSearch variable is a Boolean object that can be null (from getExhaustedSearch()). Using it directly in the condition without null checking will cause a NullPointerException when unboxing. This is a critical bug that needs to be fixed.	Medium
	Handle potential null value properly The exhaustedSearch parameter is a primitive boolean but can receive null from getExhaustedSearch(). This will cause a NullPointerException when unboxing. Change the parameter type to Boolean and add null handling. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [715-722] -private boolean isFilteredExactSearchRequireAfterANNSearch(final int filterIdsCount, final boolean exhaustedSearch, final int annResultCount) { +private boolean isFilteredExactSearchRequireAfterANNSearch(final int filterIdsCount, final Boolean exhaustedSearch, final int annResultCount) { if(filterWeight == null) { return false; } if(filterIdsCount < knnQuery.getK()) { return false; } - return knnQuery.getK() > annResultCount || exhaustedSearch; + return knnQuery.getK() > annResultCount || Boolean.TRUE.equals(exhaustedSearch); Suggestion importance[1-10]: 8 __ Why: The method signature uses primitive boolean for exhaustedSearch, but it receives a Boolean object from getExhaustedSearch() which can be null. This will cause a NullPointerException during unboxing. The suggestion correctly identifies this issue and proposes using Boolean with null-safe comparison.	Medium
General	Correct exhausted search detection logic The logic for determining exhausted search may be incorrect. According to Lucene documentation, GREATER_THAN_OR_EQUAL_TO indicates the search was exhausted or hit a limit. Consider checking if the relation equals GREATER_THAN_OR_EQUAL_TO instead. src/main/java/org/opensearch/knn/index/query/KNNWeight.java [353] -boolean exhaustedSearch = topDocs.totalHits.relation() != TotalHits.Relation.EQUAL_TO; +boolean exhaustedSearch = topDocs.totalHits.relation() == TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO; Suggestion importance[1-10]: 7 __ Why: The current logic uses != to check if the relation is not EQUAL_TO, but according to Lucene semantics, GREATER_THAN_OR_EQUAL_TO specifically indicates the search was exhausted. The suggested change makes the logic more explicit and correct, though the current implementation may work in practice if there are only two possible values.	Medium

[navneet1v]

@MrFlap what is the impact of this change? can you please add some details on that part

[MrFlap]

@navneet1v the impact of the change is that we will do exact search more aggressively on queries where filtered documents are sparse. Adding this to the PR description. Keeping it as a draft since I still need to write tests and do benchmarking on 1% case to make sure this actually closes the gap.

[github-actions]

Persistent review updated to latest commit fb4a05c

[MrFlap]

PR Code Suggestions ✨

Latest suggestions up to fb4a05c Explore these optional code suggestions:
Category **Suggestion ** Impact
Possible issue
Track exhaustedSearch regardless of explain mode

The exhaustedSearch flag is only populated when isExplain() is true, but it's used unconditionally in isExactSearchRequire. This creates a critical bug where exact search decisions depend on whether explain mode is enabled, leading to inconsistent search behavior.

src/main/java/org/opensearch/knn/index/query/KNNWeight.java [354-361]

-if (knnQuery.isExplain()) {
-    knnExplanation.addLeafResult(context.id(), annResultsCount);
-    knnExplanation.addExhaustedSearch(context.id(), exhaustedSearch);
-}
-...
+knnExplanation.addLeafResult(context.id(), annResultsCount);
+knnExplanation.addExhaustedSearch(context.id(), exhaustedSearch);
+
 if (isExactSearchRequire(context, filterCardinality, exhaustedSearch, annResultsCount)) {

Suggestion importance[1-10]: 10

High

Previous suggestions

Suggestions up to commit a71fc1f

This is just not true. We only insert to knnExplanation for the purpose of explain. That function does not alter search behavior.

[codecov]

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.75%. Comparing base (ffae669) to head (29d324e).

Files with missing lines	Patch %	Lines
...java/org/opensearch/knn/index/query/KNNWeight.java	87.50%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##               main    #3354   +/-   ##
=========================================
  Coverage     83.75%   83.75%           
- Complexity     4354     4359    +5     
=========================================
  Files           453      453           
  Lines         15759    15773   +14     
  Branches       2052     2056    +4     
=========================================
+ Hits          13199    13211   +12     
- Misses         1774     1776    +2     
  Partials        786      786           

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

Persistent review updated to latest commit 29d324e

[navneet1v]

Code looks good to me.