diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index 4f777eb085..7de82939ce 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -2601,26 +2601,19 @@ where // // Iteration preserves the pool's nearest-first order. Pruning stores results // by cache position and resolves each position to `local_id` afterward. - let sorted_cache: Vec<(f32, Option<_>)> = pool - .iter() - .map(|neighbor| { - // Filter out self loops. - let id = neighbor.id(); - if exclude(*id) { - (*neighbor.distance(), None) - } else { - (*neighbor.distance(), map.get(*id)) - } - }) - .collect(); + let mut storage = Vec::with_capacity(pool.len()); + let cache = pool.map_in( + &mut storage, + |id| if exclude(*id) { None } else { map.get(*id) }, + ); let found = prune::robust_prune( - &sorted_cache, + cache, states, degree, alpha, self.config.prune_kind(), - |neighbor, result| { + |neighbor: &M::Element<'_>, result: &M::Element<'_>| { computer.evaluate_similarity((*neighbor).reborrow(), result.reborrow()) }, ); diff --git a/diskann/src/graph/internal/prune.rs b/diskann/src/graph/internal/prune.rs index 61b2e88ad6..2ea1eb212a 100644 --- a/diskann/src/graph/internal/prune.rs +++ b/diskann/src/graph/internal/prune.rs @@ -5,11 +5,9 @@ use thiserror::Error; -use super::SortedNeighbors; - use crate::{ ANNError, error, - graph::{AdjacencyList, config::PruneKind}, + graph::{AdjacencyList, config::PruneKind, internal::SortedNeighbors}, neighbor::Neighbor, utils::{IntoUsize, VectorId}, }; @@ -97,14 +95,14 @@ pub(crate) struct State { /// Select a degree-bounded neighbor set with Vamana RobustPrune. /// -/// `sorted_cache` stores each source distance and candidate vector in ascending +/// `candidates` stores each source distance and candidate vector in ascending /// source-distance order. `None` excludes that candidate without changing /// positional alignment. `states` has one entry for each candidate position. /// /// The function writes selected candidate indexes to `states[..result]` and /// returns `result`. The caller converts those indexes to graph IDs. pub(in crate::graph) fn robust_prune( - sorted_cache: &[(f32, Option)], + candidates: SortedNeighbors<'_, Option>, states: &mut [State], degree: usize, alpha: f32, @@ -114,11 +112,6 @@ pub(in crate::graph) fn robust_prune( where D: FnMut(&V, &V) -> f32, { - debug_assert!( - sorted_cache.is_sorted_by_key(|(distance, _)| distance), - "candidate cache must be sorted by source distance" - ); - let mut current_alpha = 1.0f32; let increment_factor = alpha.min(1.2); @@ -138,11 +131,11 @@ where // // On the implementation side, we use `states` in the following way: // - // * `states[n].neighbor` is the **index** in `sorted_cache` of the `n`th **neighbor**. + // * `states[n].neighbor` is the **index** in `candidates` of the `n`th **neighbor**. // Note that a "neighbor" is a candidate that passes pruning. // // Very important: to get the index `j` in the above description, we need to - // check `sorted_cache[states[n].neighbor]`. + // check `candidates[states[n].neighbor]`. // // This indexing naturally skips candidates `j` that have not been promoted to // neighbors. @@ -152,14 +145,14 @@ where // excludes it from future consideration. // // * `states[i].last_checked` is the highest value of `n` against which the - // occlude factor for `j = sorted_cache[states[n].neighbor]` has been checked. + // occlude factor for `j = candidates[states[n].neighbor]` has been checked. // // The maximum value this should reach is `i`. // // Note that we use `states` for both "candidate" and "neighbor" tracking. let mut found = 0; while found < degree { - for (i, (neighbor_distance, neighbor)) in sorted_cache.iter().enumerate() { + for (i, neighbor) in candidates.iter().enumerate() { if found >= degree { break; } @@ -176,10 +169,11 @@ where continue; } - // Retrieval from the cache might not be perfect. + // Retrieval from the candidates might not be perfect. // - // This neighbor did not end up in the cache, then just skip it. - let neighbor = match neighbor { + // This neighbor did not end up in the candidates, then just skip it. + let neighbor_distance = neighbor.distance(); + let neighbor = match neighbor.id() { Some(n) => n, None => { debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); @@ -198,7 +192,7 @@ where let result_position = states[last_checked as usize].neighbor.into_usize(); last_checked += 1; - // If the position of this result in `sorted_cache` is greater than or equal + // If the position of this result in `candidates` is greater than or equal // to the current working position, then skip this candidate. if result_position >= i { debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); @@ -209,9 +203,9 @@ where // Otherwise, compute the distance between the result and this neighbor // and update the occlude factor. - let distance = match &sorted_cache[result_position] { - (_, Some(v)) => compute_distance(neighbor, v), - (_, None) => f32::MAX, + let distance = match candidates[result_position].id() { + Some(v) => compute_distance(neighbor, v), + None => f32::MAX, }; // Update occlude factor diff --git a/diskann/src/graph/internal/sorted_neighbors.rs b/diskann/src/graph/internal/sorted_neighbors.rs index 8993ee05a5..73b403e963 100644 --- a/diskann/src/graph/internal/sorted_neighbors.rs +++ b/diskann/src/graph/internal/sorted_neighbors.rs @@ -11,19 +11,22 @@ use crate::neighbor::{self, Neighbor}; /// A utility that asserts the contained neighbors are sorted by distance. #[derive(Debug)] -pub struct SortedNeighbors<'a, I>(&'a [Neighbor]) -where - I: Eq; - -impl<'a, I> SortedNeighbors<'a, I> -where - I: Eq + std::fmt::Debug, -{ +pub(crate) struct SortedNeighbors<'a, I>(&'a [Neighbor]); + +impl Clone for SortedNeighbors<'_, I> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for SortedNeighbors<'_, I> {} + +impl<'a, I> SortedNeighbors<'a, I> { /// Create a new `SortedNeighbors` around `neighbors` truncated to `max` length. /// /// As a by-product calling this method, `neighbors` will be resized to at most /// `max` and be sorted. - pub fn new(neighbors: &'a mut Vec>, max: usize) -> Self { + pub(crate) fn new(neighbors: &'a mut Vec>, max: usize) -> Self { // Here- we use `select_nth_unstable` to get the `position` index in the correct // location. We can then sort the prefix slice returned by that API. // @@ -42,12 +45,29 @@ where neighbors.truncate(max); Self(&*neighbors) } + + /// Apply the projection `f` to each element in `self` and store the result in `other`. + /// + /// The returned [`SortedNeighbors`] inherits the sorted property from `self`. + /// + /// # Side Effects + /// + /// This method removes all pre-existing elements from `storage`. + pub(crate) fn map_in<'b, F, J>( + self, + storage: &'b mut Vec>, + mut f: F, + ) -> SortedNeighbors<'b, J> + where + F: FnMut(&I) -> J, + { + storage.clear(); + storage.extend(self.iter().map(|n| Neighbor::new(f(n.id()), *n.distance()))); + SortedNeighbors(storage) + } } -impl Deref for SortedNeighbors<'_, I> -where - I: Eq, -{ +impl Deref for SortedNeighbors<'_, I> { type Target = [Neighbor]; fn deref(&self) -> &Self::Target { self.0 @@ -108,4 +128,46 @@ mod tests { } } } + + #[test] + fn test_map() { + let messages = ["a", "b", "c", "d", "e", "f"]; + + let mut storage = vec![Neighbor::new("foo", 1.0), Neighbor::new("bar", 0.0)]; + + { + let mut neighbors = vec![ + Neighbor::new(0usize, 5.0f32), + Neighbor::new(1, 4.0), + Neighbor::new(2, 3.0), + Neighbor::new(3, 2.0), + Neighbor::new(4, 1.0), + Neighbor::new(5, 0.0), + ]; + + let sorted = SortedNeighbors::new(&mut neighbors, 6); + let cache = sorted.map_in(&mut storage, |id: &usize| messages[*id]); + + assert_eq_verbose!( + *cache, + [ + Neighbor::new("f", 0.0f32), + Neighbor::new("e", 1.0), + Neighbor::new("d", 2.0), + Neighbor::new("c", 3.0), + Neighbor::new("b", 4.0), + Neighbor::new("a", 5.0), + ] + .as_slice() + ); + } + + // Empty + { + let mut neighbors = Vec::>::new(); + let sorted = SortedNeighbors::new(&mut neighbors, 10); + let cache = sorted.map_in(&mut storage, |id: &usize| messages[*id]); + assert!(cache.is_empty()); + } + } } diff --git a/diskann/src/test/cmp.rs b/diskann/src/test/cmp.rs index 19cc670bd9..fc6d2b5f9e 100644 --- a/diskann/src/test/cmp.rs +++ b/diskann/src/test/cmp.rs @@ -108,6 +108,15 @@ pub(crate) use assert_eq_verbose; // Implementation // //////////////////// +impl VerboseEq for &T +where + T: VerboseEq + ?Sized, +{ + fn verbose_eq(&self, other: &Self) -> ANNResult<()> { + (*self).verbose_eq(*other) + } +} + /// Display implementation for recording a field mismatch. #[derive(Debug)] pub(crate) struct Field(pub(crate) &'static str); @@ -130,7 +139,7 @@ macro_rules! impl_via_partial_eq { impl $crate::test::cmp::VerboseEq for $T { fn verbose_eq(&self, other: &Self) -> $crate::ANNResult<()> { if self != other { - Err($crate::ANNError::new(NotEq(self.clone(), other.clone()))) + Err($crate::ANNError::new(NotEq(self.to_owned(), other.to_owned()))) } else { Ok(()) } @@ -143,7 +152,7 @@ macro_rules! impl_via_partial_eq { } impl_via_partial_eq!( - u8, u16, u32, u64, i8, i16, i32, i64, usize, f32, f64, String, bool, + u8, u16, u32, u64, i8, i16, i32, i64, usize, f32, f64, String, str, bool, ); macro_rules! impl_tuple {