Enable optimistic search to memory optimized search. #2933

0ctopus13prime · 2025-10-08T05:41:06Z

Description

Enable optimistic search for memory optimized search and deprecate MultiLeafKnnCollector which has an early termination logic.

This PR has three big changes:

Now, when memory-optimized search is enabled, all queries use NativeEngineKnnVectorQuery.
KnnQuery, which only provides a ScorerSupplier and performs search within a single leaf segment (with the resulting Scorer being consumed by an external BulkScorer under the standard Lucene search flow). But optimistic search requires coordination across segments. It needs to run an initial (first-phase) search, then identify and revisit only the segments likely to contain promising results.
To support this coordinated two-phase process, NativeEngineKnnVectorQuery is a more suitable entry point than KnnQuery.
Backported Lucene components required for optimistic search, specifically:
2.1. ReentrantKnnCollectorManager
2.2. SeededMappedDISI
2.3. SeededTopDocsDISI

Related Issues

Resolves #[Issue number to be closed when this PR is merged]
#2924

Check List

[O] New functionality includes testing.
[O] New functionality has been documented.
[O] API changes companion pull request created.
[O] 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.

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

0ctopus13prime · 2025-10-08T22:20:25Z

@Vikasht34 @shatejas
All CI passed! Could you take a look at this when you have time?
Thank you

Vikasht34 · 2025-10-08T22:21:29Z

Will look into this PR :- Tomorrow Morning

0ctopus13prime · 2025-10-09T05:24:11Z

src/main/java/org/opensearch/knn/index/query/KNNWeight.java

+            );
        }

