feat: added lazy loading. - #39
Open
Giulio2002 wants to merge 8 commits into
Open
Conversation
Giulio2002
marked this pull request as draft
December 11, 2025 15:24
Giulio2002
marked this pull request as ready for review
December 11, 2025 19:21
npiesco
added a commit
to npiesco/hnsw
that referenced
this pull request
Aug 17, 2026
Two gaps, both closed on both indexes. ## Pluggable feature storage Upstream rust-cv#39 proposed keeping features outside the heap so a large index does not have to hold its whole corpus in memory. The idea is right; that implementation is not adoptable here. - Its DiskFeatureStore::get returns a reference into a single thread_local! UnsafeCell scratch buffer. Two live references from one store alias the same address. That is UB, and it silently turns distance(a, b) into distance(x, x). - Its MmapFeatureStore::get has no bounds check, so index == capacity dereferences one past the end of the mapping. - Its own tests build 8 points at M=12/M0=24, which never saturates a neighbor list and so never reaches the pruning branch at all - the one place the aliasing bug is fatal. - Its add_neighbor hunk is written against nearest-M truncation. Accepting it would revert d7b9d27 and restore the cluster-isolation bug. - Its copy_from_slice fix is already here, in both indexes, via 8d37b6b. So the trait is written here rather than taken. FeatureStore<T> is a plain borrow-based trait: get_feature returns &T, push_feature appends, feature_count reports the size. Hnsw and HnswRuntime both gain a trailing storage parameter defaulting to Vec<T>, so every existing signature is unchanged. The borrow-stability contract is load-bearing, not decorative. The diversity heuristic added in d7b9d27 holds three feature references live at once while pruning a saturated list - the target, the candidate under consideration, and each already-kept neighbor it is compared against. A store that decoded into a shared buffer would alias all three onto one address, every comparison would collapse toward distance(x, x), the heuristic would degenerate back into nearest-M truncation, and dense clusters would close themselves off again. That is precisely the failure d7b9d27 fixed, so the contract is stated explicitly on the trait and tested rather than assumed. Suitable backings: Vec<T>, a fixed mmap region, an arena, an append-only cache. Unsuitable: reused decode buffers, evicting caches, values synthesized per read. The methods are named get_feature / push_feature / feature_count rather than get / push / len for a concrete reason: a trait method named `get` on Vec<T> shadows slice::get at every call site where the trait is in scope, because trait methods on the receiver type are considered before inherent methods reached through Deref. That broke Vec<bool> indexing in the soft-delete code below. tests/feature_store.rs drives a real mmap-backed store over a real file against a 160-point four-cluster corpus, which does saturate M0=24 lists and does reach the pruning branch: - mmap and Vec storage return identical ranked results, const index - mmap and Vec storage return identical ranked results, runtime index - a dense 37-point far cluster still cannot hide the three points beside the origin when the features live in an mmap - three concurrently live handles keep distinct stable addresses, and a real distance between two of them stays non-zero That last one is the direct test of the contract: against a scratch-buffer store it fails. ## Soft-delete on the runtime index HnswRuntime had no delete surface at all, so the two indexes were not actually mirrors. Ported live_count, is_deleted, mark_delete and entry from the const index, along with both search-path changes: the zero-layer tombstone branch that traverses a deleted node without admitting it to the result heap, and the retain that drops tombstones from the final set. Traversing through tombstones rather than around them is what keeps the graph connected after deletion, and skipping only the result heap is what stops a deleted node from consuming result budget. tests/incremental_delete.rs now covers both indexes, including runtime_and_const_agree_under_identical_deletions, which applies the same deletion set to the same corpus in both and asserts the ranked lists match. ## Verification 21 tests pass with --all-features: clustered_recall 5, feature_store 4, incremental_delete 3, random 2, runtime_parity 1, serde 1, simple 3, simple_discrete 2. runtime_parity and serde passing matter most - the first shows the two indexes did not drift, the second shows the added serde bounds and the skipped PhantomData leave the snapshot format unchanged. cargo clippy --all-targets -- -D warnings is clean with default features and with --all-features, and cargo fmt --check is clean. That required fixing lints that predate this change in hnsw_const.rs, the examples and the benches; no warning was suppressed.
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
1 participant
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.
Lazy loading support
First of all, most of the LOC are just to integrate it with the examples so that we can easily test regressions and if there is any memory improvement at all. the actual library diff is very tiny in comparison. I apologise in advance and we can split it into multiple PRs. We can also just revert the examples and make the feature small
Summary
Adds
FeatureStore<T>trait who abstracts feature vector storage, enabling lazy loading from disk, mmap, or custom backends/whatever you want. The default implementation usesVec<T>for backward compatibility. This PR also addsnew_with_storageandnew_with_storage_and_paramsconstructors, updates the recall examples with mmap-based storage demonstrations, and fixes a bug incopy_from_slicewhere the destination slice wasn't properly sized.Motivation
Previously, all feature vectors were stored in a
Vec<T>in memory. For large datasets, this can exceed available RAM. TheFeatureStoretrait allows users to provide custom storage backends that could:note: that all api changes are backward compatible (i think)
There was also a small bug which I also uncovered while testing the recall examples:
In
src/hnsw/hnsw_const.rs, thenearestmethod had incorrectcopy_from_slicecalls that would panic when the destination buffer was larger than the number of results found:This bug would trigger when Searching with an amount of neighbors requested but with fewer results available (not sure if this was a bug before but I checked and seems everything work separately).
Benchmarks - Recall
This stayed identical, I am going to post the gnuplots for completeness. below are the
Vecbased plotsDiscrete:
Mmap plots:
Benchmarks - Memory
Tested with 1 million 128-dimensional f32 vectors (512 bytes per feature):
To prove the usefulness of the lazy loading, I benchmarked the results on the RAM consumption. The disk-based storage reduces memory usage by ~83%. As a matter of fact, the disk-based backend, only the graph structure remains in memory, bringing total usage down to ~216 MB.
These results are expected to scale linearly. a 100MN vector index cannot be stored on a laptop if you keep it in memory as it would be 130 GB, however, with lazy loading, you can store on 32 GB machine >100MN and if you had access to 130 GB, you could store around 0.5 Billion vectors.