diff --git a/Cargo.lock b/Cargo.lock index 5270a0cfe9..68dabeac21 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -604,7 +604,7 @@ dependencies = [ [[package]] name = "diskann-garnet" -version = "6.0.0" +version = "6.0.1" dependencies = [ "bytemuck", "crossbeam", @@ -615,6 +615,7 @@ dependencies = [ "diskann-utils", "diskann-vector", "foldhash", + "parking_lot", "rand", "strum", "thiserror 2.0.21", diff --git a/diskann-garnet/Cargo.toml b/diskann-garnet/Cargo.toml index ef05cbdd65..43f9087b2f 100644 --- a/diskann-garnet/Cargo.toml +++ b/diskann-garnet/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "diskann-garnet" -version = "6.0.0" +version = "6.0.1" edition = "2024" authors.workspace = true license.workspace = true @@ -22,6 +22,7 @@ diskann-providers.workspace = true diskann-utils.workspace = true diskann-vector.workspace = true foldhash = "0.2.0" +parking_lot = "0.12.5" rand.workspace = true strum = { version = "0.28.0", features = ["derive"] } thiserror.workspace = true diff --git a/diskann-garnet/README.md b/diskann-garnet/README.md index 7081b9b2f7..6acb41f39e 100644 --- a/diskann-garnet/README.md +++ b/diskann-garnet/README.md @@ -37,8 +37,27 @@ Redis always uses cosine distance, but many vector data sets use other metrics. - `XBIN_U8`: binary quantization of 8-bit unsigned integer (using DiskANN's spherical quantizer based on RaBitQ) -Currently there is a limit of `2^32 - 1` vectors in a single instance due to -internal IDs being `u32`. This restriction will be lifted in the future. +Currently there is a limit of `2^32 - 2` user vectors in a single instance. Internal +ID 0 is reserved for the start point, and `u32::MAX` marks ID allocation exhaustion. + +### Bulk Import + +Existing DiskANN graphs can be loaded without rebuilding them through `VADD`: + +1. Create an empty set with `XVCREATE`, supplying `QUANT_STATE` for a quantized index. +2. Load its terms with `XVIMPORT`. Terms may be imported in any order and in parallel. +3. Wait for all term imports to complete, then call `XVIMPORT key FINISH` before normal use. + +Internal IDs may also arrive in any order. Importing a high ID creates all missing +free-space-map blocks through that ID; each block covers 65,536 IDs and uses 8 KiB. +The start point is created during finalization and is excluded from FSM accounting. + +At the FFI layer, `can_import` checks eligibility. + +Ordinary operations (`card`, search, ID checks) and `FINISH` permanently disable importing. + +See the [FFI reference](docs/ffi-design.rs) and [data design](docs/data-design.md) +for API and term formats. ## Installing diff --git a/diskann-garnet/diskann-garnet.nuspec b/diskann-garnet/diskann-garnet.nuspec index 8c53121e95..65af4dde43 100644 --- a/diskann-garnet/diskann-garnet.nuspec +++ b/diskann-garnet/diskann-garnet.nuspec @@ -2,7 +2,7 @@ diskann-garnet - 6.0.0 + 6.0.1 docs/README.md Microsoft https://github.com/microsoft/DiskANN diff --git a/diskann-garnet/docs/data-design.md b/diskann-garnet/docs/data-design.md index c3a0d2c0a6..512e628fd9 100644 --- a/diskann-garnet/docs/data-design.md +++ b/diskann-garnet/docs/data-design.md @@ -30,7 +30,7 @@ The read-modify-write (rmw) method accesses a single key with a callback but all ### Keys & Contexts -Garnet keys are arbitrary byte strings (e.g. `&[u8]`). The methods described above can use whatever keys they like, with some caveats, to read and write data, however those methods also take a semi-opaque context which gives the operation a scope. For the most part this context is used for internal Garnet bookkeeping and is opaque, but the least significant 3 bits are available for diskann-garnet to use for its own scoping. Keys must be a multiple of 4 bytes in length, except for keys stored under context 3 and context 5 (`Term::Attributes` and `Term::IntMap`) which are special cased in Garnet as they key is the user-provided external ID. There are no alignment requirements on keys. +Garnet keys are arbitrary byte strings (e.g. `&[u8]`). The methods described above can use whatever keys they like, with some caveats, to read and write data, however those methods also take a semi-opaque context which gives the operation a scope. For the most part this context is used for internal Garnet bookkeeping and is opaque, but the least significant 3 bits are available for diskann-garnet to use for its own scoping. Keys must be a multiple of 4 bytes in length, except for keys stored under context 5 (`Term::IntMap`), which use the user-provided external ID. There are no alignment requirements on keys. Diskann-garnet uses these bits to distinguish between differnet kinds of index data so that the same key can be used to fetch different kinds of data. For example, vector data might be stored under the same key, the vector ID, as neighbor lists by setting different context bits for the operation. @@ -58,7 +58,7 @@ precision only index, they are the most accessed term. *Key*: Internal ID as bytes; this key is always 4 bytes. *Value*: `[u32; M + 1]` stored as bytes, where `M` is the same `M` passed to VADD. Every value in the index will have this fixed size of `(M + 1) * 4` bytes. -Neighbor lists are stored as a fixed size of `(M + 1) * mem::size_of::()`. The final entry is the true length of the neighbor list. +Neighbor lists are stored as a fixed size of `(M + 1) * mem::size_of::()`. The final entry is the true length of the neighbor list and must not exceed `M`. For example, in a graph where the `M` value is given as 16, the size of a neighbor list would be `(16 + 1) * 4 = 68` bytes long. Using fixed size lists this way means that all neighbor list allocations are the same size. @@ -75,7 +75,7 @@ In a quantized index, these vectors are the most often accessed piece of data an ### Attributes -*Key*: External ID as bytes; this is variable length byte string that the user assigned. +*Key*: Internal ID as bytes; this key is always 4 bytes. *Value*: Attributes given by the user. Variable length. When vectors are inserted by the Redis Vector Set API, an arbitrary JSON blob of attributes can be attached. These attributes are stored as a utf-8 string and read/written as a whole unit. @@ -88,13 +88,17 @@ The are several terms in the index used for internal state management of the dis #### Start Points -Currently only a single start point is supported, and it is given the internal ID of 0. Its vector data is the same as the first vector that was inserted, and it will not be returned by search, and it will not be modified during the lifetime of the index. +Currently only a single start point is supported, and it is given the internal ID of 0. During normal insertion, its vector data is the same as the first vector that was inserted. It will not be returned by search, and its vector data will not be modified during the lifetime of the index. + +For a nonempty import, finalization copies the first used imported ID's full vector, quantized vector (if present), and neighbor list to ID 0. ID 0 cannot be imported directly. Start points have no associated attributes. +ID 0 is reserved outside the free space map. The FSM never allocates, counts, visits, or reuses it; the provider maintains the start point's validity separately. + #### Metadata -Metadata is currently used for the free space map which manages used and available internal IDs and the quantizer tables. +Metadata is currently used for the free space map which manages used and available internal IDs, the quantizer tables, and import eligibility. ##### Free Space Map @@ -103,9 +107,11 @@ Metadata is currently used for the free space map which manages used and availab The free space map is used to keep track of which internal IDs are allocated and in use. Please see the [ID Mapping](#id-mapping) section for more details on why mapped IDs are used. -The free space map is a series of blocks where each block is a string of 1-bit values representing the state of the corresponding internal ID. Free IDs are represented by `0b0`, used IDs by `0b1`. Blocks are created on demand during insert when they are needed. Since each block contains 64k internal ID states, the total number of FSM blocks in the index will be at least number of active IDs / 2^16, each of which is 8kB in size. +The free space map is a contiguous sequence of blocks starting at block zero. Each block contains 65,536 one-bit ID states and occupies 8 KiB: `0b0` for free and `0b1` for used. The bit for reserved ID 0 is ignored. Blocks grow on demand during insertion or import, including intermediate blocks. -During startup, the index will scan FSM blocks in sequence to restore state. It will update the correct bits in a FSM block whenever the state of an internal ID changes. +During startup, the index scans FSM blocks in sequence to restore state. It updates the corresponding bit whenever a user vector's allocation changes. A failed expansion retains its successfully written prefix, allowing subsequent expansion to continue without erasing earlier claims. + +Importing a term keyed by an internal ID claims that ID before writing the term and advances the maximum assigned ID as needed. A failed claim writes no term; a failed term write leaves the ID claimed for a retry. Repeated claims do not increase the used count. Gaps below the maximum ID are available for reuse once reuse is enabled; quantization backfill keeps reuse disabled until it finishes. ##### Quantizer Tables @@ -114,12 +120,27 @@ During startup, the index will scan FSM blocks in sequence to restore state. It Note that for the BIN quantizer, a 1 byte flag precedes the quantizer table which indicates whether quantization backfill is complete. That byte is accounted for in the value sizing above. +Quantizer state can only be replaced before importing begins, while no user IDs are occupied in the FSM and no start-point vector exists. Setup takes the import-transition mutex before the FSM quantization barrier. Initial setup from `Disabled` requires exclusive caller access before publishing the index: ordinary operations in that state do not take the transition mutex. + +##### Import Status + +*Key*: b'_imp'; this key is always 4 bytes. +*Value*: A single byte: 0 for `Disabled`, 1 for `Ready`, or 2 for `Importing`. + +When no saved status exists, an empty index starts `Ready` if it is unquantized or has supplied quantizer state, and `Disabled` otherwise. The initial status is persisted before writing Q8's generated default table, so that table cannot enable imports on reopen. Successfully supplying quantizer state during creation publishes `Ready` after setup completes. A saved status is restored rather than inferred from the absence of a start point. + +The first ordinary operation transitions `Ready` to `Disabled`; the first accepted imported term transitions `Ready` to `Importing`. More terms can then arrive in parallel, but ordinary operations fail until all verification tasks and finalization succeed. The last successful task transitions `Importing` to `Disabled`. Empty finalization also waits for every task and does not create a start point. Failed verification or finalization leaves an active import `Importing`. + +Status transitions persist the new byte before publishing it with an atomic compare-exchange, under a transition-only mutex. A failed metadata write does not publish the new in-memory status. Ordinary operations in `Disabled` only load the atomic status; admitted term writes and verification workers do not hold the transition mutex. Callers must finish all term writes before starting finalization. + +The former Boolean format did not distinguish ready indexes from unfinished imports; finish those imports before upgrading. + #### Internal ID Mapping *Key*: Internal ID as bytes; this key is always 4 bytes in length. *Value*: External ID; this is variable length byte string that the user assigned. -Each internal ID corresponds to an external ID, which is a byte string of arbitrary length. The external IDs are stored unmodified and read/written as a whole. +Each internal ID corresponds to an external ID, which is a byte string of arbitrary length, including zero. The external IDs are stored unmodified and read/written as a whole. Lookup of an external ID will happen during post processing when we return results to Garnet. @@ -134,7 +155,7 @@ Lookup of an internal ID will happen for things such as delete. ## ID Mapping -Garnet vector set IDs are arbitrary length byte strings natively. These are quite inefficient for indexing so we map each external ID to a `u32` internal ID. This imposes a maximum on the number of vectors that are indexable of `u32::MAX - 1` (the start point always consumes an ID). +Garnet vector set IDs are arbitrary length byte strings natively. These are quite inefficient for indexing so we map each external ID to a `u32` internal ID. User vector IDs range from 1 through `u32::MAX - 1`, supporting at most 2^32 - 2 vectors. ID 0 is reserved for the start point, and `u32::MAX` is the exhausted next-ID value rather than an assignable ID. When the DiskANN algorithm performs searches and other operations it works with internal IDs only. The external IDs are only used when returning data to the user (which has no concept of the internal IDs) or when asked to perform operations like delete on specific vectors which will be identified by their external ID. In order to convert back and forth for these occasions, lookup tables must be kept for the mapping. @@ -175,4 +196,4 @@ Filtered search introduces two problems regarding attributes. The first is that It is possible to co-locate the appropriate vector term and the attribute term as a single value which may alleviate much of this problem. Because we can read partial values in the read callback, we can access the correct portion of the data (vector data would be stored first as it has a known, fixed size). However, we can alleviate only one of the two problems in an index with both full precision and quantized vector terms as we must choose whether the search operation should not see skew or whether the final reranking step should not see skew. - \ No newline at end of file + diff --git a/diskann-garnet/docs/ffi-design.rs b/diskann-garnet/docs/ffi-design.rs index 6ff12a7c4e..6cfcf03d74 100644 --- a/diskann-garnet/docs/ffi-design.rs +++ b/diskann-garnet/docs/ffi-design.rs @@ -45,6 +45,16 @@ enum InsertResult { SuccessUpdate = 3, } +/// Status returned by `finish_import`, encoded as a `u8`. +/// +/// `Success` means the task succeeded, including finalization if it was the last task. +/// `TaskFailed` permits retrying that task; `FinishFailed` is a non-retryable finalization failure. +enum ImportResult { + Success = 0, + TaskFailed = 1, + FinishFailed = 2, +} + /// Read one or more keys from Garnet. /// /// `keys` holds `key_count` keys, each 4-byte length prefixed. `value_length_hint` is the @@ -314,9 +324,11 @@ extern "C" fn remove( id_len: usize, ) -> bool; -/// Return number of vectors stored in index. +/// Return number of live vectors in index, excluding the reserved start point. /// /// Equivalent to VCARD (https://redis.io/docs/latest/commands/vcard/) can be approximate, must be fast. +/// +/// Returns `u64::MAX` on error. #[unsafe(no_mangle)] extern "C" fn card(context: u64, index_ptr: *const c_void) -> u64; @@ -381,3 +393,70 @@ extern "C" fn search_neighbors( output_distances_len: usize, overflow: *mut *mut c_void, ) -> i32; + +/// Set serialized quantizer state on an empty, quantized index, for XVCREATE. +/// +/// Pass serialized quantizer bytes. The metric and dimension must match. Importing must not have +/// begun and the index must contain no data. Initial setup on an import-disabled index requires +/// exclusive access before publishing it to concurrent commands. Success enables importing; +/// Q8's automatically generated default state does not. +/// +/// Returns true on success and false otherwise. On failure, discard the vector set: in-memory +/// state and persisted metadata may already have changed. +#[unsafe(no_mangle)] +extern "C" fn set_quant_state( + context: u64, + index_ptr: *const c_void, + state: *const u8, + state_len: usize, +) -> bool; + +/// Return whether the index accepts imported terms, both before and during import. +/// This query does not begin or disable importing. +#[unsafe(no_mangle)] +extern "C" fn can_import(context: u64, index_ptr: *const c_void) -> bool; + +/// Import a term directly into storage, for XVIMPORT. +/// +/// term_type is 0 for vectors, 1 for neighbors, 2 for quantized vectors, 3 for +/// attributes, 5 for external-to-internal ID mappings, or 6 for +/// internal-to-external ID mappings. Metadata (4) cannot be imported. IDs and +/// values must be non-null and non-empty (except for external IDs which may be +/// empty). Terms may arrive in any order and must use IDs between 1 and +/// u32::MAX - 1. Importing a high ID creates every missing FSM block through +/// that ID. ID 0 is reserved for the start point, which finalization creates +/// outside FSM accounting. External-to-internal mapping values must be a +/// nonzero native-endian u32. Neighbor terms must have `(max_degree + 1) * 4` +/// bytes and a trailing native-endian u32 count no greater than max_degree. +/// +/// The first accepted term begins importing. Ordinary operations fail until FINISH succeeds. +/// A failed term write keeps the import in progress; callers may repair or retry that term. +/// +/// Returns true on success and false otherwise. +#[unsafe(no_mangle)] +extern "C" fn import_term( + context: u64, + index_ptr: *const c_void, + term_type: u32, + id: *const u8, + id_len: usize, + value: *const u8, + value_len: usize, +) -> bool; + +/// Verify imported terms and finalize the index, for XVIMPORT ... FINISH. +/// +/// Finish all term writes before starting these tasks; no term writes may overlap verification. +/// Each task receives its own task_idx in 0..task_count. Do not repeat successful tasks. The last +/// successful task finalizes the index and enables ordinary operations. Until then, including +/// after a failed verification or finalization, ordinary operations remain blocked. An empty +/// import also waits for every task and closes without creating a start point. +/// +/// Returns an `ImportResult` discriminant. +#[unsafe(no_mangle)] +extern "C" fn finish_import( + context: u64, + index_ptr: *const c_void, + task_idx: usize, + task_count: usize, +) -> u8; diff --git a/diskann-garnet/src/dyn_index.rs b/diskann-garnet/src/dyn_index.rs index dd60d58956..1c754045f3 100644 --- a/diskann-garnet/src/dyn_index.rs +++ b/diskann-garnet/src/dyn_index.rs @@ -5,8 +5,8 @@ use crate::{ SearchResults, - garnet::{Context, GarnetId}, - provider::{DynamicQuantization, GarnetProvider}, + garnet::{Context, GarnetError, GarnetId, Term}, + provider::{DynamicQuantization, GarnetProvider, GarnetProviderError}, }; use diskann::{ ANNResult, @@ -69,7 +69,7 @@ pub(crate) trait DynIndex: Send + Sync { fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()>; /// Return an approximate count of vectors in the index - fn approximate_count(&self) -> u64; + fn approximate_count(&self, context: &Context) -> ANNResult; /// Return the maximum degree of the index graph fn max_degree(&self) -> usize; @@ -109,6 +109,26 @@ pub(crate) trait DynIndex: Send + Sync { /// message to an area (e.g. `Term::Quantized` for quantization related /// messages). fn log(&self, context: &Context, msg: &str); + + /// Set the state of the quantizer. + /// Returns true on success and false otherwise. + fn set_quant_state(&self, context: &Context, state: &[u8]) -> bool; + + /// Returns whether imported terms are accepted, without changing the import status. + fn can_import(&self, context: &Context) -> bool; + + /// Import an index term. + fn import_term(&self, context: &Context, term: Term, id: &[u8], value: &[u8]) -> bool; + + /// Finalize an import. + /// Returns task success and the optional finalization outcome. An empty import returns + /// `(false, Some(true))` so the FFI keeps start points uninitialized. + fn finish_import( + &self, + context: &Context, + task_idx: usize, + task_count: usize, + ) -> (bool, Option); } impl DynIndex for DiskANNIndex> { @@ -116,6 +136,9 @@ impl DynIndex for DiskANNIndex> { /// /// The data slice here must be aligned to `T` or this will panic. fn insert(&self, context: &Context, id: &GarnetId, data: &[u8], attrs: &[u8]) -> ANNResult<()> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.run(|_| async { let _pending = self.inner.provider().reserve_external_id(id).await?; self.inner @@ -130,6 +153,9 @@ impl DynIndex for DiskANNIndex> { } fn set_attributes(&self, context: &Context, id: &GarnetId, data: &[u8]) -> ANNResult<()> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.run(|_| async { let provider = self.inner.provider(); let _pending = provider.reserve_external_id(id).await?; @@ -140,6 +166,9 @@ impl DynIndex for DiskANNIndex> { } fn delete_attributes(&self, context: &Context, id: &GarnetId) -> ANNResult<()> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.run(|_| async { let provider = self.inner.provider(); let _pending = provider.reserve_external_id(id).await?; @@ -156,6 +185,9 @@ impl DynIndex for DiskANNIndex> { params: search::Knn, output: &mut SearchResults<'_>, ) -> ANNResult { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } let query = bytemuck::cast_slice::(data); self.search(params, &DynamicQuantization, context, query, output) } @@ -167,11 +199,19 @@ impl DynIndex for DiskANNIndex> { params: search::Knn, output: &mut SearchResults<'_>, ) -> ANNResult { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } // Look up internal ID let iid = self.inner.provider().to_internal_id(context, id)?; let data = self.inner.provider().get_full_vector(context, iid)?; - let data_bytes = bytemuck::cast_slice::(&data); - self.search_vector(context, data_bytes, params, output) + self.search( + params, + &DynamicQuantization, + context, + data.as_slice(), + output, + ) } fn filtered_search_vector( @@ -181,6 +221,9 @@ impl DynIndex for DiskANNIndex> { params: search::InlineFilterSearch, output: &mut SearchResults<'_>, ) -> ANNResult { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } let query = bytemuck::cast_slice::(data); self.search(params, &DynamicQuantization, context, query, output) } @@ -192,14 +235,25 @@ impl DynIndex for DiskANNIndex> { params: search::InlineFilterSearch, output: &mut SearchResults<'_>, ) -> ANNResult { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } // Look up internal ID let iid = self.inner.provider().to_internal_id(context, id)?; let data = self.inner.provider().get_full_vector(context, iid)?; - let data_bytes = bytemuck::cast_slice::(&data); - self.filtered_search_vector(context, data_bytes, params, output) + self.search( + params, + &DynamicQuantization, + context, + data.as_slice(), + output, + ) } fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.run(|_| async { let _pending = self.inner.provider().reserve_external_id(id).await?; self.inner @@ -214,8 +268,11 @@ impl DynIndex for DiskANNIndex> { }) } - fn approximate_count(&self) -> u64 { - self.inner.provider().max_internal_id() as u64 + fn approximate_count(&self, context: &Context) -> ANNResult { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } + Ok(self.inner.provider().total_used() as u64) } fn max_degree(&self) -> usize { @@ -223,6 +280,9 @@ impl DynIndex for DiskANNIndex> { } fn maybe_set_start_point(&self, context: &Context, data: &[u8]) -> ANNResult<()> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.inner .provider() .maybe_set_start_point(context, bytemuck::cast_slice::(data)) @@ -230,14 +290,23 @@ impl DynIndex for DiskANNIndex> { } fn internal_id_exists(&self, context: &Context, id: u32) -> bool { + if !self.inner.provider().disable_import(context) { + return false; + } self.inner.provider().vector_iid_exists(context, id) } fn external_id_exists(&self, context: &Context, id: &GarnetId) -> bool { + if !self.inner.provider().disable_import(context) { + return false; + } self.inner.provider().vector_id_exists(context, id) } fn train_quantizer(&self, context: &Context) -> bool { + if !self.inner.provider().disable_import(context) { + return false; + } self.inner.provider().train_quantizer(context) } @@ -247,6 +316,9 @@ impl DynIndex for DiskANNIndex> { task_idx: usize, task_count: usize, ) -> bool { + if !self.inner.provider().disable_import(context) { + return false; + } self.inner .provider() .backfill_quant_vectors(context, task_idx, task_count) @@ -258,14 +330,43 @@ impl DynIndex for DiskANNIndex> { count: u32, output: &mut SearchResults<'_>, ) -> bool { + if !self.inner.provider().disable_import(context) { + return false; + } self.inner.provider().random_members(context, count, output) } fn neighbors(&self, context: &Context, id: &GarnetId) -> ANNResult>> { + if !self.inner.provider().disable_import(context) { + return Err(GarnetProviderError::Garnet(GarnetError::Write).into()); + } self.inner.provider().neighbors(context, id) } fn log(&self, context: &Context, msg: &str) { self.inner.provider().log(context, msg); } + + fn set_quant_state(&self, context: &Context, state: &[u8]) -> bool { + self.inner.provider().set_quant_state(context, state) + } + + fn can_import(&self, context: &Context) -> bool { + self.inner.provider().can_import(context) + } + + fn import_term(&self, context: &Context, term: Term, id: &[u8], value: &[u8]) -> bool { + self.inner.provider().import_term(context, term, id, value) + } + + fn finish_import( + &self, + context: &Context, + task_idx: usize, + task_count: usize, + ) -> (bool, Option) { + self.inner + .provider() + .finish_import(context, task_idx, task_count) + } } diff --git a/diskann-garnet/src/ffi_tests.rs b/diskann-garnet/src/ffi_tests.rs index e039b8303c..3213f1516e 100644 --- a/diskann-garnet/src/ffi_tests.rs +++ b/diskann-garnet/src/ffi_tests.rs @@ -4,19 +4,29 @@ */ mod tests { - use std::{ffi::c_void, mem, ptr}; + use std::{ + ffi::c_void, + mem, ptr, + sync::{ + Barrier, + atomic::{AtomicBool, Ordering}, + }, + thread, + }; + use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::distance::Metric; use rand::{Rng, seq::SliceRandom}; use crate::{ - Index, InsertResult, Overflow, VectorQuantType, backfill_quant_vectors, build_quant_table, - card, check_external_id_valid, check_internal_id_valid, create_index, drop_index, - garnet::{Context, Term}, - insert, - quantization::{GarnetQuantizer, Spherical1Bit}, - remove, search_vector, set_attribute, - test_utils::Store, + ImportResult, Index, IndexState, InsertResult, Overflow, VectorQuantType, + backfill_quant_vectors, build_quant_table, can_import, card, check_external_id_valid, + check_internal_id_valid, create_index, drop_index, finish_import, + garnet::{Context, Term, WriteCallback}, + import_term, insert, + quantization::{GarnetQuantizer, MinMax8Bit, Spherical1Bit}, + remove, search_vector, set_attribute, set_quant_state, + test_utils::{STORE, Store, q8_state_with_identity_transform}, }; /// Creates an index with default test values and returns (index_ptr, Context). @@ -36,6 +46,20 @@ mod tests { store: &Store, quant_type: VectorQuantType, metric_type: i32, + ) -> (*const c_void, Context) { + create_test_index_with_write_callback( + store, + quant_type, + metric_type, + store.callbacks().write_callback(), + ) + } + + fn create_test_index_with_write_callback( + store: &Store, + quant_type: VectorQuantType, + metric_type: i32, + write_callback: WriteCallback, ) -> (*const c_void, Context) { let callbacks = store.callbacks(); let ctx = Context::new(0); @@ -56,7 +80,7 @@ mod tests { l_build, max_degree, callbacks.read_callback(), - callbacks.write_callback(), + write_callback, callbacks.delete_callback(), callbacks.rmw_callback(), callbacks.filter_callback(), @@ -224,15 +248,777 @@ mod tests { let removed = unsafe { remove(ctx.get(), index_ptr, id_bytes.as_ptr(), id_bytes.len()) }; assert!(removed); - // Currently we're not tracking deletions for cardinality, so it should still be 1 cardinality = unsafe { card(ctx.get(), index_ptr) }; - assert_eq!(cardinality, 1); + assert_eq!(cardinality, 0); unsafe { drop_index(ctx.get(), index_ptr); } } + #[test] + fn import_rejects_empty_non_external_id_fields() { + let store = Store::new(); + let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant); + let id = 1u32.to_ne_bytes(); + for term in [ + Term::Vector, + Term::Neighbors, + Term::Quantized, + Term::Attributes, + Term::Metadata, + Term::IntMap, + Term::ExtMap, + ] { + for empty_data in [ptr::null(), b"ignored".as_ptr()] { + assert_eq!( + unsafe { + import_term( + ctx.get(), + index_ptr, + term as u32, + empty_data, + 0, + id.as_ptr(), + id.len(), + ) + }, + matches!(term, Term::IntMap) + ); + assert_eq!( + unsafe { + import_term( + ctx.get(), + index_ptr, + term as u32, + id.as_ptr(), + id.len(), + empty_data, + 0, + ) + }, + matches!(term, Term::ExtMap) + ); + } + } + assert_eq!( + store.get(ctx.term(Term::IntMap).get(), b""), + Some(id.to_vec()) + ); + assert_eq!( + store.get(ctx.term(Term::ExtMap).get(), &id), + Some(Vec::new()) + ); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + + #[test] + fn can_import_tracks_eligibility() { + for quant_type in [VectorQuantType::NoQuant, VectorQuantType::Q8] { + let store = Store::new(); + let (mut index_ptr, ctx) = create_test_index(&store, quant_type); + assert_eq!( + unsafe { can_import(ctx.get(), index_ptr) }, + quant_type == VectorQuantType::NoQuant + ); + // Reopening must preserve both initial Ready and unconfigured Q8 Disabled. + unsafe { drop_index(ctx.get(), index_ptr) }; + index_ptr = create_test_index(&store, quant_type).0; + assert_eq!( + unsafe { can_import(ctx.get(), index_ptr) }, + quant_type == VectorQuantType::NoQuant + ); + let quantizer = (quant_type == VectorQuantType::Q8) + .then(|| MinMax8Bit::new(2, Metric::L2).unwrap()); + if let Some(quantizer) = &quantizer { + let state = quantizer.serialize().unwrap(); + assert!(unsafe { + set_quant_state(ctx.get(), index_ptr, state.as_ptr(), state.len()) + }); + } + + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + // Reopening and repeated eligibility queries must preserve Ready. + unsafe { drop_index(ctx.get(), index_ptr) }; + index_ptr = create_test_index(&store, quant_type).0; + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + + let id = 1u32.to_ne_bytes(); + for (internal_id, accepted, expected) in [ + (0u32, false, None), + (1, true, Some(id.as_slice())), + (0, false, Some(id.as_slice())), + ] { + let value = internal_id.to_ne_bytes(); + assert_eq!( + import_test_term(&ctx, index_ptr, Term::IntMap, &id, &value), + accepted + ); + assert_eq!( + store.get(ctx.term(Term::IntMap).get(), &id).as_deref(), + expected + ); + } + for reopen in [false, true] { + if reopen { + unsafe { drop_index(ctx.get(), index_ptr) }; + index_ptr = create_test_index(&store, quant_type).0; + } + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + if quant_type == VectorQuantType::Q8 { + let state = q8_state_with_identity_transform(2); + let key = u32::from_be_bytes(*b"_qnt").to_ne_bytes(); + let original = store.get(ctx.term(Term::Metadata).get(), &key); + assert!(!unsafe { + set_quant_state(ctx.get(), index_ptr, state.as_ptr(), state.len()) + }); + assert_eq!(store.get(ctx.term(Term::Metadata).get(), &key), original); + } + assert_eq!( + store.get(ctx.term(Term::IntMap).get(), &id), + Some(id.to_vec()) + ); + } + let attributes = b"attributes"; + assert!(import_test_term( + &ctx, + index_ptr, + Term::Attributes, + &id, + attributes + )); + if quant_type == VectorQuantType::Q8 { + let state = q8_state_with_identity_transform(2); + assert!(!unsafe { + set_quant_state(ctx.get(), index_ptr, state.as_ptr(), state.len()) + }); + } + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::TaskFailed) + ); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + unsafe { drop_index(ctx.get(), index_ptr) }; + + let (index_ptr, ctx) = create_test_index(&store, quant_type); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::TaskFailed) + ); + + // Repair the incomplete vector and retry validation on this handle. + import_single_vector( + &ctx, + index_ptr, + quantizer + .as_ref() + .map(|quantizer| quantizer as &dyn GarnetQuantizer), + ); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::Success) + ); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 1); + assert!(unsafe { + check_internal_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + assert!(unsafe { + check_external_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + + for insert_vector in [false, true] { + let store = Store::new(); + let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + if insert_vector { + assert_eq!( + insert_f32_vector(&ctx, index_ptr, 1, &[1.0, 2.0]), + InsertResult::Success + ); + } else { + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 0); + } + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + unsafe { drop_index(ctx.get(), index_ptr) }; + let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::NoQuant); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!( + unsafe { card(ctx.get(), index_ptr) }, + u64::from(insert_vector) + ); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + } + + #[test] + fn import_terms_and_finish_q8() { + let quantizer = MinMax8Bit::new(2, Metric::L2).unwrap(); + check_import_terms_and_finish(VectorQuantType::Q8, Some(&quantizer)); + } + + #[test] + fn import_terms_and_finish_noquant() { + check_import_terms_and_finish(VectorQuantType::NoQuant, None); + } + + #[test] + fn import_terms_and_finish_bin() { + let quantizer = Spherical1Bit::new(Metric::L2, 2); + let mut training_data = + rowmajor::Owned::from_element(quantizer.required_vectors(), 2, 0.0f32); + for row in 0..quantizer.required_vectors() { + training_data.row_mut(row).fill((row % 100 + 1) as f32); + } + quantizer + .train(Metric::L2, training_data.as_view()) + .unwrap(); + check_import_terms_and_finish(VectorQuantType::Bin, Some(&quantizer)); + } + + #[test] + fn import_terms_and_finish_multi_worker() { + check_import_terms_and_finish_with(VectorQuantType::NoQuant, None, |ctx, index_ptr| { + const TASK_COUNT: usize = 104; + let index = unsafe { &*index_ptr.cast::() }; + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, TASK_COUNT) }, + 0 + ); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(ctx, index_ptr); + let snapshot = STORE.with(Clone::clone); + let barrier = Barrier::new(4); + thread::scope(|scope| { + for worker in 0..4 { + let snapshot = snapshot.clone(); + let barrier = &barrier; + scope.spawn(move || { + STORE.with(|store| { + for (key, value) in snapshot { + store.insert(key, value); + } + }); + barrier.wait(); + for task_idx in (worker + 1..TASK_COUNT - 1).step_by(4) { + assert_eq!( + unsafe { + finish_import( + ctx.get(), + ptr::from_ref(index).cast(), + task_idx, + TASK_COUNT, + ) + }, + 0 + ); + } + }); + } + }); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(ctx, index_ptr); + assert_eq!( + index.state.load(Ordering::Acquire), + IndexState::NoStartPoints as usize + ); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, TASK_COUNT - 1, TASK_COUNT) }, + 0 + ); + assert_eq!( + index.state.load(Ordering::Acquire), + IndexState::Ready as usize + ); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 100); + }); + } + + #[test] + fn import_terms_and_finish_persistence_failure() { + static FAIL_WRITE: AtomicBool = AtomicBool::new(false); + static STATUS_ONLY: AtomicBool = AtomicBool::new(false); + unsafe extern "C" fn write( + context: u64, + key: *const u8, + key_len: usize, + value: *const u8, + value_len: usize, + ) -> bool { + let is_final_status = context == Context::new(0).term(Term::Metadata).get() + && unsafe { std::slice::from_raw_parts(key, key_len) } + == u32::from_be_bytes(*b"_imp").to_ne_bytes() + && unsafe { std::slice::from_raw_parts(value, value_len) } == [0]; + if (!STATUS_ONLY.load(Ordering::Acquire) || is_final_status) + && FAIL_WRITE.swap(false, Ordering::AcqRel) + { + return false; + } + unsafe { + (Store::attach().callbacks().write_callback())( + context, key, key_len, value, value_len, + ) + } + } + + // Fail either start-point persistence or the final Disabled publication. + // Neither failure is retryable on this handle, unlike task validation. + for status_only in [false, true] { + let store = Store::new(); + let (mut index_ptr, ctx) = create_test_index_with_write_callback( + &store, + VectorQuantType::NoQuant, + Metric::L2 as i32, + write, + ); + assert!(!index_ptr.is_null()); + import_single_vector(&ctx, index_ptr, None); + STATUS_ONLY.store(status_only, Ordering::Release); + FAIL_WRITE.store(true, Ordering::Release); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::FinishFailed) + ); + assert!(!FAIL_WRITE.load(Ordering::Acquire)); + for reopen in [false, true] { + if reopen { + unsafe { drop_index(ctx.get(), index_ptr) }; + index_ptr = create_test_index(&store, VectorQuantType::NoQuant).0; + } + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + } + unsafe { drop_index(ctx.get(), index_ptr) }; + } + } + + #[test] + fn import_admission_persistence_failure() { + static FAIL_WRITE: AtomicBool = AtomicBool::new(false); + unsafe extern "C" fn write( + context: u64, + key: *const u8, + key_len: usize, + value: *const u8, + value_len: usize, + ) -> bool { + !FAIL_WRITE.swap(false, Ordering::AcqRel) + && unsafe { + (Store::attach().callbacks().write_callback())( + context, key, key_len, value, value_len, + ) + } + } + + let store = Store::new(); + let (index_ptr, ctx) = create_test_index_with_write_callback( + &store, + VectorQuantType::NoQuant, + Metric::L2 as i32, + write, + ); + assert!(!index_ptr.is_null()); + let snapshot = STORE.with(Clone::clone); + FAIL_WRITE.store(true, Ordering::Release); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, u64::MAX); + assert!(!FAIL_WRITE.load(Ordering::Acquire)); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + + let id = 1u32.to_ne_bytes(); + FAIL_WRITE.store(true, Ordering::Release); + assert!(!import_test_term(&ctx, index_ptr, Term::IntMap, &id, &id)); + assert!(!FAIL_WRITE.load(Ordering::Acquire)); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + + FAIL_WRITE.store(true, Ordering::Release); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::TaskFailed) + ); + assert!(!FAIL_WRITE.load(Ordering::Acquire)); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + STORE.with(|store| { + assert_eq!(store.len(), snapshot.len()); + for entry in snapshot.iter() { + assert_eq!(store.get(entry.key()).as_deref(), Some(entry.value())); + } + }); + + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::Success) + ); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 0); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + + #[test] + fn finish_empty_import_and_reopen() { + for (quant_type, task_count) in [ + (VectorQuantType::NoQuant, 1), + (VectorQuantType::NoQuant, 4), + (VectorQuantType::Q8, 1), + (VectorQuantType::Q8, 4), + ] { + let store = Store::new(); + let (mut index_ptr, ctx) = create_test_index(&store, quant_type); + if quant_type == VectorQuantType::Q8 { + let quantizer = MinMax8Bit::new(2, Metric::L2).unwrap(); + let state = quantizer.serialize().unwrap(); + assert!(unsafe { + set_quant_state(ctx.get(), index_ptr, state.as_ptr(), state.len()) + }); + } + for task_idx in 0..task_count { + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, task_idx, task_count) }, + u8::from(ImportResult::Success) + ); + assert!( + store + .get(ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()) + .is_none() + ); + if task_idx + 1 < task_count { + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + } + } + for reopen in [false, true] { + if reopen { + unsafe { drop_index(ctx.get(), index_ptr) }; + index_ptr = create_test_index(&store, quant_type).0; + } + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 0); + assert!( + store + .get(ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()) + .is_none() + ); + } + assert_eq!( + insert_f32_vector(&ctx, index_ptr, 42, &[1.0; 2]), + InsertResult::Success + ); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 1); + let mapping = store + .get(ctx.term(Term::IntMap).get(), &42u32.to_ne_bytes()) + .unwrap(); + let internal_id = bytemuck::pod_read_unaligned::(&mapping); + assert!((1..u32::MAX).contains(&internal_id)); + let (ids, distances) = do_search(&ctx, index_ptr, &[1.0; 2], 1, None); + assert_eq!(ids, [42]); + assert_eq!(distances, [0.0]); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + } + + #[must_use] + fn import_test_term( + ctx: &Context, + index_ptr: *const c_void, + term: Term, + id: &[u8], + value: &[u8], + ) -> bool { + unsafe { + import_term( + ctx.get(), + index_ptr, + term as u32, + id.as_ptr(), + id.len(), + value.as_ptr(), + value.len(), + ) + } + } + + #[track_caller] + fn import_single_vector( + ctx: &Context, + index_ptr: *const c_void, + quantizer: Option<&dyn GarnetQuantizer>, + ) { + let id = 1u32.to_ne_bytes(); + let vector = [1.0f32, 2.0]; + assert!(import_test_term( + ctx, + index_ptr, + Term::Vector, + &id, + bytemuck::cast_slice(&vector), + )); + if let Some(quantizer) = quantizer { + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&vector, &mut quantized).unwrap(); + assert!(import_test_term( + ctx, + index_ptr, + Term::Quantized, + &id, + &quantized + )); + } + assert!(import_test_term( + ctx, + index_ptr, + Term::Neighbors, + &id, + bytemuck::cast_slice(&[0u32; 21]), + )); + assert!(import_test_term(ctx, index_ptr, Term::IntMap, &id, &id)); + assert!(import_test_term(ctx, index_ptr, Term::ExtMap, &id, &id)); + assert!(import_test_term( + ctx, + index_ptr, + Term::Attributes, + &id, + b"original" + )); + } + + fn assert_ordinary_operations_blocked(ctx: &Context, index_ptr: *const c_void) { + let snapshot = STORE.with(Clone::clone); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, u64::MAX); + assert_eq!( + insert_f32_vector(ctx, index_ptr, 42, &[3.0, 4.0]), + InsertResult::Fail + ); + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + STORE.with(|store| { + assert_eq!(store.len(), snapshot.len()); + for entry in snapshot.iter() { + assert_eq!(store.get(entry.key()).as_deref(), Some(entry.value())); + } + }); + } + + fn check_import_terms_and_finish( + quant_type: VectorQuantType, + quantizer: Option<&dyn GarnetQuantizer>, + ) { + check_import_terms_and_finish_with(quant_type, quantizer, |ctx, index_ptr| { + assert_eq!( + unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, + u8::from(ImportResult::Success) + ); + }); + } + + fn check_import_terms_and_finish_with( + quant_type: VectorQuantType, + quantizer: Option<&dyn GarnetQuantizer>, + finalize: impl Fn(&Context, *const c_void), + ) { + const VECTOR_COUNT: usize = 100; + const MAX_DEGREE: usize = 20; + const BLOCK_SIZE_IDS: u32 = 1 << 16; + + let mut internal_ids: [u32; VECTOR_COUNT] = std::array::from_fn(|index| index as u32 + 1); + internal_ids[VECTOR_COUNT - 2] = BLOCK_SIZE_IDS; + internal_ids[VECTOR_COUNT - 1] = 2 * BLOCK_SIZE_IDS + 7; + let max_id = internal_ids[VECTOR_COUNT - 1]; + + let quant_state = quantizer.map(|quantizer| quantizer.serialize().unwrap()); + let replacement_state = + (quant_type == VectorQuantType::Q8).then(|| q8_state_with_identity_transform(2)); + let terms: Vec<_> = internal_ids + .iter() + .enumerate() + .map(|(vector_index, &id)| { + let vector = [vector_index as f32 + 1.0; 2]; + let mut vector_terms = vec![(Term::Vector, bytemuck::cast_slice(&vector).to_vec())]; + if let Some(quantizer) = quantizer { + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&vector, &mut quantized).unwrap(); + vector_terms.push((Term::Quantized, quantized)); + } + let mut neighbors = [0u32; MAX_DEGREE + 1]; + let neighbor_count = if vector_index == 0 { MAX_DEGREE } else { 5 }; + for (offset, neighbor) in neighbors[..neighbor_count].iter_mut().enumerate() { + *neighbor = + internal_ids[(vector_index + VECTOR_COUNT - 1 - offset) % VECTOR_COUNT]; + } + neighbors[MAX_DEGREE] = neighbor_count as u32; + let id_bytes = id.to_ne_bytes(); + let external_id = (vector_index as u32 + 1).to_ne_bytes(); + + vector_terms.extend([ + (Term::Neighbors, bytemuck::cast_slice(&neighbors).to_vec()), + (Term::Attributes, id_bytes.to_vec()), + (Term::IntMap, id_bytes.to_vec()), + (Term::ExtMap, external_id.to_vec()), + ]); + vector_terms + }) + .collect(); + let term_count = terms[0].len(); + + for by_term in [true, false] { + let store = Store::new(); + let (index_ptr, ctx) = create_test_index(&store, quant_type); + let index = unsafe { &*index_ptr.cast::() }; + assert_eq!(index.inner.max_degree(), MAX_DEGREE); + if let Some(quant_state) = &quant_state { + let id = 1u32.to_ne_bytes(); + for (term, value) in &terms[0] { + assert!( + !import_test_term(&ctx, index_ptr, *term, &id, value), + "import accepted without preset quant state: quant_type={quant_type:?}, by_term={by_term}, term={term:?}" + ); + assert!(store.get(ctx.get() | *term as u64, &id).is_none()); + } + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, 0); + assert!(!unsafe { + check_internal_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + assert!(!unsafe { + check_external_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + assert!(unsafe { + set_quant_state( + ctx.get(), + index_ptr, + quant_state.as_ptr(), + quant_state.len(), + ) + }); + } + for rejected_id in [0u32, u32::MAX] { + let id = rejected_id.to_ne_bytes(); + let vector = [0.0f32; 2]; + let value = bytemuck::cast_slice::(&vector); + assert!(!import_test_term(&ctx, index_ptr, Term::Vector, &id, value)); + assert!(store.get(ctx.term(Term::Vector).get(), &id).is_none()); + } + + let rejected_id = (max_id + 1).to_ne_bytes(); + let mut bad_neighbors = [0u32; MAX_DEGREE + 1]; + bad_neighbors[MAX_DEGREE] = MAX_DEGREE as u32 + 1; + let value = bytemuck::cast_slice(&bad_neighbors); + assert!(!import_test_term( + &ctx, + index_ptr, + Term::Neighbors, + &rejected_id, + value + )); + assert!( + store + .get(ctx.term(Term::Neighbors).get(), &rejected_id) + .is_none() + ); + + // Ascending term-major imports grow one block at a time; descending + // vector-major imports jump directly to the final block. + for position in 0..VECTOR_COUNT * term_count { + let (vector_index, term_index) = if by_term { + (position % VECTOR_COUNT, position / VECTOR_COUNT) + } else { + ( + VECTOR_COUNT - 1 - position / term_count, + (position % term_count + usize::from(quantizer.is_some())) % term_count, + ) + }; + let (term, value) = &terms[vector_index][term_index]; + let id = if matches!(term, Term::IntMap) { + vector_index as u32 + 1 + } else { + internal_ids[vector_index] + } + .to_ne_bytes(); + assert!( + import_test_term(&ctx, index_ptr, *term, &id, value), + "import failed: by_term={by_term}, internal_id={}, term={term:?}", + internal_ids[vector_index] + ); + if position == 0 + && let Some(state) = &replacement_state + { + assert!(!unsafe { + set_quant_state(ctx.get(), index_ptr, state.as_ptr(), state.len()) + }); + assert_eq!( + store.get( + ctx.term(Term::Metadata).get(), + &u32::from_be_bytes(*b"_qnt").to_ne_bytes(), + ), + Some(quant_state.as_ref().unwrap().to_vec()) + ); + } + } + + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); + finalize(&ctx, index_ptr); + assert!(!unsafe { can_import(ctx.get(), index_ptr) }); + for rejected_id in [1u32, 101] { + let id = rejected_id.to_ne_bytes(); + for (term, value) in &terms[1] { + let term_context = ctx.get() | *term as u64; + let before = store.get(term_context, &id); + assert!( + !import_test_term(&ctx, index_ptr, *term, &id, value), + "import accepted after finish_import: by_term={by_term}, id={rejected_id}, term={term:?}" + ); + assert_eq!( + store.get(term_context, &id), + before, + "rejected import changed storage: by_term={by_term}, id={rejected_id}, term={term:?}" + ); + } + } + for unused_id in [101, BLOCK_SIZE_IDS + 1, max_id - 1] { + let id = unused_id.to_ne_bytes(); + assert!(!unsafe { + check_internal_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + assert!(!unsafe { + check_external_id_valid(ctx.get(), index_ptr, id.as_ptr(), id.len()) + }); + } + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, VECTOR_COUNT as u64); + + assert_eq!( + insert_f32_vector(&ctx, index_ptr, 101, &[101.0; 2]), + InsertResult::Success + ); + assert_eq!( + unsafe { card(ctx.get(), index_ptr) }, + VECTOR_COUNT as u64 + 1 + ); + let added_id = 101u32.to_ne_bytes(); + assert!(unsafe { + check_external_id_valid(ctx.get(), index_ptr, added_id.as_ptr(), added_id.len()) + }); + + let (ids, distances) = do_search(&ctx, index_ptr, &[50.0; 2], 5, None); + assert_eq!(ids.first(), Some(&50), "nearest vector: by_term={by_term}"); + let mut sorted_ids = ids; + sorted_ids.sort_unstable(); + assert_eq!(sorted_ids, [48, 49, 50, 51, 52]); + assert_eq!(distances.len(), 5); + assert!(distances.iter().all(|distance| distance.is_finite())); + + unsafe { + drop_index(ctx.get(), index_ptr); + } + } + } + #[test] fn update_vector_attributes() { let store = Store::new(); @@ -950,14 +1736,14 @@ mod tests { let (index_ptr, ctx) = create_test_index(&store, VectorQuantType::Bin); let index = unsafe { &*index_ptr.cast::() }; - let quantizer = Spherical1Bit::new(2); + let quantizer = Spherical1Bit::new(Metric::L2, 2); let required_vectors = quantizer.required_vectors(); let mut rng = rand::rng(); // pre-quantization phase - assert_eq!(index.inner.approximate_count(), 0); + assert_eq!(index.inner.approximate_count(&ctx).unwrap(), 0); for id in 0..required_vectors - 1 { let v = [rng.random(), rng.random()]; assert_eq!( @@ -966,7 +1752,7 @@ mod tests { ); } assert_eq!( - index.inner.approximate_count() as usize, + index.inner.approximate_count(&ctx).unwrap() as usize, required_vectors - 1 ); diff --git a/diskann-garnet/src/fsm.rs b/diskann-garnet/src/fsm.rs index c39415826d..aee9468698 100644 --- a/diskann-garnet/src/fsm.rs +++ b/diskann-garnet/src/fsm.rs @@ -19,10 +19,8 @@ //! counter which tracks the next new ID. use crossbeam::queue::ArrayQueue; -use std::sync::{ - Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard, - atomic::{AtomicBool, AtomicUsize, Ordering}, -}; +use parking_lot::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use thiserror::Error; use crate::garnet::{Callbacks, Context, GarnetError, Term}; @@ -76,13 +74,12 @@ struct IdMinter { buffer: Vec, } -/// The free space map manages an ID pool in Garnet to track which IDs are being -/// used. This is necessary to reuse deleted vectors IDs, which is needed due to -/// the small ID space of u32. +/// The free space map manages the user-vector ID pool in Garnet, including IDs +/// reclaimed after deletion and gaps left by imports. ID 0 is reserved for the +/// start point and is not tracked. /// -/// Users can call `next_id()` to get an ID which may be a newly minted ID or a -/// reused one from a previously deleted vector. Use `mark_used()` and -/// `mark_deleted()` to flag IDs as used or free respectively. +/// Use `next_id()` to allocate and mark an ID as used, `claim_id()` to claim an +/// imported ID, and `mark_free()` to release a deleted vector's ID. /// /// When quantization backfill is required, set `reuse_enabled` to `false` and /// `quantization_enabled` to true in the constructor. This will prevent all ID @@ -93,16 +90,16 @@ pub(crate) struct FreeSpaceMap { /// Garnet callbacks for reading/writing FSM keys callbacks: Callbacks, /// A flag to signal whether there are free IDs in the FSM. - /// This is set after a scan of the FSM, and is used to prevent extraneous reads - /// of FSM blocks. + /// Gaps and deletions set this flag; a scan with no free IDs clears it to + /// prevent extraneous reads of FSM blocks. has_free_ids: AtomicBool, - /// A queue of previously deleted IDs to prevent excessive reads to the FSM + /// A queue of reusable IDs to prevent excessive reads of the FSM fast_free_list: ArrayQueue, /// Controls minting new IDs and expanding the FSM blocks id_minter: RwLock, /// The total number of IDs marked used in the FSM total_used: AtomicUsize, - /// Controls when ID reuse for deleted IDs is enabled. + /// Controls when ID reuse is enabled. reuse_enabled: AtomicBool, /// Quantization backfill related parameters that must be synchronized. barrier: RwLock, @@ -126,7 +123,7 @@ impl FreeSpaceMap { quantization_enabled, }); let id_minter = RwLock::new(IdMinter { - next_id: 0, + next_id: 1, max_block: u32::MAX, buffer: vec![0u8; BLOCK_SIZE_BYTES], }); @@ -152,9 +149,8 @@ impl FreeSpaceMap { this.load_state(ctx)?; } else { // Allocate first block. - let (block_id, _, _) = this.indexes_for_id(0); - let mut id_minter = this.id_minter.write().unwrap(); - this.expand_to(&mut id_minter, ctx, block_id)?; + let mut id_minter = this.id_minter.write(); + this.expand_to(&mut id_minter, ctx, 1)?; } Ok(this) @@ -172,7 +168,7 @@ impl FreeSpaceMap { } let mut block = vec![0u8; BLOCK_SIZE_BYTES]; - let mut last_used_id = -1i64; + let mut last_used_id = 0u32; let mut total_used = 0usize; for block_id in (0..max_block_id).rev() { @@ -185,29 +181,33 @@ impl FreeSpaceMap { break; } - let mut id = block_id * BLOCK_SIZE_IDS as u32 + BLOCK_SIZE_IDS as u32 - 1; + let mut id = block_id * BLOCK_SIZE_IDS as u32 + (BLOCK_SIZE_IDS as u32 - 1); for &byte in block.iter().rev() { for bidx in (0..8).rev() { + if id == 0 { + break; + } + let used = bit_used(byte, bidx); if used { - last_used_id = last_used_id.max(id as i64); + last_used_id = last_used_id.max(id); total_used += 1; - } else if (id as i64) < last_used_id { + } else if id < last_used_id { let _ = self.fast_free_list.push(id); } - id = id.saturating_sub(1); + id -= 1; } } } - let mut id_minter = self.id_minter.write().unwrap(); + let mut id_minter = self.id_minter.write(); id_minter.max_block = max_block_id - 1; - id_minter.next_id = (last_used_id + 1) as u32; + id_minter.next_id = last_used_id + 1; - let barrier = self.barrier.get_mut().unwrap(); + let barrier = self.barrier.get_mut(); if barrier.quantization_enabled { barrier.max_id_for_backfill = id_minter.next_id.saturating_sub(1); } @@ -275,8 +275,8 @@ impl FreeSpaceMap { /// data changed or not. fn mark_id(&self, ctx: &Context, id: u32, used: bool) -> Result { { - let id_minter = self.id_minter.read().unwrap(); - if id >= id_minter.next_id { + let id_minter = self.id_minter.read(); + if id == 0 || id >= id_minter.next_id { return Err(FsmError::IdOutOfRange(id)); } } @@ -288,8 +288,8 @@ impl FreeSpaceMap { pub(crate) fn is_free(&self, ctx: &Context, id: u32) -> Result { // Don't hold the lock longer than we have to. let max_block = { - let id_minter = self.id_minter.read().unwrap(); - if id >= id_minter.next_id { + let id_minter = self.id_minter.read(); + if id == 0 || id >= id_minter.next_id { return Err(FsmError::IdOutOfRange(id)); } @@ -325,7 +325,7 @@ impl FreeSpaceMap { pub(crate) fn next_id(&self, ctx: &Context) -> Result, FsmError> { // A read barrier is acquired for the whole function. This prevents ID // minting during quantization phase changes. - let barrier = self.barrier.read().unwrap(); + let barrier = self.barrier.read(); if self.reuse_enabled.load(Ordering::Acquire) && self.has_free_ids.load(Ordering::Acquire) { // We retry reusing a freed ID until there are none or we get one and marking it used @@ -361,25 +361,54 @@ impl FreeSpaceMap { } // Mint a new ID and mark it used. - let mut id_minter = self.id_minter.write().unwrap(); + let mut id_minter = self.id_minter.write(); let id = id_minter.next_id; - id_minter.next_id += 1; + let next_id = id.checked_add(1).ok_or(FsmError::IdOutOfRange(id))?; self.expand_to(&mut id_minter, ctx, id)?; self.mark_id_unchecked(ctx, id, true)?; + id_minter.next_id = next_id; Ok(ReuseGuard::new(id, barrier)) } /// Guard writes to an existing ID without changing its allocation state. pub(crate) fn existing_id(&self, id: u32) -> ReuseGuard<'_> { - ReuseGuard::new(id, self.barrier.read().unwrap()) + ReuseGuard::new(id, self.barrier.read()) + } + + /// Claim an imported ID, creating any missing blocks and advancing the maximum ID. + /// + /// Repeated claims do not change the used count. The reserved start-point ID 0 + /// is ignored, and u32::MAX is rejected because its next ID is unrepresentable. + pub(crate) fn claim_id(&self, ctx: &Context, id: u32) -> Result<(), FsmError> { + if id == 0 { + return Ok(()); + } + let next_id = id.checked_add(1).ok_or(FsmError::IdOutOfRange(id))?; + let mut id_minter = self.id_minter.write(); + let grows = id >= id_minter.next_id; + + if grows { + self.expand_to(&mut id_minter, ctx, id)?; + } + + let _ = self.mark_id_unchecked(ctx, id, true)?; + + if grows { + if id > id_minter.next_id { + self.has_free_ids.store(true, Ordering::Release); + } + id_minter.next_id = next_id; + } + + Ok(()) } /// Return the maximum ID that has been assigned to a vector. /// /// This ID may be free if that ID has been deleted since the ID was created. pub(crate) fn max_id(&self) -> u32 { - let id_minter = self.id_minter.read().unwrap(); + let id_minter = self.id_minter.read(); id_minter.next_id.saturating_sub(1) } @@ -409,7 +438,7 @@ impl FreeSpaceMap { /// Scan the FSM blocks to fill up fast_free_list. fn refill_fast_free_list(&self, ctx: &Context) -> Result { // NOTE: We take a lock to prevent multiple refills happening simultaneously. - let _guard = self.refill_lock.lock().unwrap(); + let _guard = self.refill_lock.lock(); // If we had to wait to acquire the lock, it's possible some else refilled the list first, so check it again. if !self.fast_free_list.is_empty() { @@ -417,7 +446,7 @@ impl FreeSpaceMap { } let (max_block, next_id) = { - let id_minter = self.id_minter.read().unwrap(); + let id_minter = self.id_minter.read(); (id_minter.max_block, id_minter.next_id) }; @@ -455,7 +484,7 @@ impl FreeSpaceMap { break 'scan; } - if !bit_used(byte, bidx) { + if id != 0 && !bit_used(byte, bidx) { has_free_ids = true; self.has_free_ids.store(true, Ordering::Release); if self.fast_free_list.push(id).is_err() { @@ -474,31 +503,34 @@ impl FreeSpaceMap { Ok(has_free_ids) } - /// Ensure enough blocks exist in the FSM to hold `id`. + /// Create every missing block through the block containing `id`. + /// The initialized prefix advances only after each successful block write. fn expand_to( &self, id_minter: &mut RwLockWriteGuard, ctx: &Context, id: u32, ) -> Result<(), FsmError> { - let (block_id, _, _) = self.indexes_for_id(id); - if id_minter.max_block == u32::MAX || block_id == id_minter.max_block + 1 { + let (target_block, _, _) = self.indexes_for_id(id); + if id_minter.max_block != u32::MAX && target_block <= id_minter.max_block { + return Ok(()); + } + + let first_block = if id_minter.max_block == u32::MAX { + 0 + } else { + id_minter.max_block + 1 + }; + let ctx = ctx.term(Term::Metadata); + for block_id in first_block..=target_block { let block_key = Self::block_key(block_id); - if !self - .callbacks - .write_wid(&ctx.term(Term::Metadata), block_key, &id_minter.buffer) - { + if !self.callbacks.write_wid(&ctx, block_key, &id_minter.buffer) { return Err(FsmError::Garnet(GarnetError::Write)); } - if id_minter.max_block == u32::MAX { - id_minter.max_block = 0; - } else { - id_minter.max_block += 1; - } - } else if block_id > id_minter.max_block + 1 { - return Err(FsmError::IdOutOfRange(id)); + // Keep the successfully persisted prefix if a later write fails. + id_minter.max_block = block_id; } Ok(()) @@ -509,11 +541,10 @@ impl FreeSpaceMap { where F: FnMut(u32) -> bool, { - let max_block = { self.id_minter.read().unwrap().max_block }; + let max_block = { self.id_minter.read().max_block }; let mut block = vec![0u8; BLOCK_SIZE_BYTES]; - let mut id = 0u32; - for block_id in 0..max_block + 1 { + for block_id in 0..=max_block { let block_key = Self::block_key(block_id); if !self .callbacks @@ -522,20 +553,20 @@ impl FreeSpaceMap { return Err(FsmError::Garnet(GarnetError::Read)); } - for &byte in &block { + let first_id = block_id * BLOCK_SIZE_IDS as u32; + for (byte_idx, &byte) in block.iter().enumerate() { if byte == 0x00 { - id += 8; continue; } + let byte_id = first_id + (byte_idx * 8) as u32; for bidx in 0..8 { if bit_used(byte, bidx) { - let keep_going = f(id); - if !keep_going { + let id = byte_id + bidx as u32; + if id != 0 && !f(id) { return Ok(()); } } - id += 1; } } } @@ -545,13 +576,19 @@ impl FreeSpaceMap { /// Signal that new IDs should be quantized. pub(crate) fn enable_quantization(&self) { - // Taking the write barrier ensures that we switch quantization phases when - // 1. No `next_id` is not changing (because `next_id()` operates under a read barrier) - // 2. All vector data for any pending unquantized inserts has been written. (because - // `insert()` updates data under the same read barrier via `ReuseGuard`) - let mut guard = self.barrier.write().unwrap(); + self.enable_quantization_if(|| true); + } + + /// Prepare a quantization change while under the barrier, then enable it on success. + /// `prepare` must not acquire the FSM barrier. + pub(crate) fn enable_quantization_if(&self, prepare: impl FnOnce() -> bool) -> bool { + let mut guard = self.barrier.write(); + if !prepare() { + return false; + } guard.max_id_for_backfill = self.max_id(); guard.quantization_enabled = true; + true } /// Allow reuse of previously deleted IDs. @@ -561,7 +598,7 @@ impl FreeSpaceMap { /// Return the max ID for purposes of backfilling quantized vectors. pub(crate) fn max_id_for_backfill(&self) -> u32 { - self.barrier.read().unwrap().max_id_for_backfill + self.barrier.read().max_id_for_backfill } } @@ -582,39 +619,12 @@ fn update_status(used: bool, byte: &mut u8, bidx: usize) -> bool { #[cfg(test)] mod tests { - use std::sync::atomic::Ordering; - use crate::{ fsm::{BLOCK_SIZE_BYTES, BLOCK_SIZE_IDS, FreeSpaceMap, FsmError}, - garnet::{Context, Term}, + garnet::{Callbacks, Context, GarnetError, Term}, test_utils::Store, }; - fn verify_block(store: &Store, block_id: u32, mut f: F) - where - F: FnMut(&[u8]), - { - let ctx = Context::new(0); - let key = FreeSpaceMap::block_key(block_id); - let block = store - .get(ctx.term(Term::Metadata).get(), bytemuck::bytes_of(&key)) - .unwrap(); - f(&block); - } - - #[test] - fn create_fresh() { - let store = Store::new(); - let ctx = Context::new(0); - - // A fresh FSM should result in full block of all zeroes. - let _fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); - verify_block(&store, 0, |block| { - assert_eq!(block.len(), BLOCK_SIZE_BYTES); - assert_eq!(block.iter().filter(|b| **b != 0).count(), 0); - }); - } - #[test] fn basic_next_id() { let store = Store::new(); @@ -622,57 +632,29 @@ mod tests { let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); - // A fresh FSM should not have free IDs. - assert!(!fsm.has_free_ids.load(std::sync::atomic::Ordering::Acquire)); - - // It should return sequentials IDs starting from 0. - assert_eq!(fsm.next_id(&ctx).unwrap().id(), 0); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); } #[test] fn basic_delete() { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); - // Deleting a ID out of range should return an error. - let res = fsm.mark_free(&ctx, 0); - assert!(res.is_err()); - assert_eq!(res.err().unwrap(), FsmError::IdOutOfRange(0)); - - for _ in 0u32..64 { - let _ = fsm.next_id(&ctx).unwrap(); + assert_eq!(fsm.mark_free(&ctx, 1), Err(FsmError::IdOutOfRange(1))); + for id in 1..=64 { + assert_eq!(fsm.next_id(&ctx).unwrap().id(), id); } - verify_block(&store, 0, |block| { - for b in &block[0..8] { - assert_eq!(*b, 0b11111111); - } - for b in &block[8..BLOCK_SIZE_BYTES] { - assert_eq!(*b, 0b00000000); - } - }); - fsm.mark_free(&ctx, 37).unwrap(); fsm.mark_free(&ctx, 9).unwrap(); - - verify_block(&store, 0, |block| { - assert_eq!(block[0], 0b11111111); - assert_eq!(block[1], 0b10111111); - - for b in &block[2..4] { - assert_eq!(*b, 0b11111111); - } - assert_eq!(block[4], 0b11111011); - for b in &block[5..8] { - assert_eq!(*b, 0b11111111); - } - for b in &block[8..BLOCK_SIZE_BYTES] { - assert_eq!(*b, 0b00000000); - } - }); + fsm.mark_free(&ctx, 37).unwrap(); + assert_eq!(fsm.total_used(), 62); + for id in 1..=64 { + assert_eq!(fsm.is_free(&ctx, id).unwrap(), id == 9 || id == 37); + } + assert_eq!(fsm.is_free(&ctx, 65), Err(FsmError::IdOutOfRange(65))); } #[test] @@ -686,16 +668,13 @@ mod tests { let _ = fsm.next_id(&ctx).unwrap(); } - // Next ID should be 64 since none are free before that. - assert_eq!(fsm.next_id(&ctx).unwrap().id(), 64); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 65); // After deleting an ID, it should be returned from `next_id()`. fsm.mark_free(&ctx, 37).unwrap(); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 37); // Once all free IDs are used, fresh ones should be returned. - assert_eq!(fsm.next_id(&ctx).unwrap().id(), 65); - // And has_free_ids should now be false. - assert!(!fsm.has_free_ids.load(Ordering::Acquire)); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 66); } #[test] @@ -713,12 +692,10 @@ mod tests { // Loading FSM from store should recover all the state. let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); - assert_eq!(fsm.id_minter.read().unwrap().max_block, 0); - assert_eq!(fsm.max_id() + 1, 64); - assert!(fsm.has_free_ids.load(Ordering::Acquire)); - assert_eq!(fsm.fast_free_list.len(), 1); - assert_eq!(fsm.max_id_for_backfill(), u32::MAX); + assert_eq!(fsm.max_id(), 64); + assert_eq!(fsm.total_used(), 63); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 37); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 65); } #[test] @@ -735,18 +712,229 @@ mod tests { drop(fsm); let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), true, false).unwrap(); - assert_eq!(fsm.max_id_for_backfill(), 63); + assert_eq!(fsm.max_id_for_backfill(), 64); let next_id = fsm.next_id(&ctx).unwrap(); - assert_eq!(next_id.id(), 64); + assert_eq!(next_id.id(), 65); assert!(next_id.should_quantize()); drop(next_id); - assert_eq!(fsm.max_id_for_backfill(), 63); + assert_eq!(fsm.max_id_for_backfill(), 64); fsm.enable_reuse(); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 37); } + #[test] + fn claim_out_of_order_ids() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let block_size = BLOCK_SIZE_IDS as u32; + let high_id = 3 * block_size + 7; + + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + for id in [ + high_id, + block_size + 5, + block_size, + block_size - 1, + 1, + high_id, + ] { + fsm.claim_id(&ctx, id).unwrap(); + } + assert_eq!(fsm.max_id(), high_id); + assert_eq!(fsm.total_used(), 5); + for id in [2, block_size + 1, 2 * block_size, high_id - 1] { + assert!(fsm.is_free(&ctx, id).unwrap(), "id={id}"); + } + let mut used = Vec::new(); + fsm.visit_used(&ctx, |id| { + used.push(id); + true + }) + .unwrap(); + assert_eq!( + used, + [1, block_size - 1, block_size, block_size + 5, high_id] + ); + let empty_block = store + .get( + ctx.term(Term::Metadata).get(), + &FreeSpaceMap::block_key(2).to_ne_bytes(), + ) + .unwrap(); + assert_eq!(empty_block.len(), BLOCK_SIZE_BYTES); + assert!(empty_block.iter().all(|byte| *byte == 0)); + fsm.mark_free(&ctx, high_id).unwrap(); + assert_eq!(fsm.total_used(), 4); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), high_id); + assert_eq!(fsm.total_used(), 5); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); + drop(fsm); + + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + assert_eq!(fsm.max_id(), high_id); + assert_eq!(fsm.total_used(), 6); + used.clear(); + fsm.visit_used(&ctx, |id| { + used.push(id); + true + }) + .unwrap(); + assert_eq!( + used, + [1, 2, block_size - 1, block_size, block_size + 5, high_id] + ); + } + + #[test] + fn start_point_is_not_tracked() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + + fsm.claim_id(&ctx, 0).unwrap(); + assert_eq!(fsm.total_used(), 0); + assert_eq!(fsm.max_id(), 0); + assert_eq!(fsm.is_free(&ctx, 0), Err(FsmError::IdOutOfRange(0))); + assert_eq!(fsm.mark_free(&ctx, 0), Err(FsmError::IdOutOfRange(0))); + fsm.visit_used(&ctx, |id| panic!("unexpected used ID {id}")) + .unwrap(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + drop(fsm); + + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + assert_eq!(fsm.total_used(), 1); + fsm.mark_free(&ctx, 1).unwrap(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); + assert_eq!(fsm.total_used(), 2); + } + + #[test] + fn recovery_ignores_start_point_bit() { + let store = Store::new(); + let ctx = Context::new(0); + let mut block = vec![0u8; BLOCK_SIZE_BYTES]; + block[0] = 0b11000000; + store.set( + ctx.term(Term::Metadata).get(), + &FreeSpaceMap::block_key(0).to_ne_bytes(), + &block, + ); + + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + assert_eq!(fsm.total_used(), 1); + assert_eq!(fsm.max_id(), 1); + let mut used = Vec::new(); + fsm.visit_used(&ctx, |id| { + used.push(id); + true + }) + .unwrap(); + assert_eq!(used, [1]); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); + } + + #[test] + fn expansion_write_failure_preserves_claims() { + unsafe extern "C" fn write( + context: u64, + key: *const u8, + key_len: usize, + value: *const u8, + value_len: usize, + ) -> bool { + let key_bytes = unsafe { std::slice::from_raw_parts(key, key_len) }; + if key_bytes == FreeSpaceMap::block_key(2).to_ne_bytes() { + return false; + } + unsafe { + (Store::attach().callbacks().write_callback())( + context, key, key_len, value, value_len, + ) + } + } + + for reload in [false, true] { + let store = Store::new(); + let ctx = Context::new(0); + let callbacks = store.callbacks(); + let failing_callbacks = Callbacks::new( + callbacks.read_callback(), + write, + callbacks.delete_callback(), + callbacks.rmw_callback(), + callbacks.filter_callback(), + callbacks.log_callback(), + ); + let mut fsm = FreeSpaceMap::new(&ctx, failing_callbacks, false, true).unwrap(); + let high_id = 3 * BLOCK_SIZE_IDS as u32 + 7; + assert_eq!( + fsm.claim_id(&ctx, high_id), + Err(FsmError::Garnet(GarnetError::Write)) + ); + assert_eq!(fsm.max_id(), 0); + assert_eq!(fsm.total_used(), 0); + + let low_id = BLOCK_SIZE_IDS as u32 + 5; + fsm.claim_id(&ctx, low_id).unwrap(); + if reload { + fsm = FreeSpaceMap::new(&ctx, callbacks, false, true).unwrap(); + } else { + fsm.callbacks = callbacks; + } + fsm.claim_id(&ctx, high_id).unwrap(); + assert!(!fsm.is_free(&ctx, low_id).unwrap()); + assert!(!fsm.is_free(&ctx, high_id).unwrap()); + assert!(fsm.is_free(&ctx, 2 * BLOCK_SIZE_IDS as u32).unwrap()); + assert_eq!(fsm.total_used(), 2); + assert_eq!(fsm.max_id(), high_id); + } + } + + #[test] + fn sparse_claims_respect_backfill_barrier() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, false).unwrap(); + let high_id = 2 * BLOCK_SIZE_IDS as u32; + fsm.claim_id(&ctx, high_id).unwrap(); + fsm.enable_quantization(); + + let next_id = fsm.next_id(&ctx).unwrap(); + assert_eq!(next_id.id(), high_id + 1); + assert!(next_id.should_quantize()); + drop(next_id); + assert_eq!(fsm.max_id_for_backfill(), high_id); + fsm.enable_reuse(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + assert_eq!(fsm.total_used(), 3); + } + + #[test] + fn exhausted_id_space_does_not_wrap() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + assert_eq!( + fsm.claim_id(&ctx, u32::MAX), + Err(FsmError::IdOutOfRange(u32::MAX)) + ); + assert_eq!(fsm.max_id(), 0); + assert_eq!(fsm.total_used(), 0); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + + fsm.id_minter.write().next_id = u32::MAX; + assert!(matches!( + fsm.next_id(&ctx), + Err(FsmError::IdOutOfRange(u32::MAX)) + )); + fsm.mark_free(&ctx, 1).unwrap(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + } + #[test] fn dynamic_expansion() { let store = Store::new(); @@ -755,8 +943,47 @@ mod tests { let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); // Asking for more than BLOCK_SIZE_IDS will force another FSM block to be allocated. - for _ in 0u32..BLOCK_SIZE_IDS as u32 + 1 { - let _ = fsm.next_id(&ctx).unwrap(); + for id in 1..=BLOCK_SIZE_IDS as u32 + 1 { + assert_eq!(fsm.next_id(&ctx).unwrap().id(), id); } + assert_eq!(fsm.max_id(), BLOCK_SIZE_IDS as u32 + 1); + assert_eq!(fsm.total_used(), BLOCK_SIZE_IDS + 1); + for id in BLOCK_SIZE_IDS as u32 - 1..=BLOCK_SIZE_IDS as u32 + 1 { + assert!(!fsm.is_free(&ctx, id).unwrap()); + } + fsm.mark_free(&ctx, BLOCK_SIZE_IDS as u32).unwrap(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), BLOCK_SIZE_IDS as u32); + } + + #[test] + fn recover_sparse_claims() { + let store = Store::new(); + let ctx = Context::new(0); + let block_size = BLOCK_SIZE_IDS as u32; + let high_id = 2 * block_size + 7; + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + fsm.claim_id(&ctx, high_id).unwrap(); + fsm.claim_id(&ctx, block_size).unwrap(); + drop(fsm); + + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, false).unwrap(); + assert_eq!(fsm.max_id(), high_id); + assert_eq!(fsm.total_used(), 3); + let mut visited = Vec::new(); + fsm.visit_used(&ctx, |id| { + visited.push(id); + true + }) + .unwrap(); + assert_eq!(visited, [1, block_size, high_id]); + assert!(fsm.is_free(&ctx, block_size + 1).unwrap()); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), high_id + 1); + + fsm.enable_reuse(); + let reused = fsm.next_id(&ctx).unwrap().id(); + assert!(reused > 0 && reused < high_id && !visited.contains(&reused)); + assert_eq!(fsm.max_id(), high_id + 1); + assert_eq!(fsm.total_used(), 5); } } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 450e1998b3..3796e3f064 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -31,7 +31,7 @@ use diskann_vector::distance::Metric; use crate::{ alloc::AlignToEight, - garnet::{FilterCallback, LogCallback}, + garnet::{FilterCallback, LogCallback, Term}, provider::{GarnetProvider, GarnetProviderError}, }; use crate::{ @@ -525,6 +525,23 @@ impl From for InsertResult { } } +/// Borrows FFI bytes, accepting a null pointer only for an empty slice. +/// +/// # Safety +/// +/// For a nonzero length, a non-null pointer must satisfy `slice::from_raw_parts` +/// for the returned lifetime. +#[inline] +unsafe fn interpret_bytes_allow_empty<'a>(data: *const u8, len: usize) -> Option<&'a [u8]> { + if len == 0 { + Some(&[]) + } else if data.is_null() { + None + } else { + Some(unsafe { slice::from_raw_parts(data, len) }) + } +} + /// Insert a vector into the index. /// /// Returns a status corresponding to the `InsertResult` enum, distinguishing new inserts from @@ -547,7 +564,9 @@ pub unsafe extern "C" fn insert( let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(ctx); - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return InsertResult::Fail.into(); + }; let id = GarnetId::from(id_bytes); let v = if let Some(v) = interpret_vector(index.quant_type, &vector_data, vector_len) { @@ -656,14 +675,14 @@ pub unsafe extern "C" fn build_quant_table(context: u64, index_ptr: *const c_voi pub unsafe extern "C" fn backfill_quant_vectors( context: u64, index_ptr: *const c_void, - task_index: usize, + task_idx: usize, task_count: usize, ) -> bool { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(context); index .inner - .backfill_quant_vectors(&ctx, task_index, task_count) + .backfill_quant_vectors(&ctx, task_idx, task_count) } /// Set the attributes for a vector. @@ -686,7 +705,10 @@ pub unsafe extern "C" fn set_attribute( ) -> bool { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(context); - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return false; + }; let id = GarnetId::from(id_bytes); // Check if the vector exists @@ -939,7 +961,10 @@ pub unsafe extern "C" fn search_element( overflow: *mut *mut c_void, ) -> i32 { let index = unsafe { &*index_ptr.cast::() }; - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return -1; + }; let id = GarnetId::from(id_bytes); let ctx = Context::new(ctx); @@ -1071,7 +1096,10 @@ pub unsafe extern "C" fn remove( ) -> bool { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(ctx); - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return false; + }; let id = GarnetId::from(id_bytes); if !index.inner.external_id_exists(&ctx, &id) { @@ -1083,14 +1111,17 @@ pub unsafe extern "C" fn remove( /// Return the approximate count of vectors in the index. /// +/// Returns `u64::MAX` if disabling imports fails. +/// /// # Safety /// /// FFI #[unsafe(no_mangle)] -pub unsafe extern "C" fn card(_ctx: u64, index_ptr: *const c_void) -> u64 { +pub unsafe extern "C" fn card(ctx: u64, index_ptr: *const c_void) -> u64 { let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); - index.inner.approximate_count() + index.inner.approximate_count(&ctx).unwrap_or(u64::MAX) } /// Check if a given internal ID is a valid vector. @@ -1135,7 +1166,10 @@ pub unsafe extern "C" fn check_external_id_valid( ) -> bool { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(ctx); - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return false; + }; let id = GarnetId::from(id_bytes); index.inner.external_id_exists(&ctx, &id) @@ -1198,7 +1232,10 @@ pub unsafe extern "C" fn search_neighbors( ) -> i32 { let index = unsafe { &*index_ptr.cast::() }; let ctx = Context::new(ctx); - let id_bytes = unsafe { slice::from_raw_parts(id_data, id_len) }; + + let Some(id_bytes) = (unsafe { interpret_bytes_allow_empty(id_data, id_len) }) else { + return -1; + }; let id = GarnetId::from(id_bytes); let mut output = SearchResults::new( @@ -1232,6 +1269,179 @@ pub unsafe extern "C" fn search_neighbors( count as i32 } +/// Sets the quantizer internal state. +/// +/// This is intended to be used on an empty, quantized index to set the quantizer state to a +/// known value. This is called from the `XVCREATE` command. +/// +/// This ensures quantized vectors will be externally intelligible to another system with +/// the same quantizer state. +/// +/// Returns false if importing has begun, the index is non-empty, or setup fails. On failure, +/// discard the vector set: the quantizer and persisted metadata may already have changed. +/// +/// # Safety +/// +/// FFI +#[unsafe(no_mangle)] +pub unsafe extern "C" fn set_quant_state( + ctx: u64, + index_ptr: *const c_void, + state: *const u8, + state_len: usize, +) -> bool { + if index_ptr.is_null() { + return false; + } + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + if state.is_null() || state_len == 0 { + return false; + } + let state = unsafe { slice::from_raw_parts(state, state_len) }; + + index.inner.set_quant_state(&ctx, state) +} + +/// Returns whether the index accepts imported terms. +/// +/// # Safety +/// +/// `index_ptr` must be a live index returned by `create_index`, and `ctx` must identify its storage. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn can_import(ctx: u64, index_ptr: *const c_void) -> bool { + if index_ptr.is_null() { + return false; + } + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + + index.inner.can_import(&ctx) +} + +/// Imports a term directly in storage. +/// +/// This is intended for bulk loading of existing DiskANN graphs that are stored outside of Garnet, +/// allowing for creation of a vector set to happen without calling `VADD` on every vector. +/// Ordinary index operations will fail once importing begins until finalization succeeds. +/// +/// # Safety +/// +/// FFI +#[unsafe(no_mangle)] +pub unsafe extern "C" fn import_term( + ctx: u64, + index_ptr: *const c_void, + term_type: u32, + id: *const u8, + id_len: usize, + value: *const u8, + value_len: usize, +) -> bool { + if index_ptr.is_null() { + return false; + } + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + + let term = if let Ok(term) = Term::try_from(term_type) { + term + } else { + return false; + }; + + if (id_len == 0 && !matches!(term, Term::IntMap)) + || (value_len == 0 && !matches!(term, Term::ExtMap)) + { + return false; + } + let Some(id) = (unsafe { interpret_bytes_allow_empty(id, id_len) }) else { + return false; + }; + let Some(value) = (unsafe { interpret_bytes_allow_empty(value, value_len) }) else { + return false; + }; + + index.inner.import_term(&ctx, term, id, value) +} + +/// Return type for `finish_import()`. +/// +/// `Success` is returned if the tasks succeeds, and if the final task, if the finalization +/// succeeds. `TaskFailed` indicates that the task's portion of the work failed, but can be +/// retried. `FinishFailed` means the task finished successfully but finalization failed; the +/// task cannot be retried. +#[derive(Debug, Clone, Copy, PartialEq)] +enum ImportResult { + Success, + TaskFailed, + FinishFailed, +} + +impl From for u8 { + fn from(value: ImportResult) -> Self { + match value { + ImportResult::Success => 0, + ImportResult::TaskFailed => 1, + ImportResult::FinishFailed => 2, + } + } +} + +#[cfg(test)] +impl From for ImportResult { + fn from(value: u8) -> Self { + match value { + 0 => ImportResult::Success, + 1 => ImportResult::TaskFailed, + 2 => ImportResult::FinishFailed, + _ => panic!("bad ImportResult ({value})"), + } + } +} + +/// Finalize import. +/// +/// After terms are imported, this is used to verify the graph and enable the index. +/// +/// Returns 0 if the task succeeds and if the final task finalizes the index. Returns 1 when +/// the task failed, which can be retried. Returns 2 when the task succeeded but index +/// finalization has failed, which cannot be retried. +/// +/// # Safety +/// +/// FFI +#[unsafe(no_mangle)] +pub unsafe extern "C" fn finish_import( + ctx: u64, + index_ptr: *const c_void, + task_idx: usize, + task_count: usize, +) -> u8 { + if index_ptr.is_null() { + return ImportResult::TaskFailed.into(); + } + + let index = unsafe { &*index_ptr.cast::() }; + let ctx = Context::new(ctx); + + let (result, completed) = index.inner.finish_import(&ctx, task_idx, task_count); + if result && completed == Some(true) { + index + .state + .store(IndexState::Ready as usize, Ordering::Release); + } + + match (result, completed) { + (true, None) | (true, Some(true)) => ImportResult::Success, + (false, None) | (false, Some(false)) => ImportResult::TaskFailed, + // This variant is a special case that happens when an empty index is finalized. + (false, Some(true)) => ImportResult::Success, + (true, Some(false)) => ImportResult::FinishFailed, + } + .into() +} + #[cfg(test)] mod tests { use std::{mem, ptr}; @@ -1418,7 +1628,7 @@ mod tests { assert!(!index_ptr.is_null()); let index = unsafe { &*index_ptr.cast::() }; assert_eq!(index.quant_type, quant_type); - assert_eq!(index.inner.approximate_count(), 0); + assert_eq!(index.inner.approximate_count(&Context::new(0)).unwrap(), 0); unsafe { drop_index(0, index_ptr); diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 0e2ca5557f..39bc3d0cf6 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -31,6 +31,7 @@ use diskann_utils::{ use diskann_vector::{ DistanceFunction, PreprocessedDistanceFunction, contains::ContainsSimd, distance::Metric, }; +use parking_lot::{Condvar, Mutex, MutexGuard}; use std::{ any::TypeId, collections::HashSet, @@ -40,8 +41,8 @@ use std::{ mem, ops::{Deref, DerefMut, Range}, sync::{ - Arc, Condvar, Mutex, - atomic::{AtomicBool, AtomicU64, Ordering}, + Arc, OnceLock, + atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering}, }, }; use strum::VariantArray; @@ -65,6 +66,54 @@ use crate::{ /// bytes are the serialized quant table. const QUANT_STATE_KEY: u32 = u32::from_be_bytes(*b"_qnt"); +/// Persisted [`ImportStatus`] discriminant, stored as a single metadata byte. +const IMPORT_STATUS_KEY: u32 = u32::from_be_bytes(*b"_imp"); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +enum ImportStatus { + Disabled = 0, + Ready = 1, + Importing = 2, +} + +impl TryFrom for ImportStatus { + type Error = GarnetProviderError; + + fn try_from(value: u8) -> Result { + match value { + 0 => Ok(Self::Disabled), + 1 => Ok(Self::Ready), + 2 => Ok(Self::Importing), + _ => Err(GarnetProviderError::InvalidImportStatus(value)), + } + } +} + +struct AtomicImportStatus(AtomicU8); + +impl AtomicImportStatus { + const fn new(status: ImportStatus) -> Self { + Self(AtomicU8::new(status as u8)) + } + + fn load(&self) -> ImportStatus { + ImportStatus::try_from(self.0.load(Ordering::Acquire)) + .expect("atomic import status must contain a valid discriminant") + } + + fn compare_exchange(&self, current: ImportStatus, new: ImportStatus) -> bool { + self.0 + .compare_exchange( + current as u8, + new as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + } +} + /// Starting capacity of the pre-allocated rerank buffers. const RERANK_BUFFER_LENGTH: usize = 1024; @@ -114,6 +163,8 @@ pub(crate) enum GarnetProviderError { AllocFailed(#[from] AllocatorError), #[error("Invalid quantizer for vector data")] InvalidQuantizer, + #[error("Invalid import status: {0}")] + InvalidImportStatus(u8), #[error("Quantizer error: {0}")] Quantizer(#[from] GarnetQuantizerError), #[error("Post processing error: {0}")] @@ -133,7 +184,7 @@ struct BackfillGuard { impl Drop for BackfillGuard { fn drop(&mut self) { - let _ = self.ranges.lock().unwrap().remove(&self.range); + let _ = self.ranges.lock().remove(&self.range); self.notify.notify_all(); } } @@ -261,6 +312,12 @@ pub(crate) struct GarnetProvider { backfill_notify: Arc, /// Lock to ensure training only happens once. training_lock: Mutex<()>, + /// Signals whether imports are allowed, in progress, or disabled. + import_status: AtomicImportStatus, + /// Serializes state transition for imports. + import_transition: Mutex<()>, + /// Per job tracker for import verification + imports_completed: OnceLock>, /// Pool of pre-allocated buffers to use for neighbor lists id_buffer_pool: ObjectPool, /// Pool of pre-allocated buffers to use for IDs @@ -333,6 +390,8 @@ impl GarnetProvider { neighbor_cache.insert(0, neighbors); } + let mut has_quant_state = false; + let mut default_quant_state = None; let (quantizer, canonical_bytes, all_quantized) = match quant_type { VectorQuantType::NoQuant | VectorQuantType::XNoQuantU8 @@ -346,6 +405,7 @@ impl GarnetProvider { let quantizer = if let Some(quant_state) = callbacks.read_varsize_iid::(&context.term(Term::Metadata), QUANT_STATE_KEY) { + has_quant_state = true; quantization::MinMax8Bit::new_from_bytes(metric_type, &quant_state)? } else { if start_point_cache.contains_key(&0) { @@ -354,13 +414,7 @@ impl GarnetProvider { } let quantizer = quantization::MinMax8Bit::new(dim, metric_type)?; - if !callbacks.write_iid( - &context.term(Term::Metadata), - QUANT_STATE_KEY, - &quantizer.serialize()?, - ) { - return Err(GarnetError::Write.into()); - } + default_quant_state = Some(quantizer.serialize()?); quantizer }; @@ -380,8 +434,8 @@ impl GarnetProvider { (Some(quantizer), canonical_bytes, true) } VectorQuantType::Bin | VectorQuantType::XBinU8 | VectorQuantType::XBinI8 => { - let quantizer = - Box::new(quantization::Spherical1Bit::new(dim)) as Box; + let quantizer = Box::new(quantization::Spherical1Bit::new(metric_type, dim)) + as Box; let canonical_bytes = quantizer.bytes(); let mut all_quantized = false; @@ -391,16 +445,18 @@ impl GarnetProvider { if total_quant_state.len() <= 1 { return Err(GarnetProviderError::InvalidQuantizer); } + has_quant_state = true; all_quantized = total_quant_state[0] != 0; quantizer.deserialize(&total_quant_state[1..])?; - // Cache the saved start point, which should already exist if quantization is complete + // Cache the saved start point, which should already exist if quantization is complete, + // unless we have preset quantization state and an empty index (no start point at all). let mut qsv = Poly::broadcast(0u8, canonical_bytes, AlignToEight)?; if callbacks.read_single_iid(&context.term(Term::Quantized), 0, &mut qsv) { start_point_quant_cache.insert(0, qsv); - } else if all_quantized { + } else if all_quantized && start_point_cache.contains_key(&0) { return Err(GarnetProviderError::StartPoint); } } @@ -420,6 +476,39 @@ impl GarnetProvider { quantizer.is_none() || all_quantized, )?; + let mut persisted_status = [0u8]; + let import_status = if callbacks.read_single_iid( + &context.term(Term::Metadata), + IMPORT_STATUS_KEY, + &mut persisted_status, + ) { + ImportStatus::try_from(persisted_status[0])? + } else { + let status = if start_point_cache.is_empty() + && fsm.total_used() == 0 + && (quantizer.is_none() || has_quant_state) + { + ImportStatus::Ready + } else { + ImportStatus::Disabled + }; + if !callbacks.write_iid( + &context.term(Term::Metadata), + IMPORT_STATUS_KEY, + &[status as u8], + ) { + return Err(GarnetError::Write.into()); + } + status + }; + + // Persist quantizer state + if let Some(state) = default_quant_state + && !callbacks.write_iid(&context.term(Term::Metadata), QUANT_STATE_KEY, &state) + { + return Err(GarnetError::Write.into()); + } + Ok(Self { dim, metric_type, @@ -432,6 +521,9 @@ impl GarnetProvider { backfill_lock: Arc::new(Mutex::new(HashSet::new())), backfill_notify: Arc::new(Condvar::new()), training_lock: Mutex::new(()), + import_status: AtomicImportStatus::new(import_status), + import_transition: Mutex::new(()), + imports_completed: OnceLock::new(), id_buffer_pool, filtered_ids_pool, filtered_decisions_pool, @@ -475,12 +567,12 @@ impl GarnetProvider { return None; } - let mut ranges = self.backfill_lock.lock().unwrap(); + let mut ranges = self.backfill_lock.lock(); while ranges .iter() .any(|active| active.start < range.end && range.start < active.end) { - ranges = self.backfill_notify.wait(ranges).unwrap(); + self.backfill_notify.wait(&mut ranges); } let _ = ranges.insert(range.clone()); Some(BackfillGuard { @@ -498,6 +590,7 @@ impl GarnetProvider { context: &Context, point: &[T], ) -> Result<(), GarnetProviderError> { + let _guard = self.fsm.existing_id(0); let mut v = Poly::broadcast(0u8, self.dim * mem::size_of::(), AlignToEight)?; if self .callbacks @@ -534,13 +627,6 @@ impl GarnetProvider { } else { let neighbors = vec![0u32; self.max_degree + 1]; - // Grab the start point id, which must be zero. - let id = self.fsm.next_id(context)?; - if id.id() != 0 { - self.fsm.mark_free(context, id.id())?; - return Err(GarnetProviderError::StartPoint); - } - if !self .callbacks .write_iid(&context.term(Term::Vector), 0, point) @@ -652,6 +738,10 @@ impl GarnetProvider { } pub(crate) fn vector_iid_exists(&self, context: &Context, id: u32) -> bool { + if id == 0 { + return self.start_points_exist(); + } + !self.fsm.is_free(context, id).unwrap_or(true) } @@ -659,6 +749,10 @@ impl GarnetProvider { self.fsm.max_id() } + pub(crate) fn total_used(&self) -> usize { + self.fsm.total_used() + } + pub(crate) fn max_degree(&self) -> usize { self.max_degree } @@ -672,9 +766,8 @@ impl GarnetProvider { /// ensure that training happens only once. pub(crate) fn train_quantizer(&self, context: &Context) -> bool { // Ensure we don't kick off training twice. - let _training_guard = match self.training_lock.try_lock() { - Ok(g) => g, - Err(_) => return false, + let Some(_training_guard) = self.training_lock.try_lock() else { + return false; }; let quantizer = match &self.quantizer { @@ -690,11 +783,6 @@ impl GarnetProvider { if self .fsm .visit_used(context, |id| { - // Skip the start point. - if id == 0 { - return true; - } - if row_idx >= rows { return false; } @@ -764,6 +852,20 @@ impl GarnetProvider { } } + /// Determine ID range for a parallel task. + fn task_range(&self, task_idx: usize, task_count: usize, max_id: u32) -> Option<(u32, u32)> { + if max_id == u32::MAX || task_count == 0 || task_idx >= task_count { + return None; + } + + let id_count = max_id as usize + 1; + let work_count = id_count.div_ceil(task_count); + let start_id = work_count.saturating_mul(task_idx).min(id_count); + let end_id = start_id.saturating_add(work_count).min(id_count); + + Some((start_id as u32, end_id as u32)) + } + /// Bulk quantize previously inserted vectors. /// /// This function will be invoked on multiple threads. The total number of tasks and the ID of @@ -785,31 +887,17 @@ impl GarnetProvider { } }; - let max_id = self.fsm.max_id_for_backfill() as usize; - if max_id >= u32::MAX as usize { - // The max_id was somehow not sampled, so bail. - self.callbacks.log( - &context.term(Term::Quantized), - "Error: backfill_quant_vectors: Couldn't calculate max id to backfill. Index will operate full precision only mode.", - ); - return false; - } - - // If we have more tasks than vectors to backfill, we exit the extra tasks early. - let task_count = task_count.min(max_id + 1); - if task_idx >= task_count { + let (start_id, end_id) = if let Some((start, end)) = + self.task_range(task_idx, task_count, self.fsm.max_id_for_backfill()) + { + (start, end) + } else { self.callbacks.log( &context.term(Term::Quantized), - "Error: backfill_quant_vectors: Bad task index. Index will operate full precision only mode.", + "Error: backfill_quant_vectors: Bad task split. Index will operate full precision only mode.", ); return false; - } - - // Evenly divide the ID range from 0..max_id and determine this thread's backfill - // range. - let work_count = (max_id + 1).div_ceil(task_count); // will be >= 1 - let start_id = (work_count * task_idx) as u32; - let end_id = (work_count * (task_idx + 1)).min(max_id + 1) as u32; + }; let _backfill_guard = self.reserve_backfill_range(start_id..end_id); let mut v = vec![T::default(); self.dim]; @@ -974,6 +1062,446 @@ impl GarnetProvider { self.callbacks.log(context, msg); } + /// Set the quantizer state. + pub(crate) fn set_quant_state(&self, context: &Context, state: &[u8]) -> bool { + let Some(quantizer) = &self.quantizer else { + return false; + }; + + // Lock order: import transition, then the FSM barrier excluding vector writes. + let transition = self.import_transition.lock(); + let current = self.import_status.load(); + if current == ImportStatus::Importing { + return false; + } + let prepared = self.fsm.enable_quantization_if(|| { + if self.fsm.total_used() != 0 + || self.callbacks.exists_iid( + &context.term(Term::Vector), + 0, + self.dim * mem::size_of::(), + ) + { + return false; + } + + if quantizer.deserialize(state).is_err() { + return false; + } + let mut tmp = Vec::new(); + let total_quant_state = if quantizer.required_vectors() > 0 { + tmp.reserve_exact(state.len() + 1); + tmp.push(1); + tmp.extend_from_slice(state); + tmp.as_slice() + } else { + state + }; + + if !self.callbacks.write_iid( + &context.term(Term::Metadata), + QUANT_STATE_KEY, + total_quant_state, + ) { + return false; + } + + self.fsm.enable_reuse(); + self.all_quantized.store(true, Ordering::Release); + true + }); + prepared && self.publish_import_status(context, current, ImportStatus::Ready, &transition) + } + + /// Returns whether imports are enabled + pub(crate) fn can_import(&self, _context: &Context) -> bool { + self.import_status.load() != ImportStatus::Disabled + } + + /// Disables imports. + /// + /// Returns false on failure. + /// + /// This is called on most calls that use the index so that imports can only happen on + /// empty indices. + #[inline] + pub(crate) fn disable_import(&self, context: &Context) -> bool { + match self.import_status.load() { + ImportStatus::Disabled => true, + ImportStatus::Importing => false, + ImportStatus::Ready => { + self.transition_import_status(context, ImportStatus::Ready, ImportStatus::Disabled) + } + } + } + + /// Start or continue an import. + /// + /// Returns false if imports are disabled, and true if imports are allowed. + fn begin_import(&self, context: &Context) -> bool { + match self.import_status.load() { + ImportStatus::Importing => true, + ImportStatus::Disabled => false, + ImportStatus::Ready => { + self.transition_import_status(context, ImportStatus::Ready, ImportStatus::Importing) + } + } + } + + /// Change import status. + /// + /// Returns true if the desired `new` status is reached; false otherwise. + fn transition_import_status( + &self, + context: &Context, + current: ImportStatus, + new: ImportStatus, + ) -> bool { + let transition = self.import_transition.lock(); + let observed = self.import_status.load(); + if observed != current { + return observed == new; + } + self.publish_import_status(context, current, new, &transition) + } + + /// Persist import status and set the import_status flag. + fn publish_import_status( + &self, + context: &Context, + current: ImportStatus, + new: ImportStatus, + _transition: &MutexGuard<'_, ()>, + ) -> bool { + self.callbacks.write_iid( + &context.term(Term::Metadata), + IMPORT_STATUS_KEY, + &[new as u8], + ) && self.import_status.compare_exchange(current, new) + } + + pub(crate) fn import_term( + &self, + context: &Context, + term: Term, + id: &[u8], + value: &[u8], + ) -> bool { + match term { + Term::Vector | Term::Neighbors | Term::Quantized | Term::Attributes | Term::ExtMap => { + if id.len() == 4 { + let mut iid = 0u32; + bytemuck::bytes_of_mut(&mut iid).copy_from_slice(id); + + if iid == 0 || iid == u32::MAX { + // ID 0 is the start point; u32::MAX cannot advance the ID minter. + return false; + } + + let value_ok = match term { + Term::Vector => value.len() == self.dim * mem::size_of::(), + Term::Quantized => match &self.quantizer { + Some(quantizer) => value.len() == quantizer.bytes(), + None => false, + }, + Term::Neighbors => { + let count_offset = self.max_degree * mem::size_of::(); + value.len() == count_offset + mem::size_of::() + && bytemuck::pod_read_unaligned::(&value[count_offset..]) + as usize + <= self.max_degree + } + Term::ExtMap => true, + Term::Attributes => !value.is_empty(), + Term::Metadata | Term::IntMap => false, + }; + + if !value_ok { + return false; + } + + if !self.begin_import(context) { + return false; + } + + // Reserve the ID before storing data. + if self.fsm.claim_id(context, iid).is_err() { + return false; + } + + if !self.callbacks.write_iid(&context.term(term), iid, value) { + return false; + } + + true + } else { + false + } + } + Term::Metadata => false, + Term::IntMap => { + if value.len() != mem::size_of::() + || bytemuck::pod_read_unaligned::(value) == 0 + { + return false; + } + + if !self.begin_import(context) { + return false; + } + + let eid = GarnetId::from(id); + self.callbacks.write_eid(&context.term(term), &eid, value) + } + } + } + + /// Finalizes the import. + /// + /// Callers must finish all term writes before starting verification tasks. The last + /// successful task publishes `Disabled` only after finalization and persistence. + /// An empty import returns `(false, Some(true))` to leave start points uninitialized. + pub(crate) fn finish_import( + &self, + context: &Context, + task_idx: usize, + task_count: usize, + ) -> (bool, Option) { + let log_failure = |message: std::fmt::Arguments<'_>| { + self.callbacks.log( + context, + &format!("Error: finish import task {task_idx}/{task_count}: {message}"), + ); + }; + if !self.begin_import(context) { + log_failure(format_args!("imports not enabled")); + return (false, None); + } + + let max_id = self.fsm.max_id(); + + let (start_id, end_id) = + if let Some((start, end)) = self.task_range(task_idx, task_count, max_id) { + (start, end) + } else { + log_failure(format_args!("invalid task split")); + + return (false, None); + }; + + let imports_completed = self + .imports_completed + .get_or_init(|| (0..task_count).map(|_| AtomicBool::new(false)).collect()); + if imports_completed.len() != task_count { + log_failure(format_args!( + "task count changed; expected {}", + imports_completed.len() + )); + return (false, None); + } + let task_completed = &imports_completed[task_idx]; + if task_completed.load(Ordering::Acquire) { + return (true, None); + } + + for id in start_id.max(1)..end_id { + // For every ID that is used, check all terms exist for that ID + match self.fsm.is_free(context, id) { + Ok(true) => continue, + Ok(false) => (), + Err(error) => { + log_failure(format_args!("FSM check failed: id={id}, {error}")); + return (false, None); + } + } + + if !self.callbacks.exists_iid( + &context.term(Term::Vector), + id, + self.dim * mem::size_of::(), + ) { + log_failure(format_args!("missing vector: id={id}")); + return (false, None); + } + if self.is_quantized() + && !self.callbacks.exists_iid( + &context.term(Term::Quantized), + id, + self.quantizer().map(|q| q.bytes()).unwrap_or(0), + ) + { + log_failure(format_args!("missing quantized vector: id={id}")); + return (false, None); + } + if !self.callbacks.exists_iid( + &context.term(Term::Neighbors), + id, + (self.max_degree + 1) * mem::size_of::(), + ) { + log_failure(format_args!("missing neighbors: id={id}")); + return (false, None); + } + let eid = match self + .callbacks + .read_varsize_iid(&context.term(Term::ExtMap), id) + { + Some(v) => GarnetId::from(v.as_slice()), + None => { + log_failure(format_args!("missing external-ID mapping: id={id}")); + return (false, None); + } + }; + if !self + .callbacks + .exists_eid(&context.term(Term::IntMap), &eid, mem::size_of::()) + { + log_failure(format_args!("missing internal-ID mapping: id={id}")); + return (false, None); + } + } + + // Mark each task once and elect the finalizer under the same lock, so neither + // duplicate calls nor concurrent completions can finalize the import twice. + let import_finished = { + let _transition = self.import_transition.lock(); + if task_completed.load(Ordering::Acquire) { + return (true, None); + } + task_completed.store(true, Ordering::Release); + imports_completed + .iter() + .all(|completed| completed.load(Ordering::Acquire)) + }; + + let mut finish_result = None; + if import_finished { + if max_id == 0 { + return if self.transition_import_status( + context, + ImportStatus::Importing, + ImportStatus::Disabled, + ) { + (false, Some(true)) + } else { + log_failure(format_args!("status update failed")); + (true, Some(false)) + }; + } + + let mut first_id = u32::MAX; + if self + .fsm + .visit_used(context, |id| { + first_id = id; + false + }) + .is_err() + { + log_failure(format_args!("FSM scan failed")); + return (true, Some(false)); + } + + if first_id == u32::MAX { + log_failure(format_args!("no imported vectors")); + return (true, Some(false)); + } + + let mut v = match Poly::broadcast(0u8, self.dim * mem::size_of::(), AlignToEight) { + Ok(v) => v, + Err(error) => { + log_failure(format_args!("start point allocation failed: {error}")); + return (true, Some(false)); + } + }; + if !self + .callbacks + .read_single_iid(&context.term(Term::Vector), first_id, &mut v) + { + log_failure(format_args!("start point read failed: id={first_id}")); + return (true, Some(false)); + } + let qv = if let Some(quantizer) = &self.quantizer { + let qv_len = quantizer.bytes(); + let mut qv = match Poly::broadcast(0u8, qv_len, AlignToEight) { + Ok(v) => v, + Err(error) => { + log_failure(format_args!( + "quantized start point allocation failed: {error}" + )); + return (true, Some(false)); + } + }; + + if !self.callbacks.read_single_iid( + &context.term(Term::Quantized), + first_id, + &mut qv, + ) { + log_failure(format_args!( + "quantized start point read failed: id={first_id}" + )); + return (true, Some(false)); + } + + Some(qv) + } else { + None + }; + let mut ns = vec![0u32; self.max_degree + 1]; + if !self + .callbacks + .read_single_iid(&context.term(Term::Neighbors), first_id, &mut ns) + { + log_failure(format_args!( + "start point neighbors read failed: id={first_id}" + )); + return (true, Some(false)); + } + let ns_len = ns[self.max_degree] as usize; + + // Insert the start point + if !self.callbacks.write_iid(&context.term(Term::Vector), 0, &v) { + log_failure(format_args!("start point write failed")); + return (true, Some(false)); + } + // NOTE: unwrap will succeed because it is gated on the quantizer. + if let Some(qv) = &qv + && !self + .callbacks + .write_iid(&context.term(Term::Quantized), 0, qv) + { + log_failure(format_args!("startpoint quantized write failed")); + return (true, Some(false)); + } + if !self + .callbacks + .write_iid(&context.term(Term::Neighbors), 0, &ns) + { + log_failure(format_args!("start point neighbors write failed")); + return (true, Some(false)); + } + + self.start_point_cache.insert(0, v); + if let Some(qv) = qv { + // NOTE: unwrap will succeed because it is gated on the quantizer + self.start_point_quant_cache.insert(0, qv); + } + self.neighbor_cache.insert(0, ns[0..ns_len].to_vec()); + + let finalized = self.transition_import_status( + context, + ImportStatus::Importing, + ImportStatus::Disabled, + ); + if !finalized { + log_failure(format_args!("status update failed")); + } + finish_result = Some(finalized); + } + + (true, finish_result) + } + /// Returns the quantizer associated with the index. fn quantizer(&self) -> Option<&dyn GarnetQuantizer> { if let Some(quantizer) = &self.quantizer { @@ -1244,7 +1772,7 @@ impl SetElement<(&[T], &[u8])> for GarnetProvider { if let Some(quantizer) = &self.quantizer && !internal_id.should_quantize() && !quantizer.is_trained() - && self.fsm.total_used() > quantizer.required_vectors() + && self.fsm.total_used() >= quantizer.required_vectors() { context.set_quantizer_ready(); } @@ -1373,6 +1901,10 @@ impl Delete for GarnetProvider { context: &Self::Context, id: Self::InternalId, ) -> impl Future> + Send { + if id == 0 && self.start_points_exist() { + return future::ready(Ok(ElementStatus::Valid)); + } + let status = match self.fsm.is_free(context, id) { Ok(true) => ElementStatus::Deleted, Ok(false) => ElementStatus::Valid, @@ -2287,7 +2819,7 @@ mod tests { hash::BuildHasher, mem, ops::Range, - sync::{Arc, Mutex, atomic::Ordering, mpsc}, + sync::{Arc, atomic::Ordering, mpsc}, thread, time::{Duration, Instant}, }; @@ -2303,6 +2835,7 @@ mod tests { use diskann_providers::index::wrapped_async::DiskANNIndex; use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_vector::distance::Metric; + use parking_lot::Mutex; use rand::Rng; use crate::{ @@ -2314,7 +2847,7 @@ mod tests { }, provider::{GarnetProvider, GarnetProviderError, QUANT_STATE_KEY, RESERVATION_RETRY_LIMIT}, quantization::{GarnetQuantizer, Spherical1Bit}, - test_utils::{LOGS, Store}, + test_utils::{LOGS, Store, q8_state_with_identity_transform}, }; #[tokio::test] @@ -2348,7 +2881,7 @@ mod tests { assert_eq!(store.int_map_reads(), 0); assert!(store.get(ctx.term(Term::IntMap).get(), &id).is_none()); assert_eq!(provider.fsm.total_used(), 0); - LOGS.with(|logs| assert!(logs.lock().unwrap().is_empty())); + LOGS.with(|logs| assert!(logs.lock().is_empty())); } fn assert_waits_for_range( @@ -2381,7 +2914,7 @@ mod tests { early_result.unwrap(); } }); - assert!(provider.backfill_lock.lock().unwrap().is_empty()); + assert!(provider.backfill_lock.lock().is_empty()); } fn train_for_backfill(provider: &GarnetProvider, ctx: &Context) { @@ -2576,7 +3109,7 @@ mod tests { } fn fail_insert_operation(&self, context: u64) -> bool { - let mut control = self.control.lock().unwrap(); + let mut control = self.control.lock(); let term = context & TERM_BITMASK; if control .pause @@ -2701,6 +3234,364 @@ mod tests { true } + #[test] + fn import_admission_serializes_initial_transition() { + for (import_first, import_second) in [(false, true), (true, false), (false, false)] { + let store = Arc::new(DashMap::new()); + let first_context = ParallelContext::new(store.clone()); + let second_context = ParallelContext::new(store.clone()); + let ctx = first_context.context(); + let index = create_2d_f32_index_with_callbacks( + VectorQuantType::NoQuant, + Metric::L2, + ParallelContext::callbacks(), + &ctx, + ); + let provider = index.inner.provider(); + let ordinary_vector = [1.0f32, 2.0]; + let imported_vector = [3.0f32, 4.0]; + let operate = |context: &Context, importing: bool| { + if importing { + provider.import_term( + context, + Term::Vector, + &7u32.to_ne_bytes(), + bytemuck::cast_slice(&imported_vector), + ) + } else { + DynIndex::maybe_set_start_point( + &index, + context, + bytemuck::cast_slice(&ordinary_vector), + ) + .is_ok() + } + }; + let before = parallel_snapshot(&store); + let (paused_tx, paused_rx) = mpsc::channel(); + let (resume_tx, resume_rx) = mpsc::channel(); + first_context.control.lock().pause = + Some((Term::Metadata as u64, paused_tx, resume_rx)); + let (started_tx, started_rx) = mpsc::channel(); + let (finished_tx, finished_rx) = mpsc::channel(); + + thread::scope(|scope| { + let first = scope.spawn(|| { + if !import_first && !import_second { + // Start with a read so the two callers do not race start-point writes. + matches!(DynIndex::approximate_count(&index, &ctx), Ok(0)) + } else { + operate(&ctx, import_first) + } + }); + paused_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let second = scope.spawn(|| { + started_tx.send(()).unwrap(); + let result = operate(&second_context.context(), import_second); + finished_tx.send(result).unwrap(); + }); + started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let early_result = finished_rx.recv_timeout(Duration::from_millis(50)); + let still_eligible = provider.can_import(&ctx); + let while_paused = parallel_snapshot(&store); + resume_tx.send(()).unwrap(); + let first_result = first.join().unwrap(); + let second_result = match early_result { + Ok(result) => result, + Err(_) => finished_rx.recv_timeout(Duration::from_secs(5)).unwrap(), + }; + second.join().unwrap(); + + assert_eq!(early_result, Err(mpsc::RecvTimeoutError::Timeout)); + assert!(still_eligible); + assert_eq!(while_paused, before); + assert!(first_result); + assert_eq!(second_result, import_first == import_second); + }); + + assert_eq!(provider.can_import(&ctx), import_first); + assert_eq!(provider.fsm.total_used(), usize::from(import_first)); + assert_eq!( + parallel_get(&store, ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()), + (!import_first).then(|| bytemuck::cast_slice(&ordinary_vector).to_vec()) + ); + assert_eq!( + parallel_get(&store, ctx.term(Term::Vector).get(), &7u32.to_ne_bytes()), + import_first.then(|| bytemuck::cast_slice(&imported_vector).to_vec()) + ); + } + } + + #[test] + fn admitted_import_terms_proceed_in_parallel() { + for term in [Term::Vector, Term::IntMap] { + let store = Arc::new(DashMap::new()); + let first_context = ParallelContext::new(store.clone()); + let second_context = ParallelContext::new(store.clone()); + let ctx = first_context.context(); + let index = create_2d_f32_index_with_callbacks( + VectorQuantType::Q8, + Metric::L2, + ParallelContext::callbacks(), + &ctx, + ); + let provider = index.inner.provider(); + let quant_state = provider.quantizer().unwrap().serialize().unwrap(); + assert!(provider.set_quant_state(&ctx, &quant_state)); + let (first_id, first_value, second_id, second_value) = match term { + Term::Vector => ( + 7u32.to_ne_bytes().to_vec(), + bytemuck::cast_slice(&[1.0f32, 2.0]).to_vec(), + 8u32.to_ne_bytes().to_vec(), + bytemuck::cast_slice(&[3.0f32, 4.0]).to_vec(), + ), + Term::IntMap => ( + b"first".to_vec(), + 7u32.to_ne_bytes().to_vec(), + b"second".to_vec(), + 8u32.to_ne_bytes().to_vec(), + ), + _ => unreachable!(), + }; + let (paused_tx, paused_rx) = mpsc::channel(); + let (resume_tx, resume_rx) = mpsc::channel(); + first_context.control.lock().pause = Some((term as u64, paused_tx, resume_rx)); + let (finished_tx, finished_rx) = mpsc::channel(); + + thread::scope(|scope| { + let first = + scope.spawn(|| provider.import_term(&ctx, term, &first_id, &first_value)); + paused_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let second = scope.spawn(|| { + let context = second_context.context(); + let imported = provider.import_term(&context, term, &second_id, &second_value); + let replaced = provider.set_quant_state(&context, &quant_state); + let ordinary = DynIndex::maybe_set_start_point( + &index, + &context, + bytemuck::cast_slice(&[5.0f32, 6.0]), + ) + .is_ok(); + finished_tx.send((imported, replaced, ordinary)).unwrap(); + }); + let second_result = finished_rx.recv_timeout(Duration::from_secs(5)); + let first_while_paused = parallel_get(&store, ctx.term(term).get(), &first_id); + let second_while_paused = parallel_get(&store, ctx.term(term).get(), &second_id); + resume_tx.send(()).unwrap(); + assert!(first.join().unwrap()); + second.join().unwrap(); + + assert_eq!(second_result, Ok((true, false, false))); + assert_eq!(first_while_paused, None); + assert_eq!(second_while_paused, Some(second_value.clone())); + }); + + assert!(provider.can_import(&ctx)); + assert_eq!( + provider.fsm.total_used(), + if matches!(term, Term::Vector) { 2 } else { 0 } + ); + assert_eq!( + parallel_get(&store, ctx.term(term).get(), &first_id), + Some(first_value) + ); + assert_eq!( + parallel_get(&store, ctx.term(term).get(), &second_id), + Some(second_value) + ); + assert!( + parallel_get(&store, ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()).is_none() + ); + assert_eq!( + provider.quantizer().unwrap().serialize().unwrap(), + quant_state + ); + } + } + + #[test] + fn quant_state_replacement_serializes_first_write() { + for setter_first in [false, true] { + for importing in [false, true] { + let store = Arc::new(DashMap::new()); + let first_context = ParallelContext::new(store.clone()); + let second_context = ParallelContext::new(store.clone()); + let ctx = first_context.context(); + let index = create_2d_f32_index_with_callbacks( + VectorQuantType::Q8, + Metric::L2, + ParallelContext::callbacks(), + &ctx, + ); + let provider = index.inner.provider(); + let original = provider.quantizer().unwrap().serialize().unwrap(); + let replacement = q8_state_with_identity_transform(2); + assert!(provider.set_quant_state(&ctx, &original)); + let expected_state = if setter_first { + &replacement + } else { + &original + }; + let quantizer = + crate::quantization::MinMax8Bit::new_from_bytes(Metric::L2, expected_state) + .unwrap(); + let vector = [1.0f32, 2.0]; + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&vector, &mut quantized).unwrap(); + let write = |context: &Context| { + if importing { + provider.import_term( + context, + Term::Quantized, + &7u32.to_ne_bytes(), + &quantized, + ) + } else { + provider.maybe_set_start_point(context, &vector).is_ok() + } + }; + + // Pause before the first data write or quantizer metadata write. + let (paused_tx, paused_rx) = mpsc::channel(); + let (resume_tx, resume_rx) = mpsc::channel(); + let pause_term = if setter_first || importing { + Term::Metadata + } else { + Term::Vector + }; + first_context.control.lock().pause = + Some((pause_term as u64, paused_tx, resume_rx)); + let (started_tx, started_rx) = mpsc::channel(); + let (finished_tx, finished_rx) = mpsc::channel(); + thread::scope(|scope| { + let first = scope.spawn(|| { + if setter_first { + provider.set_quant_state(&ctx, &replacement) + } else { + write(&ctx) + } + }); + paused_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let second = scope.spawn(|| { + started_tx.send(()).unwrap(); + let context = second_context.context(); + let result = if setter_first { + write(&context) + } else { + provider.set_quant_state(&context, &replacement) + }; + finished_tx.send(result).unwrap(); + }); + started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let early_result = finished_rx.recv_timeout(Duration::from_millis(50)); + let state_while_paused = provider.quantizer().unwrap().serialize().unwrap(); + let metadata_while_paused = parallel_get( + &store, + ctx.term(Term::Metadata).get(), + bytemuck::bytes_of(&QUANT_STATE_KEY), + ); + resume_tx.send(()).unwrap(); + assert!(first.join().unwrap()); + let second_result = match early_result { + Ok(result) => result, + Err(_) => finished_rx.recv_timeout(Duration::from_secs(5)).unwrap(), + }; + second.join().unwrap(); + assert_eq!(early_result, Err(mpsc::RecvTimeoutError::Timeout)); + assert_eq!(second_result, setter_first); + assert_eq!(&state_while_paused, expected_state); + assert_eq!(metadata_while_paused, Some(original.to_vec())); + }); + + assert_eq!(provider.fsm.total_used(), usize::from(importing)); + assert_eq!( + &provider.quantizer().unwrap().serialize().unwrap(), + expected_state + ); + assert_eq!( + parallel_get( + &store, + ctx.term(Term::Metadata).get(), + bytemuck::bytes_of(&QUANT_STATE_KEY), + ), + Some(expected_state.to_vec()) + ); + let id = if importing { 7u32 } else { 0u32 }; + assert_eq!( + parallel_get(&store, ctx.term(Term::Quantized).get(), &id.to_ne_bytes()), + Some(quantized) + ); + } + } + } + + #[test] + fn import_failure_preserves_quantizer_eligibility() { + for failure_term in [Term::Metadata, Term::Quantized] { + let store = Arc::new(DashMap::new()); + let context = ParallelContext::new(store.clone()); + let ctx = context.context(); + let index = create_2d_f32_index_with_callbacks( + VectorQuantType::Q8, + Metric::L2, + ParallelContext::callbacks(), + &ctx, + ); + let provider = index.inner.provider(); + let state = provider.quantizer().unwrap().serialize().unwrap(); + let replacement = q8_state_with_identity_transform(2); + assert!(provider.set_quant_state(&ctx, &state)); + let id = 7u32.to_ne_bytes(); + let mut quantized = vec![0u8; provider.quantizer().unwrap().bytes()]; + provider + .quantizer() + .unwrap() + .compress(&[1.0, 2.0], &mut quantized) + .unwrap(); + context.control.lock().failure = Some((failure_term as u64, 0)); + assert!(!provider.import_term(&ctx, Term::Quantized, &id, &quantized)); + assert!(parallel_get(&store, ctx.term(Term::Quantized).get(), &id).is_none()); + let empty = matches!(failure_term, Term::Metadata); + assert_eq!(provider.fsm.total_used(), usize::from(!empty)); + drop(index); + + // Recover a failed admission or a claimed ID whose term write failed. + let index = create_2d_f32_index_with_callbacks( + VectorQuantType::Q8, + Metric::L2, + ParallelContext::callbacks(), + &ctx, + ); + let provider = index.inner.provider(); + assert_eq!(provider.set_quant_state(&ctx, &replacement), empty); + let expected_state = if empty { &replacement } else { &state }; + assert_eq!( + &provider.quantizer().unwrap().serialize().unwrap(), + expected_state + ); + assert_eq!( + parallel_get( + &store, + ctx.term(Term::Metadata).get(), + bytemuck::bytes_of(&QUANT_STATE_KEY) + ), + Some(expected_state.to_vec()) + ); + provider + .quantizer() + .unwrap() + .compress(&[1.0, 2.0], &mut quantized) + .unwrap(); + assert!(provider.import_term(&ctx, Term::Quantized, &id, &quantized)); + assert_eq!(provider.fsm.total_used(), 1); + assert!(!provider.set_quant_state(&ctx, &replacement)); + assert_eq!( + parallel_get(&store, ctx.term(Term::Quantized).get(), &id), + Some(quantized) + ); + } + } + fn wait_for_pending_receiver(provider: &GarnetProvider, id: &GarnetId) -> bool { let id_hash = provider.pending_external_ids.hasher().hash_one(&id[..]); let deadline = Instant::now() + Duration::from_secs(5); @@ -2764,7 +3655,7 @@ mod tests { let (first, second) = thread::scope(|scope| { let first = scope.spawn(|| { let state = ParallelContext::new(store.clone()); - *state.control.lock().unwrap() = InsertControl { + *state.control.lock() = InsertControl { failure: first_fails.then_some(failure), pause: Some((failure.0, entered_tx, resume_rx)), before_vector: None, @@ -2790,7 +3681,7 @@ mod tests { } else { (failure.0, if first_fails { failure.1 } else { 0 }) }; - *state.control.lock().unwrap() = InsertControl { + *state.control.lock() = InsertControl { failure: second_fails.then_some(second_failure), pause: None, before_vector: None, @@ -2803,7 +3694,7 @@ mod tests { bytemuck::cast_slice(&second_vector), b"second", ); - let snapshot = state.control.lock().unwrap().before_vector.take().unwrap(); + let snapshot = state.control.lock().before_vector.take().unwrap(); ((result.is_ok(), context.insert_is_update()), snapshot) }); let waiting = wait_for_pending_receiver(provider, &id); @@ -3241,7 +4132,7 @@ mod tests { error, GarnetProviderError::Garnet(GarnetError::Write) )); - assert!(provider.backfill_lock.lock().unwrap().is_empty()); + assert!(provider.backfill_lock.lock().is_empty()); assert_eq!(provider.fsm.max_id(), max_id); assert_eq!(provider.fsm.total_used(), total_used); assert_eq!( @@ -3268,6 +4159,57 @@ mod tests { } } + #[tokio::test] + async fn quantization_threshold_excludes_start_point() { + let store = Store::new(); + let ctx = Context::new(0); + let provider = GarnetProvider::::new( + 2, + VectorQuantType::Bin, + Metric::L2, + 10, + store.callbacks(), + &ctx, + ) + .unwrap(); + provider.maybe_set_start_point(&ctx, &[1.0, -1.0]).unwrap(); + assert_eq!(provider.total_used(), 0); + assert_eq!( + provider.status_by_internal_id(&ctx, 0).await.unwrap(), + diskann::provider::ElementStatus::Valid + ); + + let required = Spherical1Bit::new(Metric::L2, 2).required_vectors(); + for id in 1..required as u32 { + let external_id = GarnetId::from(&id.to_ne_bytes()[..]); + let vector = [(id % 17) as f32, (id % 31) as f32]; + provider + .set_element(&ctx, &external_id, (&vector, &[])) + .await + .unwrap() + .complete() + .await; + } + assert!(!ctx.quantizer_ready()); + let external_id = GarnetId::from(&(required as u32).to_ne_bytes()[..]); + provider + .set_element(&ctx, &external_id, (&[1.0, -1.0], &[])) + .await + .unwrap() + .complete() + .await; + assert!(ctx.quantizer_ready()); + assert_eq!(provider.total_used(), required); + assert!(provider.train_quantizer(&ctx)); + assert!(provider.backfill_quant_vectors(&ctx, 0, 1)); + assert!(provider.is_quantized()); + assert!(provider.callbacks.exists_iid( + &ctx.term(Term::Quantized), + 0, + provider.quant_vector_size(), + )); + } + fn create_2d_f32_index( quant_type: VectorQuantType, metric: Metric, @@ -3307,7 +4249,7 @@ mod tests { let ctx = Context::new(0); let index = create_2d_f32_index(VectorQuantType::Bin, Metric::L2, &store, &ctx); let provider = index.inner.provider(); - let required_vecs = Spherical1Bit::new(2).required_vectors(); + let required_vecs = Spherical1Bit::new(Metric::L2, 2).required_vectors(); let mut rng = rand::rng(); @@ -3381,7 +4323,7 @@ mod tests { let ctx = Context::new(0); let index = create_2d_f32_index(VectorQuantType::Bin, Metric::L2, &store, &ctx); let provider = index.inner.provider(); - let required_vecs = Spherical1Bit::new(2).required_vectors(); + let required_vecs = Spherical1Bit::new(Metric::L2, 2).required_vectors(); let mut rng = rand::rng(); @@ -3501,7 +4443,7 @@ mod tests { let ctx = Context::new(0); let index = create_2d_f32_index(VectorQuantType::Bin, Metric::L2, &store, &ctx); let provider = index.inner.provider(); - let required_vecs = Spherical1Bit::new(2).required_vectors(); + let required_vecs = Spherical1Bit::new(Metric::L2, 2).required_vectors(); let mut rng = rand::rng(); @@ -3633,9 +4575,22 @@ mod tests { .unwrap(); } + let state = provider.quantizer().unwrap().serialize().unwrap(); + let replacement = q8_state_with_identity_transform(2); + assert!(!provider.set_quant_state(&ctx, &replacement)); + assert_eq!(provider.quantizer().unwrap().serialize().unwrap(), state); + // Drop and re-create the index, keeping the same backing store let index = create_2d_f32_index(VectorQuantType::Q8, Metric::L2, &store, &ctx); let provider = index.inner.provider(); + assert!(!provider.set_quant_state(&ctx, &replacement)); + assert_eq!( + store.get( + ctx.term(Term::Metadata).get(), + bytemuck::bytes_of(&QUANT_STATE_KEY) + ), + Some(state.to_vec()) + ); // There should be saved quant state. assert!( diff --git a/diskann-garnet/src/quantization.rs b/diskann-garnet/src/quantization.rs index 13556352c6..3a55683d0e 100644 --- a/diskann-garnet/src/quantization.rs +++ b/diskann-garnet/src/quantization.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -use std::{num::NonZero, sync::RwLock}; +use std::num::NonZero; use diskann::utils::VectorRepr; use diskann_quantization::{ @@ -19,6 +19,7 @@ use diskann_quantization::{ }; use diskann_utils::views::rowmajor::{self, Matrix}; use diskann_vector::{DistanceFunction, PreprocessedDistanceFunction, distance::Metric}; +use parking_lot::RwLock; use thiserror::Error; use crate::provider::{GarnetDistanceComputer, GarnetQueryComputer}; @@ -86,13 +87,15 @@ pub(crate) trait DynQueryComputer: Send + Sync { /// vectors (but not thousands) for training. Quantized vectors have 1 bit per dimension plus up /// to 6 bytes of overhead. pub(crate) struct Spherical1Bit { + metric: Metric, dim: usize, inner: RwLock>>, } impl Spherical1Bit { - pub(crate) fn new(dim: usize) -> Self { + pub(crate) fn new(metric: Metric, dim: usize) -> Self { Self { + metric, dim, inner: RwLock::new(None), } @@ -109,7 +112,7 @@ impl GarnetQuantizer for Spherical1Bit { } fn is_trained(&self) -> bool { - self.inner.read().unwrap().is_some() + self.inner.read().is_some() } fn train( @@ -131,7 +134,7 @@ impl GarnetQuantizer for Spherical1Bit { ) .map_err(|e| GarnetQuantizerError::Training(Box::new(e)))?; - let mut inner = self.inner.write().unwrap(); + let mut inner = self.inner.write(); *inner = Some( spherical::iface::Impl::<1>::new(quantizer) .map_err(|e| GarnetQuantizerError::Alloc(Box::new(e)))?, @@ -141,7 +144,7 @@ impl GarnetQuantizer for Spherical1Bit { } fn compress(&self, v: &[f32], into: &mut [u8]) -> Result<(), GarnetQuantizerError> { - let guard = self.inner.read().unwrap(); + let guard = self.inner.read(); if let Some(quantizer) = &*guard { spherical::iface::Quantizer::::compress( quantizer, @@ -157,7 +160,7 @@ impl GarnetQuantizer for Spherical1Bit { } fn distance_computer(&self) -> Result { - let guard = self.inner.read().unwrap(); + let guard = self.inner.read(); if let Some(quantizer) = &*guard { let computer = quantizer .distance_computer(GlobalAllocator) @@ -169,7 +172,7 @@ impl GarnetQuantizer for Spherical1Bit { } fn query_computer(&self, query: &[f32]) -> Result { - let guard = self.inner.read().unwrap(); + let guard = self.inner.read(); if let Some(quantizer) = &*guard { let computer = quantizer .fused_query_computer( @@ -187,7 +190,7 @@ impl GarnetQuantizer for Spherical1Bit { } fn serialize(&self) -> Result, GarnetQuantizerError> { - let guard = self.inner.read().unwrap(); + let guard = self.inner.read(); if let Some(quantizer) = &*guard { quantizer .serialize(GlobalAllocator) @@ -198,14 +201,24 @@ impl GarnetQuantizer for Spherical1Bit { } fn deserialize(&self, state: &[u8]) -> Result<(), GarnetQuantizerError> { - let mut guard = self.inner.write().unwrap(); + if state.len() < 8 { + return Err(GarnetQuantizerError::UnsupportedSerialization); + } + + let mut guard = self.inner.write(); if guard.is_some() { Err(GarnetQuantizerError::UnsupportedSerialization) } else { let q = spherical::iface::Impl::<1>::try_deserialize(state, GlobalAllocator) .map_err(|e| GarnetQuantizerError::Deserialization(Box::new(e)))?; - *guard = Some(q); - Ok(()) + if as spherical::iface::Quantizer>::full_dim(&q) == self.dim && + as spherical::iface::Quantizer>::dim(&q) == self.dim && + as spherical::iface::Quantizer>::metric(&q) == self.metric { + *guard = Some(q); + Ok(()) + } else { + Err(GarnetQuantizerError::UnsupportedSerialization) + } } } } @@ -237,7 +250,8 @@ impl DynQueryComputer for iface::QueryComputer { /// quantized vector has 8 bits per dimension and 20 bytes of overhead. pub(crate) struct MinMax8Bit { metric: Metric, - inner: minmax::MinMaxQuantizer, + dim: usize, + inner: RwLock, } impl MinMax8Bit { @@ -260,7 +274,8 @@ impl MinMax8Bit { Ok(Self { metric, - inner: minmax::MinMaxQuantizer::new(transform, grid_scale), + dim: dim.get(), + inner: RwLock::new(minmax::MinMaxQuantizer::new(transform, grid_scale)), }) } @@ -270,7 +285,11 @@ impl MinMax8Bit { ) -> Result { let inner = MinMaxQuantizer::try_deserialize(bytes) .map_err(|e| GarnetQuantizerError::Deserialization(Box::new(e)))?; - Ok(Self { metric, inner }) + Ok(Self { + metric, + dim: inner.dim(), + inner: RwLock::new(inner), + }) } } @@ -280,7 +299,7 @@ impl GarnetQuantizer for MinMax8Bit { } fn bytes(&self) -> usize { - minmax::Data::<8>::canonical_bytes(self.inner.dim()) + minmax::Data::<8>::canonical_bytes(self.dim) } fn is_trained(&self) -> bool { @@ -296,9 +315,10 @@ impl GarnetQuantizer for MinMax8Bit { } fn compress(&self, v: &[f32], into: &mut [u8]) -> Result<(), GarnetQuantizerError> { - let into = minmax::DataMutRef::<8>::from_canonical_front_mut(into, self.inner.dim()) + let into = minmax::DataMutRef::<8>::from_canonical_front_mut(into, self.dim) .map_err(|e| GarnetQuantizerError::Compression(Box::new(e)))?; self.inner + .read() .compress_into(v, into) .map_err(|e| GarnetQuantizerError::Compression(Box::new(e)))?; Ok(()) @@ -308,17 +328,18 @@ impl GarnetQuantizer for MinMax8Bit { let computer = GarnetDistanceComputer::new( ::distance( self.metric, - Some(self.inner.dim()), + Some(self.dim), ), ); Ok(computer) } fn query_computer(&self, query: &[f32]) -> Result { + let inner = self.inner.read(); let computer = GarnetQueryComputer::new(MinMax8BitQueryComputer::new( - &self.inner, + &inner, query, - self.inner.dim(), + self.dim, self.metric, )?); Ok(computer) @@ -326,12 +347,25 @@ impl GarnetQuantizer for MinMax8Bit { fn serialize(&self) -> Result, GarnetQuantizerError> { self.inner + .read() .serialize(GlobalAllocator) .map_err(|e| GarnetQuantizerError::Alloc(Box::new(e))) } - fn deserialize(&self, _state: &[u8]) -> Result<(), GarnetQuantizerError> { - Err(GarnetQuantizerError::UnsupportedSerialization) + fn deserialize(&self, state: &[u8]) -> Result<(), GarnetQuantizerError> { + if state.len() < 8 { + return Err(GarnetQuantizerError::UnsupportedSerialization); + } + + let mut guard = self.inner.write(); + let quantizer = MinMaxQuantizer::try_deserialize(state) + .map_err(|error| GarnetQuantizerError::Deserialization(Box::new(error)))?; + if quantizer.dim() == self.dim && quantizer.output_dim() == self.dim { + *guard = quantizer; + Ok(()) + } else { + Err(GarnetQuantizerError::UnsupportedSerialization) + } } } @@ -385,7 +419,7 @@ mod tests { #[test] fn basic_spherical_1bit() { - let quantizer = Spherical1Bit::new(2); + let quantizer = Spherical1Bit::new(Metric::L2, 2); assert_eq!(quantizer.required_vectors(), 1000); assert_eq!(quantizer.bytes(), 1 + 6); diff --git a/diskann-garnet/src/test_utils.rs b/diskann-garnet/src/test_utils.rs index 9b6fa47932..01663f88ad 100644 --- a/diskann-garnet/src/test_utils.rs +++ b/diskann-garnet/src/test_utils.rs @@ -6,12 +6,10 @@ use crate::garnet::{Callbacks, ReadDataCallback, RmwDataCallback, TERM_BITMASK, Term}; use core::slice; use dashmap::DashMap; +use parking_lot::Mutex; use std::{ ffi::c_void, - sync::{ - Mutex, - atomic::{AtomicUsize, Ordering}, - }, + sync::atomic::{AtomicUsize, Ordering}, }; thread_local! { @@ -54,7 +52,7 @@ impl Store { pub fn clear(&self) { STORE.with(|s| s.clear()); LOGS.with(|l| { - let mut guard = l.lock().unwrap(); + let mut guard = l.lock(); guard.clear(); }); FULL_READS.with(|fr| fr.store(0, Ordering::Release)); @@ -121,12 +119,35 @@ impl Store { pub fn log(&self, context: u64, msg: &str) { LOGS.with(|l| { - let mut guard = l.lock().unwrap(); + let mut guard = l.lock(); guard.push((context, msg.to_owned())); }); } } +/// A deterministic Q8 state distinct from the default DoubleHadamard transform. +pub fn q8_state_with_identity_transform( + dim: usize, +) -> diskann_quantization::alloc::Poly<[u8], diskann_quantization::alloc::GlobalAllocator> { + use diskann_quantization::{ + algorithms::{Transform, TransformKind}, + alloc::GlobalAllocator, + minmax::MinMaxQuantizer, + num::POSITIVE_ONE_F32, + }; + + let transform = Transform::new( + TransformKind::Null, + std::num::NonZero::new(dim).unwrap(), + None, + GlobalAllocator, + ) + .unwrap(); + MinMaxQuantizer::new(transform, POSITIVE_ONE_F32) + .serialize(GlobalAllocator) + .unwrap() +} + unsafe extern "C" fn test_read( ctx: u64, count: u32,