From c2e6be62ee022574402d81464ca6a5a0a7bd0c5c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:55:11 +0000 Subject: [PATCH 1/3] refactor(graph): extract robust prune loop Keep the prune module, comments, names, preparation, output, and saturation in place; expose only the selection loop and require the existing SortedNeighbors witness. --- diskann/src/graph/index.rs | 148 ++-------------------- diskann/src/graph/internal/prune.rs | 183 +++++++++++++++++++++++++++- 2 files changed, 191 insertions(+), 140 deletions(-) diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index 7a4d09064..351b1fefb 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -2596,9 +2596,6 @@ where states.clear(); states.resize(pool.len(), prune::State::default()); - let mut current_alpha = 1.0f32; - let increment_factor = alpha.min(1.2); - // To avoid many hash lookups, we pull out just the candidates we're going to prune // into an auxiliary vector which can be accessed linearly. // @@ -2617,140 +2614,17 @@ where }) .collect(); - // For an alpha value `A`, a candidate `i` is promoted to a neighbor if for all - // ``` - // max{j < i | j is a neighbor}(occlude_factor(i, j)) - // ``` - // This process happens with multiple values of `A`. - // - // We can compute this efficiently using the following rules: - // - // 1. For a candidate `i`, start scanning `j < i`, computing occlude factors. - // 2. If we find an occlude factor greater than `A`, record that `i` has visited - // `j`, stop computing occlude factors, and move on to `i + 1`. - // 3. If we reach `j == i - 1` with the maximum occlude factor less than `A`, then - // `i` gets promoted to a neighbor. - // - // On the implementation side, we use `states` in the following way: - // - // * `states[n].neighbor` is the **index** in `pool` 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 `pool[states[n].neighbor]`. - // - // This indexing naturally skips candidates `j` that have not been promoted to - // neighbors. - // - // * `states[i].occlude_factor` is the maximum occlude factor found for a candidate - // `i`. This gets set to `f32::MAX` when `i` is promoted to a neighbor which - // excludes it from future consideration. - // - // * `states[i].last_checked` is the highest value of `n` against which the - // occlude factor for `j = pool[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 cache.iter().enumerate() { - if found >= degree { - break; - } - - // The tracking states for candidate `i`. - let prune::State { - mut occlude_factor, - mut last_checked, - .. - } = states[i]; - - // If the occlusion factor for this neighbor is too high, skip it. - if occlude_factor > current_alpha { - continue; - } - - // Retrieval from the cache might not be perfect. - // - // This neighbor did not end up in the cache, then just skip it. - let neighbor = match neighbor { - Some(n) => n, - None => { - debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); - // SAFETY: We've already checked `states[i]`. - unsafe { states.get_unchecked_mut(i) }.occlude_factor = f32::MAX; - continue; - } - }; - - // Increment `position` until we've compared with all current entries in - // `result`. - // - // When the list is empty, the loop is skipped allowing the first undeleted - // element to be added. - while last_checked as usize != found { - let result_position = states[last_checked as usize].neighbor.into_usize(); - last_checked += 1; - - // If the position of this result in `pool` is greater than or equal - // 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"); - // SAFETY: We've already checked `states[i]`. - unsafe { states.get_unchecked_mut(i) }.last_checked = last_checked; - continue; - } - - // Otherwise, compute the distance between the result and this neighbor - // and update the occlude factor. - let distance = match &cache[result_position] { - (_, Some(v)) => { - computer.evaluate_similarity((*neighbor).reborrow(), v.reborrow()) - } - (_, None) => f32::MAX, - }; - - // Update occlude factor - occlude_factor = self.config.prune_kind().update_occlude_factor( - *neighbor_distance, - distance, - occlude_factor, - current_alpha, - ); - - // Check if the most recent update to the occlusion factor removes this - // neighbor from consideration. - if occlude_factor > current_alpha { - break; - } - } - - debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); - // SAFETY: We've already checked `states[i]`. - let state = unsafe { states.get_unchecked_mut(i) }; - - state.last_checked = last_checked; - if occlude_factor > current_alpha { - state.occlude_factor = occlude_factor; - continue; - } - - // This neighbor has passed all the requirements of being a candidate. - state.occlude_factor = f32::MAX; - - // This conversion should always succeed. - states[found].neighbor = i as u16; - found += 1; - } - - // Exit if we completed the final iteration. - if current_alpha == alpha { - break; - } - // Update current alpha for the next iteration. - current_alpha = (current_alpha * increment_factor).min(alpha); - } + let found = prune::robust_prune( + pool, + &cache, + states, + degree, + alpha, + self.config.prune_kind(), + |neighbor, result| { + computer.evaluate_similarity((*neighbor).reborrow(), result.reborrow()) + }, + ); let mut guard = neighbors.resize(found); std::iter::zip(guard.iter_mut(), states.iter()).for_each(|(d, s)| { diff --git a/diskann/src/graph/internal/prune.rs b/diskann/src/graph/internal/prune.rs index 39c727e88..216e56b10 100644 --- a/diskann/src/graph/internal/prune.rs +++ b/diskann/src/graph/internal/prune.rs @@ -7,7 +7,12 @@ use thiserror::Error; use super::SortedNeighbors; -use crate::{ANNError, error, graph::AdjacencyList, neighbor::Neighbor, utils::VectorId}; +use crate::{ + ANNError, error, + graph::{AdjacencyList, config::PruneKind}, + neighbor::Neighbor, + utils::{IntoUsize, VectorId}, +}; /// Options provided to prune. See the field-level documentation for more details. /// @@ -78,8 +83,8 @@ where /// Position-wise state tracking. /// -/// Refer to the inline documentation in [`DiskANNIndex::occlude_list`] for documentation -/// on the use of these fields. +/// Refer to the inline documentation in [`robust_prune`] for documentation on the use +/// of these fields. #[derive(Debug, Clone, Copy, Default)] pub(crate) struct State { /// The occlude factor for the pool item at the corresponding index. @@ -90,6 +95,178 @@ pub(crate) struct State { pub(in crate::graph) neighbor: u16, } +/// Select positions from source-distance-sorted candidates. +/// +/// `pool` enforces nondecreasing source-distance order. `cache[i].1` contains the +/// optional value for `pool[i]`; `None` marks an excluded or unavailable candidate while +/// preserving its position. `states` must contain one +/// default-initialized entry per candidate. The caller owns allocation, ID translation, +/// and saturation. +pub(in crate::graph) fn robust_prune( + pool: &SortedNeighbors<'_, I>, + cache: &[(f32, Option)], + states: &mut [State], + degree: usize, + alpha: f32, + prune_kind: PruneKind, + mut compute_distance: D, +) -> usize +where + I: Eq, + D: FnMut(&V, &V) -> f32, +{ + assert_eq!( + pool.len(), + cache.len(), + "RobustPrune cache must have one entry per sorted candidate" + ); + assert_eq!( + cache.len(), + states.len(), + "RobustPrune state must have one entry per sorted candidate" + ); + + let mut current_alpha = 1.0f32; + let increment_factor = alpha.min(1.2); + + // For an alpha value `A`, a candidate `i` is promoted to a neighbor if for all + // ``` + // max{j < i | j is a neighbor}(occlude_factor(i, j)) + // ``` + // This process happens with multiple values of `A`. + // + // We can compute this efficiently using the following rules: + // + // 1. For a candidate `i`, start scanning `j < i`, computing occlude factors. + // 2. If we find an occlude factor greater than `A`, record that `i` has visited + // `j`, stop computing occlude factors, and move on to `i + 1`. + // 3. If we reach `j == i - 1` with the maximum occlude factor less than `A`, then + // `i` gets promoted to a neighbor. + // + // On the implementation side, we use `states` in the following way: + // + // * `states[n].neighbor` is the **index** in `pool` 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 `pool[states[n].neighbor]`. + // + // This indexing naturally skips candidates `j` that have not been promoted to + // neighbors. + // + // * `states[i].occlude_factor` is the maximum occlude factor found for a candidate + // `i`. This gets set to `f32::MAX` when `i` is promoted to a neighbor which + // excludes it from future consideration. + // + // * `states[i].last_checked` is the highest value of `n` against which the + // occlude factor for `j = pool[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)) in cache.iter().enumerate() { + if found >= degree { + break; + } + + let neighbor_distance = pool[i].distance(); + + // The tracking states for candidate `i`. + let State { + mut occlude_factor, + mut last_checked, + .. + } = states[i]; + + // If the occlusion factor for this neighbor is too high, skip it. + if occlude_factor > current_alpha { + continue; + } + + // Retrieval from the cache might not be perfect. + // + // This neighbor did not end up in the cache, then just skip it. + let neighbor = match neighbor { + Some(n) => n, + None => { + debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); + // SAFETY: We've already checked `states[i]`. + unsafe { states.get_unchecked_mut(i) }.occlude_factor = f32::MAX; + continue; + } + }; + + // Increment `position` until we've compared with all current entries in + // `result`. + // + // When the list is empty, the loop is skipped allowing the first undeleted + // element to be added. + while last_checked as usize != found { + let result_position = states[last_checked as usize].neighbor.into_usize(); + last_checked += 1; + + // If the position of this result in `pool` is greater than or equal + // 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"); + // SAFETY: We've already checked `states[i]`. + unsafe { states.get_unchecked_mut(i) }.last_checked = last_checked; + continue; + } + + // Otherwise, compute the distance between the result and this neighbor + // and update the occlude factor. + let distance = match &cache[result_position] { + (_, Some(v)) => compute_distance(neighbor, v), + (_, None) => f32::MAX, + }; + + // Update occlude factor + occlude_factor = prune_kind.update_occlude_factor( + *neighbor_distance, + distance, + occlude_factor, + current_alpha, + ); + + // Check if the most recent update to the occlusion factor removes this + // neighbor from consideration. + if occlude_factor > current_alpha { + break; + } + } + + debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); + // SAFETY: We've already checked `states[i]`. + let state = unsafe { states.get_unchecked_mut(i) }; + + state.last_checked = last_checked; + if occlude_factor > current_alpha { + state.occlude_factor = occlude_factor; + continue; + } + + // This neighbor has passed all the requirements of being a candidate. + state.occlude_factor = f32::MAX; + + // This conversion should always succeed. + states[found].neighbor = i as u16; + found += 1; + } + + // Exit if we completed the final iteration. + if current_alpha == alpha { + break; + } + // Update current alpha for the next iteration. + current_alpha = (current_alpha * increment_factor).min(alpha); + } + + found +} + #[derive(Debug, Clone, Copy, Error)] #[error("retrieval of main vector id {} failed during prune aggregation", self.0)] pub(crate) struct FailedVectorRetrieval(I) From d230a2c7043e4b4932ed9bbef259d41a0dda5d62 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:00:44 +0000 Subject: [PATCH 2/3] docs(prune): describe RobustPrune domain output --- diskann/src/graph/internal/prune.rs | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/internal/prune.rs b/diskann/src/graph/internal/prune.rs index 216e56b10..33987a8b8 100644 --- a/diskann/src/graph/internal/prune.rs +++ b/diskann/src/graph/internal/prune.rs @@ -95,13 +95,15 @@ pub(crate) struct State { pub(in crate::graph) neighbor: u16, } -/// Select positions from source-distance-sorted candidates. +/// Select a degree-bounded neighbor set with Vamana RobustPrune. /// -/// `pool` enforces nondecreasing source-distance order. `cache[i].1` contains the -/// optional value for `pool[i]`; `None` marks an excluded or unavailable candidate while -/// preserving its position. `states` must contain one -/// default-initialized entry per candidate. The caller owns allocation, ID translation, -/// and saturation. +/// `pool` contains candidates in nearest-first order from the source point. +/// `cache` contains the vector for the candidate at the same position. `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( pool: &SortedNeighbors<'_, I>, cache: &[(f32, Option)], @@ -115,17 +117,6 @@ where I: Eq, D: FnMut(&V, &V) -> f32, { - assert_eq!( - pool.len(), - cache.len(), - "RobustPrune cache must have one entry per sorted candidate" - ); - assert_eq!( - cache.len(), - states.len(), - "RobustPrune state must have one entry per sorted candidate" - ); - let mut current_alpha = 1.0f32; let increment_factor = alpha.min(1.2); From 149417b76ee6df1a69bba10e908811abe46fc086 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:02:31 +0000 Subject: [PATCH 3/3] fix(graph): use cached prune distances --- diskann/src/graph/internal/prune.rs | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/internal/prune.rs b/diskann/src/graph/internal/prune.rs index 33987a8b8..9386bfa2c 100644 --- a/diskann/src/graph/internal/prune.rs +++ b/diskann/src/graph/internal/prune.rs @@ -97,8 +97,8 @@ pub(crate) struct State { /// Select a degree-bounded neighbor set with Vamana RobustPrune. /// -/// `pool` contains candidates in nearest-first order from the source point. -/// `cache` contains the vector for the candidate at the same position. `None` +/// `pool` certifies nearest-first candidate order from the source point. +/// `cache` stores each matching source distance and candidate vector. `None` /// excludes that candidate without changing positional alignment. `states` has /// one entry for each candidate position. /// @@ -117,6 +117,12 @@ where I: Eq, D: FnMut(&V, &V) -> f32, { + debug_assert_eq!( + pool.len(), + cache.len(), + "sorted candidate pool and cache must have equal lengths" + ); + let mut current_alpha = 1.0f32; let increment_factor = alpha.min(1.2); @@ -136,11 +142,11 @@ where // // On the implementation side, we use `states` in the following way: // - // * `states[n].neighbor` is the **index** in `pool` of the `n`th **neighbor**. + // * `states[n].neighbor` is the **index** in `cache` 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 `pool[states[n].neighbor]`. + // check `cache[states[n].neighbor]`. // // This indexing naturally skips candidates `j` that have not been promoted to // neighbors. @@ -150,20 +156,18 @@ where // excludes it from future consideration. // // * `states[i].last_checked` is the highest value of `n` against which the - // occlude factor for `j = pool[states[n].neighbor]` has been checked. + // occlude factor for `j = cache[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)) in cache.iter().enumerate() { + for (i, (neighbor_distance, neighbor)) in cache.iter().enumerate() { if found >= degree { break; } - let neighbor_distance = pool[i].distance(); - // The tracking states for candidate `i`. let State { mut occlude_factor, @@ -198,8 +202,8 @@ where let result_position = states[last_checked as usize].neighbor.into_usize(); last_checked += 1; - // If the position of this result in `pool` is greater than or equal - // the current working position, then skip this candidate. + // If the position of this result in `cache` 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"); // SAFETY: We've already checked `states[i]`.