diff --git a/Cargo.lock b/Cargo.lock index 6a192ec3c..73f9b792a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -605,7 +605,7 @@ dependencies = [ [[package]] name = "diskann-garnet" -version = "6.0.1" +version = "7.0.0" dependencies = [ "bytemuck", "crossbeam", diff --git a/diskann-garnet/Cargo.toml b/diskann-garnet/Cargo.toml index 43f9087b2..9232a2d2c 100644 --- a/diskann-garnet/Cargo.toml +++ b/diskann-garnet/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "diskann-garnet" -version = "6.0.1" +version = "7.0.0" edition = "2024" authors.workspace = true license.workspace = true diff --git a/diskann-garnet/README.md b/diskann-garnet/README.md index 6acb41f39..ffd682860 100644 --- a/diskann-garnet/README.md +++ b/diskann-garnet/README.md @@ -37,8 +37,8 @@ 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 - 2` user vectors in a single instance. Internal -ID 0 is reserved for the start point, and `u32::MAX` marks ID allocation exhaustion. +An ndex supports up to `2^32 - 1` user vectors. User IDs range from 0 through +`u32::MAX - 1`; `u32::MAX` is reserved for the start point. ### Bulk Import diff --git a/diskann-garnet/diskann-garnet.nuspec b/diskann-garnet/diskann-garnet.nuspec index 65af4dde4..179626f4e 100644 --- a/diskann-garnet/diskann-garnet.nuspec +++ b/diskann-garnet/diskann-garnet.nuspec @@ -2,7 +2,7 @@ diskann-garnet - 6.0.1 + 7.0.0 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 512e628fd..f0e13a963 100644 --- a/diskann-garnet/docs/data-design.md +++ b/diskann-garnet/docs/data-design.md @@ -88,13 +88,11 @@ 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. 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. +Only a single start point is supported. Its internal ID is fixed when the provider is created and defaults to `DEFAULT_START_POINT_ID` (`u32::MAX`). 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. +Imports must supply the full vector, neighbor list, and (for quantized indexes) quantized vector at the configured start-point ID. Finalization validates these terms and populates the caches. The neighbor count must fit the configured degree, and each listed neighbor must be a valid ID. -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. +Start points have no associated attributes or ID mappings and are not counted as active vectors. #### Metadata @@ -107,11 +105,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 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. +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 the configured start-point ID is ignored. Blocks grow on demand during insertion or import, including intermediate blocks. 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. +Importing a term keyed by internal ID marks that ID as used unless it is the start point's ID, expanding the map and advancing the maximum assigned ID as needed. The start point is tracked separately from the FSM. ##### Quantizer Tables @@ -120,7 +118,7 @@ Importing a term keyed by an internal ID claims that ID before writing the term 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. +Quantizer state can only be replaced while no ordinary IDs are occupied in the FSM and no start points exist. ##### Import Status @@ -140,7 +138,7 @@ The former Boolean format did not distinguish ready indexes from unfinished impo *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, including zero. 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. Lookup of an external ID will happen during post processing when we return results to Garnet. @@ -155,7 +153,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. 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. +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. One internal ID is reserved for the start point and cannot be allocated to inserted vectors. 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. diff --git a/diskann-garnet/docs/ffi-design.rs b/diskann-garnet/docs/ffi-design.rs index 6cfcf03d7..684873a20 100644 --- a/diskann-garnet/docs/ffi-design.rs +++ b/diskann-garnet/docs/ffi-design.rs @@ -127,6 +127,8 @@ type LogCallback = unsafe extern "C" fn(context: u64, message: *const u8, messag /// reduce_dim == 0 to indicate no reduction requested. Dimensionality reduction is not /// implemented, so this parameter is currently ignored. /// +/// start_point_id is fixed for the index lifetime and must be supplied again on reopen. +/// /// metric_type is passed as a raw i32. Valid values are: /// - 0: Cosine /// - 1: InnerProduct @@ -148,6 +150,7 @@ extern "C" fn create_index( metric_type: i32, l_build: u32, max_degree: u32, + start_point_id: u32, read_callback: ReadCallback, write_callback: WriteCallback, delete_callback: DeleteCallback, @@ -422,11 +425,8 @@ extern "C" fn can_import(context: u64, index_ptr: *const c_void) -> bool; /// 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` +/// empty). Terms may arrive in any order and must use IDs between 0 and +/// u32::MAX - 1. 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. diff --git a/diskann-garnet/src/ffi_recall_tests.rs b/diskann-garnet/src/ffi_recall_tests.rs index f2d006d96..a7e0d4722 100644 --- a/diskann-garnet/src/ffi_recall_tests.rs +++ b/diskann-garnet/src/ffi_recall_tests.rs @@ -11,7 +11,7 @@ mod tests { use crate::{ Overflow, VectorQuantType, create_index, drop_index, garnet::Context, insert, - search_vector, test_utils::Store, + provider::DEFAULT_START_POINT_ID, search_vector, test_utils::Store, }; /// Helper to insert a vector with a string external ID and FP32 data. @@ -170,6 +170,7 @@ mod tests { metric as i32, l_build, max_degree, + DEFAULT_START_POINT_ID, callbacks.read_callback(), callbacks.write_callback(), callbacks.delete_callback(), diff --git a/diskann-garnet/src/ffi_tests.rs b/diskann-garnet/src/ffi_tests.rs index 3213f1516..b3f8b39fe 100644 --- a/diskann-garnet/src/ffi_tests.rs +++ b/diskann-garnet/src/ffi_tests.rs @@ -24,6 +24,7 @@ mod tests { check_internal_id_valid, create_index, drop_index, finish_import, garnet::{Context, Term, WriteCallback}, import_term, insert, + provider::DEFAULT_START_POINT_ID, quantization::{GarnetQuantizer, MinMax8Bit, Spherical1Bit}, remove, search_vector, set_attribute, set_quant_state, test_utils::{STORE, Store, q8_state_with_identity_transform}, @@ -60,6 +61,22 @@ mod tests { quant_type: VectorQuantType, metric_type: i32, write_callback: WriteCallback, + ) -> (*const c_void, Context) { + create_test_index_with_start_point( + store, + quant_type, + metric_type, + write_callback, + DEFAULT_START_POINT_ID, + ) + } + + fn create_test_index_with_start_point( + store: &Store, + quant_type: VectorQuantType, + metric_type: i32, + write_callback: WriteCallback, + start_point_id: u32, ) -> (*const c_void, Context) { let callbacks = store.callbacks(); let ctx = Context::new(0); @@ -79,6 +96,7 @@ mod tests { metric_type, l_build, max_degree, + start_point_id, callbacks.read_callback(), write_callback, callbacks.delete_callback(), @@ -103,6 +121,38 @@ mod tests { } } + #[test] + fn basic_create_index_with_start_point() { + for start_point_id in [0, 42, DEFAULT_START_POINT_ID] { + let store = Store::new(); + let create = || { + create_test_index_with_start_point( + &store, + VectorQuantType::NoQuant, + Metric::L2 as i32, + store.callbacks().write_callback(), + start_point_id, + ) + }; + let (index_ptr, ctx) = create(); + assert!(!index_ptr.is_null()); + let point = [1.0f32, 2.0]; + assert_eq!( + insert_f32_vector(&ctx, index_ptr, 1, &point), + InsertResult::Success + ); + assert_eq!( + store.get(ctx.term(Term::Vector).get(), &start_point_id.to_ne_bytes()), + Some(bytemuck::cast_slice::(&point).to_vec()), + ); + unsafe { drop_index(ctx.get(), index_ptr) }; + let (index_ptr, ctx) = create(); + assert!(!index_ptr.is_null()); + assert_eq!(do_search(&ctx, index_ptr, &point, 1, None).0, [1]); + unsafe { drop_index(ctx.get(), index_ptr) }; + } + } + #[test] fn create_index_with_invalid_metric_returns_null() { let store = Store::new(); @@ -157,12 +207,14 @@ mod tests { let store = Store::new(); let (index_ptr, ctx) = create_test_index(&store, quant_type); - for (id, vector, expected, expected_internal_id) in [ - (42, [1.0, 2.0], 1, 1u32), + for (id, vector, expected, expected_count) in [ + (42, [1.0, 2.0], 1, 1u64), (42, [2.0, 1.0], 3, 1), (43, [3.0, 4.0], 1, 2), (43, [4.0, 3.0], 3, 2), ] { + let existing_internal_id = + store.get(ctx.term(Term::IntMap).get(), bytemuck::bytes_of(&id)); assert_eq!( u8::from(insert_f32_vector(&ctx, index_ptr, id, &vector)), expected, @@ -171,18 +223,14 @@ mod tests { let internal_id = store .get(ctx.term(Term::IntMap).get(), bytemuck::bytes_of(&id)) .unwrap(); - assert_eq!( - bytemuck::pod_read_unaligned::(&internal_id), - expected_internal_id - ); + if let Some(existing_internal_id) = existing_internal_id { + assert_eq!(internal_id, existing_internal_id); + } assert_eq!( store.get(ctx.term(Term::Vector).get(), &internal_id), Some(bytemuck::cast_slice::(&vector).to_vec()) ); - assert_eq!( - unsafe { card(ctx.get(), index_ptr) }, - u64::from(expected_internal_id) - ); + assert_eq!(unsafe { card(ctx.get(), index_ptr) }, expected_count); } assert_eq!(store.int_map_reads(), 4); @@ -344,10 +392,12 @@ mod tests { assert!(unsafe { can_import(ctx.get(), index_ptr) }); let id = 1u32.to_ne_bytes(); + let zero = 0u32.to_ne_bytes(); for (internal_id, accepted, expected) in [ - (0u32, false, None), + (DEFAULT_START_POINT_ID, false, None), + (0, true, Some(zero.as_slice())), (1, true, Some(id.as_slice())), - (0, false, Some(id.as_slice())), + (DEFAULT_START_POINT_ID, false, Some(id.as_slice())), ] { let value = internal_id.to_ne_bytes(); assert_eq!( @@ -546,7 +596,6 @@ mod tests { #[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, @@ -554,51 +603,39 @@ mod tests { 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_WRITE.swap(false, Ordering::AcqRel) + && 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); + 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); + // Finalization only writes ImportStatus::Disabled; start points were imported. + 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; } - unsafe { drop_index(ctx.get(), index_ptr) }; + assert!(unsafe { can_import(ctx.get(), index_ptr) }); + assert_ordinary_operations_blocked(&ctx, index_ptr); } + unsafe { drop_index(ctx.get(), index_ptr) }; } #[test] @@ -653,6 +690,21 @@ mod tests { } }); + let start_id = DEFAULT_START_POINT_ID.to_ne_bytes(); + assert!(import_test_term( + &ctx, + index_ptr, + Term::Vector, + &start_id, + bytemuck::cast_slice(&[1.0f32, 2.0]), + )); + assert!(import_test_term( + &ctx, + index_ptr, + Term::Neighbors, + &start_id, + bytemuck::cast_slice(&[0u32; 21]), + )); assert_eq!( unsafe { finish_import(ctx.get(), index_ptr, 0, 1) }, u8::from(ImportResult::Success) @@ -663,7 +715,7 @@ mod tests { } #[test] - fn finish_empty_import_and_reopen() { + fn finish_import_requires_start_point() { for (quant_type, task_count) in [ (VectorQuantType::NoQuant, 1), (VectorQuantType::NoQuant, 4), @@ -681,17 +733,24 @@ mod tests { } for task_idx in 0..task_count { assert!(unsafe { can_import(ctx.get(), index_ptr) }); + let expected = if task_idx + 1 == task_count { + ImportResult::FinishFailed + } else { + ImportResult::Success + }; assert_eq!( unsafe { finish_import(ctx.get(), index_ptr, task_idx, task_count) }, - u8::from(ImportResult::Success) + u8::from(expected) ); assert!( store - .get(ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()) + .get( + ctx.term(Term::Vector).get(), + &DEFAULT_START_POINT_ID.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); } } @@ -700,27 +759,12 @@ mod tests { 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_ordinary_operations_blocked(&ctx, index_ptr); + let start_id = DEFAULT_START_POINT_ID.to_ne_bytes(); + assert!(!unsafe { + check_internal_id_valid(ctx.get(), index_ptr, start_id.as_ptr(), start_id.len()) + }); } - 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) }; } } @@ -753,6 +797,7 @@ mod tests { quantizer: Option<&dyn GarnetQuantizer>, ) { let id = 1u32.to_ne_bytes(); + let start_id = DEFAULT_START_POINT_ID.to_ne_bytes(); let vector = [1.0f32, 2.0]; assert!(import_test_term( ctx, @@ -761,6 +806,13 @@ mod tests { &id, bytemuck::cast_slice(&vector), )); + assert!(import_test_term( + ctx, + index_ptr, + Term::Vector, + &start_id, + bytemuck::cast_slice(&vector), + )); if let Some(quantizer) = quantizer { let mut quantized = vec![0u8; quantizer.bytes()]; quantizer.compress(&vector, &mut quantized).unwrap(); @@ -771,6 +823,13 @@ mod tests { &id, &quantized )); + assert!(import_test_term( + ctx, + index_ptr, + Term::Quantized, + &start_id, + &quantized, + )); } assert!(import_test_term( ctx, @@ -779,6 +838,16 @@ mod tests { &id, bytemuck::cast_slice(&[0u32; 21]), )); + let mut start_neighbors = [0u32; 21]; + start_neighbors[0] = 1; + start_neighbors[20] = 1; + assert!(import_test_term( + ctx, + index_ptr, + Term::Neighbors, + &start_id, + bytemuck::cast_slice(&start_neighbors), + )); 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( @@ -867,6 +936,23 @@ mod tests { .collect(); let term_count = terms[0].len(); + let start_vector = [0.5f32; 2]; + let mut start_neighbors = [0u32; MAX_DEGREE + 1]; + start_neighbors[0] = internal_ids[VECTOR_COUNT - 1]; + start_neighbors[MAX_DEGREE] = 1; + let mut start_terms = vec![ + (Term::Vector, bytemuck::cast_slice(&start_vector).to_vec()), + ( + Term::Neighbors, + bytemuck::cast_slice(&start_neighbors).to_vec(), + ), + ]; + if let Some(quantizer) = quantizer { + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&start_vector, &mut quantized).unwrap(); + start_terms.push((Term::Quantized, quantized)); + } + for by_term in [true, false] { let store = Store::new(); let (index_ptr, ctx) = create_test_index(&store, quant_type); @@ -897,13 +983,16 @@ mod tests { ) }); } - 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 id = DEFAULT_START_POINT_ID.to_ne_bytes(); + let value = b"{}"; + assert!(!import_test_term( + &ctx, + index_ptr, + Term::Attributes, + &id, + value + )); + assert!(store.get(ctx.term(Term::Attributes).get(), &id).is_none()); let rejected_id = (max_id + 1).to_ne_bytes(); let mut bad_neighbors = [0u32; MAX_DEGREE + 1]; @@ -961,10 +1050,21 @@ mod tests { } } + let start_id = DEFAULT_START_POINT_ID.to_ne_bytes(); + for (term, value) in &start_terms { + assert!(import_test_term(&ctx, index_ptr, *term, &start_id, value)); + } + 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 (term, value) in &start_terms { + assert_eq!( + store.get(ctx.get() | *term as u64, &start_id).as_ref(), + Some(value) + ); + } for rejected_id in [1u32, 101] { let id = rejected_id.to_ne_bytes(); for (term, value) in &terms[1] { @@ -1216,8 +1316,7 @@ mod tests { let vector_bytes = bytemuck::cast_slice(&vector); let vector_len = 2; - // First insert will get ID=1 - let iid = 1u32; + let iid = 0u32; let iid_bytes = bytemuck::bytes_of(&iid); // Check internal ID does not exist @@ -1778,7 +1877,7 @@ mod tests { ); // previous insert is unquantized; inserted before training - let iid = required_vectors as u32; + let iid = required_vectors as u32 - 1; assert!( store .get(ctx.term(Term::Quantized).get(), bytemuck::bytes_of(&iid)) @@ -1786,7 +1885,7 @@ mod tests { ); // latest insert is quantized - let iid = required_vectors as u32 + 1; // +1 to account for the start vector + let iid = required_vectors as u32; let qv = store .get(ctx.term(Term::Quantized).get(), bytemuck::bytes_of(&iid)) .expect("missing quant vector"); @@ -1797,7 +1896,7 @@ mod tests { unsafe { backfill_quant_vectors(ctx.get(), index_ptr, 0, 1) }; // all previous inserts are now quantized - for iid in 1..=required_vectors as u32 { + for iid in (0..=required_vectors as u32).chain([DEFAULT_START_POINT_ID]) { let qv = store .get(ctx.term(Term::Quantized).get(), bytemuck::bytes_of(&iid)) .expect("missing quant vector"); diff --git a/diskann-garnet/src/fsm.rs b/diskann-garnet/src/fsm.rs index aee946869..04e5efdc1 100644 --- a/diskann-garnet/src/fsm.rs +++ b/diskann-garnet/src/fsm.rs @@ -75,8 +75,8 @@ struct IdMinter { } /// 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. +/// reclaimed after deletion and gaps left by imports. The configured start-point +/// ID is reserved and is not tracked. /// /// 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. @@ -89,6 +89,7 @@ struct IdMinter { pub(crate) struct FreeSpaceMap { /// Garnet callbacks for reading/writing FSM keys callbacks: Callbacks, + reserved_id: u32, /// A flag to signal whether there are free IDs in the FSM. /// Gaps and deletions set this flag; a scan with no free IDs clears it to /// prevent extraneous reads of FSM blocks. @@ -111,6 +112,7 @@ impl FreeSpaceMap { pub(crate) fn new( ctx: &Context, callbacks: Callbacks, + reserved_id: u32, quantization_enabled: bool, reuse_enabled: bool, ) -> Result { @@ -123,7 +125,7 @@ impl FreeSpaceMap { quantization_enabled, }); let id_minter = RwLock::new(IdMinter { - next_id: 1, + next_id: 0, max_block: u32::MAX, buffer: vec![0u8; BLOCK_SIZE_BYTES], }); @@ -131,6 +133,7 @@ impl FreeSpaceMap { let mut this = Self { callbacks, + reserved_id, has_free_ids, fast_free_list, id_minter, @@ -150,7 +153,7 @@ impl FreeSpaceMap { } else { // Allocate first block. let mut id_minter = this.id_minter.write(); - this.expand_to(&mut id_minter, ctx, 1)?; + this.expand_to(&mut id_minter, ctx, 0)?; } Ok(this) @@ -168,7 +171,7 @@ impl FreeSpaceMap { } let mut block = vec![0u8; BLOCK_SIZE_BYTES]; - let mut last_used_id = 0u32; + let mut last_used_id = -1i64; let mut total_used = 0usize; for block_id in (0..max_block_id).rev() { @@ -185,19 +188,17 @@ impl FreeSpaceMap { 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); - total_used += 1; - } else if id < last_used_id { - let _ = self.fast_free_list.push(id); + if id != self.reserved_id { + let used = bit_used(byte, bidx); + if used { + last_used_id = last_used_id.max(id as i64); + total_used += 1; + } else if (id as i64) < last_used_id { + let _ = self.fast_free_list.push(id); + } } - id -= 1; + id = id.saturating_sub(1); } } } @@ -205,7 +206,7 @@ impl FreeSpaceMap { let mut id_minter = self.id_minter.write(); id_minter.max_block = max_block_id - 1; - id_minter.next_id = last_used_id + 1; + id_minter.next_id = (last_used_id + 1) as u32; let barrier = self.barrier.get_mut(); if barrier.quantization_enabled { @@ -276,7 +277,7 @@ impl FreeSpaceMap { fn mark_id(&self, ctx: &Context, id: u32, used: bool) -> Result { { let id_minter = self.id_minter.read(); - if id == 0 || id >= id_minter.next_id { + if id == self.reserved_id || id >= id_minter.next_id { return Err(FsmError::IdOutOfRange(id)); } } @@ -289,7 +290,7 @@ impl FreeSpaceMap { // Don't hold the lock longer than we have to. let max_block = { let id_minter = self.id_minter.read(); - if id == 0 || id >= id_minter.next_id { + if id == self.reserved_id || id >= id_minter.next_id { return Err(FsmError::IdOutOfRange(id)); } @@ -362,7 +363,10 @@ impl FreeSpaceMap { // Mint a new ID and mark it used. let mut id_minter = self.id_minter.write(); - let id = id_minter.next_id; + let mut id = id_minter.next_id; + if id == self.reserved_id { + id = id.checked_add(1).ok_or(FsmError::IdOutOfRange(id))?; + } 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)?; @@ -378,11 +382,11 @@ impl FreeSpaceMap { /// 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. + /// Repeated claims do not change the used count. The configured start-point ID + /// and u32::MAX are rejected; u32::MAX cannot represent its next ID. pub(crate) fn claim_id(&self, ctx: &Context, id: u32) -> Result<(), FsmError> { - if id == 0 { - return Ok(()); + if id == self.reserved_id { + return Err(FsmError::IdOutOfRange(id)); } let next_id = id.checked_add(1).ok_or(FsmError::IdOutOfRange(id))?; let mut id_minter = self.id_minter.write(); @@ -484,7 +488,7 @@ impl FreeSpaceMap { break 'scan; } - if id != 0 && !bit_used(byte, bidx) { + if id != self.reserved_id && !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() { @@ -563,7 +567,7 @@ impl FreeSpaceMap { for bidx in 0..8 { if bit_used(byte, bidx) { let id = byte_id + bidx as u32; - if id != 0 && !f(id) { + if id != self.reserved_id && !f(id) { return Ok(()); } } @@ -630,17 +634,33 @@ mod tests { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); } + #[test] + fn claim_zero_in_empty_map() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), u32::MAX, false, true).unwrap(); + + fsm.claim_id(&ctx, 0).unwrap(); + // Second call should be idempotent + fsm.claim_id(&ctx, 0).unwrap(); + assert!(!fsm.is_free(&ctx, 0).unwrap()); + assert_eq!(fsm.total_used(), 1); + assert_eq!(fsm.max_id(), 0); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); + assert_eq!(fsm.total_used(), 2); + } + #[test] fn basic_delete() { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.mark_free(&ctx, 1), Err(FsmError::IdOutOfRange(1))); for id in 1..=64 { @@ -662,7 +682,7 @@ mod tests { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); for _ in 0u32..64 { let _ = fsm.next_id(&ctx).unwrap(); @@ -682,7 +702,7 @@ mod tests { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); for _ in 0u32..64 { let _ = fsm.next_id(&ctx).unwrap(); @@ -691,7 +711,7 @@ mod tests { fsm.mark_free(&ctx, 37).unwrap(); // Loading FSM from store should recover all the state. - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.max_id(), 64); assert_eq!(fsm.total_used(), 63); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 37); @@ -702,7 +722,7 @@ mod tests { fn backfill_recovery() { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, false).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, false).unwrap(); for _ in 0..64 { let _ = fsm.next_id(&ctx).unwrap(); @@ -711,7 +731,7 @@ mod tests { fsm.enable_quantization(); drop(fsm); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), true, false).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, true, false).unwrap(); assert_eq!(fsm.max_id_for_backfill(), 64); let next_id = fsm.next_id(&ctx).unwrap(); @@ -728,7 +748,7 @@ mod tests { 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 fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); let block_size = BLOCK_SIZE_IDS as u32; let high_id = 3 * block_size + 7; @@ -773,7 +793,7 @@ mod tests { assert_eq!(fsm.next_id(&ctx).unwrap().id(), 2); drop(fsm); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.max_id(), high_id); assert_eq!(fsm.total_used(), 6); used.clear(); @@ -792,9 +812,9 @@ mod tests { 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(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); - fsm.claim_id(&ctx, 0).unwrap(); + assert_eq!(fsm.claim_id(&ctx, 0), Err(FsmError::IdOutOfRange(0))); assert_eq!(fsm.total_used(), 0); assert_eq!(fsm.max_id(), 0); assert_eq!(fsm.is_free(&ctx, 0), Err(FsmError::IdOutOfRange(0))); @@ -804,7 +824,7 @@ mod tests { assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); drop(fsm); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.total_used(), 1); fsm.mark_free(&ctx, 1).unwrap(); assert_eq!(fsm.next_id(&ctx).unwrap().id(), 1); @@ -824,7 +844,7 @@ mod tests { &block, ); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!(fsm.total_used(), 1); assert_eq!(fsm.max_id(), 1); let mut used = Vec::new(); @@ -869,7 +889,7 @@ mod tests { callbacks.filter_callback(), callbacks.log_callback(), ); - let mut fsm = FreeSpaceMap::new(&ctx, failing_callbacks, false, true).unwrap(); + let mut fsm = FreeSpaceMap::new(&ctx, failing_callbacks, 0, false, true).unwrap(); let high_id = 3 * BLOCK_SIZE_IDS as u32 + 7; assert_eq!( fsm.claim_id(&ctx, high_id), @@ -881,7 +901,7 @@ mod tests { 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(); + fsm = FreeSpaceMap::new(&ctx, callbacks, 0, false, true).unwrap(); } else { fsm.callbacks = callbacks; } @@ -898,7 +918,7 @@ mod tests { 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 fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, false).unwrap(); let high_id = 2 * BLOCK_SIZE_IDS as u32; fsm.claim_id(&ctx, high_id).unwrap(); fsm.enable_quantization(); @@ -917,7 +937,7 @@ mod tests { 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(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); assert_eq!( fsm.claim_id(&ctx, u32::MAX), Err(FsmError::IdOutOfRange(u32::MAX)) @@ -940,7 +960,7 @@ mod tests { let store = Store::new(); let ctx = Context::new(0); - let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), false, true).unwrap(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, true).unwrap(); // Asking for more than BLOCK_SIZE_IDS will force another FSM block to be allocated. for id in 1..=BLOCK_SIZE_IDS as u32 + 1 { @@ -961,13 +981,13 @@ mod tests { 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(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, 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(); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), 0, false, false).unwrap(); assert_eq!(fsm.max_id(), high_id); assert_eq!(fsm.total_used(), 3); let mut visited = Vec::new(); @@ -986,4 +1006,85 @@ mod tests { assert_eq!(fsm.max_id(), high_id + 1); assert_eq!(fsm.total_used(), 5); } + + #[test] + fn reserved_id_is_not_tracked() { + for reserved_id in [0, 2, u32::MAX] { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), reserved_id, false, true).unwrap(); + + assert_eq!(fsm.total_used(), 0); + assert_eq!(fsm.id_minter.read().max_block, 0); + assert_eq!( + fsm.claim_id(&ctx, reserved_id), + Err(FsmError::IdOutOfRange(reserved_id)) + ); + assert_eq!( + fsm.mark_free(&ctx, reserved_id), + Err(FsmError::IdOutOfRange(reserved_id)) + ); + + let expected_ids: Vec = (0..6).filter(|&id| id != reserved_id).take(5).collect(); + for &id in &expected_ids { + assert_eq!(fsm.next_id(&ctx).unwrap().id(), id); + } + assert_eq!(fsm.total_used(), expected_ids.len()); + assert_eq!( + fsm.is_free(&ctx, reserved_id), + Err(FsmError::IdOutOfRange(reserved_id)) + ); + + let mut visited_ids = Vec::new(); + fsm.visit_used(&ctx, |id| { + visited_ids.push(id); + true + }) + .unwrap(); + assert_eq!(visited_ids, expected_ids); + + let deleted_id = expected_ids[1]; + fsm.mark_free(&ctx, deleted_id).unwrap(); + while fsm.fast_free_list.pop().is_some() {} + assert_eq!(fsm.next_id(&ctx).unwrap().id(), deleted_id); + assert_eq!(fsm.total_used(), expected_ids.len()); + + fsm.mark_free(&ctx, deleted_id).unwrap(); + drop(fsm); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), reserved_id, false, true).unwrap(); + assert_eq!(fsm.total_used(), expected_ids.len() - 1); + assert_eq!(fsm.next_id(&ctx).unwrap().id(), deleted_id); + assert_eq!( + fsm.next_id(&ctx).unwrap().id(), + expected_ids.last().unwrap() + 1 + ); + assert_eq!(fsm.total_used(), expected_ids.len() + 1); + + if reserved_id < BLOCK_SIZE_IDS as u32 { + let block = store + .get( + ctx.term(Term::Metadata).get(), + &FreeSpaceMap::block_key(0).to_ne_bytes(), + ) + .unwrap(); + let (_, byte_index, bit_index) = fsm.indexes_for_id(reserved_id); + assert!(!super::bit_used(block[byte_index], bit_index)); + } + } + } + + #[test] + fn reserved_max_id_does_not_overflow() { + let store = Store::new(); + let ctx = Context::new(0); + let fsm = FreeSpaceMap::new(&ctx, store.callbacks(), u32::MAX, false, true).unwrap(); + fsm.id_minter.write().next_id = u32::MAX; + + assert!(matches!( + fsm.next_id(&ctx), + Err(FsmError::IdOutOfRange(u32::MAX)) + )); + assert_eq!(fsm.total_used(), 0); + assert_eq!(fsm.id_minter.read().max_block, 0); + } } diff --git a/diskann-garnet/src/lib.rs b/diskann-garnet/src/lib.rs index 0fb40ed98..2f0f2cc78 100644 --- a/diskann-garnet/src/lib.rs +++ b/diskann-garnet/src/lib.rs @@ -239,12 +239,17 @@ impl SearchOutputBuffer for SearchResults<'_> { /// Helper generic function to create the correct type-erased `Arc`. /// This also returns a bool indicating whether quantization is needed. +#[expect( + clippy::too_many_arguments, + reason = "Mirrors native index creation parameters" +)] fn create_index_impl( quant_type: VectorQuantType, config: config::Config, dim: usize, metric_type: Metric, max_degree: usize, + start_point_id: u32, callbacks: Callbacks, context: Context, ) -> Result<(Arc, bool), GarnetProviderError> { @@ -253,6 +258,7 @@ fn create_index_impl( quant_type, metric_type, max_degree, + start_point_id, callbacks, &context, )?; @@ -293,9 +299,12 @@ fn create_index_impl( /// flag controls whether supplemental control is needed from Garnet to manage quantizers /// which require training and backfill. /// +/// The same immutable `start_point_id` must be supplied when reopening the index. +/// /// # Safety /// -/// FFI +/// Callbacks must remain valid for the index lifetime, and `quantization_needed` +/// must point to writable memory for a bool. #[unsafe(no_mangle)] pub unsafe extern "C" fn create_index( ctx: u64, @@ -305,6 +314,7 @@ pub unsafe extern "C" fn create_index( metric_type: i32, l_build: u32, max_degree: u32, + start_point_id: u32, read_callback: ReadCallback, write_callback: WriteCallback, delete_callback: DeleteCallback, @@ -354,6 +364,7 @@ pub unsafe extern "C" fn create_index( dim as usize, metric_type, max_degree as usize, + start_point_id, callbacks, context, ) { @@ -370,6 +381,7 @@ pub unsafe extern "C" fn create_index( dim as usize, metric_type, max_degree as usize, + start_point_id, callbacks, context, ) { @@ -386,6 +398,7 @@ pub unsafe extern "C" fn create_index( dim as usize, metric_type, max_degree as usize, + start_point_id, callbacks, context, ) { @@ -1457,6 +1470,7 @@ mod tests { use crate::{ Index, IndexState, PolyCow, SearchResults, VectorQuantType, drop_index, garnet::{Context, GarnetId, Term}, + provider::DEFAULT_START_POINT_ID, test_utils::Store, }; @@ -1531,6 +1545,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), @@ -1617,6 +1632,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), @@ -1650,6 +1666,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), @@ -1710,6 +1727,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), @@ -1780,6 +1798,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), @@ -1972,6 +1991,7 @@ mod tests { Metric::L2.into(), 10, 8, + DEFAULT_START_POINT_ID, store.callbacks().read_callback(), store.callbacks().write_callback(), store.callbacks().delete_callback(), diff --git a/diskann-garnet/src/provider.rs b/diskann-garnet/src/provider.rs index 39bc3d0cf..9ff52a52a 100644 --- a/diskann-garnet/src/provider.rs +++ b/diskann-garnet/src/provider.rs @@ -59,6 +59,9 @@ use crate::{ }, }; +#[cfg(test)] +pub(crate) const DEFAULT_START_POINT_ID: u32 = u32::MAX; + /// Quantization state and table are stored under this key in Garnet under the metadata /// term. /// @@ -297,6 +300,8 @@ pub(crate) struct GarnetProvider { /// Note: Unlike DiskANN, this is the true maximum. Neighbors can never /// exceed this degree. max_degree: usize, + /// The internal ID the start point will use + start_point_id: u32, /// Garnet storage engine callbacks callbacks: Callbacks, pending_external_ids: DashMap, foldhash::fast::RandomState>, @@ -346,6 +351,7 @@ impl GarnetProvider { quant_type: VectorQuantType, metric_type: Metric, max_degree: usize, + start_point_id: u32, callbacks: Callbacks, context: &Context, ) -> Result { @@ -375,19 +381,36 @@ impl GarnetProvider { let neighbor_cache = DashMap::with_capacity_and_hasher(1, foldhash::fast::RandomState::default()); + let mut persisted_status = [0u8]; + let saved_import_status = if callbacks.read_single_iid( + &context.term(Term::Metadata), + IMPORT_STATUS_KEY, + &mut persisted_status, + ) { + Some(ImportStatus::try_from(persisted_status[0])?) + } else { + None + }; + // Try to read the start point from Garnet let mut v = Poly::broadcast(0u8, dim * mem::size_of::(), AlignToEight)?; - if callbacks.read_single_iid(&context.term(Term::Vector), 0, &mut v) { + if matches!(saved_import_status, None | Some(ImportStatus::Disabled)) + && callbacks.read_single_iid(&context.term(Term::Vector), start_point_id, &mut v) + { let mut neighbors = vec![0u32; max_degree + 1]; - if !callbacks.read_single_iid(&context.term(Term::Neighbors), 0, &mut neighbors) { + if !callbacks.read_single_iid( + &context.term(Term::Neighbors), + start_point_id, + &mut neighbors, + ) { return Err(GarnetError::Read.into()); } - start_point_cache.insert(0, v); + start_point_cache.insert(start_point_id, v); let len = neighbors[max_degree] as usize; neighbors.truncate(len); - neighbor_cache.insert(0, neighbors); + neighbor_cache.insert(start_point_id, neighbors); } let mut has_quant_state = false; @@ -408,7 +431,7 @@ impl GarnetProvider { has_quant_state = true; quantization::MinMax8Bit::new_from_bytes(metric_type, &quant_state)? } else { - if start_point_cache.contains_key(&0) { + if start_point_cache.contains_key(&start_point_id) { // If we have a start point, we should have had a quantizer. return Err(GarnetProviderError::InvalidQuantizer); } @@ -427,8 +450,12 @@ impl GarnetProvider { // it exists. 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); + if callbacks.read_single_iid( + &context.term(Term::Quantized), + start_point_id, + &mut qsv, + ) { + start_point_quant_cache.insert(start_point_id, qsv); } (Some(quantizer), canonical_bytes, true) @@ -454,9 +481,13 @@ impl GarnetProvider { // 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 && start_point_cache.contains_key(&0) { + if callbacks.read_single_iid( + &context.term(Term::Quantized), + start_point_id, + &mut qsv, + ) { + start_point_quant_cache.insert(start_point_id, qsv); + } else if all_quantized && start_point_cache.contains_key(&start_point_id) { return Err(GarnetProviderError::StartPoint); } } @@ -470,19 +501,15 @@ impl GarnetProvider { let fsm: FreeSpaceMap = FreeSpaceMap::new( context, callbacks, + start_point_id, quantizer .as_ref() .is_some_and(|quantizer| quantizer.is_trained()), 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])? + let import_status = if let Some(status) = saved_import_status { + status } else { let status = if start_point_cache.is_empty() && fsm.total_used() == 0 @@ -513,6 +540,7 @@ impl GarnetProvider { dim, metric_type, max_degree, + start_point_id, callbacks, pending_external_ids: DashMap::with_hasher(foldhash::fast::RandomState::default()), quantizer, @@ -590,18 +618,19 @@ impl GarnetProvider { context: &Context, point: &[T], ) -> Result<(), GarnetProviderError> { - let _guard = self.fsm.existing_id(0); + let _guard = self.fsm.existing_id(self.start_point_id); let mut v = Poly::broadcast(0u8, self.dim * mem::size_of::(), AlignToEight)?; if self .callbacks - .read_single_iid(&context.term(Term::Vector), 0, &mut v) + .read_single_iid(&context.term(Term::Vector), self.start_point_id, &mut v) { // Garnet already has a start point, so use that instead of `point` let mut neighbors = vec![0u32; self.max_degree + 1]; - if !self - .callbacks - .read_single_iid(&context.term(Term::Neighbors), 0, &mut neighbors) - { + if !self.callbacks.read_single_iid( + &context.term(Term::Neighbors), + self.start_point_id, + &mut neighbors, + ) { return Err(GarnetError::Read.into()); } @@ -609,27 +638,30 @@ impl GarnetProvider { && let Some(quantizer) = self.quantizer() { let mut qpoint = vec![0u8; quantizer.bytes()]; - if !self - .callbacks - .read_single_iid(&context.term(Term::Quantized), 0, &mut qpoint) - { + if !self.callbacks.read_single_iid( + &context.term(Term::Quantized), + self.start_point_id, + &mut qpoint, + ) { return Err(GarnetError::Read.into()); } - self.start_point_quant_cache - .insert(0, Poly::from_iter(qpoint.iter().copied(), AlignToEight)?); + self.start_point_quant_cache.insert( + self.start_point_id, + Poly::from_iter(qpoint.iter().copied(), AlignToEight)?, + ); } - self.start_point_cache.insert(0, v); + self.start_point_cache.insert(self.start_point_id, v); let len = neighbors[self.max_degree] as usize; neighbors.truncate(len); - self.neighbor_cache.insert(0, neighbors); + self.neighbor_cache.insert(self.start_point_id, neighbors); } else { let neighbors = vec![0u32; self.max_degree + 1]; if !self .callbacks - .write_iid(&context.term(Term::Vector), 0, point) + .write_iid(&context.term(Term::Vector), self.start_point_id, point) { return Err(GarnetError::Write.into()); } @@ -644,40 +676,45 @@ impl GarnetProvider { T::as_f32(point).map_err(|e| GarnetQuantizerError::Compression(Box::new(e)))?; quantizer.compress(&point_f32, &mut qpoint)?; - if !self - .callbacks - .write_iid(&context.term(Term::Quantized), 0, &qpoint) - { + if !self.callbacks.write_iid( + &context.term(Term::Quantized), + self.start_point_id, + &qpoint, + ) { return Err(GarnetError::Write.into()); } - self.start_point_quant_cache - .insert(0, Poly::from_iter(qpoint.iter().copied(), AlignToEight)?); + self.start_point_quant_cache.insert( + self.start_point_id, + Poly::from_iter(qpoint.iter().copied(), AlignToEight)?, + ); } - if !self - .callbacks - .write_iid(&context.term(Term::Neighbors), 0, &neighbors) - { + if !self.callbacks.write_iid( + &context.term(Term::Neighbors), + self.start_point_id, + &neighbors, + ) { return Err(GarnetError::Write.into()); } self.start_point_cache.insert( - 0, + self.start_point_id, Poly::from_iter( bytemuck::cast_slice::(point).iter().copied(), AlignToEight, )?, ); self.neighbor_cache - .insert(0, Vec::with_capacity(self.max_degree + 1)); + .insert(self.start_point_id, Vec::with_capacity(self.max_degree + 1)); } Ok(()) } pub(crate) fn start_points_exist(&self) -> bool { - self.start_point_cache.get(&0).is_some() && self.neighbor_cache.get(&0).is_some() + self.start_point_cache.contains_key(&self.start_point_id) + && self.neighbor_cache.contains_key(&self.start_point_id) } pub(crate) fn set_attributes( @@ -738,11 +775,11 @@ impl GarnetProvider { } pub(crate) fn vector_iid_exists(&self, context: &Context, id: u32) -> bool { - if id == 0 { - return self.start_points_exist(); + if id == self.start_point_id { + self.start_points_exist() + } else { + !self.fsm.is_free(context, id).unwrap_or(true) } - - !self.fsm.is_free(context, id).unwrap_or(true) } pub(crate) fn max_internal_id(&self) -> u32 { @@ -904,6 +941,10 @@ impl GarnetProvider { let mut f = vec![0f32; self.dim]; let mut q = vec![0u8; quantizer.bytes()]; for id in start_id..end_id { + // Skip the start point as it will be done by the last worker + if id == self.start_point_id { + continue; + } if !self .callbacks .read_single_iid(&context.term(Term::Vector), id, &mut v) @@ -933,13 +974,15 @@ impl GarnetProvider { if backfill_finished { // The final thread to finish backfilling will add the quantized started points. - if let Some(v) = self.start_point_cache.get(&0) + if let Some(v) = self.start_point_cache.get(&self.start_point_id) && let Ok(v_f32) = T::as_f32(bytemuck::cast_slice::(&v)) && quantizer.compress(&v_f32, &mut q).is_ok() { - let _ = self - .callbacks - .write_iid(&context.term(Term::Quantized), 0, &q); + let _ = self.callbacks.write_iid( + &context.term(Term::Quantized), + self.start_point_id, + &q, + ); // set the cache let point = if let Ok(p) = Poly::from_iter(q.iter().copied(), AlignToEight) { @@ -951,7 +994,8 @@ impl GarnetProvider { ); return false; }; - self.start_point_quant_cache.insert(0, point); + self.start_point_quant_cache + .insert(self.start_point_id, point); } // Now that all vectors have quant vectors associated, unlock ID reuse. @@ -1044,7 +1088,7 @@ impl GarnetProvider { let d = ::distance(self.metric_type, Some(self.dim)); let mut result = Vec::with_capacity(self.max_degree); for &nbr_id in neighbors.iter() { - if nbr_id == 0 { + if nbr_id == self.start_point_id { // Skip the start point continue; } @@ -1076,11 +1120,19 @@ impl GarnetProvider { } 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::(), - ) + || [ + (Term::Vector, self.dim * mem::size_of::()), + ( + Term::Neighbors, + (self.max_degree + 1) * mem::size_of::(), + ), + (Term::Quantized, quantizer.bytes()), + ] + .into_iter() + .any(|(term, bytes)| { + self.callbacks + .exists_iid(&context.term(term), self.start_point_id, bytes) + }) { return false; } @@ -1193,8 +1245,10 @@ impl GarnetProvider { 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. + let is_start_point = iid == self.start_point_id; + if (is_start_point && matches!(term, Term::Attributes | Term::ExtMap)) + || (!is_start_point && iid == u32::MAX) + { return false; } @@ -1224,8 +1278,8 @@ impl GarnetProvider { return false; } - // Reserve the ID before storing data. - if self.fsm.claim_id(context, iid).is_err() { + // Reserve ordinary IDs before storing data; the start point stays outside the FSM. + if !is_start_point && self.fsm.claim_id(context, iid).is_err() { return false; } @@ -1241,7 +1295,7 @@ impl GarnetProvider { Term::Metadata => false, Term::IntMap => { if value.len() != mem::size_of::() - || bytemuck::pod_read_unaligned::(value) == 0 + || bytemuck::pod_read_unaligned::(value) == self.start_point_id { return false; } @@ -1258,9 +1312,10 @@ impl GarnetProvider { /// 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. + /// Callers must finish all term writes before starting verification tasks. + /// The last successful task publishes `ImportStatus::Disabled` only after + /// finalization and persistence. Start points must be supplied even when no + /// ordinary vectors were imported. pub(crate) fn finish_import( &self, context: &Context, @@ -1279,6 +1334,7 @@ impl GarnetProvider { } let max_id = self.fsm.max_id(); + let empty_import = self.fsm.total_used() == 0; let (start_id, end_id) = if let Some((start, end)) = self.task_range(task_idx, task_count, max_id) { @@ -1304,7 +1360,11 @@ impl GarnetProvider { return (true, None); } - for id in start_id.max(1)..end_id { + let ids = if empty_import { 0..0 } else { start_id..end_id }; + for id in ids { + if id == self.start_point_id { + continue; + } // For every ID that is used, check all terms exist for that ID match self.fsm.is_free(context, id) { Ok(true) => continue, @@ -1375,37 +1435,6 @@ impl GarnetProvider { 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) => { @@ -1413,11 +1442,15 @@ impl GarnetProvider { 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}")); + if !self.callbacks.read_single_iid( + &context.term(Term::Vector), + self.start_point_id, + &mut v, + ) { + log_failure(format_args!( + "start point read failed: id={}", + self.start_point_id + )); return (true, Some(false)); } let qv = if let Some(quantizer) = &self.quantizer { @@ -1434,11 +1467,12 @@ impl GarnetProvider { if !self.callbacks.read_single_iid( &context.term(Term::Quantized), - first_id, + self.start_point_id, &mut qv, ) { log_failure(format_args!( - "quantized start point read failed: id={first_id}" + "quantized start point read failed: id={}", + self.start_point_id )); return (true, Some(false)); } @@ -1448,45 +1482,38 @@ impl GarnetProvider { None }; let mut ns = vec![0u32; self.max_degree + 1]; - if !self - .callbacks - .read_single_iid(&context.term(Term::Neighbors), first_id, &mut ns) - { + if !self.callbacks.read_single_iid( + &context.term(Term::Neighbors), + self.start_point_id, + &mut ns, + ) { log_failure(format_args!( - "start point neighbors read failed: id={first_id}" + "start point neighbors read failed: id={}", + self.start_point_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")); + if ns_len > self.max_degree { + log_failure(format_args!("start point neighbor count exceeds degree")); return (true, Some(false)); } - // NOTE: unwrap will succeed because it is gated on the quantizer. - if let Some(qv) = &qv - && !self + for term in [Term::Attributes, Term::ExtMap] { + if 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)); + .exists_iid(&context.term(term), self.start_point_id, 0) + { + log_failure(format_args!("unexpected start point term: {term:?}")); + return (true, Some(false)); + } } - self.start_point_cache.insert(0, v); + self.start_point_cache.insert(self.start_point_id, 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.start_point_quant_cache.insert(self.start_point_id, qv); } - self.neighbor_cache.insert(0, ns[0..ns_len].to_vec()); + self.neighbor_cache + .insert(self.start_point_id, ns[0..ns_len].to_vec()); let finalized = self.transition_import_status( context, @@ -1518,9 +1545,7 @@ impl GarnetProvider { pub(crate) fn quantization_needed(&self) -> bool { if let Some(quantizer) = &self.quantizer { - !self.is_quantized() - && quantizer.is_trained() - && self.max_internal_id() as usize > quantizer.required_vectors() + !self.is_quantized() && quantizer.is_trained() } else { false } @@ -1533,7 +1558,7 @@ impl GarnetProvider { ) -> Result, GarnetProviderError> { let mut v = vec![T::default(); self.dim]; - if iid == 0 { + if iid == self.start_point_id { let guard = if let Some(r) = self.start_point_cache.get(&iid) { r } else { @@ -1561,7 +1586,7 @@ impl GarnetProvider { ) -> bool { let mut guard = neighbors.resize(self.max_degree + 1); - if iid == 0 + if iid == self.start_point_id && let Some(cached) = self.neighbor_cache.get(&iid) { guard[0..cached.len()].copy_from_slice(&cached); @@ -1601,7 +1626,7 @@ impl GarnetProvider { (self.max_degree + 1) * mem::size_of::(), |data: &mut [u32]| { data.copy_from_slice(&guard); - if iid == 0 { + if iid == self.start_point_id { self.neighbor_cache.insert(iid, neighbors.to_vec()); } }, @@ -1642,7 +1667,7 @@ impl GarnetProvider { data[max_degree] = len as u32; } - if iid == 0 + if iid == self.start_point_id && let Some(mut ns) = self.neighbor_cache.get_mut(&iid) { ns.clear(); @@ -1901,10 +1926,13 @@ 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)); + if id == self.start_point_id { + return future::ready(if self.start_points_exist() { + Ok(ElementStatus::Valid) + } else { + Err(GarnetProviderError::StartPoint) + }); } - let status = match self.fsm.is_free(context, id) { Ok(true) => ElementStatus::Deleted, Ok(false) => ElementStatus::Valid, @@ -1941,8 +1969,6 @@ pub(crate) struct DynamicAccessor<'a, T: VectorRepr> { } impl<'a, T: VectorRepr> DynamicAccessor<'a, T> { - const START_ID: u32 = 0; - pub(crate) fn new( provider: &'a GarnetProvider, context: &'a Context, @@ -1985,17 +2011,25 @@ impl<'a, T: VectorRepr> DynamicAccessor<'a, T> { }) } - /// Return the distance to the start point (the point with `ID == 0`). + /// Return the distance to the start point. fn start_point_distance(&mut self) -> Result { if self.quantized && let Some(_quantizer) = self.provider.quantizer() { - match self.provider.start_point_quant_cache.get(&Self::START_ID) { + match self + .provider + .start_point_quant_cache + .get(&self.provider.start_point_id) + { Some(guard) => Ok(self.computer.evaluate_similarity(&*guard)), None => Err(GarnetProviderError::Garnet(GarnetError::Read)), } } else { - match self.provider.start_point_cache.get(&Self::START_ID) { + match self + .provider + .start_point_cache + .get(&self.provider.start_point_id) + { Some(guard) => Ok(self.computer.evaluate_similarity(&*guard)), None => Err(GarnetProviderError::Garnet(GarnetError::Read)), } @@ -2036,7 +2070,7 @@ impl HasId for DynamicAccessor<'_, T> { impl SearchAccessor for DynamicAccessor<'_, T> { fn starting_points(&self) -> impl Future>> + Send { let points = if self.provider.start_points_exist() { - vec![Self::START_ID] + vec![self.provider.start_point_id] } else { vec![] }; @@ -2047,7 +2081,8 @@ impl SearchAccessor for DynamicAccessor<'_, T> { &self, ) -> impl Future bool + Send + Sync + 'static>> + Send { - future::ready(Ok(move |id| id != Self::START_ID)) + let start_point_id = self.provider.start_point_id; + future::ready(Ok(move |id| id != start_point_id)) } fn start_point_distances(&mut self, mut f: F) -> impl Future> + Send @@ -2062,7 +2097,7 @@ impl SearchAccessor for DynamicAccessor<'_, T> { let result = match self.start_point_distance() { Ok(dist) => { - f(Self::START_ID, dist); + f(self.provider.start_point_id, dist); Ok(()) } Err(err) => Err(ANNError::from(err)), @@ -2091,7 +2126,7 @@ impl SearchAccessor for DynamicAccessor<'_, T> { .get_neighbors(self.context, nl_id, &mut id_buffer); self.filtered_ids.clear(); for id in id_buffer.iter().copied().filter(|id| pred.eval_mut(id)) { - if id == Self::START_ID { + if id == self.provider.start_point_id { let dist = match self.start_point_distance() { Ok(dist) => dist, Err(err) => return future::ready(Err(ANNError::from(err))), @@ -2316,7 +2351,7 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { let result = match self.start_point_distance() { Ok(dist) => { - f(glue::Decision::reject(Self::START_ID), dist); + f(glue::Decision::reject(self.provider.start_point_id), dist); Ok(()) } Err(err) => Err(ANNError::from(err)), @@ -2349,7 +2384,7 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { self.filtered_ids.clear(); for id in id_buffer.iter().copied().filter(|id| pred.eval_mut(id)) { - if id == Self::START_ID { + if id == self.provider.start_point_id { let dist = match self.start_point_distance() { Ok(dist) => dist, Err(err) => return future::ready(Err(ANNError::from(err))), @@ -2423,7 +2458,7 @@ impl FilteredAccessor for DynamicAccessor<'_, T> { self.filtered_ids.clear(); for id in id_buffer.iter().copied() { - if id != Self::START_ID && pred.eval(&id) { + if id != self.provider.start_point_id && pred.eval(&id) { self.filtered_ids.push(4); self.filtered_ids.push(id); } @@ -2591,7 +2626,7 @@ where self.filtered_ids.clear(); for id in itr { - if id == 0 { + if id == self.provider.start_point_id { if self.quantized && let Entry::Vacant(e) = self.set.entry(id) { @@ -2830,7 +2865,7 @@ mod tests { config::{self, defaults::GRAPH_SLACK_FACTOR}, search, }, - provider::{DataProvider, Delete, Guard, SetElement}, + provider::{DataProvider, Delete, ElementStatus, Guard, SetElement}, }; use diskann_providers::index::wrapped_async::DiskANNIndex; use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; @@ -2845,7 +2880,10 @@ mod tests { Callbacks, Context, GarnetError, GarnetId, ReadDataCallback, RmwDataCallback, TERM_BITMASK, Term, WriteCallback, }, - provider::{GarnetProvider, GarnetProviderError, QUANT_STATE_KEY, RESERVATION_RETRY_LIMIT}, + provider::{ + DEFAULT_START_POINT_ID, GarnetProvider, GarnetProviderError, QUANT_STATE_KEY, + RESERVATION_RETRY_LIMIT, + }, quantization::{GarnetQuantizer, Spherical1Bit}, test_utils::{LOGS, Store, q8_state_with_identity_transform}, }; @@ -2859,6 +2897,7 @@ mod tests { VectorQuantType::NoQuant, Metric::L2, 10, + DEFAULT_START_POINT_ID, store.callbacks(), &ctx, ) @@ -3036,7 +3075,7 @@ mod tests { #[test] fn backfill_waits_for_overlapping_updates() { - for (range, should_wait) in [(1..2, true), (2..3, false)] { + for (offset, should_wait) in [(0, true), (1, false)] { let store = Arc::new(DashMap::new()); let state = ParallelContext::new(store); let ctx = state.context(); @@ -3052,7 +3091,8 @@ mod tests { provider.maybe_set_start_point(&ctx, &original).unwrap(); DynIndex::insert(&index, &ctx, &id, bytemuck::cast_slice(&original), &[]).unwrap(); train_for_backfill(provider, &ctx); - assert_waits_for_range(provider, range, should_wait, || { + let range_start = provider.to_internal_id(&ctx, &id).unwrap() + offset; + assert_waits_for_range(provider, range_start..range_start + 1, should_wait, || { assert!(provider.backfill_quant_vectors(&ctx, 0, 1)); }); assert!(provider.all_quantized.load(Ordering::Acquire)); @@ -3312,7 +3352,11 @@ mod tests { 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()), + parallel_get( + &store, + ctx.term(Term::Vector).get(), + &provider.start_point_id.to_ne_bytes(), + ), (!import_first).then(|| bytemuck::cast_slice(&ordinary_vector).to_vec()) ); assert_eq!( @@ -3400,7 +3444,12 @@ mod tests { Some(second_value) ); assert!( - parallel_get(&store, ctx.term(Term::Vector).get(), &0u32.to_ne_bytes()).is_none() + parallel_get( + &store, + ctx.term(Term::Vector).get(), + &provider.start_point_id.to_ne_bytes(), + ) + .is_none() ); assert_eq!( provider.quantizer().unwrap().serialize().unwrap(), @@ -3516,7 +3565,11 @@ mod tests { ), Some(expected_state.to_vec()) ); - let id = if importing { 7u32 } else { 0u32 }; + let id = if importing { + 7u32 + } else { + provider.start_point_id + }; assert_eq!( parallel_get(&store, ctx.term(Term::Quantized).get(), &id.to_ne_bytes()), Some(quantized) @@ -3826,6 +3879,7 @@ mod tests { VectorQuantType::NoQuant, Metric::L2, 10, + DEFAULT_START_POINT_ID, store.callbacks(), &ctx, ) @@ -3885,6 +3939,7 @@ mod tests { VectorQuantType::NoQuant, Metric::L2, 10, + DEFAULT_START_POINT_ID, store.callbacks(), &ctx, ) @@ -4034,9 +4089,16 @@ mod tests { ] { let store = Store::new(); let ctx = Context::new(0); - let mut provider = - GarnetProvider::::new(2, quant_type, Metric::L2, 10, store.callbacks(), &ctx) - .unwrap(); + let mut provider = GarnetProvider::::new( + 2, + quant_type, + Metric::L2, + 10, + DEFAULT_START_POINT_ID, + store.callbacks(), + &ctx, + ) + .unwrap(); let id = GarnetId::from(bytemuck::bytes_of(&42u32)); let original = [0.0f32, 1.0]; provider.maybe_set_start_point(&ctx, &original).unwrap(); @@ -4168,6 +4230,7 @@ mod tests { VectorQuantType::Bin, Metric::L2, 10, + DEFAULT_START_POINT_ID, store.callbacks(), &ctx, ) @@ -4175,7 +4238,10 @@ mod tests { 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(), + provider + .status_by_internal_id(&ctx, provider.start_point_id) + .await + .unwrap(), diskann::provider::ElementStatus::Valid ); @@ -4205,11 +4271,183 @@ mod tests { assert!(provider.is_quantized()); assert!(provider.callbacks.exists_iid( &ctx.term(Term::Quantized), - 0, + provider.start_point_id, provider.quant_vector_size(), )); } + #[test] + fn import_start_point_terms() { + for start_point_id in [DEFAULT_START_POINT_ID, 0, 2] { + for quant_type in [VectorQuantType::NoQuant, VectorQuantType::Q8] { + let store = Store::new(); + let ctx = Context::new(0); + let create_provider = || { + GarnetProvider::::new( + 2, + quant_type, + Metric::L2, + 10, + start_point_id, + store.callbacks(), + &ctx, + ) + .unwrap() + }; + let provider = create_provider(); + assert_eq!( + store.get( + ctx.term(Term::Metadata).get(), + &super::IMPORT_STATUS_KEY.to_ne_bytes(), + ), + Some(vec![u8::from(quant_type == VectorQuantType::NoQuant)]) + ); + if let Some(quantizer) = provider.quantizer() { + assert!(provider.set_quant_state(&ctx, &quantizer.serialize().unwrap())); + } + + let id_bytes = start_point_id.to_ne_bytes(); + let point = [1.0f32, 2.0]; + assert!(provider.import_term( + &ctx, + Term::Vector, + &id_bytes, + bytemuck::cast_slice(&point) + )); + drop(provider); + let provider = create_provider(); + assert!(provider.can_import(&ctx)); + assert!(!provider.start_points_exist()); + if let Some(quantizer) = provider.quantizer() { + let state = quantizer.serialize().unwrap(); + assert!(!provider.set_quant_state(&ctx, &q8_state_with_identity_transform(2))); + assert_eq!(quantizer.serialize().unwrap(), state); + } + assert!(provider.import_term( + &ctx, + Term::Neighbors, + &id_bytes, + bytemuck::cast_slice(&[0u32; 11]) + )); + if let Some(quantizer) = provider.quantizer() { + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&point, &mut quantized).unwrap(); + assert!(provider.import_term(&ctx, Term::Quantized, &id_bytes, &quantized)); + } else { + assert!(!provider.import_term(&ctx, Term::Quantized, &id_bytes, &[0u8])); + } + + assert!(!provider.import_term(&ctx, Term::Attributes, &id_bytes, b"attrs")); + assert!(!provider.import_term(&ctx, Term::ExtMap, &id_bytes, b"external")); + assert!(!provider.import_term(&ctx, Term::IntMap, b"external", &id_bytes)); + for term in [Term::Attributes, Term::ExtMap] { + assert!(store.get(ctx.term(term).get(), &id_bytes).is_none()); + } + assert!( + store + .get(ctx.term(Term::IntMap).get(), b"external") + .is_none() + ); + assert_eq!(provider.fsm.total_used(), 0); + assert_eq!(provider.fsm.max_id(), 0); + assert!(!provider.start_points_exist()); + assert_eq!(provider.finish_import(&ctx, 0, 1), (true, Some(true))); + assert!(provider.start_points_exist()); + assert_eq!( + provider.get_full_vector(&ctx, start_point_id).unwrap(), + point + ); + assert_eq!(provider.fsm.total_used(), 0); + } + } + } + + #[test] + fn finish_import_requires_start_point_terms() { + for missing_term in 0..3 { + let store = Store::new(); + let ctx = Context::new(0); + let provider = GarnetProvider::::new( + 2, + VectorQuantType::Q8, + Metric::L2, + 10, + DEFAULT_START_POINT_ID, + store.callbacks(), + &ctx, + ) + .unwrap(); + let quantizer = provider.quantizer().unwrap(); + assert!(provider.set_quant_state(&ctx, &quantizer.serialize().unwrap())); + let point = [1.0f32, 2.0]; + let neighbors = [0u32; 11]; + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&point, &mut quantized).unwrap(); + for (term_index, (term, value)) in [ + (Term::Vector, bytemuck::cast_slice(&point)), + (Term::Neighbors, bytemuck::cast_slice(&neighbors)), + (Term::Quantized, quantized.as_slice()), + ] + .into_iter() + .enumerate() + { + if term_index != missing_term { + assert!(provider.import_term( + &ctx, + term, + &DEFAULT_START_POINT_ID.to_ne_bytes(), + value + )); + } + } + + assert_eq!(provider.finish_import(&ctx, 0, 1), (true, Some(false))); + assert!(!provider.start_points_exist()); + } + } + + #[test] + fn finish_import_rejects_invalid_start_point_neighbors() { + for (neighbor_count, neighbor) in [(11, 0), (1, 0), (1, DEFAULT_START_POINT_ID)] { + let store = Store::new(); + let ctx = Context::new(0); + let provider = GarnetProvider::::new( + 2, + VectorQuantType::NoQuant, + Metric::L2, + 10, + DEFAULT_START_POINT_ID, + store.callbacks(), + &ctx, + ) + .unwrap(); + let mut neighbors = [0u32; 11]; + neighbors[0] = neighbor; + neighbors[10] = neighbor_count; + let id = DEFAULT_START_POINT_ID.to_ne_bytes(); + assert!(provider.import_term( + &ctx, + Term::Vector, + &id, + bytemuck::cast_slice(&[1.0f32, 2.0]) + )); + let imported = + provider.import_term(&ctx, Term::Neighbors, &id, bytemuck::cast_slice(&neighbors)); + assert_eq!(imported, neighbor_count <= 10); + if !imported { + // Exercise finalization's check for malformed data already in storage. + store.set( + ctx.term(Term::Neighbors).get(), + &id, + bytemuck::cast_slice(&neighbors), + ); + } + + assert_eq!(provider.finish_import(&ctx, 0, 1), (true, Some(false))); + assert!(!provider.start_points_exist()); + } + } + fn create_2d_f32_index( quant_type: VectorQuantType, metric: Metric, @@ -4225,8 +4463,16 @@ mod tests { callbacks: Callbacks, ctx: &Context, ) -> DiskANNIndex> { - let provider = - GarnetProvider::::new(2, quant_type, metric, 10, callbacks, ctx).unwrap(); + let provider = GarnetProvider::::new( + 2, + quant_type, + metric, + 10, + DEFAULT_START_POINT_ID, + callbacks, + ctx, + ) + .unwrap(); let config = config::Builder::new( (10.0 / GRAPH_SLACK_FACTOR) as usize, @@ -4240,6 +4486,83 @@ mod tests { DiskANNIndex::new_with_current_thread_runtime(config, provider) } + #[tokio::test] + async fn start_point_is_frozen_and_untracked() { + for start_point_id in [DEFAULT_START_POINT_ID, 0, 2] { + for quant_type in [VectorQuantType::NoQuant, VectorQuantType::Q8] { + let store = Store::new(); + let ctx = Context::new(0); + let create_provider = || { + GarnetProvider::::new( + 2, + quant_type, + Metric::L2, + 10, + start_point_id, + store.callbacks(), + &ctx, + ) + .unwrap() + }; + let provider = create_provider(); + let point = [1.0f32, 2.0]; + + assert!(!provider.start_points_exist()); + assert!(!provider.vector_iid_exists(&ctx, start_point_id)); + assert!(matches!( + provider.status_by_internal_id(&ctx, start_point_id).await, + Err(GarnetProviderError::StartPoint) + )); + assert!(provider.disable_import(&ctx)); + provider.maybe_set_start_point(&ctx, &point).unwrap(); + assert!(provider.start_points_exist()); + assert!(provider.vector_iid_exists(&ctx, start_point_id)); + assert_eq!(provider.fsm.total_used(), 0); + assert_eq!(provider.max_internal_id(), 0); + assert!(provider.fsm.is_free(&ctx, start_point_id).is_err()); + assert_eq!( + provider + .status_by_internal_id(&ctx, start_point_id) + .await + .unwrap(), + ElementStatus::Valid + ); + + provider.maybe_set_start_point(&ctx, &[9.0, 9.0]).unwrap(); + assert_eq!( + provider.get_full_vector(&ctx, start_point_id).unwrap(), + point + ); + drop(provider); + + let provider = create_provider(); + assert!(provider.start_points_exist()); + assert_eq!(provider.total_used(), 0); + assert_eq!( + provider.get_full_vector(&ctx, start_point_id).unwrap(), + point + ); + assert!(provider.to_external_id(&ctx, start_point_id).is_err()); + for term in [Term::Attributes, Term::ExtMap, Term::IntMap] { + assert!( + store + .get(ctx.term(term).get(), bytemuck::bytes_of(&start_point_id)) + .is_none() + ); + } + if let Some(quantizer) = provider.quantizer() { + let mut quantized = vec![0u8; quantizer.bytes()]; + quantizer.compress(&point, &mut quantized).unwrap(); + let cached = provider + .start_point_quant_cache + .get(&start_point_id) + .unwrap(); + assert_eq!(&**cached, quantized.as_slice()); + } + } + } + } + /// Test that restarts during phase one quant bootstrap work. /// Phase one is all index activity before the index has the required /// number of vectors to begin quantization. @@ -4276,7 +4599,7 @@ mod tests { assert!(!provider.is_quantized()); let max_id = provider.max_internal_id(); - assert_eq!(max_id, last_inserted_id + 1); + assert_eq!(max_id, last_inserted_id); // There should be no saved quant state. assert!( @@ -4354,7 +4677,7 @@ mod tests { // is_quantized won't be true until backfill is complete assert!(!provider.is_quantized()); let max_id = provider.max_internal_id(); - assert_eq!(max_id, last_inserted_id + 1); + assert_eq!(max_id, last_inserted_id); // There should be saved quant state. assert!( @@ -4379,7 +4702,7 @@ mod tests { assert!(!provider.is_quantized()); let max_id = provider.max_internal_id(); - assert_eq!(max_id, last_inserted_id + 1); + assert_eq!(max_id, last_inserted_id); // Quant should be needed now, since backfill has never run assert!(provider.quantization_needed()); diff --git a/diskann-garnet/src/test_utils.rs b/diskann-garnet/src/test_utils.rs index 01663f88a..fe75aae50 100644 --- a/diskann-garnet/src/test_utils.rs +++ b/diskann-garnet/src/test_utils.rs @@ -304,7 +304,7 @@ mod tests { #[test] fn garnet_provider_with_store_callbacks() { - use crate::provider::GarnetProvider; + use crate::provider::{DEFAULT_START_POINT_ID, GarnetProvider}; use diskann_vector::distance::Metric; let store: Store = Store::new(); @@ -319,6 +319,7 @@ mod tests { VectorQuantType::NoQuant, Metric::L2, max_degree, + DEFAULT_START_POINT_ID, callbacks, &ctx, ); @@ -338,6 +339,7 @@ mod tests { VectorQuantType::NoQuant, Metric::L2, max_degree, + DEFAULT_START_POINT_ID, callbacks, &ctx, );