leaves12138 commented on code in PR #62: URL: https://github.com/apache/paimon-vector-index/pull/62#discussion_r3650891288
########## 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, Review Comment: The sharded-fit calculation omits the KMeans peak that runs immediately before the graph allocations. `kmeans_train` always materializes a `train_data` copy (up to `shard_count * 256 * dimension` floats), plus its assignments, centroid buffers, and SGEMM score matrix, but none of those bytes are included here. At the supported limits this can be roughly another 64 MiB just for the copied training vectors (`64 * 256 * 1024 * 4`), before the other scratch buffers. Consequently `graph_build_shard_count` can select a shard count whose estimate is within `diskann.memory-budget-bytes` while the actual build exceeds that budget. Please include the KMeans training peak (and overlap it with the already-live fixed allocations) when deciding whether a sharded build fits. -- 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]
