leaves12138 commented on code in PR #62: URL: https://github.com/apache/paimon-vector-index/pull/62#discussion_r3650864883
########## 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) { + let processed = self.preprocess_vectors(data, n); + if let Some(sample) = + bounded_pq_training_sample(&processed, n, self.d, self.build_params.seed) + { + self.pq.train(&sample, DISKANN_MAX_PQ_TRAINING_VECTORS); + } else { + self.pq.train(&processed, n); + } + } + + pub fn add(&mut self, data: &[f32], ids: &[i64]) { + self.ids.extend_from_slice(ids); + self.vectors + .extend_from_slice(self.preprocess_vectors(data, ids.len()).as_ref()); + } + + pub fn estimate_build_memory_bytes(&self) -> io::Result<usize> { + let n = self.ids.len(); + let workers = rayon::current_num_threads().max(1); + let raw_vectors = checked_bytes(self.vectors.len(), size_of::<f32>(), "raw vectors")?; + let row_ids = checked_bytes(n, size_of::<i64>(), "row IDs")?; + let row_id_encoding_scratch = row_id_encoding_scratch_bytes(n)?; + let pq_codes = checked_bytes(n, self.pq.code_size(), "PQ codes")?; + let pq_codebook = checked_bytes(self.pq.centroids.len(), size_of::<f32>(), "PQ codebook")?; + let pq_build_distances = if self.build_params.build_distance + == DiskAnnBuildDistance::ProductQuantized + { + self.pq + .m + .checked_mul(self.pq.ksub) + .and_then(|value| value.checked_mul(self.pq.ksub)) + .and_then(|value| value.checked_mul(size_of::<f32>())) + .ok_or_else(|| invalid_input("DiskANN PQ build-distance table size overflows"))? + } else { + 0 + }; + let row_id_order = checked_bytes(n, size_of::<u32>(), "row-ID order")?; + let adjacency_index = checked_bytes( + n, + DISKANN_ADJACENCY_LOCATOR_NODE_BYTES, + "adjacency locators", + )? + .checked_add(checked_bytes( + n.div_ceil(DISKANN_ADJACENCY_LOCATOR_BLOCK_NODES), + size_of::<u64>(), + "adjacency locator block offsets", + )?) + .ok_or_else(|| invalid_input("DiskANN adjacency index size overflows usize"))?; + let vamana = estimate_vamana_memory_bytes( + n, + self.build_params.max_degree, + self.build_params.build_search_list_size, + workers, + ) + .ok_or_else(|| invalid_input("DiskANN Vamana memory estimate overflows usize"))?; + let graph_stage_peak = vamana.build_peak_bytes.max(vamana.remap_peak_bytes); + + [ + raw_vectors, + row_ids, + row_id_encoding_scratch, + pq_codes, + pq_codebook, + pq_build_distances, + row_id_order, + adjacency_index, + graph_stage_peak, + ] + .into_iter() + .try_fold(0usize, |total, value| { + total + .checked_add(value) + .ok_or_else(|| invalid_input("DiskANN memory estimate overflows usize")) + }) + } + + fn graph_build_shard_count(&self) -> io::Result<usize> { + let estimated = self.estimate_build_memory_bytes()?; + if estimated <= self.build_params.memory_budget_bytes { + return Ok(1); + } + let n = self.ids.len(); + let workers = rayon::current_num_threads().max(1); + let vamana = estimate_vamana_memory_bytes( + n, + self.build_params.max_degree, + self.build_params.build_search_list_size, + workers, + ) + .ok_or_else(|| invalid_input("DiskANN Vamana memory estimate overflows usize"))?; + let graph_peak = vamana.build_peak_bytes.max(vamana.remap_peak_bytes); + let fixed_bytes = estimated + .checked_sub(graph_peak) + .ok_or_else(|| invalid_input("DiskANN fixed memory estimate underflows"))?; + let max_shards = 64.min(n / 2); + for shard_count in 2..=max_shards { + let Some(sharded_graph) = estimate_sharded_vamana_memory_bytes( + n, + self.d, + self.build_params.max_degree.min(n.saturating_sub(1)), + shard_count, + ) else { + continue; + }; + if fixed_bytes + .checked_add(sharded_graph) + .is_some_and(|peak| peak <= self.build_params.memory_budget_bytes) + { + return Ok(shard_count); + } + } + Err(invalid_input(format!( + "DiskANN estimated build memory {} exceeds memory budget {}; overlapping sharded build also cannot fit", + estimated, self.build_params.memory_budget_bytes + ))) + } + + pub(crate) fn validate_for_write(&self) -> io::Result<()> { + validate_diskann_format_configuration(self.d, self.pq.m, self.pq.nbits, self.build_params)?; + if self.build_params.memory_budget_bytes == 0 { + return Err(invalid_input( + "DiskANN memory budget must be greater than zero", + )); + } + if self.ids.is_empty() || u32::try_from(self.ids.len()).is_err() { + return Err(invalid_input( + "DiskANN vector count must be between 1 and u32::MAX", + )); + } + let expected_vectors = self + .ids + .len() + .checked_mul(self.d) + .ok_or_else(|| invalid_input("DiskANN vector shape overflows usize"))?; + if self.vectors.len() != expected_vectors { + return Err(invalid_input(format!( + "DiskANN vector length {} does not match {} row IDs * dimension {}", + self.vectors.len(), + self.ids.len(), + self.d + ))); + } + if let Some(offset) = self.vectors.iter().position(|value| !value.is_finite()) { + return Err(invalid_input(format!( + "DiskANN vector data contains a non-finite value at offset {}", + offset + ))); + } + if self.build_params.raw_vector_encoding == DiskAnnRawVectorEncoding::F16 { + if let Some(offset) = self + .vectors + .iter() + .position(|&value| !half::f16::from_f32(value).is_finite()) + { + return Err(invalid_input(format!( + "DiskANN vector data at offset {} is outside the finite f16 range", + offset + ))); + } + } + + let expected_ksub = 1usize + .checked_shl(self.pq.nbits as u32) + .ok_or_else(|| invalid_input("DiskANN PQ centroid count overflows usize"))?; + let expected_centroids = self + .d + .checked_mul(expected_ksub) + .ok_or_else(|| invalid_input("DiskANN PQ codebook shape overflows usize"))?; + if self.pq.d != self.d + || self.pq.ksub != expected_ksub + || self.pq.centroids.len() != expected_centroids + || !self.pq.has_valid_layout() + { + return Err(invalid_input("DiskANN PQ codebook shape is invalid")); + } + if let Some(offset) = self + .pq + .centroids + .iter() + .position(|value| !value.is_finite()) + { + return Err(invalid_input(format!( + "DiskANN PQ codebook contains a non-finite value at offset {}", + offset + ))); + } + Ok(()) + } + + pub(crate) fn prepare_build(&self) -> io::Result<PreparedDiskAnn> { + self.validate_for_write()?; + let graph_shards = self.graph_build_shard_count()?; + + let pq_started = Instant::now(); + let mut pq_codes = vec![0u8; self.ids.len() * self.pq.code_size()]; + self.pq + .encode_batch(&self.vectors, self.ids.len(), &mut pq_codes); + let pq_encoding = pq_started.elapsed(); + let (mut graph, vamana_stats) = if graph_shards > 1 { Review Comment: Could we preserve the selected build-distance mode when the memory budget triggers sharding? `build_sharded_with_stats` builds each local graph through `build_sequential_with_metric`, which always uses full-precision distances, so both an explicit `diskann.build-distance=product_quantized` and the FastBuild/Balanced preset are silently ignored whenever `graph_shards > 1`. I reproduced this by building the same four-shard graph with the enum toggled; the graphs are identical. This also makes the budget model inconsistent: it accounts for the PQ build-distance table although the sharded execution instead allocates full-precision local vectors. Please add a PQ-guided sharded path, or reject/explicitly resolve this combination before shard selection, with coverage for both build-distance modes. ########## core/src/index.rs: ########## @@ -188,9 +254,243 @@ impl VectorIndexConfig { Self::IvfFlat { nlist, .. } | Self::IvfPq { nlist, .. } | Self::IvfRq { nlist, .. } - | Self::IvfHnswFlat { nlist, .. } - | Self::IvfHnswSq { nlist, .. } => *nlist, + | Self::IvfSq { nlist, .. } => *nlist, + Self::DiskAnn { .. } => 1, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ResolvedVectorIndexConfig { + pub index_type: IndexType, + pub dimension: usize, + pub nlist: usize, + pub metric: MetricType, + pub pq_m: Option<usize>, + pub pq_bits: Option<usize>, + pub rq_bits: Option<usize>, + pub use_opq: bool, + pub diskann_build: Option<DiskAnnBuildParams>, +} + +impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { + fn from(config: &VectorIndexConfig) -> Self { + match config { + VectorIndexConfig::IvfFlat { + dimension, + nlist, + metric, + } + | VectorIndexConfig::IvfSq { + dimension, + nlist, + metric, + } => Self { + index_type: config.index_type(), + dimension: *dimension, + nlist: *nlist, + metric: *metric, + pq_m: None, + pq_bits: None, + rq_bits: None, + use_opq: false, + diskann_build: None, + }, + VectorIndexConfig::IvfPq { + dimension, + nlist, + m, + metric, + use_opq, + } => Self { + index_type: IndexType::IvfPq, + dimension: *dimension, + nlist: *nlist, + metric: *metric, + pq_m: Some(*m), + pq_bits: Some(8), + rq_bits: None, + use_opq: *use_opq, + diskann_build: None, + }, + VectorIndexConfig::IvfRq { + dimension, + nlist, + bits, + metric, + } => Self { + index_type: IndexType::IvfRq, + dimension: *dimension, + nlist: *nlist, + metric: *metric, + pq_m: None, + pq_bits: None, + rq_bits: Some(*bits), + use_opq: false, + diskann_build: None, + }, + VectorIndexConfig::DiskAnn { + dimension, + metric, + pq_m, + pq_bits, + build, + } => Self { + index_type: IndexType::DiskAnn, + dimension: *dimension, + nlist: 1, + metric: *metric, + pq_m: Some(*pq_m), + pq_bits: Some(*pq_bits), + rq_bits: None, + use_opq: false, + diskann_build: Some(*build), + }, + } + } +} + +#[derive(Debug, Clone)] +pub struct VectorIndexBuildPlan { + pub config: VectorIndexConfig, + pub expected_vector_count: Option<usize>, + pub objective: TuningObjective, +} + +impl VectorIndexBuildPlan { + pub fn from_options(options: &HashMap<String, String>) -> io::Result<Self> { + let mut options = ConfigOptions::new(options)?; + let index_type = parse_index_type_option(&options.required("index.type")?)?; + let dimension = parse_usize_option("dimension", &options.required("dimension")?)?; + let expected_vector_count = options + .optional("expected-vector-count") + .map(|value| parse_usize_option("expected-vector-count", &value)) + .transpose()?; + if expected_vector_count == Some(0) { + return Err(invalid_input( + "expected-vector-count must be greater than 0", + )); + } + let metric = parse_metric_option(&options.required("metric")?)?; + let target_recall = options + .optional("target-recall") + .map(|value| parse_f32_option("target-recall", &value)) + .transpose()?; + if target_recall.is_some_and(|recall| !recall.is_finite() || !(0.0..=1.0).contains(&recall)) + { + return Err(invalid_input("target-recall must be finite and in [0, 1]")); + } + let max_bytes_per_vector = options Review Comment: Could we validate `max-bytes-per-vector` as an actual maximum before returning the plan? It is currently accepted for every index type without a final size check: IVF-FLAT/IVF-SQ ignore it, while DiskANN uses it only to choose PQ/raw-vector encodings even though the format always stores the raw vector, adjacency, row ID, and PQ code. For example, `index.type=diskann`, `dimension=128`, and `max-bytes-per-vector=32` succeeds and selects F16, but the raw vector alone is already 256 bytes per row; `ivf_flat` with a one-byte limit also succeeds. This silently violates the requested storage objective. Please estimate the complete persisted per-vector cost and reject unsatisfiable combinations, or narrow/rename the option and reject index types for which it is not enforced. -- 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]
