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

Merged
weiyaoluo (SeliMeli) merged 3 commits into from
Aug 17, 2026
Merged

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
6bf2139
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
153 changes: 13 additions & 140 deletions diskann/src/graph/index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2596,15 +2596,12 @@ 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.
//
// During the pruning phase, we store results by their relative position in the
// cache, and only resolve to their `local_id` at the end.
let cache: Vec<(f32, Option<_>)> = pool
// 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.
Expand All @@ -2617,140 +2614,16 @@ 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(
&sorted_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
174 changes: 171 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,169 @@ pub(crate) struct State {
pub(in crate::graph) neighbor: u16,
}

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

// 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 `sorted_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 `sorted_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 = sorted_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 sorted_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 `sorted_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 &sorted_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