vector-store: diskann basic queries - #535
Conversation
There was a problem hiding this comment.
Pull request overview
Adds baseline DiskANN support for vector insertion, removal, ANN search, and counting.
Changes:
- Implements a single-task DiskANN actor and ID mapping.
- Adds search validation, metric handling, and count support.
- Adds unit and integration tests.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
crates/vector-store/src/vs_index/diskann.rs |
Implements core DiskANN operations. |
crates/vector-store/tests/integration/diskann.rs |
Adds DiskANN integration coverage. |
crates/vector-store/tests/integration/main.rs |
Registers the DiskANN tests. |
| if let Err(err) = self | ||
| .index | ||
| .insert(&FullPrecision, &DefaultContext, &id, embedding.as_slice()) | ||
| .await |
There was a problem hiding this comment.
MAX_POINTS in general is a problem, should it be configurable, or a constant or something else entirely? It controls the maximum number of points that a DiskANNIndex can keep in memory. Diskann explicitly returns an Err on overflow. Unlike usearch, the memory can't be resized, that's why it's good to decide on a correct approach earlier rather than later.
| VsIndex::RemoveVector { primary_id, .. } => { | ||
| state.remove_vector(primary_id).await; | ||
| } |
There was a problem hiding this comment.
Should I write additional tests or keep them consistent with usearch?
1c4caaf to
6ae9367
Compare
1c1fd23 to
617ff5f
Compare
ewienik
left a comment
ewienik
left a comment
There was a problem hiding this comment.
the current solution uses a ExternalId <-> Epoch mapping
The epoch is the 16 oldest bits of the PrimaryId
because diskann requires u32 external ids we map the 32 LSBs
to the epoch number and resolve this number when to recreate the PrimaryId
This doesnt fix the fact that only 2^32-1 ids can be distributed
but that's the limitation of diskann
the solution currently operates on a sigle tokio worker
I understand that limitation for 32bits is from memory DiskAnn provider. Using PrimaryId::idx() as an ExternalId seems a bad idea - Epoch is a part of the PrimaryId. In the same operation wee can remove old PrimaryId, increase only Epoch and insert a epoch-changed-only PrimaryId - as we are proceeding concurrently we should take care of the Epoch in the ExternalId.
Consider using a mapping between PrimaryId and internal 32 bit counter as ExternalId. After removing some PrimaryId (and associated ExternalId) you can store unused ExternalId in a cache.
Consider using BTreeMap instead of HashMap -> with a large number of vectors realloc of HashMap takes significant time.
617ff5f to
f1a0dda
Compare
|
Changes:
|
f1a0dda to
3112dce
Compare
|
microsoft/DiskANN#1206 was just merged into DiskANN. It not easy to really understand what's going on but this might allow us rely on diskANN internal id mapping instead of our own. Apart from that it also provides improvements towards safety of concurrent operations. It has some open follow ups so it might be a bit early to adopt this right now, but it's worth keeping an eye on this. Relevant docs: //! An in-memory provider for the DiskANN graph index.
//!
//! This type supports the following:
//!
//! * Arbitrary external IDs for store data (provided they satisfy [`Id`].
//! * Support for concurrent insertions, deletions, and searches.
//! * Specialized implementations of [`glue::SearchAccessor::expand_beam`] enabling full
//! inlining of distance kernels.
//!
//! Known areas for future work:
//!
//! * Insert and delete protection: The [`DiskANNIndex`](diskann::graph::DiskANNIndex) doesn't
//! support ergonomic insert or delete guards to protect slots during insert or delete
//! operations. This leaves open a situation where an item can be inserted and during
//! the insertion algorithm, it is deleted, and then re-inserted.
//!
//! This can cause some issue within the main indexing algorithms which assume the inserted
//! ID is present but requires upstream changes to properly fix.
//!
//! * Failed insert rollback: again, this needs some upstream changes to full support.
//!
//! * Quantization + reranking: The current version of this index targets just a single
//! data-store and is planned to be addressed in the near future.
//!
//! * Lack of save/load support: The index is currently ephemeral, but there are plans to
//! address this gap.
|
3112dce to
3eb537b
Compare
the current solution uses a PrimaryId <-> u32 ExternalId DiskANN unfortunately supports u32 ids with DefaultProvider the solution currently operates on a sigle tokio worker
index.search() outputs a vector of neighbours that need distances to be added to them, hence the State::search() helper State::ann is similar to usearch thanks to this, with the difference being no partitions in diskann as of yet. added space_type, l_default and beam_width to DiskannParams added to unit test
these tests are the working subset of usearch tests basically everything but filtered ann(since it's beyond scope) and http server responsiveness since it uses usearch simulator
3eb537b to
2886672
Compare
I think this is a good catch. I'd rather go with this new code as it seems it allows to use our PrimaryId instead of problematic mapping with u32. Do you think it is possible?
We've similar situation with usearch - there is implemented specific Permit structure to support concurrent operations. Additionally we guarantee that every delete/insert operation will use different PartitionId (Epoch variable) - so we should be safe. I think removing our mapping PrimaryId <-> u32 has a lot of benefits and I don't see issues with a new inmem provider. Can you make a try to update dependencies? |
I can either do this right now, or we can wait until crate version 0.56.0 is released, but ETA is unknown. |
So better seems to do it right now to avoid problems with id mapping. |
|
microsoft/DiskANN#1303 waiting on this, because clippy fails |
You can temporarily use your github PR as a patched reference in Cargo.toml |
This PR focuses on a simple one-by-one actor loop implementation of:
The aim is to establish a correctly behaving baseline to build upon.
Things that are currently beyond the scope of this PR:
Ref: VECTOR-714