diff --git a/Cargo.toml b/Cargo.toml index 71e9bde7e9..5cf3daa427 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,6 +87,7 @@ criterion = "0.5.1" crossbeam-queue = { version = "0.3.12", default-features = false } crossbeam-utils = { version = "0.8.23", default-features = false } dashmap = "6.1.0" +erased-serde = "0.4.11" futures-util = { version = "0.3.31", default-features = false } half = "2.6.0" hashbrown = "0.16.0" diff --git a/diskann-benchmark-runner/Cargo.toml b/diskann-benchmark-runner/Cargo.toml index b381c8ae2d..7b4cebe0b3 100644 --- a/diskann-benchmark-runner/Cargo.toml +++ b/diskann-benchmark-runner/Cargo.toml @@ -10,7 +10,7 @@ edition = "2024" [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } -erased-serde = "0.4.11" +erased-serde = { workspace = true } half = { workspace = true } indicatif = "0.18.3" serde = { workspace = true, features = ["derive"] } diff --git a/diskann-inmem/integration/jsons/graph/product/cosine-baseline.json b/diskann-inmem/integration/jsons/graph/product/cosine-baseline.json index 681808b2ae..a411bcc4d4 100644 --- a/diskann-inmem/integration/jsons/graph/product/cosine-baseline.json +++ b/diskann-inmem/integration/jsons/graph/product/cosine-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 39095, "distance": 549876, "get_neighbors": 298634, - "get_vector": 2008258, + "get_vector": 2032786, "query_distance": 1724191, "set_neighbors": 11168, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/product/ip-baseline.json b/diskann-inmem/integration/jsons/graph/product/ip-baseline.json index bbeacb85d3..b297b495b0 100644 --- a/diskann-inmem/integration/jsons/graph/product/ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/product/ip-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 128225, "distance": 5596404, "get_neighbors": 384814, - "get_vector": 2467136, + "get_vector": 3131891, "query_distance": 1545792, "set_neighbors": 41655, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/product/l2-baseline.json b/diskann-inmem/integration/jsons/graph/product/l2-baseline.json index 58eae3a798..ef2b4550a4 100644 --- a/diskann-inmem/integration/jsons/graph/product/l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/product/l2-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 37939, "distance": 533446, "get_neighbors": 297940, - "get_vector": 1985381, + "get_vector": 2008439, "query_distance": 1702322, "set_neighbors": 11098, "set_vector": 10000 @@ -171,7 +171,7 @@ "append_neighbors": 37939, "distance": 533446, "get_neighbors": 297940, - "get_vector": 1985381, + "get_vector": 2008439, "query_distance": 1702322, "set_neighbors": 11098, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json index df05e5f862..5eb27bb7e9 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/scalar/eight-bit-l2-baseline.json @@ -50,7 +50,7 @@ "append_neighbors": 100311, "distance": 3063537, "get_neighbors": 357204, - "get_vector": 2824455, + "get_vector": 3132525, "query_distance": 2259492, "set_neighbors": 24670, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json index 81d7d4133b..fec59a0639 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/scalar/four-bit-ip-baseline.json @@ -50,7 +50,7 @@ "append_neighbors": 128219, "distance": 5554287, "get_neighbors": 380443, - "get_vector": 3110524, + "get_vector": 3775405, "query_distance": 2193419, "set_neighbors": 41661, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json index a4e231ca1e..fa38a15c47 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/scalar/one-bit-l2-baseline.json @@ -50,7 +50,7 @@ "append_neighbors": 125092, "distance": 4903759, "get_neighbors": 401105, - "get_vector": 3459145, + "get_vector": 3990928, "query_distance": 2651349, "set_neighbors": 35323, "set_vector": 10000 @@ -173,7 +173,7 @@ "append_neighbors": 125092, "distance": 4903759, "get_neighbors": 401105, - "get_vector": 3459145, + "get_vector": 3990928, "query_distance": 2651349, "set_neighbors": 35323, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json index 0b2d7f568e..0443958789 100644 --- a/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json +++ b/diskann-inmem/integration/jsons/graph/scalar/two-bit-cosine-baseline.json @@ -50,7 +50,7 @@ "append_neighbors": 107423, "distance": 3639726, "get_neighbors": 368584, - "get_vector": 3001191, + "get_vector": 3380724, "query_distance": 2360497, "set_neighbors": 28073, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine-baseline.json index dcc0686f51..43276b2cd5 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/four-bit-cosine-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 100733, "distance": 3083799, "get_neighbors": 358175, - "get_vector": 2811083, + "get_vector": 3124613, "query_distance": 2240111, "set_neighbors": 24930, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/four-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/four-bit-ip-baseline.json index 73651ee1c9..75e58e2c8e 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/four-bit-ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/four-bit-ip-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 128260, "distance": 5569637, "get_neighbors": 380639, - "get_vector": 3119103, + "get_vector": 3783123, "query_distance": 2202704, "set_neighbors": 41620, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/four-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/four-bit-l2-baseline.json index 5c186c5337..50c8ac295a 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/four-bit-l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/four-bit-l2-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 100455, "distance": 3077148, "get_neighbors": 357944, - "get_vector": 2803941, + "get_vector": 3115938, "query_distance": 2234455, "set_neighbors": 24857, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine-baseline.json index 127a152e9f..7609b13993 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/one-bit-cosine-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 104911, "distance": 3391469, "get_neighbors": 371616, - "get_vector": 3179260, + "get_vector": 3517759, "query_distance": 2574056, "set_neighbors": 26119, "set_vector": 10000 @@ -171,7 +171,7 @@ "append_neighbors": 104911, "distance": 3391469, "get_neighbors": 371616, - "get_vector": 3179260, + "get_vector": 3517759, "query_distance": 2574056, "set_neighbors": 26119, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/one-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/one-bit-ip-baseline.json index 9327a602a4..f911c402a5 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/one-bit-ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/one-bit-ip-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 128385, "distance": 5555853, "get_neighbors": 380952, - "get_vector": 3252844, + "get_vector": 3914239, "query_distance": 2338882, "set_neighbors": 41495, "set_vector": 10000 @@ -171,7 +171,7 @@ "append_neighbors": 128385, "distance": 5555853, "get_neighbors": 380952, - "get_vector": 3252844, + "get_vector": 3914239, "query_distance": 2338882, "set_neighbors": 41495, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/one-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/one-bit-l2-baseline.json index bf89867931..47a995d612 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/one-bit-l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/one-bit-l2-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 104879, "distance": 3387869, "get_neighbors": 370460, - "get_vector": 3161112, + "get_vector": 3500913, "query_distance": 2555730, "set_neighbors": 26181, "set_vector": 10000 @@ -171,7 +171,7 @@ "append_neighbors": 104879, "distance": 3387869, "get_neighbors": 370460, - "get_vector": 3161112, + "get_vector": 3500913, "query_distance": 2555730, "set_neighbors": 26181, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine-baseline.json index bbde79c7ed..90062231ac 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/two-bit-cosine-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 103313, "distance": 3222189, "get_neighbors": 362992, - "get_vector": 2907863, + "get_vector": 3236282, "query_distance": 2319765, "set_neighbors": 25639, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/two-bit-ip-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/two-bit-ip-baseline.json index ddb80c697f..d9d8b62058 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/two-bit-ip-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/two-bit-ip-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 128278, "distance": 5549845, "get_neighbors": 380582, - "get_vector": 3182824, + "get_vector": 3846466, "query_distance": 2266878, "set_neighbors": 41602, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/graph/spherical/two-bit-l2-baseline.json b/diskann-inmem/integration/jsons/graph/spherical/two-bit-l2-baseline.json index e5b1c7678f..ddbfc257fa 100644 --- a/diskann-inmem/integration/jsons/graph/spherical/two-bit-l2-baseline.json +++ b/diskann-inmem/integration/jsons/graph/spherical/two-bit-l2-baseline.json @@ -49,7 +49,7 @@ "append_neighbors": 102671, "distance": 3195226, "get_neighbors": 362395, - "get_vector": 2896757, + "get_vector": 3222110, "query_distance": 2311680, "set_neighbors": 25493, "set_vector": 10000 diff --git a/diskann-inmem/integration/jsons/integration-baseline.json b/diskann-inmem/integration/jsons/integration-baseline.json index 4e2b6f53c7..f7d371e3dd 100644 --- a/diskann-inmem/integration/jsons/integration-baseline.json +++ b/diskann-inmem/integration/jsons/integration-baseline.json @@ -48,7 +48,7 @@ "append_neighbors": 96949, "distance": 2867876, "get_neighbors": 352139, - "get_vector": 2781513, + "get_vector": 3067092, "query_distance": 2240744, "set_neighbors": 23599, "set_vector": 10000 @@ -169,7 +169,7 @@ "append_neighbors": 96949, "distance": 2867876, "get_neighbors": 352139, - "get_vector": 2781513, + "get_vector": 3067092, "query_distance": 2240744, "set_neighbors": 23599, "set_vector": 10000 @@ -290,7 +290,7 @@ "append_neighbors": 96949, "distance": 2867876, "get_neighbors": 352139, - "get_vector": 2781513, + "get_vector": 3067092, "query_distance": 2240744, "set_neighbors": 23599, "set_vector": 10000 @@ -414,7 +414,7 @@ "append_neighbors": 97055, "distance": 2867292, "get_neighbors": 352087, - "get_vector": 2778779, + "get_vector": 3064106, "query_distance": 2238420, "set_neighbors": 23587, "set_vector": 10000 diff --git a/diskann-inmem/src/counters.rs b/diskann-inmem/src/counters.rs index 5c10954268..882ad4adda 100644 --- a/diskann-inmem/src/counters.rs +++ b/diskann-inmem/src/counters.rs @@ -43,6 +43,7 @@ mod inner { pub(crate) fn query_distance(&mut self, _i: u64) {} pub(crate) fn distance_ref(&self, _i: u64) {} pub(crate) fn get_vector(&mut self, _i: u64) {} + pub(crate) fn get_vector_ref(&self, _i: u64) {} pub(crate) fn set_vector(&mut self, _i: u64) {} pub(crate) fn get_neighbors(&mut self, _i: u64) {} pub(crate) fn set_neighbors(&mut self, _i: u64) {} @@ -135,6 +136,10 @@ mod inner { *self.get_vector.get_mut() += i; } + pub(crate) fn get_vector_ref(&self, i: u64) { + self.get_vector.fetch_add(i, Relaxed); + } + pub(crate) fn set_vector(&mut self, i: u64) { self.set_vector += i; } diff --git a/diskann-inmem/src/provider.rs b/diskann-inmem/src/provider.rs index a5b9185f93..e67cfc886d 100644 --- a/diskann-inmem/src/provider.rs +++ b/diskann-inmem/src/provider.rs @@ -415,7 +415,6 @@ impl glue::SearchAccessor for SearchAccessor<'_> { #[derive(Debug)] pub struct PruneAccessor<'a> { prune: Box, - keys: hashbrown::HashMap>, neighbors: &'a Neighbors, counters: LocalCounters<'a>, } @@ -428,7 +427,6 @@ impl<'a> PruneAccessor<'a> { ) -> Self { Self { prune, - keys: hashbrown::HashMap::new(), neighbors, counters, } @@ -456,18 +454,18 @@ impl<'a> Distance<'a> { /// An opaque element-ref for [`PruneAccessor`]. #[derive(Debug, Clone, Copy)] #[repr(transparent)] -pub struct ElementRef(repr::PruneKey); +pub struct ElementRef<'a>(&'a [u8]); -impl<'a> diskann_utils::Reborrow<'a> for ElementRef { - type Target = ElementRef; +impl<'a> diskann_utils::Reborrow<'a> for ElementRef<'_> { + type Target = ElementRef<'a>; fn reborrow(&'a self) -> Self { - *self + Self(self.0) } } -impl diskann_vector::DistanceFunction for Distance<'_> { +impl diskann_vector::DistanceFunction, ElementRef<'_>, f32> for Distance<'_> { #[inline] - fn evaluate_similarity(&self, x: ElementRef, y: ElementRef) -> f32 { + fn evaluate_similarity(&self, x: ElementRef<'_>, y: ElementRef<'_>) -> f32 { self.counters.distance_ref(1); self.prune.evaluate(x.0, y.0) } @@ -483,7 +481,7 @@ impl glue::PruneAccessor for PruneAccessor<'_> { where Self: 'a; - type ElementRef<'a> = ElementRef; + type ElementRef<'a> = ElementRef<'a>; type View<'a> = &'a Self @@ -501,16 +499,11 @@ impl glue::PruneAccessor for PruneAccessor<'_> { async fn fill<'a, Itr>( &'a mut self, - itr: Itr, + _itr: Itr, ) -> ANNResult<(Self::View<'a>, Self::Distance<'a>)> where Itr: ExactSizeIterator + Clone + Send + Sync, { - self.keys.clear(); - self.keys.extend(itr.map(|i| (i, None))); - let count: usize = self.prune.prepare(self.keys.iter_mut())?; - self.counters.get_vector(count as u64); - Ok((self, Distance::new(&*self.prune, self.counters.fork()))) } } @@ -573,14 +566,15 @@ impl provider::NeighborAccessorMut for PruneAccessor<'_> { } impl workingset::View for &PruneAccessor<'_> { - type ElementRef<'a> = ElementRef; + type ElementRef<'a> = ElementRef<'a>; type Element<'a> - = ElementRef + = ElementRef<'a> where Self: 'a; - fn get(&self, id: u32) -> Option { - self.keys.get(&id)?.map(ElementRef) + fn get(&self, id: u32) -> Option> { + self.counters.get_vector_ref(1); + self.prune.get(id).map(ElementRef) } } diff --git a/diskann-inmem/src/repr/full.rs b/diskann-inmem/src/repr/full.rs index df59567f00..04a7330c78 100644 --- a/diskann-inmem/src/repr/full.rs +++ b/diskann-inmem/src/repr/full.rs @@ -1148,27 +1148,24 @@ mod tests { prune: &mut repr::internal::intrusive::Prune<'_, repr::test::TestDistance>, ids: &[u32], ) { - let mut items: HashMap> = - ids.iter().map(|id| (*id, None)).collect(); - - let processed = <_ as repr::Prune>::prepare(prune, items.iter_mut()).unwrap(); - assert_eq!(processed, items.values().filter(|i| i.is_some()).count()); - - // Ensure that `prepare` agrees with `points`. - for (k, v) in items.iter() { - match v { - Some(_) => assert!(points.contains_key(k)), - None => assert!(!points.contains_key(k)), + // Ensure that distances agree. + for &k0 in ids.iter() { + let v0 = <_ as repr::Prune>::get(prune, k0); + if !points.contains_key(&k0) { + assert!(v0.is_none(), "Failed for key {}", k0); + continue; } - } - fn filter((k, v): (&u32, &Option)) -> Option<(u32, repr::PruneKey)> { - v.map(|v| (*k, v)) - } + let v0 = v0.unwrap(); + for &k1 in ids.iter() { + let v1 = <_ as repr::Prune>::get(prune, k1); + if !points.contains_key(&k1) { + assert!(v1.is_none(), "Failed for key {}", k1); + continue; + } + + let v1 = v1.unwrap(); - // Ensure that distances agree. - for (k0, v0) in items.iter().filter_map(filter) { - for (k1, v1) in items.iter().filter_map(filter) { // Manually implement `TestDistance`. let expected = points[&k0] + points[&k1]; let got = <_ as repr::Prune>::evaluate(prune, v0, v1); @@ -1303,13 +1300,10 @@ mod tests { // Prune { - let mut prune = ::make_prune(&full).unwrap(); - let mut points: HashMap> = - [(internal_query, None), (start_id, None)] - .into_iter() - .collect(); - prune.prepare(points.iter_mut()).unwrap(); - let got = prune.evaluate(points[&internal_query].unwrap(), points[&start_id].unwrap()); + let prune = ::make_prune(&full).unwrap(); + let query = prune.get(internal_query).unwrap(); + let start = prune.get(start_id).unwrap(); + let got = prune.evaluate(query, start); assert!( approx_eq(expected, got), "{ctx} - expected {expected}, got {got}" diff --git a/diskann-inmem/src/repr/internal/intrusive.rs b/diskann-inmem/src/repr/internal/intrusive.rs index 1d0fa3ed2a..555249ab09 100644 --- a/diskann-inmem/src/repr/internal/intrusive.rs +++ b/diskann-inmem/src/repr/internal/intrusive.rs @@ -182,17 +182,6 @@ pub(in crate::repr) struct Prune<'a, D> { /// The distance computation. Ensure this is only provided with slices of length /// `reader.bytes()`. distance: D, - - /// Buffered pointers to slices in `reader`. These - /// - /// * Must come from slices yielded by `reader`. - /// - /// * Must have a length of exactly `reader.bytes()`. - /// - /// * Have no lifetime because they need to be outlived `reader`. We only materialize - /// them as slices in internal methods taking `&[mut] self` where we know the reader - /// is valid. - buffer: Vec<*const u8>, } impl<'a, D> Prune<'a, D> { @@ -203,11 +192,7 @@ impl<'a, D> Prune<'a, D> { /// [`store::intrusive::Reader::bytes`] will be passed to `distance`. However, unsafe /// code may not rely on this. pub(in crate::repr) fn new(reader: store::intrusive::Reader<'a>, distance: D) -> Self { - Self { - reader, - distance, - buffer: Vec::new(), - } + Self { reader, distance } } /// Wrap `self` in a `Box`. @@ -225,44 +210,11 @@ impl repr::Prune for Prune<'_, D> where D: RawDistance, { - fn prepare( - &mut self, - items: hashbrown::hash_map::IterMut<'_, u32, Option>, - ) -> ANNResult { - let mut counter = repr::PruneKey::counter(); - self.buffer.clear(); - self.buffer.reserve(items.len()); - - for (id, key) in items { - if let Some(v) = self.reader.read(id.into_usize()) { - self.buffer.push(v.as_ptr()); - - *key = Some(counter); - - // Potential overflow issue - but it's exceedingly unlikely that - // someone will provide a prune list exceeding `u16::MAX`. - // - // In addition, `diskann` limits this bound as well. - counter = counter.increment()?; - } - } - - Ok(counter.index()) + fn get(&self, id: u32) -> Option<&[u8]> { + self.reader.read(id.into_usize()) } - fn evaluate(&self, a: repr::PruneKey, b: repr::PruneKey) -> f32 { - let len = self.reader.bytes().value(); - - // SAFETY: The only way a pointer ends up in `self.buffer` is from `self.prepare`, - // where they are obtained from slices returned from `self.reader`. - // - // Further, `self.reader` must be retained for the lifetime of `self`. Thus, they - // point to valid data of length `len` and are protected from concurrent mutation. - let a = unsafe { std::slice::from_raw_parts(self.buffer[a.index()], len) }; - - // SAFETY: See above. - let b = unsafe { std::slice::from_raw_parts(self.buffer[b.index()], len) }; - + fn evaluate(&self, a: &[u8], b: &[u8]) -> f32 { #[expect( clippy::expect_used, reason = "diskann currently does not support fallible prune" diff --git a/diskann-inmem/src/repr/mod.rs b/diskann-inmem/src/repr/mod.rs index ce725894a1..5c5dbd4084 100644 --- a/diskann-inmem/src/repr/mod.rs +++ b/diskann-inmem/src/repr/mod.rs @@ -24,10 +24,7 @@ mod internal_docs { //! relocation. } -use std::num::NonZeroU16; - use diskann::{ANNResult, neighbor::Neighbor}; -use thiserror::Error; use crate::{ counters::LocalCounters, @@ -242,67 +239,39 @@ pub(crate) trait PostProcess: Send + Sync + std::fmt::Debug { /// Trait-object-based implementation for [`diskann::graph::glue::PruneAccessor`]. /// -/// [`Self::prepare`] assigns a [`PruneKey`] to each retrieved entry, allowing implementations -/// to buffer data in a representation suitable for pruning. A prune session consists of one -/// call to [`Self::prepare`] followed by calls to [`Self::evaluate`] using the keys it -/// produced. +/// Currently, we return raw slices to the underlying data. Limitations in the prune contract +/// prevent us from doing any thing fancier. +/// +/// ## Musings +/// +/// Returning raw slices prevents us from doing something like hybrid pruning for PQ. At +/// least not without some unpleasantness in `evaluate`. We accept this limitation for now +/// since the performance of returning a raw slice is too tempting to give up. /// -/// [`Self::prepare`] can be called multiple times for a single [`Prune`] object. -/// Implementations may assume that each call starts a new prune session and do not need to -/// buffer entries from all calls. +/// A slightly better interface may be to use something like: +/// ```text +/// struct ElementRef<'a> { +/// ptr: NonNull<()>, +/// _lifetime: PhantomData<&'a ()>, +/// } +/// ``` +/// allowing us to store something slightly slimmer. However, because +/// [`diskann::graph::glue::PruneAccessor::fill`] is a public method and +/// [`diskann_vector::DistanceFunction`] is a safe trait with a safe method, we cannot +/// expose these raw pointers soundly as someone in safe code could mix and match pointers. +/// +/// At least slices have their length tagged. Our internal code does not yet do unsafe things +/// assuming the lengths of two slices are the same. It would be nice to get there though. pub(crate) trait Prune: Send + Sync + std::fmt::Debug { - /// Prepare items for pruning. - /// - /// The `items` iterator contains all the internal IDs that need to be pruned as keys. - /// Each value in `items` is an output slot that should be populated with a [`PruneKey`] - /// if retrieval succeeds. Implementations are responsible for providing unique - /// [`PruneKey`]s and may assume these slots are initialized to `None`. + /// The the data associated with `id`. /// - /// Returns the total number of internal IDs that were successfully buffered. - fn prepare( - &mut self, - items: hashbrown::hash_map::IterMut<'_, u32, Option>, - ) -> ANNResult; + /// Return `None` if `id` is unreadable (including out-of-bounds). + fn get(&self, id: u32) -> Option<&[u8]>; - /// Compute the distance between the elements referenced by the two [`PruneKey`]s. + /// Compute the distance between the two raw elements. /// - /// Implementations may assume that `a` and `b` were produced by the most recent call to - /// [`Self::prepare`] and are allowed to panic if this is violated. This property may + /// Implementations may assume that `a` and `b` were produced by a call to + /// [`Self::get`] and are allowed to panic if this is violated. This property may /// **not** be relied on for `unsafe` code. - fn evaluate(&self, a: PruneKey, b: PruneKey) -> f32; -} - -/// A [`Prune`]-local miniature internal ID. -#[derive(Debug, Clone, Copy)] -pub(crate) struct PruneKey(NonZeroU16); - -impl PruneKey { - const ONE: Self = Self(NonZeroU16::new(1).unwrap()); - - /// Return the initial [`PruneKey`] in a sequence. - pub(crate) fn counter() -> Self { - Self::ONE - } - - /// Increment `self` by 1. - /// - /// Overflows if called approximately 2^16 times. - pub(crate) fn increment(self) -> Result { - match self.0.checked_add(1) { - Some(v) => Ok(Self(v)), - None => Err(Overflow), - } - } - - /// Return the zero-based index represented by this key. - pub(crate) fn index(self) -> usize { - usize::from(self.0.get()) - 1 - } + fn evaluate(&self, a: &[u8], b: &[u8]) -> f32; } - -/// Incrementing a [`PruneKey`] overflowed. -#[derive(Debug, Error)] -#[error("prune list exceeded u16::MAX")] -pub(crate) struct Overflow; - -diskann::convert_error!(Overflow); diff --git a/diskann-inmem/src/repr/test/mod.rs b/diskann-inmem/src/repr/test/mod.rs index d07b79d0de..95be8ee16a 100644 --- a/diskann-inmem/src/repr/test/mod.rs +++ b/diskann-inmem/src/repr/test/mod.rs @@ -287,33 +287,36 @@ pub(super) fn test_prune( ids: &[SlotId], ctx: &dyn std::fmt::Display, ) { - let num_present_ids = ids - .iter() - .filter(|&&slot_id| distances.contains_key(&(slot_id, slot_id))) - .count(); - - let mut items: HashMap> = - ids.iter().map(|slot_id| (slot_id.value(), None)).collect(); - - let count = accessor.prepare(items.iter_mut()).unwrap(); - assert_eq!(count, num_present_ids, "{ctx}"); - let mut visited = 0; for slot_id0 in ids.iter() { - if let Some(key0) = items[&slot_id0.value()] { - for slot_id1 in ids.iter() { - if let Some(key1) = items[&slot_id1.value()] { - let d = accessor.evaluate(key0, key1); - let expected = distances[&(*slot_id0, *slot_id1)]; - assert_eq!( - d, expected, - "failed for {} x {} -- {}", - slot_id0, slot_id1, ctx - ); - - visited += 1; - } + let v0 = { + let v = accessor.get(slot_id0.value()); + if !distances.contains_key(&(*slot_id0, *slot_id0)) { + assert!(v.is_none(), "failed on {}", slot_id0); + continue; } + v.unwrap() + }; + + for slot_id1 in ids.iter() { + let v1 = { + let v = accessor.get(slot_id1.value()); + if !distances.contains_key(&(*slot_id1, *slot_id1)) { + assert!(v.is_none(), "failed on {}", slot_id1); + continue; + } + v.unwrap() + }; + + let d = accessor.evaluate(v0, v1); + let expected = distances[&(*slot_id0, *slot_id1)]; + assert_eq!( + d, expected, + "failed for {} x {} -- {}", + slot_id0, slot_id1, ctx + ); + + visited += 1; } } diff --git a/diskann-inmem/src/store/intrusive.rs b/diskann-inmem/src/store/intrusive.rs index 590b54261c..f08f022780 100644 --- a/diskann-inmem/src/store/intrusive.rs +++ b/diskann-inmem/src/store/intrusive.rs @@ -389,8 +389,12 @@ impl<'a> Reader<'a> { } /// Return the number of bytes for each entry. + #[expect( + unused, + reason = "this is a reasonable method to have - it's just not used" + )] pub(crate) fn bytes(&self) -> Bytes { - self.bytes_plus_tag().unchecked_sub(AtomicTag::SIZE) + self.unpadded.unchecked_sub(AtomicTag::SIZE) } /// Return the number of bytes plus the atomic tag.