Add InsertResult::SuccessUpdate

diskann-garnet/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "diskann-garnet"
-version = "5.0.3"
+version = "6.0.0"
 edition = "2024"
 authors.workspace = true
 license.workspace = true

diskann-garnet/diskann-garnet.nuspec

@@ -2,7 +2,7 @@
 <package>
   <metadata>
     <id>diskann-garnet</id>
-    <version>5.0.3</version>
+    <version>6.0.0</version>
     <readme>docs/README.md</readme>
     <authors>Microsoft</authors>
     <projectUrl>https://github.com/microsoft/DiskANN</projectUrl>

diskann-garnet/docs/data-design.md

@@ -36,7 +36,7 @@ Diskann-garnet uses these bits to distinguish between differnet kinds of index d
 
 ### Key Data Prefixing
 
-In order to reduce allocations in the data access path in Garnet, Garnet needs some place to scribble state into during operations. It uses a single byte immediately preceding the first key byte for this purpose. This means that any key pointer given to Garnet access methods must contain valid space preceding the real key. For this reason, key data pointers are `*mut u8` and not `*const u8` and care must be taken to ensure the memory preceding that pointer is valid. In diskann-garnet, we precede the key data with at least 4 bytes of scratch space.
+Read callbacks receive keys with four-byte length prefixes. Write, delete, and RMW callbacks receive the key bytes and length separately.
 
 ## Term Types
 

diskann-garnet/docs/ffi-design.rs

