Reuse visited lists across batched Python queries - #676
Open
ShreyPandit wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR reduces visited-list pool locking in Python HNSW batch queries by leasing one
VisitedListper activeParallelForworker instead of acquiring and returning one for every query row.For a batch with
Qrows andWactive workers, the number of mutex-protected pool operations falls from roughly2Qto at most2 min(Q, W). Each row still receives a fresh visited generation, and the HNSW traversal anddistance calculations are unchanged.
Implementation
The patch adds a non-copyable RAII
VisitedListLease. Pool acquisition performs the first existing generationreset,
next()performs one reset before every subsequent row, and destruction returns the exact list to theoriginating pool on both normal and exceptional exit.
The level-zero search implementation is factored so it can consume caller-owned visited scratch. The existing C++
searchKnnNoExceptionspath retains its normal acquire/release behavior. Only PythonIndex.knn_queryuses the newentry point, with one lazily acquired lease per worker. Both normalized and non-normalized query paths use the same
lease protocol.
The patch does not change the graph, distance arithmetic, result ordering, index format, construction path,
BFIndex, or public Python signatures. The number of simultaneously live visited lists remains bounded by thenumber of active workers.
Performance
The A/B benchmark used 30,000 deterministic standard-normal float32 base vectors, 2,000 held-out queries,
dimension 64,
M=24,ef_construction=120,ef_search=48,k=10, and 16 query workers. For each of three buildseeds, baseline and candidate loaded the same baseline-produced serialized index. Each treatment ran in a fresh
process with two warmups and 31 retained full-batch timings.
The median paired gain is +4.73%. The mean paired ratio is
1.0486x, with a two-sided 95% paired-t interval of[1.0414x, 1.0558x]over the three build seeds. Peak RSS ratios were 0.9960, 0.9998, and 0.9966, so no memoryincrease was observed.
A separate two-round matrix measured one- and two-row calls within 0.6% of baseline and found improvements from
eight rows upward in the tested configurations. The performance claim is limited to batched Python queries and the
tested environment; it is not a universal throughput claim.
Correctness and tests
test_updatesandtest_updates update: passed in both builds.BFIndextest: passed; all five Python examples passed.BruteforceSearchmisaligned-load finding reproduces unchanged on theunmodified base. The new focused test passes under the combined sanitizers.
The added tests cover generation wrap, exception return, l2/ip/cosine, normalized queries, filters, deletions,
repeated use, concurrent batches, and the serial/threaded scheduling boundary.
Compatibility
The implementation is C++11-compatible and preserves the current exception-enabled and no-exception APIs. Index
serialization is unchanged. This PR targets
developatdf5814a61d6ac129fff6e5c27aa3ab99474aefac.AI assistance
OpenAI Codex
gpt-5.6-solwithxhighreasoning produced and screened the initial candidate. OpenAI Codexgpt-5.6-solwithmaxreasoning was used to port, review, and validate the final change.I reviewed and understand every submitted change.