leaves12138 commented on code in PR #62: URL: https://github.com/apache/paimon-vector-index/pull/62#discussion_r3651750164
########## core/src/diskann.rs: ########## @@ -0,0 +1,1039 @@ +// 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::kmeans::KMeansConfig; +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(()) +} + +pub(crate) fn validate_diskann_training_budget( + dimension: usize, + metric: MetricType, + pq_m: usize, + pq_bits: usize, + memory_budget_bytes: usize, +) -> io::Result<()> { + let minimum_training_vectors = 1usize << pq_bits; + pq_training_plan_with_sample_buffers( + dimension, + pq_m, + minimum_training_vectors, + minimum_training_vectors, + memory_budget_bytes, + usize::from(metric == MetricType::Cosine) + 1, + ) + .map(|_| ()) + .map_err(|error| { + invalid_input(format!( + "DiskANN memory budget cannot fit minimum PQ training: {error}" + )) + }) +} + +pub(crate) fn diskann_training_sample_limit( + dimension: usize, + metric: MetricType, + pq_m: usize, + pq_bits: usize, + memory_budget_bytes: usize, +) -> io::Result<usize> { + pq_training_plan_with_sample_buffers( + dimension, + pq_m, + 1usize << pq_bits, + DISKANN_MAX_PQ_TRAINING_VECTORS, + memory_budget_bytes, + usize::from(metric == MetricType::Cosine) + 1, + ) + .map(|plan| plan.sample_count) +} + +#[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 plan = pq_training_plan( Review Comment: The high-level trainer reservoir is now budgeted, but the public low-level `DiskAnnIndex::train` path still calls `pq_training_plan` with one sample buffer. For cosine training with downsampling, this method simultaneously retains the sampled buffer returned by `bounded_pq_training_sample` and the normalized owned buffer returned by `preprocess_vectors`. With `d=1024`, `m=256`, 50,000 vectors, and a 128 MiB budget, the current plan selects 28,111 samples while the same peak model with two coexisting sample buffers selects only 14,379. Could this path pass the metric-aware sample-buffer count as well, and return an error rather than silently falling back to a one-vector plan when the budget is infeasible? ########## core/src/diskann_io.rs: ########## @@ -0,0 +1,5419 @@ +// 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. + +//! DiskANN v1 persistence and readers. +//! +//! The normative byte layout is documented in the repository's +//! [storage-format specification](https://github.com/apache/paimon-vector-index/blob/main/core/STORAGE_FORMAT.md#diskann-v1). + +use crate::diskann::{ + validate_diskann_format_configuration, DiskAnnBuildParams, DiskAnnBuildStats, DiskAnnIndex, + DiskAnnRawVectorEncoding, DiskAnnStorageLayout, PreparedDiskAnn, + DISKANN_ADJACENCY_LOCATOR_BLOCK_NODES as ADJACENCY_LOCATOR_BLOCK_NODES, + DISKANN_ADJACENCY_LOCATOR_NODE_BYTES, +}; +use crate::diskann_search::{DiskAnnQueryScratch, DiskAnnSearchStats}; +use crate::distance::MetricType; +use crate::io::{ReadRequest, SeekRead, SeekReadCapabilities, SeekWrite}; +use crate::pq::ProductQuantizer; +use crate::read_options::{ + DeploymentProfile, ReadPlan, ResolvedVectorIndexReaderOptions, VectorIndexReadPlan, + VectorIndexReaderOptions, +}; +use rayon::prelude::*; +use std::collections::{HashMap, HashSet}; +use std::io; +use std::ops::{Index, IndexMut}; +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::{Arc, Condvar, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +pub const DISKANN_MAGIC: u32 = 0x4E4E4144; // "DANN" +pub const DISKANN_VERSION: u32 = 1; +pub const DISKANN_HEADER_SIZE: usize = 256; +pub const DISKANN_PAGE_SIZE: u32 = 4096; +const FLAG_BFS_LAYOUT: u32 = 1 << 0; +const FLAG_SEPARATE_ADJACENCY_AND_VECTORS: u32 = 1 << 1; +const FLAG_ADAPTIVE_ADJACENCY: u32 = 1 << 2; +const FLAG_PQ_CODES: u32 = 1 << 3; +const FLAG_ROW_ID_ORDER: u32 = 1 << 4; +const FLAG_INTERLEAVED_ADJACENCY_AND_VECTORS: u32 = 1 << 5; +pub const DISKANN_REQUIRED_FLAGS: u32 = + FLAG_BFS_LAYOUT | FLAG_ADAPTIVE_ADJACENCY | FLAG_PQ_CODES | FLAG_ROW_ID_ORDER; +const DISKANN_SUPPORTED_FLAGS: u32 = DISKANN_REQUIRED_FLAGS + | FLAG_SEPARATE_ADJACENCY_AND_VECTORS + | FLAG_INTERLEAVED_ADJACENCY_AND_VECTORS; +const SECTION_COUNT: usize = 7; +const ADJACENCY_LOCATOR_SIZE: u32 = DISKANN_ADJACENCY_LOCATOR_NODE_BYTES as u32; +const ADJACENCY_LOCATOR_ENCODING: u32 = 3; +const ADJACENCY_LOCATOR_BLOCK_BASE_SIZE: usize = size_of::<u64>(); +const ADJACENCY_RAW_U32_FLAG: u16 = 1 << 15; +const ADJACENCY_DEGREE_MASK: u16 = ADJACENCY_RAW_U32_FLAG - 1; +const ROW_ID_SECTION_HEADER_SIZE: usize = 32; +const ROW_ID_ENCODING_RAW_I64: u32 = 0; +const ROW_ID_ENCODING_FOR_BITPACK: u32 = 1; +const PQ_CODEBOOK_MAGIC: u32 = 0x3151_5044; // "DPQ1" +const PQ_CODEBOOK_VERSION: u32 = 1; +const PQ_CODEBOOK_HEADER_SIZE: usize = 32; +const DISKANN_WRITE_BUFFER_SIZE: usize = 1024 * 1024; +const DISKANN_RESIDENT_DECODE_BUFFER_SIZE: usize = 1024 * 1024; +const DISKANN_ADJACENCY_PRELOAD_ALIGNMENT: usize = 64 * 1024; +const AUTO_PROFILE_MEMORY_LATENCY_THRESHOLD: Duration = Duration::from_micros(50); +const AUTO_PROFILE_LOCAL_LATENCY_THRESHOLD: Duration = Duration::from_micros(750); +const AUTO_PROFILE_REMOTE_LATENCY_THRESHOLD: Duration = Duration::from_millis(3); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SectionRange { + pub offset: u64, + pub length: u64, +} + +impl SectionRange { + pub const fn new(offset: u64, length: u64) -> Self { + Self { offset, length } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DiskAnnSections { + pub codebook: SectionRange, + pub row_ids: SectionRange, + pub pq_codes: SectionRange, + pub row_id_order: SectionRange, + pub adjacency_index: SectionRange, + pub adjacency: SectionRange, + pub vectors: SectionRange, +} + +impl DiskAnnSections { + fn from_array(sections: [SectionRange; SECTION_COUNT]) -> Self { + Self { + codebook: sections[0], + row_ids: sections[1], + pq_codes: sections[2], + row_id_order: sections[3], + adjacency_index: sections[4], + adjacency: sections[5], + vectors: sections[6], + } + } + + fn as_array(self) -> [SectionRange; SECTION_COUNT] { + [ + self.codebook, + self.row_ids, + self.pq_codes, + self.row_id_order, + self.adjacency_index, + self.adjacency, + self.vectors, + ] + } +} + +impl Index<usize> for DiskAnnSections { + type Output = SectionRange; + + fn index(&self, index: usize) -> &Self::Output { + match index { + 0 => &self.codebook, + 1 => &self.row_ids, + 2 => &self.pq_codes, + 3 => &self.row_id_order, + 4 => &self.adjacency_index, + 5 => &self.adjacency, + 6 => &self.vectors, + _ => panic!("DiskANN section index {index} is out of range"), + } + } +} + +impl IndexMut<usize> for DiskAnnSections { + fn index_mut(&mut self, index: usize) -> &mut Self::Output { + match index { + 0 => &mut self.codebook, + 1 => &mut self.row_ids, + 2 => &mut self.pq_codes, + 3 => &mut self.row_id_order, + 4 => &mut self.adjacency_index, + 5 => &mut self.adjacency, + 6 => &mut self.vectors, + _ => panic!("DiskANN section index {index} is out of range"), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct AdjacencyLocator { + pub page_index: u32, + pub byte_offset: u16, + degree_and_flags: u16, +} + +impl AdjacencyLocator { + fn new( + page_index: u32, + byte_offset: u16, + degree: usize, + encoding: AdjacencyListEncoding, + ) -> io::Result<Self> { + let degree = u16::try_from(degree) + .map_err(|_| invalid_input("DiskANN adjacency degree exceeds u16"))?; + if degree > ADJACENCY_DEGREE_MASK { + return Err(invalid_input( + "DiskANN adjacency degree exceeds locator capacity", + )); + } + let encoding_flag = match encoding { + AdjacencyListEncoding::DeltaVarint => 0, + AdjacencyListEncoding::RawU32 => ADJACENCY_RAW_U32_FLAG, + }; + Ok(Self { + page_index, + byte_offset, + degree_and_flags: degree | encoding_flag, + }) + } + + pub(crate) fn degree(self) -> usize { + usize::from(self.degree_and_flags & ADJACENCY_DEGREE_MASK) + } + + pub(crate) fn encoding(self) -> AdjacencyListEncoding { + if self.degree_and_flags & ADJACENCY_RAW_U32_FLAG == 0 { + AdjacencyListEncoding::DeltaVarint + } else { + AdjacencyListEncoding::RawU32 + } + } +} + +#[derive(Debug)] +struct AdjacencyIndex { + block_offsets: Box<[u64]>, + relative_offsets: Box<[u16]>, + degree_and_flags: Box<[u16]>, +} + +impl AdjacencyIndex { + #[cfg(test)] + fn from_locators(locators: &[AdjacencyLocator]) -> io::Result<Self> { + let block_count = locators.len().div_ceil(ADJACENCY_LOCATOR_BLOCK_NODES); + let mut block_offsets = Vec::new(); + block_offsets + .try_reserve_exact(block_count) + .map_err(|_| invalid_data("DiskANN adjacency block-offset allocation failed"))?; + let mut relative_offsets = Vec::new(); + relative_offsets + .try_reserve_exact(locators.len()) + .map_err(|_| invalid_data("DiskANN adjacency relative-offset allocation failed"))?; + let mut degree_and_flags = Vec::new(); + degree_and_flags + .try_reserve_exact(locators.len()) + .map_err(|_| invalid_data("DiskANN adjacency metadata allocation failed"))?; + + for (node, locator) in locators.iter().copied().enumerate() { + let absolute_offset = adjacency_locator_absolute_offset(locator)?; + if node.is_multiple_of(ADJACENCY_LOCATOR_BLOCK_NODES) { + block_offsets.push(absolute_offset); + } + let block_offset = *block_offsets + .last() + .expect("each adjacency locator belongs to a block"); + let relative_offset = absolute_offset + .checked_sub(block_offset) + .and_then(|offset| u16::try_from(offset).ok()) + .ok_or_else(|| { + invalid_data("DiskANN adjacency locator exceeds its block offset range") + })?; + relative_offsets.push(relative_offset); + degree_and_flags.push(locator.degree_and_flags); + } + Ok(Self { + block_offsets: block_offsets.into_boxed_slice(), + relative_offsets: relative_offsets.into_boxed_slice(), + degree_and_flags: degree_and_flags.into_boxed_slice(), + }) + } + + fn len(&self) -> usize { + self.relative_offsets.len() + } + + fn locator(&self, node: usize) -> Option<AdjacencyLocator> { + let relative_offset = u64::from(*self.relative_offsets.get(node)?); + let degree_and_flags = *self.degree_and_flags.get(node)?; + let block_offset = *self + .block_offsets + .get(node / ADJACENCY_LOCATOR_BLOCK_NODES)?; + let absolute_offset = block_offset.checked_add(relative_offset)?; + let page_index = u32::try_from(absolute_offset / u64::from(DISKANN_PAGE_SIZE)).ok()?; + let byte_offset = u16::try_from(absolute_offset % u64::from(DISKANN_PAGE_SIZE)).ok()?; + Some(AdjacencyLocator { + page_index, + byte_offset, + degree_and_flags, + }) + } + + fn partition_point(&self, mut predicate: impl FnMut(AdjacencyLocator) -> bool) -> usize { + let mut left = 0; + let mut right = self.len(); + while left < right { + let middle = left + (right - left) / 2; + let locator = self + .locator(middle) + .expect("validated DiskANN adjacency index"); + if predicate(locator) { + left = middle + 1; + } else { + right = middle; + } + } + left + } +} + +fn adjacency_locator_absolute_offset(locator: AdjacencyLocator) -> io::Result<u64> { + u64::from(locator.page_index) + .checked_mul(u64::from(DISKANN_PAGE_SIZE)) + .and_then(|offset| offset.checked_add(u64::from(locator.byte_offset))) + .ok_or_else(|| invalid_data("DiskANN adjacency locator offset overflows")) +} + +fn adjacency_index_serialized_len(vector_count: usize) -> io::Result<u64> { + let block_count = vector_count.div_ceil(ADJACENCY_LOCATOR_BLOCK_NODES); + let block_bytes = block_count + .checked_mul(ADJACENCY_LOCATOR_BLOCK_BASE_SIZE) + .ok_or_else(|| invalid_input("DiskANN adjacency block-offset size overflows usize"))?; + let locator_bytes = vector_count + .checked_mul(ADJACENCY_LOCATOR_SIZE as usize) + .ok_or_else(|| invalid_input("DiskANN adjacency locator size overflows usize"))?; + u64::try_from( + block_bytes + .checked_add(locator_bytes) + .ok_or_else(|| invalid_input("DiskANN adjacency index size overflows usize"))?, + ) + .map_err(|_| invalid_input("DiskANN adjacency index size exceeds u64")) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AdjacencyListEncoding { + DeltaVarint, + RawU32, +} + +fn encode_adjacency_list( + neighbors: &[u32], + encoded: &mut Vec<u8>, +) -> io::Result<AdjacencyListEncoding> { + encoded.clear(); + let (encoding, encoded_len) = plan_adjacency_list(neighbors)?; + encoded + .try_reserve_exact(encoded_len) + .map_err(|_| invalid_input("DiskANN adjacency allocation failed"))?; + match encoding { + AdjacencyListEncoding::DeltaVarint => { + let mut previous = 0u32; + for &neighbor in neighbors { + append_u32_varint(encoded, neighbor - previous); + previous = neighbor; + } + } + AdjacencyListEncoding::RawU32 => { + for &neighbor in neighbors { + encoded.extend_from_slice(&neighbor.to_le_bytes()); + } + } + } + Ok(encoding) +} + +fn plan_adjacency_list(neighbors: &[u32]) -> io::Result<(AdjacencyListEncoding, usize)> { + let raw_len = neighbors + .len() + .checked_mul(size_of::<u32>()) + .ok_or_else(|| invalid_input("DiskANN adjacency list size overflows usize"))?; + let delta_len = adjacency_delta_varint_len(neighbors) + .ok_or_else(|| invalid_input("DiskANN adjacency neighbors must be strictly increasing"))?; + if neighbors.is_empty() || delta_len < raw_len { + return Ok((AdjacencyListEncoding::DeltaVarint, delta_len)); + } + Ok((AdjacencyListEncoding::RawU32, raw_len)) +} + +fn adjacency_delta_varint_len(neighbors: &[u32]) -> Option<usize> { + let mut previous = 0u32; + let mut encoded_len = 0usize; + for (index, &neighbor) in neighbors.iter().enumerate() { + if index != 0 && neighbor <= previous { + return None; + } + encoded_len = encoded_len.checked_add(u32_varint_len(neighbor - previous))?; + previous = neighbor; + } + Some(encoded_len) +} + +fn u32_varint_len(value: u32) -> usize { + let significant_bits = (u32::BITS - value.leading_zeros()).max(1); + significant_bits.div_ceil(7) as usize +} + +pub(crate) fn decode_adjacency_list( + bytes: &[u8], + degree: usize, + encoding: AdjacencyListEncoding, + neighbors: &mut Vec<u32>, +) -> io::Result<usize> { + neighbors.clear(); + neighbors + .try_reserve(degree) + .map_err(|_| invalid_data("DiskANN adjacency decode allocation failed"))?; + match encoding { + AdjacencyListEncoding::DeltaVarint => { + let mut position = 0usize; + let mut previous = 0u32; + for _ in 0..degree { + let delta = read_u32_varint(bytes, &mut position)?; + let neighbor = previous + .checked_add(delta) + .ok_or_else(|| invalid_data("DiskANN adjacency delta overflows u32"))?; + neighbors.push(neighbor); + previous = neighbor; + } + Ok(position) + } + AdjacencyListEncoding::RawU32 => { + let encoded_len = degree + .checked_mul(size_of::<u32>()) + .ok_or_else(|| invalid_data("DiskANN raw adjacency size overflows usize"))?; + let encoded = bytes + .get(..encoded_len) + .ok_or_else(|| invalid_data("DiskANN raw adjacency list is truncated"))?; + neighbors.extend(encoded.chunks_exact(4).map(|value| { + u32::from_le_bytes(value.try_into().expect("fixed adjacency neighbor")) + })); + Ok(encoded_len) + } + } +} + +fn append_u32_varint(encoded: &mut Vec<u8>, mut value: u32) { + while value >= 0x80 { + encoded.push((value as u8 & 0x7f) | 0x80); + value >>= 7; + } + encoded.push(value as u8); +} + +fn read_u32_varint(bytes: &[u8], position: &mut usize) -> io::Result<u32> { + let mut value = 0u32; + let start = *position; + for shift in (0..=28).step_by(7) { + let byte = *bytes + .get(*position) + .ok_or_else(|| invalid_data("DiskANN adjacency varint is truncated"))?; + *position += 1; + if shift == 28 && byte > 0x0f { + return Err(invalid_data("DiskANN adjacency varint exceeds u32")); + } + value |= u32::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + if *position - start > 1 && byte == 0 { + return Err(invalid_data("DiskANN adjacency varint is not canonical")); + } + return Ok(value); + } + } + Err(invalid_data("DiskANN adjacency varint exceeds five bytes")) +} + +#[derive(Debug, Clone, PartialEq)] +pub struct DiskAnnHeader { + pub flags: u32, + pub dimension: u32, + pub metric: u32, + pub vector_count: u64, + pub entry_node: u32, + pub max_degree: u32, + pub build_search_list_size: u32, + pub alpha: f32, + pub seed: u64, + pub pq_m: u32, + pub pq_bits: u32, + pub page_size: u32, + pub adjacency_locator_size: u32, + pub adjacency_locator_encoding: u32, + pub raw_vector_encoding: u32, + pub vector_record_size: u32, + pub file_len: u64, + pub sections: DiskAnnSections, +} + +pub struct DiskAnnIndexReader<R: SeekRead> { + reader: R, + pub header: DiskAnnHeader, + resident: Option<Arc<DiskAnnResidentData>>, + options: ResolvedVectorIndexReaderOptions, + read_capabilities: SeekReadCapabilities, + effective_read_tier: DeploymentProfile, + random_read_latency: Duration, + hot_adjacency: Arc<[u8]>, + row_id_order: Arc<Mutex<RowIdOrderState>>, + pub(crate) query_scratch: Box<DiskAnnQueryScratch>, + pub(crate) last_search_stats: DiskAnnSearchStats, + pub(crate) batch_workers: Vec<DiskAnnIndexReader<R>>, + pub(crate) calibrated_l_search: Option<usize>, +} + +struct DiskAnnResidentData { + pq: ProductQuantizer, + row_ids: RowIdStorage, + pq_codes: Vec<u8>, + adjacency_index: AdjacencyIndex, + adjacency_validation: AdjacencyValidationCache, + adjacency_cache: SharedWindowCache, + raw_vector_cache: SharedWindowCache, +} + +#[derive(Clone, Copy, Default)] +struct OffsetLruLink { + older: Option<u64>, + newer: Option<u64>, +} + +#[derive(Default)] +pub(crate) struct OffsetLru { + links: HashMap<u64, OffsetLruLink>, + oldest: Option<u64>, + newest: Option<u64>, +} + +impl OffsetLru { + pub(crate) fn touch(&mut self, offset: u64) { + if self.newest == Some(offset) { + return; + } + let already_present = match self.links.entry(offset) { + std::collections::hash_map::Entry::Occupied(_) => true, + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(OffsetLruLink::default()); + false + } + }; + if already_present { + self.detach(offset); + } + let older = self.newest; + if let Some(older) = older { + self.links + .get_mut(&older) + .expect("DiskANN LRU newest offset must exist") + .newer = Some(offset); + } else { + self.oldest = Some(offset); + } + let link = self + .links + .get_mut(&offset) + .expect("DiskANN LRU touched offset must exist"); + link.older = older; + link.newer = None; + self.newest = Some(offset); + } + + pub(crate) fn remove(&mut self, offset: u64) { + if self.links.contains_key(&offset) { + self.detach(offset); + self.links.remove(&offset); + } + } + + pub(crate) fn pop_oldest(&mut self) -> Option<u64> { + let offset = self.oldest?; + self.remove(offset); + Some(offset) + } + + pub(crate) fn clear(&mut self) { + self.links.clear(); + self.oldest = None; + self.newest = None; + } + + #[cfg(test)] + pub(crate) fn is_empty(&self) -> bool { + self.links.is_empty() + } + + #[cfg(test)] + pub(crate) fn len(&self) -> usize { + self.links.len() + } + + #[cfg(test)] + pub(crate) fn oldest_offsets(&self) -> Vec<u64> { + let mut offsets = Vec::with_capacity(self.links.len()); + let mut current = self.oldest; + while let Some(offset) = current { + offsets.push(offset); + current = self.links.get(&offset).and_then(|link| link.newer); + } + debug_assert_eq!(offsets.len(), self.links.len()); + offsets + } + + fn detach(&mut self, offset: u64) { + let link = *self + .links + .get(&offset) + .expect("DiskANN LRU detached offset must exist"); + if let Some(older) = link.older { + self.links + .get_mut(&older) + .expect("DiskANN LRU older offset must exist") + .newer = link.newer; + } else { + self.oldest = link.newer; + } + if let Some(newer) = link.newer { + self.links + .get_mut(&newer) + .expect("DiskANN LRU newer offset must exist") + .older = link.older; + } else { + self.newest = link.older; + } + } +} + +pub(crate) enum SharedWindowCacheLookup { + Hit(Arc<Vec<u8>>), + Reserved, + Loading, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) struct CacheLockMetrics { + pub(crate) acquisitions: usize, + pub(crate) wait_nanos: u64, +} + +struct SharedWindowCacheState { + entries: HashMap<u64, Arc<Vec<u8>>>, + loading: HashSet<u64>, + recency: OffsetLru, + retained_bytes: usize, +} + +struct SharedWindowCacheShard { + capacity_bytes: AtomicUsize, + state: Mutex<SharedWindowCacheState>, + waiters: Condvar, +} + +const SHARED_WINDOW_CACHE_SHARDS: usize = 16; +const MAX_SHARED_READ_WINDOW_BYTES: usize = 64 * 1024; + +pub(crate) struct SharedWindowCache { + shards: Box<[SharedWindowCacheShard]>, +} + +impl SharedWindowCache { + fn new(capacity_bytes: usize) -> Self { + Self::new_with_max_shards(capacity_bytes, SHARED_WINDOW_CACHE_SHARDS) + } + + fn new_with_max_shards(capacity_bytes: usize, max_shards: usize) -> Self { + let max_shards = max_shards.max(1); + let shard_count = if capacity_bytes >= max_shards * MAX_SHARED_READ_WINDOW_BYTES { + max_shards + } else { + 1 + }; + let base_capacity = capacity_bytes / shard_count; + let remainder = capacity_bytes % shard_count; + let shards = (0..shard_count) + .map(|shard| SharedWindowCacheShard { + capacity_bytes: AtomicUsize::new(base_capacity + usize::from(shard < remainder)), + state: Mutex::new(SharedWindowCacheState { + entries: HashMap::new(), + loading: HashSet::new(), + recency: OffsetLru::default(), + retained_bytes: 0, + }), + waiters: Condvar::new(), + }) + .collect::<Vec<_>>() + .into_boxed_slice(); + Self { shards } + } + + fn shard_index(&self, offset: u64) -> usize { + let page = offset / u64::from(DISKANN_PAGE_SIZE); + let mut mixed = page.wrapping_add(0x9e37_79b9_7f4a_7c15); + mixed = (mixed ^ (mixed >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + mixed = (mixed ^ (mixed >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + ((mixed ^ (mixed >> 31)) as usize) % self.shards.len() + } + + fn shard(&self, offset: u64) -> &SharedWindowCacheShard { + &self.shards[self.shard_index(offset)] + } + + #[cfg(test)] + fn shard_count(&self) -> usize { + self.shards.len() + } + + #[cfg(test)] + fn total_capacity(&self) -> usize { + self.shards + .iter() + .map(|shard| shard.capacity_bytes.load(AtomicOrdering::Relaxed)) + .sum() + } + + fn add_lock_metrics(total: &mut CacheLockMetrics, metrics: CacheLockMetrics) { + total.acquisitions = total.acquisitions.saturating_add(metrics.acquisitions); + total.wait_nanos = total.wait_nanos.saturating_add(metrics.wait_nanos); + } + + fn set_total_capacity(&self, capacity_bytes: usize) -> io::Result<()> { + let base_capacity = capacity_bytes / self.shards.len(); + let remainder = capacity_bytes % self.shards.len(); + for (shard_index, shard) in self.shards.iter().enumerate() { + let shard_capacity = base_capacity + usize::from(shard_index < remainder); + shard + .capacity_bytes + .store(shard_capacity, AtomicOrdering::Relaxed); + let (mut state, _) = Self::lock_state(shard)?; + while state.retained_bytes > shard_capacity { + let Some(oldest) = state.recency.pop_oldest() else { + break; + }; + if let Some(evicted) = state.entries.remove(&oldest) { + state.retained_bytes = state.retained_bytes.saturating_sub(evicted.capacity()); + } + } + } + Ok(()) + } + + fn lock_state( + shard: &SharedWindowCacheShard, + ) -> io::Result<(MutexGuard<'_, SharedWindowCacheState>, CacheLockMetrics)> { + let started = Instant::now(); + let state = shard + .state + .lock() + .map_err(|_| invalid_data("DiskANN shared window cache state is poisoned"))?; + Ok(( + state, + CacheLockMetrics { + acquisitions: 1, + wait_nanos: u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX), + }, + )) + } + + fn remove_loading(&self, offsets: &[u64], shard_index: usize) -> io::Result<CacheLockMetrics> { + let shard = &self.shards[shard_index]; + let (mut state, metrics) = Self::lock_state(shard)?; + for offset in offsets { + if self.shard_index(*offset) == shard_index { + state.loading.remove(offset); + } + } + shard.waiters.notify_all(); + Ok(metrics) + } + + pub(crate) fn lookup_or_reserve( + &self, + offset: u64, + length: usize, + ) -> io::Result<(SharedWindowCacheLookup, CacheLockMetrics)> { + let shard = self.shard(offset); + let (mut state, metrics) = Self::lock_state(shard)?; + if let Some(payload) = state.entries.get(&offset).cloned() { + if payload.len() == length { + state.recency.touch(offset); + return Ok((SharedWindowCacheLookup::Hit(payload), metrics)); + } + state.entries.remove(&offset); + state.recency.remove(offset); + state.retained_bytes = state.retained_bytes.saturating_sub(payload.capacity()); + } + if state.loading.contains(&offset) { + return Ok((SharedWindowCacheLookup::Loading, metrics)); + } + state.loading.insert(offset); + Ok((SharedWindowCacheLookup::Reserved, metrics)) + } + + pub(crate) fn publish( + &self, + offset: u64, + payload: Arc<Vec<u8>>, + ) -> io::Result<(usize, CacheLockMetrics)> { + let shard = self.shard(offset); + let (mut state, metrics) = Self::lock_state(shard)?; + state.loading.remove(&offset); + if let Some(previous) = state.entries.insert(offset, Arc::clone(&payload)) { + state.retained_bytes = state.retained_bytes.saturating_sub(previous.capacity()); + } + state.retained_bytes = state.retained_bytes.saturating_add(payload.capacity()); + state.recency.touch(offset); + let mut evictions = 0usize; + let capacity_bytes = shard.capacity_bytes.load(AtomicOrdering::Relaxed); + while state.retained_bytes > capacity_bytes { + let Some(oldest) = state.recency.pop_oldest() else { + break; + }; + if let Some(evicted) = state.entries.remove(&oldest) { + state.retained_bytes = state.retained_bytes.saturating_sub(evicted.capacity()); + evictions = evictions.saturating_add(1); + } + } + shard.waiters.notify_all(); + Ok((evictions, metrics)) + } + + pub(crate) fn cancel(&self, offsets: &[u64]) -> io::Result<CacheLockMetrics> { + let mut metrics = CacheLockMetrics::default(); + for shard_index in 0..self.shards.len() { + if offsets + .iter() + .any(|offset| self.shard_index(*offset) == shard_index) + { + Self::add_lock_metrics(&mut metrics, self.remove_loading(offsets, shard_index)?); + } + } + Ok(metrics) + } + + pub(crate) fn wait_for( + &self, + offset: u64, + length: usize, + ) -> io::Result<(Option<Arc<Vec<u8>>>, CacheLockMetrics)> { + let shard = self.shard(offset); + let (mut state, metrics) = Self::lock_state(shard)?; + while state.loading.contains(&offset) { + state = shard + .waiters + .wait(state) + .map_err(|_| invalid_data("DiskANN shared window cache state is poisoned"))?; + } + let payload = state + .entries + .get(&offset) + .filter(|payload| payload.len() == length) + .cloned(); + if payload.is_some() { + state.recency.touch(offset); + } + Ok((payload, metrics)) + } +} + +const ADJACENCY_PAGE_UNVALIDATED: u8 = 0; +const ADJACENCY_PAGE_VALIDATING: u8 = 1; +const ADJACENCY_PAGE_VALID: u8 = 2; +const ADJACENCY_PAGE_INVALID: u8 = 3; + +#[derive(Clone)] +struct CachedValidationError { + kind: io::ErrorKind, + message: String, +} + +impl CachedValidationError { + fn from_error(error: &io::Error) -> Self { + Self { + kind: error.kind(), + message: error.to_string(), + } + } + + fn to_error(&self) -> io::Error { + io::Error::new(self.kind, self.message.clone()) + } +} + +struct AdjacencyValidationCache { + states: Box<[AtomicU8]>, + errors: Mutex<HashMap<usize, CachedValidationError>>, + wait_lock: Mutex<()>, + waiters: Condvar, +} + +impl AdjacencyValidationCache { + fn new(page_count: usize) -> io::Result<Self> { + let mut states = Vec::new(); + states + .try_reserve_exact(page_count) + .map_err(|_| invalid_data("DiskANN adjacency validation cache allocation failed"))?; + states.extend((0..page_count).map(|_| AtomicU8::new(ADJACENCY_PAGE_UNVALIDATED))); + Ok(Self { + states: states.into_boxed_slice(), + errors: Mutex::new(HashMap::new()), + wait_lock: Mutex::new(()), + waiters: Condvar::new(), + }) + } + + fn get_or_validate( + &self, + page_index: usize, + validate: impl FnOnce() -> io::Result<()>, + ) -> io::Result<()> { + let state = self + .states + .get(page_index) + .ok_or_else(|| invalid_data("DiskANN adjacency page is out of range"))?; + loop { + match state.load(AtomicOrdering::Acquire) { + ADJACENCY_PAGE_UNVALIDATED => { + if state + .compare_exchange( + ADJACENCY_PAGE_UNVALIDATED, + ADJACENCY_PAGE_VALIDATING, + AtomicOrdering::AcqRel, + AtomicOrdering::Acquire, + ) + .is_err() + { + continue; + } + let mut claim = AdjacencyValidationClaim { + state, + wait_lock: &self.wait_lock, + waiters: &self.waiters, + published: false, + }; + return match validate() { + Ok(()) => { + claim.publish(ADJACENCY_PAGE_VALID); + Ok(()) + } + Err(error) => { + self.errors + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(page_index, CachedValidationError::from_error(&error)); + claim.publish(ADJACENCY_PAGE_INVALID); + Err(error) + } + }; + } + ADJACENCY_PAGE_VALIDATING => { + let guard = self + .wait_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let _guard = self + .waiters + .wait_while(guard, |_| { + state.load(AtomicOrdering::Acquire) == ADJACENCY_PAGE_VALIDATING + }) + .unwrap_or_else(std::sync::PoisonError::into_inner); + } + ADJACENCY_PAGE_VALID => return Ok(()), + ADJACENCY_PAGE_INVALID => { + let errors = self + .errors + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + return Err(errors.get(&page_index).map_or_else( + || invalid_data("DiskANN adjacency page failed validation"), + CachedValidationError::to_error, + )); + } + _ => return Err(invalid_data("invalid DiskANN adjacency validation state")), + } + } + } +} + +struct AdjacencyValidationClaim<'a> { + state: &'a AtomicU8, + wait_lock: &'a Mutex<()>, + waiters: &'a Condvar, + published: bool, +} + +impl AdjacencyValidationClaim<'_> { + fn publish(&mut self, state: u8) { + let _guard = self + .wait_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + self.state.store(state, AtomicOrdering::Release); + self.published = true; + self.waiters.notify_all(); + } +} + +impl Drop for AdjacencyValidationClaim<'_> { + fn drop(&mut self) { + if !self.published { + let _guard = self + .wait_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + self.state + .store(ADJACENCY_PAGE_UNVALIDATED, AtomicOrdering::Release); + self.waiters.notify_all(); + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum RowIdStorage { + Raw(Vec<i64>), + ForBitPacked { + base: i64, + bit_width: u8, + count: usize, + payload: Vec<u8>, + }, +} + +impl RowIdStorage { + #[cfg(test)] + fn encode(row_ids: Vec<i64>) -> io::Result<Self> { + Self::encode_from_fn(row_ids.len(), |node| row_ids[node]) + } + + fn encode_from_fn(count: usize, row_id_at: impl Fn(usize) -> i64) -> io::Result<Self> { + if count == 0 { + return Err(invalid_input("DiskANN row IDs must not be empty")); + } + let mut base = row_id_at(0); + let mut maximum = base; + for node in 1..count { + let row_id = row_id_at(node); + base = base.min(row_id); + maximum = maximum.max(row_id); + } + let span = u64::try_from(maximum as i128 - base as i128) + .expect("the difference between two i64 values fits in u64"); + let bit_width = if span == 0 { + 0 + } else { + (u64::BITS - span.leading_zeros()) as u8 + }; + if bit_width == u64::BITS as u8 { + let mut row_ids = Vec::new(); + row_ids + .try_reserve_exact(count) + .map_err(|_| invalid_input("DiskANN raw row-ID allocation failed"))?; + row_ids.extend((0..count).map(row_id_at)); + return Ok(Self::Raw(row_ids)); + } + + let payload_len = packed_row_id_payload_len(count, bit_width)?; + let mut payload = Vec::new(); + payload + .try_reserve_exact(payload_len) + .map_err(|_| invalid_input("DiskANN packed row-ID allocation failed"))?; + payload.resize(payload_len, 0); + for node in 0..count { + let row_id = row_id_at(node); + let delta = u64::try_from(row_id as i128 - base as i128) + .expect("the row ID is not below the selected base"); + pack_row_id_delta(&mut payload, node, bit_width, delta)?; + } + Ok(Self::ForBitPacked { + base, + bit_width, + count, + payload, + }) + } + + fn len(&self) -> usize { + match self { + Self::Raw(row_ids) => row_ids.len(), + Self::ForBitPacked { count, .. } => *count, + } + } + + #[cfg(test)] + fn bit_width(&self) -> u8 { + match self { + Self::Raw(_) => u64::BITS as u8, + Self::ForBitPacked { bit_width, .. } => *bit_width, + } + } + + fn get(&self, node: usize) -> Option<i64> { + match self { + Self::Raw(row_ids) => row_ids.get(node).copied(), + Self::ForBitPacked { + base, + bit_width, + count, + payload, + } => { + if node >= *count { + return None; + } + let delta = unpack_row_id_delta(payload, node, *bit_width)?; + i64::try_from(*base as i128 + delta as i128).ok() + } + } + } + + fn try_for_each( + &self, + mut visitor: impl FnMut(usize, i64) -> io::Result<()>, + ) -> io::Result<()> { + match self { + Self::Raw(row_ids) => { + for (node, &row_id) in row_ids.iter().enumerate() { + visitor(node, row_id)?; + } + } + Self::ForBitPacked { + base, + bit_width, + count, + payload, + } => { + let mut bit_offset = 0usize; + for node in 0..*count { + let delta = unpack_row_id_delta_at_bit_offset(payload, bit_offset, *bit_width) + .ok_or_else(|| { + invalid_data("DiskANN packed row-ID payload is truncated") + })?; + let row_id = i64::try_from(*base as i128 + delta as i128) + .map_err(|_| invalid_data("DiskANN packed row ID overflows i64"))?; + visitor(node, row_id)?; + bit_offset = bit_offset + .checked_add(*bit_width as usize) + .ok_or_else(|| invalid_data("DiskANN packed row-ID offset overflows"))?; + } + } + } + Ok(()) + } + + fn payload_len(&self) -> io::Result<usize> { + match self { + Self::Raw(row_ids) => row_ids + .len() + .checked_mul(size_of::<i64>()) + .ok_or_else(|| invalid_input("DiskANN raw row-ID length overflows usize")), + Self::ForBitPacked { payload, .. } => Ok(payload.len()), + } + } + + fn serialized_len(&self) -> io::Result<usize> { + ROW_ID_SECTION_HEADER_SIZE + .checked_add(self.payload_len()?) + .ok_or_else(|| invalid_input("DiskANN row-ID section length overflows usize")) + } + + fn section_header(&self) -> [u8; ROW_ID_SECTION_HEADER_SIZE] { + let mut bytes = [0u8; ROW_ID_SECTION_HEADER_SIZE]; + match self { + Self::Raw(row_ids) => { + put_u32(&mut bytes, 0, ROW_ID_ENCODING_RAW_I64); + put_u32(&mut bytes, 4, u64::BITS); + put_u64(&mut bytes, 8, row_ids.len() as u64); + } + Self::ForBitPacked { + base, + bit_width, + count, + .. + } => { + put_u32(&mut bytes, 0, ROW_ID_ENCODING_FOR_BITPACK); + put_u32(&mut bytes, 4, *bit_width as u32); + put_u64(&mut bytes, 8, *count as u64); + put_u64(&mut bytes, 16, *base as u64); + } + } + bytes + } +} + +fn packed_row_id_payload_len(count: usize, bit_width: u8) -> io::Result<usize> { + count + .checked_mul(bit_width as usize) + .and_then(|bits| bits.checked_add(7)) + .map(|bits| bits / 8) + .ok_or_else(|| invalid_input("DiskANN packed row-ID length overflows usize")) +} + +fn raw_row_id_section_len(count: usize) -> io::Result<usize> { + count + .checked_mul(size_of::<i64>()) + .and_then(|payload| payload.checked_add(ROW_ID_SECTION_HEADER_SIZE)) + .ok_or_else(|| invalid_input("DiskANN raw row-ID section length overflows usize")) +} + +fn pack_row_id_delta(payload: &mut [u8], node: usize, bit_width: u8, delta: u64) -> io::Result<()> { + if bit_width == 0 { + return Ok(()); + } + let bit_offset = node + .checked_mul(bit_width as usize) + .ok_or_else(|| invalid_input("DiskANN packed row-ID offset overflows usize"))?; + let byte_offset = bit_offset / 8; + let shift = bit_offset % 8; + let byte_count = (shift + bit_width as usize).div_ceil(8); + let encoded = (delta as u128) << shift; + let destination = payload + .get_mut(byte_offset..byte_offset + byte_count) + .ok_or_else(|| invalid_input("DiskANN packed row-ID payload is truncated"))?; + for (index, byte) in destination.iter_mut().enumerate() { + *byte |= (encoded >> (index * 8)) as u8; + } + Ok(()) +} + +fn unpack_row_id_delta(payload: &[u8], node: usize, bit_width: u8) -> Option<u64> { + let bit_offset = node.checked_mul(bit_width as usize)?; + unpack_row_id_delta_at_bit_offset(payload, bit_offset, bit_width) +} + +fn unpack_row_id_delta_at_bit_offset( + payload: &[u8], + bit_offset: usize, + bit_width: u8, +) -> Option<u64> { + if bit_width == 0 { + return Some(0); + } + let byte_offset = bit_offset / 8; + let shift = bit_offset % 8; + let byte_count = (shift + bit_width as usize).div_ceil(8); + let source = payload.get(byte_offset..byte_offset + byte_count)?; + let encoded = source + .iter() + .enumerate() + .fold(0u128, |value, (index, &byte)| { + value | ((byte as u128) << (index * 8)) + }); + let mask = (1u128 << bit_width) - 1; + Some(((encoded >> shift) & mask) as u64) +} + +#[derive(Debug, Clone, Copy)] +struct RowIdSectionHeader { + encoding: u32, + bit_width: u8, + count: usize, + base: i64, +} + +fn decode_row_id_section_header( + bytes: &[u8], + section_len: usize, + expected_count: usize, +) -> io::Result<RowIdSectionHeader> { + if bytes.len() < ROW_ID_SECTION_HEADER_SIZE { + return Err(invalid_data("DiskANN row-ID section header is truncated")); + } + if bytes[24..ROW_ID_SECTION_HEADER_SIZE] + .iter() + .any(|&byte| byte != 0) + { + return Err(invalid_data( + "DiskANN row-ID section reserved bytes must be zero", + )); + } + let count = usize::try_from(get_u64(bytes, 8)) + .map_err(|_| invalid_data("DiskANN row-ID count exceeds usize"))?; + if count != expected_count { + return Err(invalid_data("DiskANN row-ID count does not match header")); + } + let encoding = get_u32(bytes, 0); + let width = get_u32(bytes, 4); + let base = get_u64(bytes, 16) as i64; + let (bit_width, payload_len) = match encoding { + ROW_ID_ENCODING_RAW_I64 => { + if width != u64::BITS || base != 0 { + return Err(invalid_data("invalid DiskANN raw row-ID metadata")); + } + ( + u64::BITS as u8, + count + .checked_mul(size_of::<i64>()) + .ok_or_else(|| invalid_data("DiskANN raw row-ID length overflows usize"))?, + ) + } + ROW_ID_ENCODING_FOR_BITPACK => { + let bit_width = u8::try_from(width) + .map_err(|_| invalid_data("invalid DiskANN FOR row-ID bit width"))?; + if bit_width >= u64::BITS as u8 { + return Err(invalid_data("invalid DiskANN FOR row-ID bit width")); + } + let payload_len = packed_row_id_payload_len(count, bit_width) + .map_err(|_| invalid_data("DiskANN packed row-ID length overflows usize"))?; + (bit_width, payload_len) + } + _ => return Err(invalid_data("unsupported DiskANN row-ID encoding")), + }; + if ROW_ID_SECTION_HEADER_SIZE.checked_add(payload_len) != Some(section_len) { + return Err(invalid_data("invalid DiskANN row-ID payload length")); + } + Ok(RowIdSectionHeader { + encoding, + bit_width, + count, + base, + }) +} + +fn validate_row_id_storage(storage: &RowIdStorage) -> io::Result<()> { + if let RowIdStorage::ForBitPacked { + bit_width, + count, + payload, + .. + } = storage + { + let used_bits = count + .checked_mul(*bit_width as usize) + .ok_or_else(|| invalid_data("DiskANN packed row-ID length overflows usize"))?; + let tail_bits = used_bits % 8; + if tail_bits != 0 + && payload + .last() + .is_some_and(|&byte| byte & !((1u8 << tail_bits) - 1) != 0) + { + return Err(invalid_data("DiskANN packed row-ID tail bits must be zero")); + } + } + storage.try_for_each(|_, _| Ok(())) +} + +#[cfg(test)] +fn encode_row_id_section(storage: &RowIdStorage) -> io::Result<Vec<u8>> { + let section_len = storage.serialized_len()?; + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(section_len) + .map_err(|_| invalid_input("DiskANN row-ID section allocation failed"))?; + bytes.extend_from_slice(&storage.section_header()); + match storage { + RowIdStorage::Raw(row_ids) => row_ids + .iter() + .for_each(|row_id| bytes.extend_from_slice(&row_id.to_le_bytes())), + RowIdStorage::ForBitPacked { payload, .. } => bytes.extend_from_slice(payload), + } + Ok(bytes) +} + +#[cfg(test)] +fn decode_row_id_section(bytes: &[u8], expected_count: usize) -> io::Result<RowIdStorage> { + let header = decode_row_id_section_header(bytes, bytes.len(), expected_count)?; + let payload = &bytes[ROW_ID_SECTION_HEADER_SIZE..]; + let storage = match header.encoding { + ROW_ID_ENCODING_RAW_I64 => { + let mut row_ids = Vec::new(); + row_ids + .try_reserve_exact(header.count) + .map_err(|_| invalid_data("DiskANN raw row-ID allocation failed"))?; + row_ids.extend(payload.chunks_exact(8).map(|value| { + i64::from_le_bytes(value.try_into().expect("validated eight-byte row ID")) + })); + RowIdStorage::Raw(row_ids) + } + ROW_ID_ENCODING_FOR_BITPACK => { + let mut packed = Vec::new(); + packed + .try_reserve_exact(payload.len()) + .map_err(|_| invalid_data("DiskANN packed row-ID allocation failed"))?; + packed.extend_from_slice(payload); + RowIdStorage::ForBitPacked { + base: header.base, + bit_width: header.bit_width, + count: header.count, + payload: packed, + } + } + _ => unreachable!("row-ID encoding was validated"), + }; + validate_row_id_storage(&storage)?; + Ok(storage) +} + +#[derive(Default)] +enum RowIdOrderState { + #[default] + NotLoaded, + Loaded(Arc<[u32]>), + UnavailableByBudget, +} + +fn classify_read_tier(random_read_latency: Duration) -> DeploymentProfile { + if random_read_latency < AUTO_PROFILE_MEMORY_LATENCY_THRESHOLD { + DeploymentProfile::Memory + } else if random_read_latency < AUTO_PROFILE_LOCAL_LATENCY_THRESHOLD { + DeploymentProfile::LocalStorage + } else if random_read_latency < AUTO_PROFILE_REMOTE_LATENCY_THRESHOLD { + DeploymentProfile::RemoteStorage + } else { + DeploymentProfile::ObjectStore + } +} + +impl<R: SeekRead> DiskAnnIndexReader<R> { + pub fn open(reader: R) -> io::Result<Self> { + Self::open_with_options(reader, VectorIndexReaderOptions::default()) + } + + pub fn open_with_options(mut reader: R, options: VectorIndexReaderOptions) -> io::Result<Self> { + let read_capabilities = reader.read_capabilities(); + let mut bytes = [0u8; DISKANN_HEADER_SIZE]; + let header_read_started = Instant::now(); + reader + .pread(&mut [ReadRequest::new(0, &mut bytes)]) + .map_err(|error| map_read_error(error, "header"))?; + let measured_header_read_latency = header_read_started.elapsed(); + let header = DiskAnnHeader::decode(&bytes)?; + let random_read_latency = if read_capabilities.estimated_random_read_latency_nanos > 0 { + Duration::from_nanos(read_capabilities.estimated_random_read_latency_nanos) + } else { + measured_header_read_latency.max(Duration::from_nanos(1)) + }; + let effective_read_tier = classify_read_tier(random_read_latency); + let options = options.resolve_cache_budgets( + effective_read_tier, + resident_steady_bytes(&header)?, + usize::try_from(header.sections.adjacency.length).unwrap_or(usize::MAX), + usize::try_from(header.sections.vectors.length).unwrap_or(usize::MAX), + ); + Ok(Self { + reader, + header, + resident: None, + options, + read_capabilities, + effective_read_tier, + random_read_latency, + hot_adjacency: Arc::from([]), + row_id_order: Arc::new(Mutex::new(RowIdOrderState::NotLoaded)), + query_scratch: Box::default(), + last_search_stats: DiskAnnSearchStats::default(), + batch_workers: Vec::new(), + calibrated_l_search: None, + }) + } + + pub fn ensure_resident(&mut self) -> io::Result<()> { + if self.resident.is_some() { + return Ok(()); + } + let peak_bytes = resident_peak_bytes(&self.header)?; + if peak_bytes > self.options.max_resident_bytes { + return Err(invalid_data(format!( + "DiskANN resident warmup requires {} bytes, exceeding reader budget {}", + peak_bytes, self.options.max_resident_bytes + ))); + } + + let (mut pq, row_ids, pq_codes, adjacency_index) = + read_resident_sections(&mut self.reader, &self.header)?; + pq.try_rebuild_norms_cache() + .map_err(|_| invalid_data("DiskANN PQ norms allocation failed"))?; + validate_pq_code_padding(&self.header, &pq_codes)?; + let adjacency_validation = + AdjacencyValidationCache::new(adjacency_page_count(&self.header)?)?; + self.resident = Some(Arc::new(DiskAnnResidentData { + pq, + row_ids, + pq_codes, + adjacency_index, + adjacency_validation, + adjacency_cache: SharedWindowCache::new(self.options.adjacency_cache_bytes), + raw_vector_cache: SharedWindowCache::new(self.options.raw_vector_cache_bytes), + })); + Ok(()) + } + + pub fn pq(&self) -> io::Result<&ProductQuantizer> { + Ok(&self.resident()?.pq) + } + + pub fn row_id(&self, node: usize) -> io::Result<i64> { + self.resident()? + .row_ids + .get(node) + .ok_or_else(|| invalid_data("DiskANN row-ID node is out of range")) + } + + pub fn row_id_count(&self) -> io::Result<usize> { + Ok(self.resident()?.row_ids.len()) + } + + pub(crate) fn try_for_each_row_id( + &self, + visitor: impl FnMut(usize, i64) -> io::Result<()>, + ) -> io::Result<()> { + self.resident()?.row_ids.try_for_each(visitor) + } + + pub fn pq_codes(&self) -> io::Result<&[u8]> { + Ok(&self.resident()?.pq_codes) + } + + pub(crate) fn adjacency_locator(&self, node: usize) -> io::Result<AdjacencyLocator> { + self.resident()? + .adjacency_index + .locator(node) + .ok_or_else(|| invalid_data("DiskANN adjacency index is truncated")) + } + + pub(crate) fn adjacency_cache(&self) -> io::Result<&SharedWindowCache> { + Ok(&self.resident()?.adjacency_cache) + } + + pub(crate) fn raw_vector_cache(&self) -> io::Result<&SharedWindowCache> { + Ok(&self.resident()?.raw_vector_cache) + } + + fn resize_shared_cache_budgets(&self, total_bytes: usize) -> io::Result<()> { + let desired_adjacency = self.options.adjacency_cache_bytes; + let desired_raw = self.options.raw_vector_cache_bytes; + let desired_total = desired_adjacency.saturating_add(desired_raw); + let total_bytes = total_bytes.min(desired_total); + let adjacency_bytes = if desired_total == 0 { + 0 + } else { + usize::try_from( + (total_bytes as u128 * desired_adjacency as u128) / desired_total as u128, + ) + .unwrap_or(total_bytes) + }; + let raw_bytes = total_bytes.saturating_sub(adjacency_bytes); + let resident = self.resident()?; + resident + .adjacency_cache + .set_total_capacity(adjacency_bytes)?; + resident.raw_vector_cache.set_total_capacity(raw_bytes) + } + + fn loaded_row_id_order_bytes(&self) -> io::Result<usize> { + let state = self + .row_id_order + .lock() + .map_err(|_| invalid_data("DiskANN row-ID lookup state is poisoned"))?; + Ok(match &*state { + RowIdOrderState::Loaded(order) => order + .len() + .checked_mul(size_of::<u32>()) + .ok_or_else(|| invalid_data("DiskANN row-ID order size overflows usize"))?, + RowIdOrderState::NotLoaded | RowIdOrderState::UnavailableByBudget => 0, + }) + } + + pub(crate) fn ensure_row_id_order(&mut self) -> io::Result<Option<Arc<[u32]>>> { + self.ensure_resident()?; + let mut state = self + .row_id_order + .lock() + .map_err(|_| invalid_data("DiskANN row-ID lookup state is poisoned"))?; + match &*state { + RowIdOrderState::Loaded(order) => return Ok(Some(order.clone())), + RowIdOrderState::UnavailableByBudget => return Ok(None), + RowIdOrderState::NotLoaded => {} + } + let peak_bytes = row_id_order_peak_bytes(&self.header, self.hot_adjacency.len())?; + if peak_bytes > self.options.max_resident_bytes { + *state = RowIdOrderState::UnavailableByBudget; + return Ok(None); + } + // Reserve the decode peak before allocating the immutable lookup. + // Cache hits remain lock-free with respect to this budget operation; + // only cache publication reads the adjusted capacity. + self.resize_shared_cache_budgets(self.options.max_resident_bytes - peak_bytes)?; + let order = read_u32_section( + &mut self.reader, + self.header.sections.row_id_order, + "row-ID order", + )?; + validate_row_id_order(&self.resident()?.row_ids, &order)?; + let order: Arc<[u32]> = Arc::from(order); + *state = RowIdOrderState::Loaded(order.clone()); + let steady_with_order = resident_steady_bytes(&self.header)? + .checked_add(self.hot_adjacency.len()) + .and_then(|bytes| { + order + .len() + .checked_mul(size_of::<u32>()) + .and_then(|order_bytes| bytes.checked_add(order_bytes)) + }) + .ok_or_else(|| invalid_data("DiskANN filtered resident size overflows usize"))?; + self.resize_shared_cache_budgets( + self.options + .max_resident_bytes + .saturating_sub(steady_with_order), + )?; + Ok(Some(order)) + } + + pub fn optimize_for_search(&mut self) -> io::Result<()> { + self.ensure_resident()?; + if self.options.adjacency_preload_bytes == 0 || !self.hot_adjacency.is_empty() { + return Ok(()); + } + let adjacency = self.header.sections.adjacency; + let requested_len = self + .options + .adjacency_preload_bytes + .min(adjacency.length as usize); + let requested_len = requested_len + .div_ceil(DISKANN_ADJACENCY_PRELOAD_ALIGNMENT) + .saturating_mul(DISKANN_ADJACENCY_PRELOAD_ALIGNMENT) + .min(adjacency.length as usize); + let row_id_order_bytes = self.loaded_row_id_order_bytes()?; + let available_bytes = self + .options + .max_resident_bytes + .saturating_sub(resident_steady_bytes(&self.header)?) + .saturating_sub(row_id_order_bytes); + let mut preload_len = requested_len.min(available_bytes); + if preload_len < adjacency.length as usize { + preload_len = preload_len / DISKANN_ADJACENCY_PRELOAD_ALIGNMENT + * DISKANN_ADJACENCY_PRELOAD_ALIGNMENT; + } + if preload_len == 0 { + self.resize_shared_cache_budgets(available_bytes)?; + return Ok(()); + } + self.resize_shared_cache_budgets(available_bytes.saturating_sub(preload_len))?; + let mut payload = vec![0u8; preload_len]; + self.reader + .pread(&mut [ReadRequest::new(adjacency.offset, &mut payload)]) + .map_err(|error| map_read_error(error, "adjacency preload"))?; + let payload: Arc<[u8]> = Arc::from(payload); + let resident = self + .resident + .as_ref() + .expect("resident sections were loaded before adjacency preload"); + payload + .par_chunks_exact(DISKANN_PAGE_SIZE as usize) + .enumerate() + .try_for_each(|(page_index, page)| { + resident + .adjacency_validation + .get_or_validate(page_index, || { + validate_adjacency_page_payload( + &self.header, + &resident.adjacency_index, + page_index, + page, + ) + }) + })?; + self.hot_adjacency = payload; + Ok(()) + } + + #[cfg(test)] + pub(crate) fn effective_read_tier(&self) -> DeploymentProfile { + self.effective_read_tier + } + + #[cfg(test)] + pub(crate) fn random_read_latency(&self) -> Duration { + self.random_read_latency + } + + pub(crate) const fn options(&self) -> ResolvedVectorIndexReaderOptions { + self.options + } + + pub fn read_capabilities(&self) -> SeekReadCapabilities { + self.read_capabilities + } + + pub fn vector_read_plan(&self) -> VectorIndexReadPlan { + let plan = self.read_plan(); + VectorIndexReadPlan { + random_read_latency_nanos: u64::try_from(self.random_read_latency.as_nanos()) + .unwrap_or(u64::MAX), + window_bytes: plan.window_bytes, + max_ranges_per_read: self.read_capabilities.max_ranges_per_pread, + graph_beam_width: plan.graph_beam_width, + filtered_graph_beam_width: plan.filtered_graph_beam_width, + adjacency_preload_bytes: self.options.adjacency_preload_bytes, + adjacency_cache_bytes: self.options.adjacency_cache_bytes, Review Comment: `vector_read_plan` reports the originally resolved cache budgets, not the effective capacities after lazy state consumes memory. `ensure_row_id_order` can shrink both shared caches through `resize_shared_cache_budgets`, but `self.options` remains unchanged. I extended `diskann_row_id_order_reserves_budget_from_shared_caches`: after loading the row-ID order, the public plan still reported 32,768 cache bytes while the actual combined cache capacity was 10,060 bytes. Since this API is documented as the concrete read plan, could it report the current hot-adjacency size and shared-cache capacities instead of the initial desired values? ########## include/paimon_vindex.hpp: ########## @@ -355,35 +407,74 @@ class Reader { } ~Reader() { + std::lock_guard<std::mutex> lock(native_handle_mutex_); if (handle_) paimon_vindex_reader_free(handle_); } Metadata metadata() const { + std::lock_guard<std::mutex> lock(native_handle_mutex_); PaimonVindexMetadata raw; - check(paimon_vindex_reader_metadata(handle_, &raw)); + check(paimon_vindex_reader_metadata(require_open(), &raw)); Metadata result; result.index_type = raw.index_type; result.dimension = raw.dimension; result.nlist = raw.nlist; result.metric = raw.metric; result.total_vectors = raw.total_vectors; result.pq_m = raw.pq_m; - result.hnsw_m = raw.hnsw_m; - result.hnsw_ef_construction = raw.hnsw_ef_construction; - result.hnsw_max_level = raw.hnsw_max_level; + result.pq_bits = raw.pq_bits; + result.rq_bits = raw.rq_bits; + result.diskann_max_degree = raw.diskann_max_degree; + result.diskann_build_search_list_size = raw.diskann_build_search_list_size; + result.diskann_alpha = raw.diskann_alpha; return result; } void optimize_for_search() { - check(paimon_vindex_reader_optimize_for_search(handle_)); + std::lock_guard<std::mutex> lock(native_handle_mutex_); + check(paimon_vindex_reader_optimize_for_search(require_open())); + } + + void warmup_queries( + const float* queries, size_t query_count, size_t l_search = 0) { + std::lock_guard<std::mutex> lock(native_handle_mutex_); Review Comment: Serializing the C++ handle fixes the cross-thread race, but a plain `std::mutex` turns callback reentry into a permanent deadlock. I reproduced this with an IVF-FLAT `InputFile::read_ranges_fn` that calls `reader.metadata()` during `reader.search()`; the callback blocks reacquiring `native_handle_mutex_`, and the process timed out after five seconds. Java and Python now explicitly track the owning thread and reject this pattern. Could the C++ wrapper do the same rather than blocking same-thread reentry? ########## core/src/ivfsq_io.rs: ########## @@ -0,0 +1,1312 @@ +// 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. + +//! Stable v1 storage and positional-I/O search for IVF-SQ8. + +use crate::distance::{preprocess_vectors, MetricType}; +use crate::index_io_util::{ + bounded_ivf_payload_batch_end, bounded_ivf_stream_chunk_rows, bytes_to_f32_vec, + checked_list_bytes, checked_list_offset, checked_section_size, decode_delta_varint_ids, + decode_roaring_filter, encode_delta_varint_ids, ivf_payload_is_oversized, + pread_batched_payloads, read_delta_varint_ids_at, u64_to_i64, usize_to_i32, usize_to_i64, + validate_positive_i32, validate_reserved_zero, validate_search_inputs, write_f32_slice, + write_i32_le, write_i64_le, write_u32_le, +}; +use crate::io::{ReadRequest, SeekRead, SeekWrite}; +use crate::ivfpq::RowIdFilter; +use crate::ivfsq::IVFSQIndex; +use crate::kmeans; +use crate::sq::ScalarQuantizer; +use crate::topk::TopKHeap; +use rayon::prelude::*; +use std::io; +use std::mem::size_of; + +pub const IVF_SQ_MAGIC: u32 = 0x49565351; // "IVSQ" +pub const IVF_SQ_VERSION: u32 = 1; +pub const IVF_SQ_HEADER_SIZE: usize = 64; +pub const IVF_SQ_BITS: u32 = 8; +const FLAG_DELTA_IDS: u32 = 1 << 0; +const FLAG_BLOCKED_CODES: u32 = 1 << 1; +const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS | FLAG_BLOCKED_CODES; +const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS; +pub(crate) const IVF_SQ_SCAN_BLOCK_SIZE: usize = 32; + +pub fn write_ivfsq_index(index: &IVFSQIndex, out: &mut dyn SeekWrite) -> io::Result<()> { + validate_index_shape(index)?; + let total_vectors = index.ids.iter().try_fold(0i64, |sum, ids| { + let count = usize_to_i64(ids.len(), "total vector count")?; + sum.checked_add(count).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "total vector count exceeds i64 length limit", + ) + }) + })?; + let sorted_lists = (0..index.nlist) + .map(|list_id| build_sorted_sq_list_metadata(index, list_id)) + .collect::<io::Result<Vec<_>>>()?; + + write_u32_le(out, IVF_SQ_MAGIC)?; + write_u32_le(out, IVF_SQ_VERSION)?; + write_i32_le(out, usize_to_i32(index.d, "dimension")?)?; + write_i32_le(out, usize_to_i32(index.nlist, "nlist")?)?; + write_u32_le(out, index.metric as u32)?; + write_i64_le(out, total_vectors)?; + write_u32_le(out, IVF_SQ_BITS)?; + write_u32_le(out, REQUIRED_FLAGS)?; + let (sq_min, sq_max) = sq_global_bounds(&index.sq.mins, &index.sq.maxs); + out.write_all(&sq_min.to_le_bytes())?; + out.write_all(&sq_max.to_le_bytes())?; + out.write_all(&[0u8; 20])?; + + write_f32_slice(out, &index.sq.mins)?; + write_f32_slice(out, &index.sq.maxs)?; + for sq in &index.list_sqs { + write_f32_slice(out, &sq.mins)?; + write_f32_slice(out, &sq.maxs)?; + } + write_f32_slice(out, &index.quantizer_centroids)?; + + let offset_table_size = index.nlist.checked_mul(16).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "IVF-SQ offset table size overflow", + ) + })?; + let data_start = out + .pos() + .checked_add(offset_table_size as u64) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "IVF-SQ data offset overflow") + })?; + let mut list_offsets = vec![0i64; index.nlist]; + let mut list_counts = vec![0i32; index.nlist]; + let mut list_id_bytes_lens = vec![0i32; index.nlist]; + let mut current_offset = data_start; + + for (list_id, list) in sorted_lists.iter().enumerate() { + list_offsets[list_id] = u64_to_i64(current_offset, "list offset")?; + list_counts[list_id] = usize_to_i32(list.order.len(), "list count")?; + if !list.order.is_empty() { + list_id_bytes_lens[list_id] = usize_to_i32(list.id_bytes.len(), "delta ID section")?; + current_offset = current_offset + .checked_add(list_payload_len( + list.order.len(), + index.code_size(), + list.id_bytes.len(), + )? as u64) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "IVF-SQ list offset overflow") + })?; + } + } + + for list_id in 0..index.nlist { + write_i64_le(out, list_offsets[list_id])?; + write_i32_le(out, list_counts[list_id])?; + write_i32_le(out, list_id_bytes_lens[list_id])?; + } + for (list_id, list) in sorted_lists.iter().enumerate() { + if list.order.is_empty() { + continue; + } + let codes = block_sorted_sq_codes( + &index.codes[list_id], + &list.order, + index.d, + IVF_SQ_SCAN_BLOCK_SIZE, + ); + out.write_all(&codes)?; + write_i64_le(out, list.base_id)?; + write_i32_le(out, usize_to_i32(list.id_bytes.len(), "delta ID section")?)?; + out.write_all(&list.id_bytes)?; + } + Ok(()) +} + +pub struct IVFSQIndexReader<R: SeekRead> { + reader: R, + pub d: usize, + pub nlist: usize, + pub metric: MetricType, + pub total_vectors: i64, + pub sq: ScalarQuantizer, + pub list_sqs: Vec<ScalarQuantizer>, + pub quantizer_centroids: Vec<f32>, + pub list_offsets: Vec<i64>, + pub list_counts: Vec<i32>, + pub list_id_bytes_lens: Vec<i32>, + loaded: bool, +} + +impl<R: SeekRead> IVFSQIndexReader<R> { + pub fn open(mut reader: R) -> io::Result<Self> { + let mut header = [0u8; IVF_SQ_HEADER_SIZE]; + reader.pread(&mut [ReadRequest::new(0, &mut header)])?; + Self::open_with_header(reader, header) + } + + pub(crate) fn open_with_header( + mut reader: R, + header: [u8; IVF_SQ_HEADER_SIZE], + ) -> io::Result<Self> { + let read_u32 = + |offset: usize| u32::from_le_bytes(header[offset..offset + 4].try_into().unwrap()); + let read_i32 = + |offset: usize| i32::from_le_bytes(header[offset..offset + 4].try_into().unwrap()); + let read_i64 = + |offset: usize| i64::from_le_bytes(header[offset..offset + 8].try_into().unwrap()); + let read_f32 = + |offset: usize| f32::from_le_bytes(header[offset..offset + 4].try_into().unwrap()); + + let magic = read_u32(0); + if magic != IVF_SQ_MAGIC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Invalid IVF-SQ magic: 0x{magic:08X}"), + )); + } + let version = read_u32(4); + if version != IVF_SQ_VERSION { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported IVF-SQ version: {version}"), + )); + } + let d = validate_positive_i32(read_i32(8), "d")? as usize; + let nlist = validate_positive_i32(read_i32(12), "nlist")? as usize; + let metric_code = read_u32(16); + let metric = MetricType::from_code(metric_code).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("Unknown metric type: {metric_code}"), + ) + })?; + let total_vectors = read_i64(20); + if total_vectors < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ total vector count must be non-negative", + )); + } + let bits = read_u32(28); + if bits != IVF_SQ_BITS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported IVF-SQ bit width: {bits}"), + )); + } + let flags = read_u32(32); + let sq_min_summary = read_f32(36); + let sq_max_summary = read_f32(40); + validate_reserved_zero(&header[44..64], "IVF-SQ")?; + let unknown_flags = flags & !SUPPORTED_FLAGS; + if unknown_flags != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Unsupported IVF-SQ flags: 0x{unknown_flags:08X}"), + )); + } + if flags & REQUIRED_FLAGS != REQUIRED_FLAGS { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ v1 requires delta-varint IDs and 32-row blocked codes", + )); + } + + let bounds_values = checked_section_size(nlist + 1, d)? + .checked_mul(2) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ bounds size overflow") + })?; + let bounds_bytes = bounds_values.checked_mul(4).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ bounds byte length overflow", + ) + })?; + let centroid_values = checked_section_size(nlist, d)?; + let centroid_bytes = centroid_values.checked_mul(4).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ centroid byte length overflow", + ) + })?; + let offset_table_bytes = nlist.checked_mul(16).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ offset table byte length overflow", + ) + })?; + let metadata_bytes = bounds_bytes + .checked_add(centroid_bytes) + .and_then(|size| size.checked_add(offset_table_bytes)) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ metadata size overflow") + })?; + let mut metadata = vec![0u8; metadata_bytes]; + reader.pread(&mut [ReadRequest::new(IVF_SQ_HEADER_SIZE as u64, &mut metadata)])?; + let (sq, list_sqs, mut position) = { + let mut position = 0usize; + let mut next_f32_section = |count: usize| -> io::Result<Vec<f32>> { + let byte_len = count.checked_mul(4).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ f32 size overflow") + })?; + let end = position.checked_add(byte_len).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ metadata offset overflow", + ) + })?; + let values = bytes_to_f32_vec(&metadata[position..end])?; + position = end; + Ok(values) + }; + + let mins = next_f32_section(d)?; + let maxs = next_f32_section(d)?; + validate_sq_bounds(d, &mins, &maxs)?; + let (sq_min, sq_max) = sq_global_bounds(&mins, &maxs); + if sq_min.to_bits() != sq_min_summary.to_bits() + || sq_max.to_bits() != sq_max_summary.to_bits() + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ bounds summary does not match global SQ bounds", + )); + } + let sq = ScalarQuantizer::with_dimension_bounds(d, mins, maxs); + let mut list_sqs = Vec::with_capacity(nlist); + for _ in 0..nlist { + let mins = next_f32_section(d)?; + let maxs = next_f32_section(d)?; + validate_sq_bounds(d, &mins, &maxs)?; + list_sqs.push(ScalarQuantizer::with_dimension_bounds(d, mins, maxs)); + } + (sq, list_sqs, position) + }; + + let quantizer_centroids = bytes_to_f32_vec(&metadata[position..position + centroid_bytes])?; + position += centroid_bytes; + let offset_table = &metadata[position..]; + let mut list_offsets = vec![0; nlist]; + let mut list_counts = vec![0; nlist]; + let mut list_id_bytes_lens = vec![0; nlist]; + let mut actual_total = 0i64; + for (list_id, entry) in offset_table.chunks_exact(16).enumerate() { + list_offsets[list_id] = i64::from_le_bytes(entry[0..8].try_into().unwrap()); + let count = i32::from_le_bytes(entry[8..12].try_into().unwrap()); + let id_bytes_len = i32::from_le_bytes(entry[12..16].try_into().unwrap()); + if count < 0 || id_bytes_len < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("negative IVF-SQ list metadata at list {list_id}"), + )); + } + if count > 0 && id_bytes_len == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("missing delta ID bytes for non-empty IVF-SQ list {list_id}"), + )); + } + actual_total = actual_total.checked_add(count as i64).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ vector count overflow") + })?; + list_counts[list_id] = count; + list_id_bytes_lens[list_id] = id_bytes_len; + } + if actual_total != total_vectors { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "IVF-SQ header vector count {total_vectors} does not match list total {actual_total}" + ), + )); + } + + Ok(Self { + reader, + d, + nlist, + metric, + total_vectors, + sq, + list_sqs, + quantizer_centroids, + list_offsets, + list_counts, + list_id_bytes_lens, + loaded: true, + }) + } + + pub fn ensure_loaded(&mut self) -> io::Result<()> { + debug_assert!(self.loaded); + Ok(()) + } + + pub fn optimize_for_search(&mut self) -> io::Result<()> { + self.ensure_loaded() + } + + pub fn read_inverted_list(&mut self, list_id: usize) -> io::Result<(Vec<i64>, Vec<u8>)> { + let mut lists = self.read_inverted_lists(&[list_id])?; + let list = lists.pop().expect("one requested list has one result"); + Ok((list.ids, list.codes)) + } + + pub fn read_inverted_lists(&mut self, list_ids: &[usize]) -> io::Result<Vec<SqListData>> { + self.ensure_loaded()?; + let mut results = (0..list_ids.len()).map(|_| None).collect::<Vec<_>>(); + let mut metas = Vec::new(); + let mut payloads = Vec::new(); + for (input_index, &list_id) in list_ids.iter().enumerate() { + if list_id >= self.nlist { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("list_id {list_id} out of range (nlist={})", self.nlist), + )); + } + let count = self.list_counts[list_id] as usize; + if count == 0 { + results[input_index] = Some(SqListData { + list_id, + ids: Vec::new(), + codes: Vec::new(), + }); + continue; + } + let id_bytes_len = self.list_id_bytes_lens[list_id] as usize; + let payload_len = list_payload_len(count, self.d, id_bytes_len)?; + metas.push(BatchedListRead { + input_index, + list_id, + count, + id_bytes_len, + offset: checked_list_offset(self.list_offsets[list_id], list_id)?, + }); + payloads.push(vec![0u8; payload_len]); + } + + if !metas.is_empty() { + let offsets = metas.iter().map(|meta| meta.offset).collect::<Vec<_>>(); + pread_batched_payloads(&mut self.reader, &offsets, &mut payloads)?; + for (meta, payload) in metas.into_iter().zip(payloads) { + let (ids, codes) = + decode_list_payload(payload, meta.count, meta.id_bytes_len, self.d)?; + results[meta.input_index] = Some(SqListData { + list_id: meta.list_id, + ids, + codes, + }); + } + } + results + .into_iter() + .map(|result| { + result.ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "missing batched IVF-SQ list read result", + ) + }) + }) + .collect() + } + + fn batch_read_end(&self, list_ids: &[usize]) -> io::Result<usize> { + let payload_lengths = list_ids + .iter() + .map(|&list_id| self.list_payload_len(list_id)) + .collect::<io::Result<Vec<_>>>()?; + bounded_ivf_payload_batch_end( + &payload_lengths, + self.reader.read_capabilities().max_ranges_per_pread, + ) + } + + fn list_payload_len(&self, list_id: usize) -> io::Result<usize> { + if list_id >= self.nlist { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("list_id {list_id} out of range (nlist={})", self.nlist), + )); + } + let count = self.list_counts[list_id] as usize; + if count == 0 { + Ok(0) + } else { + list_payload_len(count, self.d, self.list_id_bytes_lens[list_id] as usize) + } + } + + fn for_each_streamed_list_chunk( + &mut self, + list_id: usize, + mut consume: impl FnMut(&[i64], &[u8]), + ) -> io::Result<()> { + self.ensure_loaded()?; + let count = self.list_counts[list_id] as usize; + let list_offset = checked_list_offset(self.list_offsets[list_id], list_id)?; + let code_bytes = checked_list_bytes(count, self.d)?; + let id_offset = list_offset.checked_add(code_bytes as u64).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ ID offset overflow") + })?; + let ids = read_delta_varint_ids_at( + &mut self.reader, + id_offset, + count, + self.list_id_bytes_lens[list_id] as usize, + "IVF-SQ", + )?; + let retained_id_bytes = ids.len().checked_mul(size_of::<i64>()).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "IVF-SQ decoded ID size overflow", + ) + })?; + let mut row_start = 0usize; + while row_start < count { + let chunk_rows = bounded_ivf_stream_chunk_rows( + count - row_start, + self.d, + retained_id_bytes, + IVF_SQ_SCAN_BLOCK_SIZE, + )?; + let chunk_bytes = chunk_rows.checked_mul(self.d).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ chunk size overflow") + })?; + let chunk_offset = list_offset + .checked_add(row_start.checked_mul(self.d).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ chunk offset overflow") + })? as u64) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "IVF-SQ chunk offset overflow") + })?; + let mut codes = vec![0u8; chunk_bytes]; + self.reader + .pread(&mut [ReadRequest::new(chunk_offset, &mut codes)])?; + let row_end = row_start + chunk_rows; + consume(&ids[row_start..row_end], &codes); + row_start = row_end; + } + Ok(()) + } + + pub fn search( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + ) -> io::Result<(Vec<i64>, Vec<f32>)> { + self.search_with_filter(query, k, nprobe, None) + } + + pub fn search_with_filter( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, + ) -> io::Result<(Vec<i64>, Vec<f32>)> { + self.ensure_loaded()?; + validate_search_inputs(query, 1, self.d, k, nprobe)?; + let query = preprocess_vectors(query, 1, self.d, self.metric); + let (probe_indices, _) = kmeans::find_topk( + &query, + &self.quantizer_centroids, + self.nlist, + self.d, + nprobe, + ); + let mut heap = TopKHeap::new(k); + let d = self.d; + let metric = self.metric; + let mut batch_start = 0usize; + while batch_start < probe_indices.len() { + let first_list = probe_indices[batch_start]; + if ivf_payload_is_oversized(self.list_payload_len(first_list)?) { + let centroid = + self.quantizer_centroids[first_list * d..(first_list + 1) * d].to_vec(); + let sq = self.list_sqs.get(first_list).unwrap_or(&self.sq).clone(); + let mut scratch = SqScanScratch::default(); + self.for_each_streamed_list_chunk(first_list, |ids, codes| { + scan_sq_rows( + &query, + ids, + codes, + ¢roid, + &sq, + metric, + filter, + &mut scratch, + &mut heap, + ); + })?; + batch_start += 1; + continue; + } + let count = self.batch_read_end(&probe_indices[batch_start..])?.max(1); + let batch_end = (batch_start + count).min(probe_indices.len()); + let lists = self.read_inverted_lists(&probe_indices[batch_start..batch_end])?; + let centroids = &self.quantizer_centroids; + let list_sqs = &self.list_sqs; + let global_sq = &self.sq; + let candidate_count = lists.iter().map(|list| list.ids.len()).sum::<usize>(); + if candidate_count >= PARALLEL_SQ_SCAN_MIN_CANDIDATES { + let per_list_results = lists + .par_iter() + .map_init(SqScanScratch::default, |scratch, list| { + let mut local_heap = TopKHeap::new(k); + let list_id = list.list_id; + scan_sq_list( + &query, + list, + ¢roids[list_id * d..(list_id + 1) * d], + list_sqs.get(list_id).unwrap_or(global_sq), + metric, + filter, + scratch, + &mut local_heap, + ); + local_heap.into_sorted() + }) + .collect::<Vec<_>>(); + for results in per_list_results { + for (distance, row_id) in results { + heap.push(distance, row_id); + } + } + } else { + let mut scratch = SqScanScratch::default(); + for list in &lists { + let list_id = list.list_id; + scan_sq_list( + &query, + list, + ¢roids[list_id * d..(list_id + 1) * d], + list_sqs.get(list_id).unwrap_or(global_sq), + metric, + filter, + &mut scratch, + &mut heap, + ); + } + } + batch_start = batch_end; + } + Ok(padded_results(heap, k)) + } + + pub fn search_with_roaring_filter( + &mut self, + query: &[f32], + k: usize, + nprobe: usize, + roaring_filter_bytes: &[u8], + ) -> io::Result<(Vec<i64>, Vec<f32>)> { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.search_with_filter(query, k, nprobe, Some(&filter)) + } +} + +pub fn search_batch_ivfsq_reader<R: SeekRead>( + reader: &mut IVFSQIndexReader<R>, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, +) -> io::Result<(Vec<i64>, Vec<f32>)> { + search_batch_ivfsq_reader_filter(reader, queries, nq, k, nprobe, None) +} + +pub fn search_batch_ivfsq_reader_filter<R: SeekRead>( + reader: &mut IVFSQIndexReader<R>, + queries: &[f32], + nq: usize, + k: usize, + nprobe: usize, + filter: Option<&dyn RowIdFilter>, +) -> io::Result<(Vec<i64>, Vec<f32>)> { + reader.ensure_loaded()?; + validate_search_inputs(queries, nq, reader.d, k, nprobe)?; + let processed = preprocess_vectors(queries, nq, reader.d, reader.metric); + let (all_probe_indices, _) = kmeans::find_topk_batch( + &processed, + nq, + &reader.quantizer_centroids, + reader.nlist, + reader.d, + nprobe, + ); + let mut seen = vec![false; reader.nlist]; + let mut unique_lists = Vec::new(); + for list_ids in &all_probe_indices { + for &list_id in list_ids { + if !seen[list_id] { + seen[list_id] = true; + unique_lists.push(list_id); + } + } + } + let mut list_to_queries = vec![Vec::new(); reader.nlist]; + for (query_index, list_ids) in all_probe_indices.iter().enumerate() { + for &list_id in list_ids { + list_to_queries[list_id].push(query_index); + } + } + let d = reader.d; + let metric = reader.metric; + let mut heaps = (0..nq).map(|_| TopKHeap::new(k)).collect::<Vec<_>>(); + let mut stream_scratches = (0..nq) Review Comment: Streaming the oversized list bounds the I/O payload, but batch search retains one chunk-sized distance array for every query. I added a temporary unit test that scanned a 65,536-row chunk with 64 `SqScanScratch` instances; after the loop they retained 16,777,216 bytes (`64 * 65,536 * sizeof(f32)`) rather than one reusable 256 KiB buffer. Production chunk sizes can be much larger, so a large query batch can still consume GiBs despite the 64 MiB payload bound. Could queries be processed in bounded groups, or use scratch that is released/reused instead of one retained buffer per query? The transposed 8-bit IVF-PQ oversized-list path has the same per-query distance-scratch shape. -- 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]
