leaves12138 commented on code in PR #62: URL: https://github.com/apache/paimon-vector-index/pull/62#discussion_r3651293862
########## core/src/vamana.rs: ########## @@ -0,0 +1,2738 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::diskann::DiskAnnBuildParams; +use crate::distance::{ + fvec_distance, fvec_l2sqr, fvec_l2sqr_four, fvec_l2sqr_scaled_exceeds, MetricType, +}; +use crate::kmeans::{self, KMeansConfig}; +use crate::pq::ProductQuantizer; +use crate::sparse_table::{estimated_memory_bytes as sparse_table_memory_bytes, SparseTable}; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; +use std::cmp::{Ordering, Reverse}; +use std::collections::{BinaryHeap, VecDeque}; +use std::io; +use std::ops::Index; +use std::sync::{Mutex, RwLock}; +use std::time::{Duration, Instant}; + +const PARALLEL_ADJACENCY_NODES_PER_SHARD: usize = 256; +const PARALLEL_BUILD_BATCH_NODES_PER_WORKER: usize = 8; +const CONNECTIVITY_SOURCE_SAMPLE_SIZE: usize = 64; +const SPARSE_BUILD_VISITED_MIN_MEMORY_SAVINGS: usize = 16; + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ScoredNode { + pub id: u32, + pub distance: f32, +} + +impl Eq for ScoredNode {} + +impl PartialOrd for ScoredNode { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +impl Ord for ScoredNode { + fn cmp(&self, other: &Self) -> Ordering { + scored_node_order(self, other) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct VamanaGraph { + pub entry_node: u32, + pub(crate) adjacency: CompactAdjacency, +} + +pub(crate) struct VamanaMemoryEstimate { + pub(crate) build_peak_bytes: usize, + pub(crate) remap_peak_bytes: usize, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct ConnectivityRepairStats { + full_reachability_traversals: usize, + source_distance_evaluations: usize, + edges_added: usize, +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct VamanaBuildStats { + pub(crate) initialization: Duration, + pub(crate) pass_one: Duration, + pub(crate) pass_two: Duration, + pub(crate) connectivity_repair: Duration, +} + +pub(crate) fn estimate_vamana_memory_bytes( + node_count: usize, + max_degree: usize, + search_list_size: usize, + workers: usize, +) -> Option<VamanaMemoryEstimate> { + let edge_bytes = max_degree.checked_mul(size_of::<u32>())?; + let builder_edges = node_count.checked_mul(edge_bytes)?; + let builder_degrees = node_count.checked_mul(size_of::<u16>())?; + let builder_shards = node_count + .div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD) + .checked_mul(size_of::<RwLock<AdjacencyShard>>())?; + let builder_graph = builder_edges + .checked_add(builder_degrees)? + .checked_add(builder_shards)?; + let build_order = node_count.checked_mul(size_of::<usize>())?; + let expected_visited = search_list_size + .checked_mul(max_degree)? + .checked_add(1)? + .min(node_count); + let dense_worker_states = node_count + .checked_mul(size_of::<u8>())? + .checked_add(expected_visited.checked_mul(size_of::<u32>())?)?; + let sparse_worker_states = sparse_table_memory_bytes(expected_visited, size_of::<u8>())?; + let worker_states = if sparse_worker_states + .checked_mul(SPARSE_BUILD_VISITED_MIN_MEMORY_SAVINGS) + .is_some_and(|threshold| threshold < dense_worker_states) + { + sparse_worker_states + } else { + dense_worker_states + }; + let worker_candidates = search_list_size + .checked_mul(3)? + .checked_mul(size_of::<ScoredNode>())?; + let prune_candidates = search_list_size.checked_add(max_degree)?; + let worker_prune = + prune_candidates.checked_mul(size_of::<ScoredNode>().checked_add(size_of::<u32>())?)?; + let worker_candidate_ids = search_list_size.checked_mul(size_of::<u32>())?; + let worker_neighbors = max_degree.checked_mul(2 * size_of::<u32>())?; + let worker_scratch = workers.max(1).checked_mul( + worker_states + .checked_add(worker_candidates)? + .checked_add(worker_prune)? + .checked_add(worker_candidate_ids)? + .checked_add(worker_neighbors)?, + )?; + let reverse_edge_batch = workers + .max(1) + .checked_mul(PARALLEL_BUILD_BATCH_NODES_PER_WORKER)? + .checked_mul(max_degree)? + .checked_mul(size_of::<(u32, u32)>())?; + let build_peak_bytes = builder_graph + .checked_add(build_order)? + .checked_add(worker_scratch)? + .checked_add(reverse_edge_batch)?; + + let final_shards = node_count + .div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD) + .checked_mul(size_of::<AdjacencyShard>())?; + let compact_graph = builder_edges + .checked_add(builder_degrees)? + .checked_add(final_shards)?; + let permutations = node_count.checked_mul(2 * size_of::<u32>())?; + let permutation_visited = node_count.checked_mul(size_of::<bool>())?; + let permutation_scratch = edge_bytes.checked_add(size_of::<u16>())?; + let remap_peak_bytes = compact_graph + .checked_add(permutations)? + .checked_add(permutation_visited)? + .checked_add(permutation_scratch)?; + Some(VamanaMemoryEstimate { + build_peak_bytes, + remap_peak_bytes, + }) +} + +pub(crate) fn estimate_sharded_vamana_memory_bytes( + node_count: usize, + dimension: usize, + max_degree: usize, + shard_count: usize, +) -> Option<usize> { + if node_count == 0 || shard_count < 2 { + return None; + } + let edge_bytes = max_degree.checked_mul(size_of::<u32>())?; + let compact_graph = node_count.checked_mul(edge_bytes.checked_add(size_of::<u16>())?)?; + let assignments = node_count.checked_mul(2 * size_of::<usize>())?; + let local_count = overlapping_shard_capacity(node_count, shard_count)?; + let memberships = shard_count + .checked_mul(local_count)? + .checked_mul(size_of::<u32>())? + .checked_add(shard_count.checked_mul(size_of::<Vec<u32>>())?)?; + let centroids = shard_count + .checked_mul(dimension)? + .checked_mul(size_of::<f32>())?; + let local_vectors = local_count + .checked_mul(dimension)? + .checked_mul(size_of::<f32>())?; + let local_ids = local_count.checked_mul(size_of::<u32>())?; + // Sequential local construction briefly holds nested and compact + // adjacency plus its order/visited vectors. + let local_graph = local_count.checked_mul( + edge_bytes + .checked_mul(2)? + .checked_add(size_of::<Vec<u32>>())? + .checked_add(size_of::<usize>())? + .checked_add(2 * size_of::<bool>())?, + )?; + let build_peak = [ + compact_graph, + assignments, + memberships, + centroids, + local_vectors, + local_ids, + local_graph, + ] + .into_iter() + .try_fold(0usize, |total, value| total.checked_add(value))?; + let remap_peak = estimate_vamana_memory_bytes(node_count, max_degree, max_degree.max(1), 1)? + .remap_peak_bytes; + Some(build_peak.max(remap_peak)) +} + +impl VamanaGraph { + pub fn from_adjacency(entry_node: u32, adjacency: Vec<Vec<u32>>) -> Self { + let max_degree = adjacency.iter().map(Vec::len).max().unwrap_or(0); + Self { + entry_node, + adjacency: CompactAdjacency::from_nested(adjacency, max_degree), + } + } + + pub fn build( + vectors: &[f32], + count: usize, + dimension: usize, + params: DiskAnnBuildParams, + ) -> io::Result<Self> { + Self::build_with_stats(vectors, count, dimension, MetricType::L2, params) + .map(|(graph, _)| graph) + } + + pub(crate) fn build_with_stats( + vectors: &[f32], + count: usize, + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + ) -> io::Result<(Self, VamanaBuildStats)> { + Self::build_with_search_distance( + vectors, + count, + dimension, + metric, + params, + BuildSearchDistance::FullPrecision { + vectors, + dimension, + metric, + }, + ) + } + + pub(crate) fn build_with_pq_stats( + vectors: &[f32], + pq: &ProductQuantizer, + pq_codes: &[u8], + count: usize, + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + ) -> io::Result<(Self, VamanaBuildStats)> { + validate_build_inputs(vectors, count, dimension, params)?; + if pq.d != dimension + || !matches!(pq.nbits, 4 | 8) + || !pq.has_valid_layout() + || pq_codes.len() != count.saturating_mul(pq.code_size()) + { + return Err(invalid_input( + "Vamana PQ-guided build received an invalid codebook or code buffer", + )); + } + let distance_started = Instant::now(); + let distance = PqBuildDistance::new(pq, pq_codes, count, metric)?; + let distance_initialization = distance_started.elapsed(); + let (graph, mut stats) = Self::build_with_search_distance( + vectors, + count, + dimension, + metric, + params, + BuildSearchDistance::ProductQuantized(distance), + )?; + stats.initialization = stats.initialization.saturating_add(distance_initialization); + Ok((graph, stats)) + } + + pub(crate) fn build_sharded_with_stats( + vectors: &[f32], + count: usize, + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + shard_count: usize, + ) -> io::Result<(Self, VamanaBuildStats)> { + validate_build_inputs(vectors, count, dimension, params)?; + if shard_count < 2 || shard_count > count { + return Err(invalid_input( + "Vamana shard count must be between 2 and the vector count", + )); + } + let initialization_started = Instant::now(); + let cluster_config = KMeansConfig { + niter: 8, + nredo: 1, + max_points_per_centroid: 256, + seed: params.seed, + balance_factor: 0.1, + }; + let centroids = + kmeans::kmeans_train(&cluster_config, vectors, count, dimension, shard_count); + let mut assignments = (0..count) + .into_par_iter() + .map(|node| { + nearest_two_centroids( + &vectors[node * dimension..(node + 1) * dimension], + ¢roids, + shard_count, + dimension, + ) + }) + .collect::<Vec<_>>(); + let membership_capacity = overlapping_shard_capacity(count, shard_count) + .ok_or_else(|| invalid_input("Vamana overlapping-shard capacity overflows"))?; + let mut memberships = (0..shard_count) + .map(|_| Vec::with_capacity(membership_capacity)) + .collect::<Vec<_>>(); + for (node, [first, second]) in assignments.iter().copied().enumerate() { + memberships[first].push(node as u32); + if second != first { + memberships[second].push(node as u32); + } + } + rebalance_overlapping_shards( + vectors, + dimension, + ¢roids, + &mut assignments, + &mut memberships, + )?; + let initialization = initialization_started.elapsed(); + + let entry_node = centroid_entry(vectors, count, dimension, metric) as u32; + let degree = params.max_degree.min(count.saturating_sub(1)); + let mut graph = Self { + entry_node, + adjacency: CompactAdjacency::empty(count, degree), + }; + let mut mapped = Vec::with_capacity(degree); + let mut pass_one = Duration::ZERO; + let mut pass_two = Duration::ZERO; + for (shard, members) in memberships.iter().enumerate() { + if members.len() < 2 { + continue; + } + let local_started = Instant::now(); + let mut local_vectors = Vec::new(); + local_vectors + .try_reserve_exact(members.len().saturating_mul(dimension)) + .map_err(|_| invalid_input("Vamana shard vector allocation failed"))?; + for &node in members { + let node = node as usize; + local_vectors.extend_from_slice(&vectors[node * dimension..(node + 1) * dimension]); + } + let local_degree = params.max_degree.min(members.len() - 1); + let local_params = DiskAnnBuildParams { + max_degree: local_degree, + build_search_list_size: params + .build_search_list_size + .min(members.len()) + .max(local_degree), + seed: derived_seed(params.seed, shard as u64), + ..params + }; + let local_graph = Self::build_sequential_with_metric( + &local_vectors, + members.len(), + dimension, + metric, + local_params, + )?; + pass_one = pass_one.saturating_add(local_started.elapsed()); + let merge_started = Instant::now(); + for (local_node, &global_node) in members.iter().enumerate() { + mapped.clear(); + mapped.extend( + local_graph.adjacency[local_node] + .iter() + .map(|&neighbor| members[neighbor as usize]), + ); + let selected = robust_prune_candidates( + vectors, + dimension, + global_node as usize, + &mapped, + &graph.adjacency[global_node as usize], + count, + degree, + params.alpha, + metric, + ); + graph.adjacency.replace(global_node as usize, &selected); + } + pass_two = pass_two.saturating_add(merge_started.elapsed()); + } + let connectivity_started = Instant::now(); + graph.repair_connectivity(vectors, dimension, degree, metric)?; + let connectivity_repair = connectivity_started.elapsed(); + graph.validate(degree)?; + if !graph.is_fully_reachable() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "sharded Vamana graph is not fully reachable from its entry node", + )); + } + Ok(( + graph, + VamanaBuildStats { + initialization, + pass_one, + pass_two, + connectivity_repair, + }, + )) + } + + fn build_with_search_distance( + vectors: &[f32], + count: usize, + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + search_distance: BuildSearchDistance<'_>, + ) -> io::Result<(Self, VamanaBuildStats)> { + validate_build_inputs(vectors, count, dimension, params)?; + let entry_node = centroid_entry(vectors, count, dimension, metric) as u32; + let degree = params.max_degree.min(count.saturating_sub(1)); + let initialization_started = Instant::now(); + let adjacency = ParallelAdjacency::new_random(count, degree, params.seed); + let initialization = initialization_started.elapsed(); + let mut rng = StdRng::seed_from_u64(derived_seed(params.seed, u64::MAX)); + let builder = ParallelVamanaBuilder { + vectors, + dimension, + metric, + entry_node, + adjacency, + search_distance, + }; + let pass_one_started = Instant::now(); + builder.run_pass(params, 1.0, &mut rng); + let pass_one = pass_one_started.elapsed(); + let pass_two_started = Instant::now(); + builder.run_pass(params, params.alpha, &mut rng); + let mut graph = builder.finish()?; + let pass_two = pass_two_started.elapsed(); + let connectivity_started = Instant::now(); + graph.repair_connectivity(vectors, dimension, params.max_degree, metric)?; + let connectivity_repair = connectivity_started.elapsed(); + graph.validate(params.max_degree)?; + if !graph.is_fully_reachable() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "parallel Vamana graph is not fully reachable from its entry node", + )); + } + Ok(( + graph, + VamanaBuildStats { + initialization, + pass_one, + pass_two, + connectivity_repair, + }, + )) + } + + pub fn build_sequential( + vectors: &[f32], + count: usize, + dimension: usize, + params: DiskAnnBuildParams, + ) -> io::Result<Self> { + Self::build_sequential_with_metric(vectors, count, dimension, MetricType::L2, params) + } + + fn build_sequential_with_metric( + vectors: &[f32], + count: usize, + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + ) -> io::Result<Self> { + validate_build_inputs(vectors, count, dimension, params)?; + let entry_node = centroid_entry(vectors, count, dimension, metric) as u32; + let mut rng = StdRng::seed_from_u64(params.seed); + let degree = params.max_degree.min(count.saturating_sub(1)); + let adjacency = (0..count) + .map(|node| random_neighbors(&mut rng, count, node, degree)) + .collect::<Vec<_>>(); + let mut graph = Self { + entry_node, + adjacency: CompactAdjacency::from_nested(adjacency, degree), + }; + + graph.run_sequential_pass(vectors, dimension, metric, params, 1.0, &mut rng); + graph.run_sequential_pass(vectors, dimension, metric, params, params.alpha, &mut rng); + graph.repair_connectivity(vectors, dimension, params.max_degree, metric)?; + graph.validate(params.max_degree)?; + if !graph.is_fully_reachable() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Vamana graph is not fully reachable from its entry node", + )); + } + Ok(graph) + } + + fn run_sequential_pass( + &mut self, + vectors: &[f32], + dimension: usize, + metric: MetricType, + params: DiskAnnBuildParams, + alpha: f32, + rng: &mut StdRng, + ) { + let mut order = (0..self.adjacency.len()).collect::<Vec<_>>(); + order.shuffle(rng); + for node in order { + let query_start = node * dimension; + let candidates = self.greedy_search_with_metric( + vectors, + dimension, + metric, + &vectors[query_start..query_start + dimension], + params.build_search_list_size, + ); + let candidate_ids = candidates + .into_iter() + .map(|candidate| candidate.id) + .collect::<Vec<_>>(); + let selected = self.robust_prune_with_metric( + vectors, + dimension, + metric, + node, + &candidate_ids, + params.max_degree, + alpha, + ); + self.adjacency.replace(node, &selected); + for neighbor in selected { + self.insert_reverse_edge( + vectors, + dimension, + metric, + neighbor as usize, + node as u32, + params.max_degree, + alpha, + ); + } + } + } + + fn insert_reverse_edge( + &mut self, + vectors: &[f32], + dimension: usize, + metric: MetricType, + node: usize, + neighbor: u32, + max_degree: usize, + alpha: f32, + ) { + if self.adjacency[node].contains(&neighbor) { + return; + } + if self.adjacency[node].len() < max_degree { + self.adjacency.push(node, neighbor); + return; + } + let mut candidates = self.adjacency[node].to_vec(); + candidates.push(neighbor); + let selected = self.robust_prune_with_metric( + vectors, + dimension, + metric, + node, + &candidates, + max_degree, + alpha, + ); + self.adjacency.replace(node, &selected); + } + + pub fn greedy_search( + &self, + vectors: &[f32], + dimension: usize, + query: &[f32], + search_list_size: usize, + ) -> Vec<ScoredNode> { + self.greedy_search_with_metric(vectors, dimension, MetricType::L2, query, search_list_size) + } + + fn greedy_search_with_metric( + &self, + vectors: &[f32], + dimension: usize, + metric: MetricType, + query: &[f32], + search_list_size: usize, + ) -> Vec<ScoredNode> { + if search_list_size == 0 || self.adjacency.is_empty() { + return Vec::new(); + } + + let mut visited = vec![false; self.adjacency.len()]; + let mut expanded = vec![false; self.adjacency.len()]; + let entry = self.entry_node as usize; + visited[entry] = true; + let mut results = vec![ScoredNode { + id: self.entry_node, + distance: node_distance(vectors, dimension, entry, query, metric), + }]; + + loop { + results.sort_by(scored_node_order); + results.truncate(search_list_size); + let Some(current) = results + .iter() + .find(|node| !expanded[node.id as usize]) + .copied() + else { + break; + }; + expanded[current.id as usize] = true; + + for &neighbor in &self.adjacency[current.id as usize] { + let neighbor = neighbor as usize; + if neighbor >= self.adjacency.len() || visited[neighbor] { + continue; + } + visited[neighbor] = true; + results.push(ScoredNode { + id: neighbor as u32, + distance: node_distance(vectors, dimension, neighbor, query, metric), + }); + } + } + + results.sort_by(scored_node_order); + results.truncate(search_list_size); + results + } + + #[allow(clippy::too_many_arguments)] + pub fn robust_prune( + &self, + vectors: &[f32], + dimension: usize, + node: usize, + candidates: &[u32], + max_degree: usize, + alpha: f32, + ) -> Vec<u32> { + self.robust_prune_with_metric( + vectors, + dimension, + MetricType::L2, + node, + candidates, + max_degree, + alpha, + ) + } + + #[allow(clippy::too_many_arguments)] + fn robust_prune_with_metric( + &self, + vectors: &[f32], + dimension: usize, + metric: MetricType, + node: usize, + candidates: &[u32], + max_degree: usize, + alpha: f32, + ) -> Vec<u32> { + robust_prune_candidates( + vectors, + dimension, + node, + candidates, + &self.adjacency[node], + self.adjacency.len(), + max_degree, + alpha, + metric, + ) + } + + pub fn is_fully_reachable(&self) -> bool { + if self.adjacency.is_empty() || self.entry_node as usize >= self.adjacency.len() { + return false; + } + let mut visited = vec![false; self.adjacency.len()]; + let mut queue = VecDeque::from([self.entry_node as usize]); + visited[self.entry_node as usize] = true; + while let Some(node) = queue.pop_front() { + for &neighbor in &self.adjacency[node] { + let neighbor = neighbor as usize; + if neighbor < visited.len() && !visited[neighbor] { + visited[neighbor] = true; + queue.push_back(neighbor); + } + } + } + visited.into_iter().all(|value| value) + } + + fn repair_connectivity( + &mut self, + vectors: &[f32], + dimension: usize, + max_degree: usize, + metric: MetricType, + ) -> io::Result<()> { + self.repair_connectivity_with_stats(vectors, dimension, max_degree, metric) + .map(|_| ()) + } + + fn repair_connectivity_with_stats( + &mut self, + vectors: &[f32], + dimension: usize, + max_degree: usize, + metric: MetricType, + ) -> io::Result<ConnectivityRepairStats> { + let (mut visited, mut parent) = self.reachability_tree(); + let mut stats = ConnectivityRepairStats { + full_reachability_traversals: 1, + ..ConnectivityRepairStats::default() + }; + let mut eligible_sources = visited + .iter() + .enumerate() + .filter_map(|(node, &reachable)| { + (reachable && self.is_repair_source(node, max_degree, &parent)).then_some(node) + }) + .collect::<Vec<_>>(); + let mut queue = VecDeque::new(); + let mut newly_reachable = Vec::new(); + let mut target_cursor = 0usize; + + loop { + while target_cursor < visited.len() && visited[target_cursor] { + target_cursor += 1; + } + if target_cursor == visited.len() { + return Ok(stats); + } + let target = target_cursor; + let target_vector = &vectors[target * dimension..target * dimension + dimension]; + let (source_index, source) = self.select_repair_source( + vectors, + dimension, + target, + target_vector, + max_degree, + &parent, + &eligible_sources, + &mut stats, + metric, + )?; + eligible_sources.swap_remove(source_index); + + if self.adjacency[source].len() == max_degree { + let removable = self.adjacency[source] + .iter() + .enumerate() + .filter(|(_, neighbor)| parent[**neighbor as usize] != Some(source)) + .map(|(slot, &neighbor)| { + ( + slot, + node_distance( + vectors, + dimension, + neighbor as usize, + &vectors[source * dimension..source * dimension + dimension], + metric, + ), + ) + }) + .max_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }) + .map(|(slot, _)| slot) + .expect("eligible full source has a non-tree edge"); + self.adjacency.swap_remove(source, removable); + } + self.adjacency.push(source, target as u32); + stats.edges_added += 1; + + visited[target] = true; + parent[target] = Some(source); + queue.push_back(target); + newly_reachable.clear(); + newly_reachable.push(target); + while let Some(node) = queue.pop_front() { + for &neighbor in &self.adjacency[node] { + let neighbor = neighbor as usize; + if !visited[neighbor] { + visited[neighbor] = true; + parent[neighbor] = Some(node); + queue.push_back(neighbor); + newly_reachable.push(neighbor); + } + } + } + if self.is_repair_source(source, max_degree, &parent) { + eligible_sources.push(source); + } + for &node in &newly_reachable { + if self.is_repair_source(node, max_degree, &parent) { + eligible_sources.push(node); + } + } + } + } + + #[allow(clippy::too_many_arguments)] + fn select_repair_source( + &self, + vectors: &[f32], + dimension: usize, + target: usize, + target_vector: &[f32], + max_degree: usize, + parent: &[Option<usize>], + eligible_sources: &[usize], + stats: &mut ConnectivityRepairStats, + metric: MetricType, + ) -> io::Result<(usize, usize)> { + if eligible_sources.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "Vamana connectivity repair found no replaceable edge", + )); + } + let sample_size = CONNECTIVITY_SOURCE_SAMPLE_SIZE.min(eligible_sources.len()); + let start = + derived_seed(target as u64, stats.edges_added as u64) as usize % eligible_sources.len(); + let mut best: Option<(usize, usize, f32)> = None; + for offset in 0..sample_size { + let index = (start + offset) % eligible_sources.len(); + let source = eligible_sources[index]; + if !self.is_repair_source(source, max_degree, parent) { + continue; + } + stats.source_distance_evaluations += 1; + let distance = node_distance(vectors, dimension, source, target_vector, metric); + let candidate = (index, source, distance); + if best.is_none_or(|(_, best_source, best_distance)| { + distance + .total_cmp(&best_distance) + .then_with(|| source.cmp(&best_source)) + .is_lt() + }) { + best = Some(candidate); + } + } + if let Some((index, source, _)) = best { + return Ok((index, source)); + } + eligible_sources + .iter() + .enumerate() + .find_map(|(index, &source)| { + self.is_repair_source(source, max_degree, parent) + .then_some((index, source)) + }) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "Vamana connectivity repair found no replaceable edge", + ) + }) + } + + fn is_repair_source(&self, source: usize, max_degree: usize, parent: &[Option<usize>]) -> bool { + self.adjacency[source].len() < max_degree + || self.adjacency[source] + .iter() + .any(|neighbor| parent[*neighbor as usize] != Some(source)) + } + + fn reachability_tree(&self) -> (Vec<bool>, Vec<Option<usize>>) { + let mut visited = vec![false; self.adjacency.len()]; + let mut parent = vec![None; self.adjacency.len()]; + let entry = self.entry_node as usize; + let mut queue = VecDeque::from([entry]); + visited[entry] = true; + while let Some(node) = queue.pop_front() { + for &neighbor in &self.adjacency[node] { + let neighbor = neighbor as usize; + if !visited[neighbor] { + visited[neighbor] = true; + parent[neighbor] = Some(node); + queue.push_back(neighbor); + } + } + } + (visited, parent) + } + + fn validate(&self, max_degree: usize) -> io::Result<()> { + for (node, neighbors) in self.adjacency.iter().enumerate() { + if neighbors.len() > max_degree { + return Err(invalid_input(format!( + "Vamana node {} degree {} exceeds maximum {}", + node, + neighbors.len(), + max_degree + ))); + } + let mut sorted = neighbors.to_vec(); + sorted.sort_unstable(); + if sorted.windows(2).any(|pair| pair[0] == pair[1]) { + return Err(invalid_input(format!( + "Vamana node {} contains duplicate neighbors", + node + ))); + } + if sorted.iter().any(|&neighbor| { + neighbor as usize >= self.adjacency.len() || neighbor as usize == node + }) { + return Err(invalid_input(format!( + "Vamana node {} contains invalid neighbor", + node + ))); + } + } + Ok(()) + } +} + +fn overlapping_shard_capacity(node_count: usize, shard_count: usize) -> Option<usize> { + if shard_count < 2 { + return None; + } + if shard_count == 2 { + return Some(node_count); + } + // For k > 2, ceil(2N / (k - 1)) leaves enough aggregate capacity even + // when one shard cannot accept a node's second, distinct membership. This + // lets deterministic overflow repair preserve exactly two memberships + // without an unbounded degenerate cluster. + node_count + .checked_mul(2)? + .checked_add(shard_count - 2) + .map(|value| value / (shard_count - 1)) +} + +fn rebalance_overlapping_shards( + vectors: &[f32], + dimension: usize, + centroids: &[f32], + assignments: &mut [[usize; 2]], + memberships: &mut [Vec<u32>], +) -> io::Result<()> { + let shard_count = memberships.len(); + let capacity = overlapping_shard_capacity(assignments.len(), shard_count) + .ok_or_else(|| invalid_input("Vamana overlapping-shard capacity overflows"))?; + for shard in 0..shard_count { + while memberships[shard].len() > capacity { + let node = memberships[shard] + .pop() + .expect("an overflowing Vamana shard is non-empty") as usize; + let slot = if assignments[node][0] == shard { + 0 + } else { + debug_assert_eq!(assignments[node][1], shard); + 1 + }; + let other = assignments[node][1 - slot]; + let vector = &vectors[node * dimension..(node + 1) * dimension]; + let mut replacement = None; + let mut replacement_distance = f32::INFINITY; + for candidate in 0..shard_count { + if candidate == other || memberships[candidate].len() >= capacity { + continue; + } + let start = candidate * dimension; + let distance = fvec_l2sqr(vector, ¢roids[start..start + dimension]); + if distance < replacement_distance + || (distance == replacement_distance + && replacement.is_none_or(|current| candidate < current)) + { + replacement = Some(candidate); + replacement_distance = distance; + } + } + let replacement = replacement.ok_or_else(|| { + invalid_input("Vamana overlapping shards cannot be capacity-balanced") + })?; + assignments[node][slot] = replacement; + memberships[replacement].push(node as u32); + } + } + debug_assert!(memberships.iter().all(|members| members.len() <= capacity)); + Ok(()) +} + +struct ParallelVamanaBuilder<'a> { + vectors: &'a [f32], + dimension: usize, + metric: MetricType, + entry_node: u32, + adjacency: ParallelAdjacency, + search_distance: BuildSearchDistance<'a>, +} + +enum BuildSearchDistance<'a> { + FullPrecision { + vectors: &'a [f32], + dimension: usize, + metric: MetricType, + }, + ProductQuantized(PqBuildDistance<'a>), +} + +impl BuildSearchDistance<'_> { + #[inline] + fn between(&self, left: usize, right: usize) -> f32 { + match self { + Self::FullPrecision { + vectors, + dimension, + metric, + } => distance_between(vectors, *dimension, left, right, *metric), + Self::ProductQuantized(distance) => distance.between(left, right), + } + } +} + +struct PqBuildDistance<'a> { + codes: &'a [u8], + code_size: usize, + m: usize, + ksub: usize, + nbits: usize, + centroid_distances: Vec<f32>, +} + +impl<'a> PqBuildDistance<'a> { + fn new( + pq: &ProductQuantizer, + codes: &'a [u8], + count: usize, + metric: MetricType, + ) -> io::Result<Self> { + let table_len = pq + .m + .checked_mul(pq.ksub) + .and_then(|value| value.checked_mul(pq.ksub)) + .ok_or_else(|| invalid_input("Vamana PQ build-distance table size overflows usize"))?; + let mut centroid_distances = Vec::new(); + centroid_distances + .try_reserve_exact(table_len) + .map_err(|_| invalid_input("Vamana PQ build-distance table allocation failed"))?; + centroid_distances.resize(table_len, 0.0); + for sub in 0..pq.m { + let chunk_dim = pq.chunk_dim(sub); + let sub_base = pq.centroid_chunk_base(sub); + let table_base = sub * pq.ksub * pq.ksub; + for left in 0..pq.ksub { + let left_start = sub_base + left * chunk_dim; + for right in left..pq.ksub { + let right_start = sub_base + right * chunk_dim; + let distance = fvec_distance( + &pq.centroids[left_start..left_start + chunk_dim], + &pq.centroids[right_start..right_start + chunk_dim], + metric, + ); + centroid_distances[table_base + left * pq.ksub + right] = distance; + centroid_distances[table_base + right * pq.ksub + left] = distance; + } + } + } + let code_size = pq.code_size(); + if codes.len() != count.saturating_mul(code_size) { + return Err(invalid_input( + "Vamana PQ code buffer does not match vector count", + )); + } + Ok(Self { + codes, + code_size, + m: pq.m, + ksub: pq.ksub, + nbits: pq.nbits, + centroid_distances, + }) + } + + #[inline] + fn code(&self, node: usize, sub: usize) -> usize { + let codes = &self.codes[node * self.code_size..(node + 1) * self.code_size]; + if self.nbits == 4 { + let byte = codes[sub / 2]; + usize::from(if sub.is_multiple_of(2) { + byte & 0x0f + } else { + byte >> 4 + }) + } else { + usize::from(codes[sub]) + } + } + + #[inline] + fn between(&self, left: usize, right: usize) -> f32 { + let mut distance = 0.0; + for sub in 0..self.m { + let left_code = self.code(left, sub); + let right_code = self.code(right, sub); + distance += self.centroid_distances + [sub * self.ksub * self.ksub + left_code * self.ksub + right_code]; + } + distance + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AdjacencyShard { + start_node: usize, + node_count: usize, + slots: Vec<u32>, + degrees: Vec<u16>, +} + +impl AdjacencyShard { + fn neighbors(&self, node: usize, max_degree: usize) -> &[u32] { + let local_node = node - self.start_node; + let start = local_node * max_degree; + &self.slots[start..start + self.degrees[local_node] as usize] + } + + fn neighbors_mut(&mut self, node: usize, max_degree: usize) -> &mut [u32] { + let local_node = node - self.start_node; + let start = local_node * max_degree; + let degree = self.degrees[local_node] as usize; + &mut self.slots[start..start + degree] + } + + fn replace(&mut self, node: usize, max_degree: usize, neighbors: &[u32]) { + debug_assert!(neighbors.len() <= max_degree); + let local_node = node - self.start_node; + let start = local_node * max_degree; + self.slots[start..start + neighbors.len()].copy_from_slice(neighbors); + self.degrees[local_node] = neighbors.len() as u16; + } + + fn push(&mut self, node: usize, max_degree: usize, neighbor: u32) { + let local_node = node - self.start_node; + let degree = self.degrees[local_node] as usize; + assert!(degree < max_degree, "compact adjacency node is full"); + self.slots[local_node * max_degree + degree] = neighbor; + self.degrees[local_node] += 1; + } + + fn swap_remove(&mut self, node: usize, max_degree: usize, slot: usize) { + let local_node = node - self.start_node; + let degree = self.degrees[local_node] as usize; + assert!(slot < degree, "compact adjacency removal slot is invalid"); + let start = local_node * max_degree; + self.slots[start + slot] = self.slots[start + degree - 1]; + self.degrees[local_node] -= 1; + } + + fn swap_node_with_buffer( + &mut self, + node: usize, + max_degree: usize, + slots: &mut [u32], + degree: &mut u16, + ) { + let local_node = node - self.start_node; + let start = local_node * max_degree; + for (stored, buffered) in self.slots[start..start + max_degree].iter_mut().zip(slots) { + std::mem::swap(stored, buffered); + } + std::mem::swap(&mut self.degrees[local_node], degree); + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CompactAdjacency { + shards: Vec<AdjacencyShard>, + node_count: usize, + max_degree: usize, +} + +impl CompactAdjacency { + fn empty(node_count: usize, max_degree: usize) -> Self { + let mut shards = + Vec::with_capacity(node_count.div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD)); + for start_node in (0..node_count).step_by(PARALLEL_ADJACENCY_NODES_PER_SHARD) { + let shard_node_count = PARALLEL_ADJACENCY_NODES_PER_SHARD.min(node_count - start_node); + shards.push(AdjacencyShard { + start_node, + node_count: shard_node_count, + slots: vec![0; shard_node_count * max_degree], + degrees: vec![0; shard_node_count], + }); + } + Self { + shards, + node_count, + max_degree, + } + } + + fn from_nested(adjacency: Vec<Vec<u32>>, max_degree: usize) -> Self { + let node_count = adjacency.len(); + let mut shards = + Vec::with_capacity(node_count.div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD)); + for start_node in (0..node_count).step_by(PARALLEL_ADJACENCY_NODES_PER_SHARD) { + let shard_node_count = PARALLEL_ADJACENCY_NODES_PER_SHARD.min(node_count - start_node); + let mut slots = vec![0; shard_node_count * max_degree]; + let mut degrees = vec![0; shard_node_count]; + for local_node in 0..shard_node_count { + let neighbors = &adjacency[start_node + local_node]; + assert!(neighbors.len() <= max_degree); + let slot_start = local_node * max_degree; + slots[slot_start..slot_start + neighbors.len()].copy_from_slice(neighbors); + degrees[local_node] = u16::try_from(neighbors.len()) + .expect("Vamana degree exceeds compact adjacency metadata"); + } + shards.push(AdjacencyShard { + start_node, + node_count: shard_node_count, + slots, + degrees, + }); + } + Self { + shards, + node_count, + max_degree, + } + } + + pub(crate) fn len(&self) -> usize { + self.node_count + } + + fn is_empty(&self) -> bool { + self.node_count == 0 + } + + pub(crate) fn neighbors(&self, node: usize) -> &[u32] { + self.shards[ParallelAdjacency::shard_index(node)].neighbors(node, self.max_degree) + } + + pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = &[u32]> { + (0..self.node_count).map(|node| self.neighbors(node)) + } + + fn replace(&mut self, node: usize, neighbors: &[u32]) { + self.shards[ParallelAdjacency::shard_index(node)].replace(node, self.max_degree, neighbors); + } + + fn push(&mut self, node: usize, neighbor: u32) { + self.shards[ParallelAdjacency::shard_index(node)].push(node, self.max_degree, neighbor); + } + + fn swap_remove(&mut self, node: usize, slot: usize) { + self.shards[ParallelAdjacency::shard_index(node)].swap_remove(node, self.max_degree, slot); + } + + pub(crate) fn permute_and_map_neighbors(&mut self, old_to_new: &[u32]) { + assert_eq!(old_to_new.len(), self.node_count); + let mut visited = vec![false; self.node_count]; + let mut slots = vec![0; self.max_degree]; + for start in 0..self.node_count { + if visited[start] { + continue; + } + slots.copy_from_slice(self.node_slots(start)); + let mut degree = self.node_degree(start); + let mut current = start; + loop { + visited[current] = true; + let destination = old_to_new[current] as usize; + assert!(destination < self.node_count); + self.swap_node_with_buffer(destination, &mut slots, &mut degree); + current = destination; + if current == start { + break; + } + assert!( + !visited[current], + "DiskANN locality mapping is not a permutation" + ); + } + } + for node in 0..self.node_count { + let neighbors = self.neighbors_mut(node); + for neighbor in neighbors.iter_mut() { + *neighbor = old_to_new[*neighbor as usize]; + } + neighbors.sort_unstable(); + } + } + + fn node_slots(&self, node: usize) -> &[u32] { + let shard = &self.shards[ParallelAdjacency::shard_index(node)]; + let local_node = node - shard.start_node; + let start = local_node * self.max_degree; + &shard.slots[start..start + self.max_degree] + } + + fn node_degree(&self, node: usize) -> u16 { + let shard = &self.shards[ParallelAdjacency::shard_index(node)]; + shard.degrees[node - shard.start_node] + } + + fn neighbors_mut(&mut self, node: usize) -> &mut [u32] { + self.shards[ParallelAdjacency::shard_index(node)].neighbors_mut(node, self.max_degree) + } + + fn swap_node_with_buffer(&mut self, node: usize, slots: &mut [u32], degree: &mut u16) { + self.shards[ParallelAdjacency::shard_index(node)].swap_node_with_buffer( + node, + self.max_degree, + slots, + degree, + ); + } +} + +impl Index<usize> for CompactAdjacency { + type Output = [u32]; + + fn index(&self, index: usize) -> &Self::Output { + self.neighbors(index) + } +} + +struct ParallelAdjacency { + shards: Vec<RwLock<AdjacencyShard>>, + node_count: usize, + max_degree: usize, +} + +impl ParallelAdjacency { + fn new_random(node_count: usize, max_degree: usize, seed: u64) -> Self { + assert!(max_degree <= node_count.saturating_sub(1)); + let shard_count = node_count.div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD); + let shards = (0..shard_count) + .into_par_iter() + .map(|shard_index| { + let start_node = shard_index * PARALLEL_ADJACENCY_NODES_PER_SHARD; + let shard_node_count = + PARALLEL_ADJACENCY_NODES_PER_SHARD.min(node_count - start_node); + let mut slots = vec![0; shard_node_count * max_degree]; + let degrees = vec![max_degree as u16; shard_node_count]; + let mut rng = StdRng::seed_from_u64(derived_seed(seed, shard_index as u64)); + let mut swaps = SparseTable::with_capacity(max_degree); + for local_node in 0..shard_node_count { + let node = start_node + local_node; + let slot_start = local_node * max_degree; + sample_random_neighbors_into( + &mut rng, + node_count, + node, + &mut swaps, + &mut slots[slot_start..slot_start + max_degree], + ); + } + RwLock::new(AdjacencyShard { + start_node, + node_count: shard_node_count, + slots, + degrees, + }) + }) + .collect(); + Self { + shards, + node_count, + max_degree, + } + } + + #[cfg(test)] + fn new( + node_count: usize, + max_degree: usize, + mut neighbors_for: impl FnMut(usize) -> Vec<u32>, + ) -> Self { + let mut shards = + Vec::with_capacity(node_count.div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD)); + for start_node in (0..node_count).step_by(PARALLEL_ADJACENCY_NODES_PER_SHARD) { + let shard_node_count = PARALLEL_ADJACENCY_NODES_PER_SHARD.min(node_count - start_node); + let mut slots = vec![0; shard_node_count * max_degree]; + let mut degrees = vec![0; shard_node_count]; + for local_node in 0..shard_node_count { + let neighbors = neighbors_for(start_node + local_node); + assert!( + neighbors.len() <= max_degree, + "parallel Vamana initializer exceeds maximum degree" + ); + let slot_start = local_node * max_degree; + slots[slot_start..slot_start + neighbors.len()].copy_from_slice(&neighbors); + degrees[local_node] = neighbors.len() as u16; + } + shards.push(RwLock::new(AdjacencyShard { + start_node, + node_count: shard_node_count, + slots, + degrees, + })); + } + Self { + shards, + node_count, + max_degree, + } + } + + fn node_count(&self) -> usize { + self.node_count + } + + fn shard_index(node: usize) -> usize { + node / PARALLEL_ADJACENCY_NODES_PER_SHARD + } + + fn copy_neighbors(&self, node: usize, target: &mut Vec<u32>) { + let shard = self.shards[Self::shard_index(node)] + .read() + .expect("parallel Vamana adjacency shard lock poisoned"); + target.extend_from_slice(shard.neighbors(node, self.max_degree)); + } + + #[cfg(test)] + fn replace(&self, node: usize, neighbors: &[u32]) { + let mut shard = self.shards[Self::shard_index(node)] + .write() + .expect("parallel Vamana adjacency shard lock poisoned"); + shard.replace(node, self.max_degree, neighbors); + } + + fn update_from_buffer( + &self, + node: usize, + replacement: &mut Vec<u32>, + update: impl FnOnce(&[u32], &mut Vec<u32>), + ) { + let mut shard = self.shards[Self::shard_index(node)] + .write() + .expect("parallel Vamana adjacency shard lock poisoned"); + update(shard.neighbors(node, self.max_degree), replacement); + shard.replace(node, self.max_degree, replacement); + } + + fn into_adjacency(self) -> io::Result<CompactAdjacency> { + let mut shards = Vec::with_capacity(self.shards.len()); + for shard in self.shards { + let shard = shard.into_inner().map_err(|_| { + io::Error::other("parallel Vamana adjacency shard lock poisoned during finish") + })?; + shards.push(shard); + } + Ok(CompactAdjacency { + shards, + node_count: self.node_count, + max_degree: self.max_degree, + }) + } +} + +enum BuildVisitStates { + Dense { + states: Vec<u8>, + touched_nodes: Vec<u32>, + }, + Sparse(SparseTable<u8>), +} + +struct GreedySearchScratch { + visit_states: BuildVisitStates, + results: BinaryHeap<ScoredNode>, + frontier: BinaryHeap<Reverse<ScoredNode>>, + neighbor_buffer: Vec<u32>, + candidate_ids: Vec<u32>, + prune_unique: Vec<u32>, + prune_pool: Vec<ScoredNode>, + prune_selected: Vec<u32>, + #[cfg(test)] + peak_retained: usize, + #[cfg(test)] + peak_frontier: usize, +} + +impl GreedySearchScratch { + fn new(node_count: usize, max_degree: usize, search_list_size: usize) -> Self { + let expected_visited = search_list_size + .saturating_mul(max_degree) + .saturating_add(1) + .min(node_count); + let dense_bytes = node_count + .saturating_mul(size_of::<u8>()) + .saturating_add(expected_visited.saturating_mul(size_of::<u32>())); + let sparse_bytes = + sparse_table_memory_bytes(expected_visited, size_of::<u8>()).unwrap_or(usize::MAX); + let visit_states = if sparse_bytes + .checked_mul(SPARSE_BUILD_VISITED_MIN_MEMORY_SAVINGS) + .is_some_and(|threshold| threshold < dense_bytes) + { + BuildVisitStates::Sparse(SparseTable::with_capacity(expected_visited)) + } else { + BuildVisitStates::Dense { + states: vec![0; node_count], + touched_nodes: Vec::with_capacity(expected_visited), + } + }; + Self { + visit_states, + results: BinaryHeap::new(), + frontier: BinaryHeap::new(), + neighbor_buffer: Vec::with_capacity(max_degree), + candidate_ids: Vec::with_capacity(search_list_size), + prune_unique: Vec::with_capacity(search_list_size.saturating_add(max_degree)), + prune_pool: Vec::with_capacity(search_list_size.saturating_add(max_degree)), + prune_selected: Vec::with_capacity(max_degree), + #[cfg(test)] + peak_retained: 0, + #[cfg(test)] + peak_frontier: 0, + } + } + + fn begin_search(&mut self) { + match &mut self.visit_states { + BuildVisitStates::Dense { + states, + touched_nodes, + } => { + for node in touched_nodes.drain(..) { + states[node as usize] = 0; + } + } + BuildVisitStates::Sparse(states) => states.clear(), + } + self.results.clear(); + self.frontier.clear(); + self.neighbor_buffer.clear(); + #[cfg(test)] + { + self.peak_retained = 0; + self.peak_frontier = 0; + } + } + + fn is_visited(&self, node: usize) -> bool { + match &self.visit_states { + BuildVisitStates::Dense { states, .. } => states[node] != 0, + BuildVisitStates::Sparse(states) => states.get(node as u32).is_some(), + } + } + + fn mark_visited(&mut self, node: usize) { + match &mut self.visit_states { + BuildVisitStates::Dense { + states, + touched_nodes, + } => { + if states[node] == 0 { + touched_nodes.push(node as u32); + } + states[node] = 1; + } + BuildVisitStates::Sparse(states) => { + states.insert(node as u32, 1); + } + } + } + + #[cfg(test)] + fn is_expanded(&self, node: usize) -> bool { + match &self.visit_states { + BuildVisitStates::Dense { states, .. } => states[node] == 3, + BuildVisitStates::Sparse(states) => states.get(node as u32) == Some(&3), + } + } + + fn is_retained_unexpanded(&self, node: usize) -> bool { + match &self.visit_states { + BuildVisitStates::Dense { states, .. } => states[node] == 2, + BuildVisitStates::Sparse(states) => states.get(node as u32) == Some(&2), + } + } + + fn mark_expanded(&mut self, node: usize) { + match &mut self.visit_states { + BuildVisitStates::Dense { states, .. } => states[node] = 3, + BuildVisitStates::Sparse(states) => { + states.insert(node as u32, 3); + } + } + } + + fn mark_retained(&mut self, node: usize) { + match &mut self.visit_states { + BuildVisitStates::Dense { states, .. } => states[node] = 2, + BuildVisitStates::Sparse(states) => { + states.insert(node as u32, 2); + } + } + } + + #[cfg(test)] + fn uses_sparse_states(&self) -> bool { + matches!(self.visit_states, BuildVisitStates::Sparse(_)) + } + + fn insert_candidate(&mut self, candidate: ScoredNode, search_list_size: usize) { + self.mark_visited(candidate.id as usize); + if search_list_size == 0 { + return; + } + if self.results.len() == search_list_size + && self.results.peek().is_some_and(|worst| candidate >= *worst) + { + return; + } + if self.results.len() == search_list_size { + let evicted = self + .results + .pop() + .expect("full result heap has a worst node"); + self.mark_visited(evicted.id as usize); + } + self.results.push(candidate); + self.frontier.push(Reverse(candidate)); + self.mark_retained(candidate.id as usize); + if self.frontier.len() > search_list_size.saturating_mul(2) { + let visit_states = &self.visit_states; + self.frontier + .retain(|Reverse(candidate)| match visit_states { + BuildVisitStates::Dense { states, .. } => states[candidate.id as usize] == 2, + BuildVisitStates::Sparse(states) => states.get(candidate.id) == Some(&2), + }); + } + #[cfg(test)] + { + self.peak_retained = self.peak_retained.max(self.results.len()); + self.peak_frontier = self.peak_frontier.max(self.frontier.len()); + } + } + + fn pop_nearest_unexpanded(&mut self) -> Option<ScoredNode> { + while let Some(Reverse(candidate)) = self.frontier.pop() { + if self.is_retained_unexpanded(candidate.id as usize) { + return Some(candidate); + } + } + None + } + + #[cfg(test)] + fn peak_retained_len(&self) -> usize { + self.peak_retained + } + + #[cfg(test)] + fn peak_frontier_len(&self) -> usize { + self.peak_frontier + } +} + +impl ParallelVamanaBuilder<'_> { + fn run_pass(&self, params: DiskAnnBuildParams, alpha: f32, rng: &mut StdRng) { + let mut order = (0..self.adjacency.node_count()).collect::<Vec<_>>(); + order.shuffle(rng); + let worker_count = rayon::current_num_threads().max(1); + let scratches = (0..worker_count) + .map(|_| { + Mutex::new(GreedySearchScratch::new( + self.adjacency.node_count(), + params.max_degree, + params.build_search_list_size, + )) + }) + .collect::<Vec<_>>(); + let batch_size = worker_count + .saturating_mul(PARALLEL_BUILD_BATCH_NODES_PER_WORKER) + .max(1); + let reverse_edges = Mutex::new(Vec::with_capacity( + batch_size.saturating_mul(params.max_degree), + )); + for batch in order.chunks(batch_size) { + batch.par_iter().for_each(|&node| { Review Comment: The seeded parallel build is nondeterministic even when the Rayon configuration is fixed. I reproduced this by building the same 512-vector graph repeatedly with the same data, parameters, `diskann.seed`, and the same 8-thread pool; the resulting `VamanaGraph`s differ. The workers read and mutate the shared adjacency concurrently within a batch, so search/update visibility depends on scheduling (and `batch_size` additionally changes with the host thread count). This makes the persisted DiskANN graph depend on runtime interleaving rather than only the input and configured seed. Could we make each batch read a stable graph snapshot and apply updates in a deterministic order, or otherwise preserve the seed's reproducibility contract? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
