Skip to content

PiPNN 1/6: extract shared RobustPrune #1315

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
weiyaoluo (SeliMeli) wants to merge 3 commits into from
+186 −140
Open

PiPNN 1/6: extract shared RobustPrune #1315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c2e6be6
refactor(graph): extract robust prune loop
SeliMeli Aug 7, 2026
d230a2c
docs(prune): describe RobustPrune domain output
SeliMeli Aug 10, 2026
149417b
fix(graph): use cached prune distances
SeliMeli Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 11 additions & 137 deletions diskann/src/graph/index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
//
Expand All @@ -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)| {
Expand Down
178 changes: 175 additions & 3 deletions diskann/src/graph/internal/prune.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down Expand Up @@ -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.
Expand All @@ -90,6 +95,173 @@ pub(crate) struct State {
pub(in crate::graph) neighbor: u16,
}

/// Select a degree-bounded neighbor set with Vamana RobustPrune.
///
/// `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.
///
/// 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<I, V, D>(
pool: &SortedNeighbors<'_, I>,
cache: &[(f32, Option<V>)],
states: &mut [State],
degree: usize,
alpha: f32,
prune_kind: PruneKind,
mut compute_distance: D,
) -> usize
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);

// 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 `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 `cache[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 = 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_distance, neighbor)) in cache.iter().enumerate() {
if found >= degree {
break;
}

// 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 `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]`.
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>(I)
Expand Down
Loading