-        /*


This logic has been moved to approximateSearch

Vikasht34 · 2025-10-10T18:47:31Z

src/main/java/org/opensearch/knn/index/query/PerLeafResult.java

+     * An immutable, empty {@link BitSet} implementation used to represent
+     * the absence of filter bits without incurring null checks or allocations.
+     */
+    public static final BitSet MATCH_ALL_BIT_SET = new BitSet() {


Curious why we are not uysing Lucene's MathAllBits or MatcNoBits?

I could not use both, as it's sub class of Bits while we need BitSet in here. 😵‍💫
Since optimistic search will call approximateSearch twice, we need to keep BitSet for reusing.

Vikasht34 · 2025-10-10T18:49:00Z

src/main/java/org/opensearch/knn/index/query/PerLeafResult.java

+        }
+
+        @Override
+        public int length() {


Are we doing any iteration on Bitset , this gooona break if we are doing?

Only place using this one is when we're getting siblings in nested case.
In there, we don't do iteration.

shatejas · 2025-10-13T00:35:36Z

when memory-optimized search is enabled, all queries use NativeEngineKnnVectorQuery.

There are caveats to doing this,

This will completely bypass the slicing logic, consuming more CPU for concurrent segment search. The concurrency will use all threads available for search. Allowing a shard to use all cores of CPU impacts other parts of the system and is deviating from existing behavior for cases where rescoring is not used
This impacts the total hits count, moving it to the shard level will always show total hits as k, where as with segment level will add up each segment results in total count. This is a behavior change and can impact cases where users rely on total hit count
This now returns k results on shard level, while its not a huge concern, it can affect results of single shard - single segment case when k < size for non-rescoring cases

I think 1 is a concern which needs discussion, 2 is manageable with some extra logic to keep behavior consistent, 3 isn't a big deal

navneet1v · 2025-10-13T08:06:14Z

when memory-optimized search is enabled, all queries use NativeEngineKnnVectorQuery.

There are caveats to doing this,

This will completely bypass the slicing logic, consuming more CPU for concurrent segment search. The concurrency will use all threads available for search. Allowing a shard to use all cores of CPU impacts other parts of the system and is deviating from existing behavior for cases where rescoring is not used

This impacts the total hits count, moving it to the shard level will always show total hits as k, where as with segment level will add up each segment results in total count. This is a behavior change and can impact cases where users rely on total hit count

This now returns k results on shard level, while its not a huge concern, it can affect results of single shard - single segment case when k < size for non-rescoring cases

I think 1 is a concern which needs discussion, 2 is manageable with some extra logic to keep behavior consistent, 3 isn't a big deal

@shatejas from a user perspective all of this is a breaking change if we are making Lucene on Faiss default. So this needs to be documented in the docs to clearly callout the behavior and how to mitigate this. Along with this, we should ensure that older indices are still on the same non memory optimized based search so that upgrades are seamless.

We have already seen GH issues in part where changes in range of cosine scores lead to issues with users. #2561

...java/org/opensearch/knn/index/query/memoryoptsearch/optimistic/Optimistic2ndSearchUtils.java

src/main/java/org/opensearch/lucene/ReentrantKnnCollectorManager.java

navneet1v · 2025-10-13T08:12:20Z

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java

+    @Setter
+    private KnnCollectorManager optimistic2ndKnnCollectorManager;
+
+    public static class OptimisticKnnCollectorManager implements KnnCollectorManager {


can we move this to a separate file?

sure, will update in the next rev.

navneet1v · 2025-10-13T08:14:04Z

src/main/java/org/opensearch/knn/index/query/PerLeafResult.java

-    public PerLeafResult(final Bits filterBits, final TopDocs result) {
-        this.filterBits = filterBits == null ? new Bits.MatchAllBits(0) : filterBits;
+    // Indicates whether this result was produced via exact or approximate search.
+    private final SearchMode searchMode;


what is the use of this parameter?

Optimistic search would do deep dive HNSW search with acquired top k results as seeds. And this will not be needed if the results acquired via exact search. Hence having search mode here, and let it bypass the second search if possible.

The cases when results will be acquired from exact search is the filters case right?

Yes! Whenever running exact search for whatever reason, it will set the mode as EXACT_SEARCH.
If that's the case, then we should not run 2nd search as it's pointless.

navneet1v · 2025-10-13T08:14:14Z

src/main/java/org/opensearch/knn/index/query/PerLeafResult.java

+     * An immutable, empty {@link BitSet} implementation used to represent
+     * the absence of filter bits without incurring null checks or allocations.
+     */
+    public static final BitSet MATCH_ALL_BIT_SET = new BitSet() {


navneet1v · 2025-10-13T08:14:37Z

src/main/java/org/opensearch/knn/index/query/KNNWeight.java

+     * @return a {@link TopDocs} object containing the top {@code k} approximate search results
+     * @throws IOException if an error occurs while reading index data or accessing vector fields
+     */
+    public TopDocs approximateSearch(


why we are making this function public?

We need this particular function in optimistic second search. Otherwise, if using searchLeaf, then we will end up building filter bitset twice.

navneet1v · 2025-10-13T08:17:31Z

We are brining in a lot of code from Lucene. Please mention the source of the code for better maintainability.

One way I would think is to move the class to org.opensearch.lucene to ensure that we know these classes are merely copie/inspired from Lucene.

Vikasht34

Looks Good !! Clean Code and Very Concise !! Thanks

navneet1v · 2025-10-15T07:43:52Z

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

+    /**
+     * A special flag used for testing purposes that forces execution of the second (exact) search
+     * in optimistic search mode, regardless of the results returned by the first approximate search.
+     * <p>
+     * This flag should never be enabled in production; it is intended for testing and debugging only.
+     */
+    private static final boolean FORCE_REENTER_TESTING;


lets remove this in next iteration of PR.

This is needed for testing the second reentrance in optimistic search.
It's bit tricky to make the data to ensure there's at least one segment whose min topk score is greater than the minimum score in the merged top-k results

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

navneet1v

Overall looks good to me.

I am not seeing how this new memory optimized search is getting integrated with the explain api.

Also let put the destination branch as main branch.

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

src/main/java/org/opensearch/lucene/OptimisticKnnCollectorManager.java

navneet1v · 2025-10-15T08:00:42Z

src/main/java/org/opensearch/knn/index/query/PerLeafResult.java

-    public PerLeafResult(final Bits filterBits, final TopDocs result) {
-        this.filterBits = filterBits == null ? new Bits.MatchAllBits(0) : filterBits;
+    // Indicates whether this result was produced via exact or approximate search.
+    private final SearchMode searchMode;


The cases when results will be acquired from exact search is the filters case right?

navneet1v · 2025-10-15T08:04:27Z

build.gradle

+    // Forcing optimistic search for testing
+    systemProperty 'mem_opt_srch.force_reenter', 'true'


do we need this ?

If you want to run search for memory optimized cases then lets create another gradle task and also a new CI that runs that task.

We need this to force it to run 2nd search in optimistic.
Since the 2nd search will kick off only if there's segment whose min score > the min score in merged results, it was tricky for me to make the data.

build.gradle

The base branch was changed.

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java

shatejas

Looks good overall, some minor comments

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

shatejas · 2025-10-22T17:30:34Z

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

+            final PerLeafResult perLeafResult = perLeafResults.get(i);
+            final TopDocs perLeaf = perLeafResults.get(i).getResult();
+            if (perLeaf.scoreDocs.length > 0 && perLeafResult.getSearchMode() == PerLeafResult.SearchMode.APPROXIMATE_SEARCH) {
+                if (FORCE_REENTER_TESTING || perLeaf.scoreDocs[perLeaf.scoreDocs.length - 1].score >= minTopKScore) {


nit: Its not ideal to have a testing flag in final code, can't we create a test case where we mock scores?

Since whether if it kicks off the second search will be entirely depending on data. Which is hard to be made, and this was the only solution I found so far to always enforce it to enter 2nd search in optimistic.

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

navneet1v · 2025-10-22T20:48:42Z

@0ctopus13prime any reason why MacOS builds are failing?

0ctopus13prime · 2025-10-22T21:45:47Z

@navneet1v
It's timed out

2025-10-22T19:33:55.9624590Z REPRODUCE WITH: ./gradlew ':integTest' --tests 'org.opensearch.knn.index.FaissIT.testEndToEnd_whenDoRadiusSearch_whenDistanceThreshold_whenMethodIsHNSWFlat_thenSucceed' -Dtests.seed=FD2D022924B4F513 -Dtests.security.manager=false -Dtests.locale=jmc-Latn-TZ -Dtests.timezone=Etc/Zulu -Druntime.java=21
2025-10-22T19:33:55.9640570Z FaissIT > classMethod FAILED
2025-10-22T19:33:55.9643660Z 
2025-10-22T19:33:55.9644380Z     java.lang.Exception: Suite timeout exceeded (>= 1200000 msec).

…pensearch-project#2904) Signed-off-by: Dooyong Kim <[email protected]>

Signed-off-by: Dooyong Kim <[email protected]>

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java

Signed-off-by: Dooyong Kim <[email protected]>

0ctopus13prime self-assigned this Oct 8, 2025

0ctopus13prime requested review from VijayanB, Vikasht34, heemin32, jmazanec15, junqiu-lei, luyuncheng, martin-gaievski, naveentatikonda, navneet1v, ryanbogan, shatejas and vamshin as code owners October 8, 2025 05:41

0ctopus13prime added the skip-changelog label Oct 8, 2025

0ctopus13prime commented Oct 8, 2025

View reviewed changes

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java Show resolved Hide resolved

0ctopus13prime force-pushed the optimistic-srch branch 2 times, most recently from 5b006af to 1e09bf0 Compare October 9, 2025 05:06

0ctopus13prime commented Oct 9, 2025

View reviewed changes

Vikasht34 reviewed Oct 11, 2025

View reviewed changes

navneet1v reviewed Oct 13, 2025

View reviewed changes

Vikasht34 previously approved these changes Oct 14, 2025

View reviewed changes

0ctopus13prime force-pushed the optimistic-srch branch 3 times, most recently from e625f04 to 0cc0b54 Compare October 15, 2025 07:00

navneet1v reviewed Oct 15, 2025

View reviewed changes

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java Outdated Show resolved Hide resolved

navneet1v reviewed Oct 15, 2025

View reviewed changes

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java Outdated Show resolved Hide resolved

navneet1v reviewed Oct 15, 2025

View reviewed changes

build.gradle Show resolved Hide resolved

0ctopus13prime changed the base branch from feature/fp16-faiss-bulk to main October 15, 2025 17:22

0ctopus13prime force-pushed the optimistic-srch branch 3 times, most recently from ea03c31 to 21aa301 Compare October 16, 2025 03:05

shatejas reviewed Oct 16, 2025

View reviewed changes

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java Outdated Show resolved Hide resolved

src/main/java/org/opensearch/knn/index/query/memoryoptsearch/MemoryOptimizedKNNWeight.java Outdated Show resolved Hide resolved

0ctopus13prime force-pushed the optimistic-srch branch 4 times, most recently from 23325ed to 0f44ba7 Compare October 21, 2025 18:24

0ctopus13prime added the backport 3.3 label Oct 21, 2025

shatejas previously approved these changes Oct 22, 2025

View reviewed changes

0ctopus13prime dismissed shatejas’s stale review via 4410b34 October 22, 2025 18:55

0ctopus13prime force-pushed the optimistic-srch branch from 0f44ba7 to 4410b34 Compare October 22, 2025 18:55

0ctopus13prime force-pushed the optimistic-srch branch from 4410b34 to 0800030 Compare October 22, 2025 21:50

0ctopus13prime added 2 commits October 23, 2025 13:28

Added MMapByteVectorValues for FP16 native scoring in LuceneOnFaiss. (o…

f86d8f2

…pensearch-project#2904) Signed-off-by: Dooyong Kim <[email protected]>

Enable optimistic search for LuceneOnFaiss.

6197e6a

Signed-off-by: Dooyong Kim <[email protected]>

0ctopus13prime force-pushed the optimistic-srch branch from 0800030 to a43fa59 Compare October 23, 2025 20:28

shatejas reviewed Oct 23, 2025

View reviewed changes

src/main/java/org/opensearch/knn/index/query/nativelib/NativeEngineKnnVectorQuery.java Show resolved Hide resolved

Added debug logging, measure execution times of 2nd search.

347347f

Signed-off-by: Dooyong Kim <[email protected]>

0ctopus13prime force-pushed the optimistic-srch branch from a43fa59 to 347347f Compare October 23, 2025 21:57

+                          );
                       }
-                      /*

		// Forcing optimistic search for testing
		systemProperty 'mem_opt_srch.force_reenter', 'true'

Enable optimistic search to memory optimized search. #2933

Are you sure you want to change the base?

Enable optimistic search to memory optimized search. #2933

Uh oh!

Conversation

0ctopus13prime commented Oct 8, 2025

Description

Related Issues

Check List

Uh oh!

Uh oh!

0ctopus13prime commented Oct 8, 2025

Uh oh!

Vikasht34 commented Oct 8, 2025

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!

shatejas commented Oct 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

navneet1v commented Oct 13, 2025

Uh oh!

Uh oh!

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!

navneet1v commented Oct 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Vikasht34 left a comment

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!

Uh oh!

Uh oh!

navneet1v left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

Uh oh!

Uh oh!

shatejas commented Oct 13, 2025 •

edited

Loading

navneet1v commented Oct 13, 2025 •

edited

Loading

navneet1v left a comment •

edited

Loading

0ctopus13prime Oct 22, 2025 •

edited

Loading