leaves12138 commented on code in PR #62: URL: https://github.com/apache/paimon-vector-index/pull/62#discussion_r3650898732
########## core/src/diskann.rs: ########## @@ -0,0 +1,798 @@ +// 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::distance::{preprocess_vectors, MetricType}; +use crate::pq::ProductQuantizer; +use crate::vamana::{ + estimate_sharded_vamana_memory_bytes, estimate_vamana_memory_bytes, VamanaGraph, +}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use std::borrow::Cow; +use std::collections::VecDeque; +use std::io; +use std::time::{Duration, Instant}; + +pub(crate) const DISKANN_ADJACENCY_LOCATOR_NODE_BYTES: usize = 4; +pub(crate) const DISKANN_ADJACENCY_LOCATOR_BLOCK_NODES: usize = 16; +/// Match the proven DiskANN training bound: more samples materially increase +/// memory and training time without consistently improving the codebook. +pub const DISKANN_MAX_PQ_TRAINING_VECTORS: usize = 50_000; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DiskAnnStorageLayout { + /// Keep compressed adjacency pages and dense raw-vector records in separate sections. + Compact, + /// Store each raw vector immediately before its compressed adjacency list. + Interleaved, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u32)] +pub enum DiskAnnRawVectorEncoding { + /// Preserve indexed vectors and final distances as little-endian `f32`. + F32 = 1, + /// Store little-endian IEEE 754 binary16 values for approximate final reranking. + F16 = 2, +} + +impl DiskAnnRawVectorEncoding { + pub(crate) const fn element_size(self) -> usize { + match self { + Self::F32 => size_of::<f32>(), + Self::F16 => size_of::<u16>(), + } + } + + pub(crate) const fn from_code(code: u32) -> Option<Self> { + match code { + 1 => Some(Self::F32), + 2 => Some(Self::F16), + _ => None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DiskAnnBuildDistance { + /// Use full-precision distances for graph traversal and robust pruning. + FullPrecision, + /// Use PQ distances for graph traversal and full precision for robust pruning. + ProductQuantized, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct DiskAnnBuildParams { + pub max_degree: usize, + pub build_search_list_size: usize, + pub alpha: f32, + pub seed: u64, + pub memory_budget_bytes: usize, + pub storage_layout: DiskAnnStorageLayout, + pub raw_vector_encoding: DiskAnnRawVectorEncoding, + pub build_distance: DiskAnnBuildDistance, +} + +impl Default for DiskAnnBuildParams { + fn default() -> Self { + Self { + max_degree: 64, + build_search_list_size: 100, + alpha: 1.2, + seed: 42, + memory_budget_bytes: 8 * 1024 * 1024 * 1024, + storage_layout: DiskAnnStorageLayout::Compact, + raw_vector_encoding: DiskAnnRawVectorEncoding::F16, + build_distance: DiskAnnBuildDistance::ProductQuantized, + } + } +} + +pub(crate) fn validate_diskann_format_configuration( + dimension: usize, + pq_m: usize, + pq_bits: usize, + build: DiskAnnBuildParams, +) -> io::Result<()> { + if dimension == 0 { + return Err(invalid_input("DiskANN dimension must be greater than 0")); + } + if dimension > 1024 { + return Err(invalid_input("DiskANN v1 dimension must be at most 1024")); + } + if pq_m == 0 { + return Err(invalid_input("DiskANN pq.m must be greater than 0")); + } + if pq_m > dimension { + return Err(invalid_input(format!( + "DiskANN pq.m {} must not exceed dimension {}", + pq_m, dimension + ))); + } + if !matches!(pq_bits, 4 | 8) { + return Err(invalid_input("DiskANN pq.bits must be 4 or 8")); + } + if build.max_degree == 0 { + return Err(invalid_input( + "DiskANN maximum degree must be greater than 0", + )); + } + if build.max_degree > 1023 { + return Err(invalid_input(format!( + "DiskANN adjacency list size {} exceeds the v1 1023-neighbor page limit", + build.max_degree.saturating_mul(size_of::<u32>()) + ))); + } + if build.build_search_list_size < build.max_degree { + return Err(invalid_input(format!( + "DiskANN build search-list size {} must be at least maximum degree {}", + build.build_search_list_size, build.max_degree + ))); + } + if u32::try_from(build.build_search_list_size).is_err() { + return Err(invalid_input("DiskANN build search-list size exceeds u32")); + } + if !build.alpha.is_finite() || build.alpha < 1.0 { + return Err(invalid_input("DiskANN alpha must be at least 1 and finite")); + } + let interleaved_record_bytes = dimension + .checked_mul(build.raw_vector_encoding.element_size()) + .and_then(|vector_bytes| { + build + .max_degree + .checked_mul(size_of::<u32>()) + .and_then(|adjacency_bytes| vector_bytes.checked_add(adjacency_bytes)) + }); + if build.storage_layout == DiskAnnStorageLayout::Interleaved + && interleaved_record_bytes.is_none_or(|record_bytes| record_bytes > 4096) + { + return Err(invalid_input( + "DiskANN interleaved raw vector and maximum adjacency list must fit in one page", + )); + } + Ok(()) +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct DiskAnnBuildStats { + /// One for the normal parallel build; greater than one when the memory + /// budget selected overlapping shard construction. + pub graph_shards: usize, + pub total: Duration, + pub pq_encoding: Duration, + pub vamana_initialization: Duration, + pub vamana_pass_one: Duration, + pub vamana_pass_two: Duration, + pub connectivity_repair: Duration, + pub locality_remap: Duration, + pub resident_serialization: Duration, + pub adjacency_serialization: Duration, + pub vector_serialization: Duration, +} + +impl DiskAnnBuildStats { + pub fn accounted_duration(self) -> Duration { + [ + self.pq_encoding, + self.vamana_initialization, + self.vamana_pass_one, + self.vamana_pass_two, + self.connectivity_repair, + self.locality_remap, + self.resident_serialization, + self.adjacency_serialization, + self.vector_serialization, + ] + .into_iter() + .sum() + } +} + +pub struct DiskAnnIndex { + pub d: usize, + pub metric: MetricType, + pub pq: ProductQuantizer, + pub build_params: DiskAnnBuildParams, + pub ids: Vec<i64>, + pub vectors: Vec<f32>, +} + +impl DiskAnnIndex { + pub fn new( + d: usize, + metric: MetricType, + pq_m: usize, + build_params: DiskAnnBuildParams, + ) -> Self { + Self::with_pq_bits(d, metric, pq_m, 8, build_params) + } + + pub fn with_pq_bits( + d: usize, + metric: MetricType, + pq_m: usize, + pq_bits: usize, + build_params: DiskAnnBuildParams, + ) -> Self { + Self { + d, + metric, + pq: ProductQuantizer::with_nbits_balanced(d, pq_m, pq_bits), + build_params, + ids: Vec::new(), + vectors: Vec::new(), + } + } + + pub fn train(&mut self, data: &[f32], n: usize) { Review Comment: `diskann.memory-budget-bytes` is not consulted during PQ training, even though this phase can be substantially larger than the later graph estimate. `ProductQuantizer::train_hot_start` trains all `m` sub-quantizers with a Rayon `par_iter`; each active KMeans task copies its sub-vector training data again and may allocate a 16 MiB SGEMM assignment matrix. With the supported 50,000 x 1,024 training sample and 64 concurrent sub-quantizers, the score matrices alone can reach about 1 GiB, on top of the retained ~195 MiB sample and the per-task data copies. A caller can therefore set a 256 MiB build budget, exceed it by several times during `finish_training`, and only have the budget considered later during graph serialization. Please make PQ-training concurrency/sample scratch honor the build budget (or reject budgets that cannot cover the training peak) rather than treating only Vamana construction as budgeted. -- 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]