@@ -36,11 +36,13 @@ enum VectorQuantType {
 /// Status returned by `insert`, encoded as a `u8`.
 ///
 /// `SuccessStartTraining` signals that the insert crossed the threshold at which the quantizer
-/// can be trained, and that Garnet should call `build_quant_table`.
+/// can be trained, and that Garnet should call `build_quant_table`. Otherwise, `Success` is returned for a newly
+/// inserted vector, and `SuccessUpdate` for an updated one.
 enum InsertResult {
     Fail = 0,
     Success = 1,
     SuccessStartTraining = 2,
+    SuccessUpdate = 3,
 }
 
 /// Read one or more keys from Garnet.
@@ -155,8 +157,8 @@ extern "C" fn drop_index(context: u64, index_ptr: *const c_void);
 
 /// Insert a vector into an index.
 ///
-/// Returns an `InsertResult` discriminant. `Fail` may result from the vector already being in
-/// the index, or from writes failing.
+/// Returns an `InsertResult` discriminant, distinguishing inserts from updates unless quantizer
+/// training becomes ready. `Fail` may result from invalid input or writes failing.
 ///
 /// vector_len is a count of elements, not bytes; the element type follows from the index's
 /// `quant_type`. The pointer need not be aligned for that element type.

diskann-garnet/src/dyn_index.rs

@@ -116,26 +116,37 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     ///
     /// 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<()> {
-        self.insert(
-            &DynamicQuantization,
-            context,
-            id,
-            (bytemuck::cast_slice::<u8, T>(data), attrs),
-        )
+        self.run(|_| async {
+            let _pending = self.inner.provider().reserve_external_id(id).await?;
+            self.inner
+                .insert(
+                    &DynamicQuantization,
+                    context,
+                    id,
+                    (bytemuck::cast_slice::<u8, T>(data), attrs),
+                )
+                .await
+        })
     }
 
     fn set_attributes(&self, context: &Context, id: &GarnetId, data: &[u8]) -> ANNResult<()> {
-        self.inner
-            .provider()
-            .set_attributes(context, id, data)
-            .map_err(|e| e.into())
+        self.run(|_| async {
+            let provider = self.inner.provider();
+            let _pending = provider.reserve_external_id(id).await?;
+            provider
+                .set_attributes(context, id, data)
+                .map_err(|error| error.into())
+        })
     }
 
     fn delete_attributes(&self, context: &Context, id: &GarnetId) -> ANNResult<()> {
-        self.inner
-            .provider()
-            .delete_attributes(context, id)
-            .map_err(|e| e.into())
+        self.run(|_| async {
+            let provider = self.inner.provider();
+            let _pending = provider.reserve_external_id(id).await?;
+            provider
+                .delete_attributes(context, id)
+                .map_err(|error| error.into())
+        })
     }
 
     fn search_vector(
@@ -189,13 +200,18 @@ impl<T: VectorRepr> DynIndex for DiskANNIndex<GarnetProvider<T>> {
     }
 
     fn remove(&self, context: &Context, id: &GarnetId) -> ANNResult<()> {
-        self.inplace_delete(
-            DynamicQuantization,
-            context,
-            id,
-            3,
-            InplaceDeleteMethod::TwoHopAndOneHop,
-        )
+        self.run(|_| async {
+            let _pending = self.inner.provider().reserve_external_id(id).await?;
+            self.inner
+                .inplace_delete(
+                    DynamicQuantization,
+                    context,
+                    id,
+                    3,
+                    InplaceDeleteMethod::TwoHopAndOneHop,
+                )
+                .await
+        })
     }
 
     fn approximate_count(&self) -> u64 {

diskann-garnet/src/ffi_recall_tests.rs

@@ -125,9 +125,9 @@ mod tests {
         let mut result_ids = Vec::with_capacity(count);
         let mut offset = 0;
         for _ in 0..count {
-            let mut id_len = 0u32;
-            bytemuck::bytes_of_mut(&mut id_len)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u32>()]);
+            let id_len = bytemuck::pod_read_unaligned::<u32>(
+                &output_id_buffer[offset..offset + mem::size_of::<u32>()],
+            );
             offset += mem::size_of::<u32>();
             let id_str = std::str::from_utf8(&output_id_buffer[offset..offset + id_len as usize])
                 .expect("id should be valid utf8");

diskann-garnet/src/ffi_tests.rs

@@ -123,6 +123,64 @@ mod tests {
         }
     }
 
+    #[test]
+    fn insert_distinguishes_updates() {
+        for quant_type in [
+            VectorQuantType::NoQuant,
+            VectorQuantType::Bin,
+            VectorQuantType::Q8,
+        ] {
+            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),
+                (42, [2.0, 1.0], 3, 1),
+                (43, [3.0, 4.0], 1, 2),
+                (43, [4.0, 3.0], 3, 2),
+            ] {
+                assert_eq!(
+                    u8::from(insert_f32_vector(&ctx, index_ptr, id, &vector)),
+                    expected,
+                    "unexpected insert status for {quant_type:?} and ID {id}"
+                );
+                let internal_id = store
+                    .get(ctx.term(Term::IntMap).get(), bytemuck::bytes_of(&id))
+                    .unwrap();
+                assert_eq!(
+                    bytemuck::pod_read_unaligned::<u32>(&internal_id),
+                    expected_internal_id
+                );
+                assert_eq!(
+                    store.get(ctx.term(Term::Vector).get(), &internal_id),
+                    Some(bytemuck::cast_slice::<f32, u8>(&vector).to_vec())
+                );
+                assert_eq!(
+                    unsafe { card(ctx.get(), index_ptr) },
+                    u64::from(expected_internal_id)
+                );
+            }
+            assert_eq!(store.int_map_reads(), 4);
+
+            let (ids, distances) = do_search(&ctx, index_ptr, &[2.0, 1.0], 3, None);
+            assert_eq!(ids, [42, 43]);
+            assert_eq!(distances[0], 0.0);
+
+            let id_bytes = bytemuck::bytes_of(&42u32);
+            assert!(unsafe { remove(ctx.get(), index_ptr, id_bytes.as_ptr(), id_bytes.len()) });
+            store.clear_read_counts();
+            assert_eq!(
+                insert_f32_vector(&ctx, index_ptr, 42, &[5.0, 6.0]),
+                InsertResult::Success
+            );
+            assert_eq!(store.int_map_reads(), 1);
+
+            unsafe {
+                drop_index(ctx.get(), index_ptr);
+            }
+        }
+    }
+
     #[test]
     fn add_check_and_remove_vector() {
         let store = Store::new();
@@ -493,16 +551,16 @@ mod tests {
         let mut output_ids = vec![];
         let mut offset = 0;
         for _ in 0..(count as usize) {
-            let mut id_len = 0u32;
-            bytemuck::bytes_of_mut(&mut id_len)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u32>()]);
+            let id_len = bytemuck::pod_read_unaligned::<u32>(
+                &output_id_buffer[offset..offset + mem::size_of::<u32>()],
+            );
             offset += mem::size_of::<u32>();
 
             assert_eq!(id_len, mem::size_of::<u64>() as u32);
 
-            let mut id = 0u64;
-            bytemuck::bytes_of_mut(&mut id)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u64>()]);
+            let id = bytemuck::pod_read_unaligned::<u64>(
+                &output_id_buffer[offset..offset + mem::size_of::<u64>()],
+            );
             offset += mem::size_of::<u64>();
 
             output_ids.push(id);
@@ -603,16 +661,16 @@ mod tests {
         let mut output_ids = vec![];
         let mut offset = 0;
         for _ in 0..(count as usize) {
-            let mut id_len = 0u32;
-            bytemuck::bytes_of_mut(&mut id_len)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u32>()]);
+            let id_len = bytemuck::pod_read_unaligned::<u32>(
+                &output_id_buffer[offset..offset + mem::size_of::<u32>()],
+            );
             offset += mem::size_of::<u32>();
 
             assert_eq!(id_len, mem::size_of::<u64>() as u32);
 
-            let mut id = 0u64;
-            bytemuck::bytes_of_mut(&mut id)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u64>()]);
+            let id = bytemuck::pod_read_unaligned::<u64>(
+                &output_id_buffer[offset..offset + mem::size_of::<u64>()],
+            );
             offset += mem::size_of::<u64>();
 
             output_ids.push(id);
@@ -772,13 +830,13 @@ mod tests {
         let mut ids = vec![];
         let mut offset = 0;
         for _ in 0..count {
-            let mut id_len = 0u32;
-            bytemuck::bytes_of_mut(&mut id_len)
-                .copy_from_slice(&output_id_buffer[offset..offset + mem::size_of::<u32>()]);
+            let id_len = bytemuck::pod_read_unaligned::<u32>(
+                &output_id_buffer[offset..offset + mem::size_of::<u32>()],
+            );
             offset += mem::size_of::<u32>();
-            let mut id = 0u32;
-            bytemuck::bytes_of_mut(&mut id)
-                .copy_from_slice(&output_id_buffer[offset..offset + id_len as usize]);
+            let id = bytemuck::pod_read_unaligned::<u32>(
+                &output_id_buffer[offset..offset + id_len as usize],
+            );
             offset += id_len as usize;
             ids.push(id);
         }

diskann-garnet/src/fsm.rs

@@ -41,7 +41,7 @@ pub(crate) enum FsmError {
     IdOutOfRange(u32),
 }
 
-/// Guard returned by `next_id()` to ensure correctness when reuse is enabled.
+/// Guard protecting vector writes against quantization phase changes.
 pub(crate) struct ReuseGuard<'a> {
     id: u32,
     barrier: RwLockReadGuard<'a, Barrier>,
@@ -59,6 +59,10 @@ impl<'a> ReuseGuard<'a> {
     pub(crate) fn should_quantize(&self) -> bool {
         self.barrier.quantization_enabled
     }
+
+    pub(crate) fn max_id_for_backfill(&self) -> u32 {
+        self.barrier.max_id_for_backfill
+    }
 }
 
 struct Barrier {
@@ -366,6 +370,11 @@ impl FreeSpaceMap {
         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())
+    }
+
     /// 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.