leaves12138 commented on code in PR #62:
URL:
https://github.com/apache/paimon-vector-index/pull/62#discussion_r3651235713
##########
python/paimon_vindex/__init__.py:
##########
@@ -47,24 +62,44 @@ class VectorIndexMetadata:
metric: str
total_vectors: int
pq_m: Optional[int] = None
- hnsw_m: Optional[int] = None
- hnsw_ef_construction: Optional[int] = None
- hnsw_max_level: Optional[int] = None
+ pq_bits: Optional[int] = None
+ rq_bits: Optional[int] = None
+ diskann_max_degree: Optional[int] = None
+ diskann_build_search_list_size: Optional[int] = None
+ diskann_alpha: Optional[float] = None
@dataclass(frozen=True)
class SearchParams:
top_k: int
- nprobe: int
- ef_search: int = 0
- query_bits: int = 0
+ search_width: SearchWidth = SearchWidth.AUTO
+ width: int = 0
+
+ @classmethod
+ def automatic(cls, top_k: int):
+ return cls(top_k=top_k)
+
+ @classmethod
+ def ivf(cls, top_k: int, nprobe: int):
+ return cls(
+ top_k=top_k,
+ search_width=SearchWidth.IVF_NPROBE,
+ width=nprobe,
+ )
+
+ @classmethod
+ def diskann(cls, top_k: int, l_search: int):
+ return cls(
+ top_k=top_k,
+ search_width=SearchWidth.DISKANN_L_SEARCH,
+ width=l_search,
+ )
def to_ffi(self):
Review Comment:
Could we validate the search parameters before converting them to the C ABI?
`ctypes.c_size_t` silently wraps Python integers, so `SearchParams.diskann(5,
-1).to_ffi().width` becomes `usize::MAX` (and oversized positive values wrap
modulo `SIZE_MAX`). In an actual search this turns a caller error into an
effectively exhaustive DiskANN search; `SearchParams.ivf(5, -1)` has the same
behavior for `nprobe`. The Java/JNI path rejects negative widths. A
`__post_init__` check for positive `top_k` and algorithm-specific width, plus
the platform upper bound, would keep the Python API consistent and avoid
silently changing the requested search.
##########
core/src/diskann_io.rs:
##########
@@ -0,0 +1,5315 @@
+// 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::{
+ ReadPlan, ResolvedVectorIndexReaderOptions, StorageProfile,
VectorIndexReaderOptions,
+};
+use rayon::prelude::*;
+use std::collections::{HashMap, HashSet};
+use std::io;
+use std::ops::{Index, IndexMut};
+use std::sync::atomic::{AtomicU8, 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);
+const AUTO_PROFILE_PROBE_COUNT: usize = 3;
+const AUTO_PROFILE_PROBE_BYTES: usize = 4 * 1024;
+
+#[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,
+ automatic_cache_budgets: bool,
+ read_capabilities: SeekReadCapabilities,
+ effective_storage_profile: StorageProfile,
+ auto_storage_profile_probe_time: Option<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: usize,
+ 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: 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()
+ }
+
+ 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 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;
+ while state.retained_bytes > shard.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_auto_storage_profile(random_read_latency: Duration) ->
StorageProfile {
+ if random_read_latency < AUTO_PROFILE_MEMORY_LATENCY_THRESHOLD {
+ StorageProfile::Memory
+ } else if random_read_latency < AUTO_PROFILE_LOCAL_LATENCY_THRESHOLD {
+ StorageProfile::LocalStorage
+ } else if random_read_latency < AUTO_PROFILE_REMOTE_LATENCY_THRESHOLD {
+ StorageProfile::RemoteStorage
+ } else {
+ StorageProfile::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];
+ reader
+ .pread(&mut [ReadRequest::new(0, &mut bytes)])
+ .map_err(|error| map_read_error(error, "header"))?;
+ let header = DiskAnnHeader::decode(&bytes)?;
+ let effective_storage_profile = options.storage_profile;
+ let automatic_cache_budgets = options.uses_automatic_cache_budgets();
+ let options = options.resolve_cache_budgets(
+ 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,
+ automatic_cache_budgets,
+ read_capabilities,
+ effective_storage_profile,
+ auto_storage_profile_probe_time: None,
+ 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
+ )));
+ }
+
+ if self.effective_storage_profile == StorageProfile::Auto {
+ if let Some(
+ profile @ (StorageProfile::Memory
+ | StorageProfile::LocalStorage
+ | StorageProfile::RemoteStorage
+ | StorageProfile::ObjectStore),
+ ) = self.reader.preferred_storage_profile()
+ {
+ self.effective_storage_profile = profile;
+ }
+ }
+ if self.effective_storage_profile == StorageProfile::Auto {
+ let latency = self.measure_random_read_latency()?;
+ self.effective_storage_profile =
classify_auto_storage_profile(latency);
+ self.auto_storage_profile_probe_time = Some(latency);
+ }
+ if self.automatic_cache_budgets {
+ self.options = VectorIndexReaderOptions::new(
+ self.effective_storage_profile,
+ self.options.max_resident_bytes,
+ )
+ .resolve_cache_budgets(
+ resident_steady_bytes(&self.header)?,
+
usize::try_from(self.header.sections.adjacency.length).unwrap_or(usize::MAX),
+
usize::try_from(self.header.sections.vectors.length).unwrap_or(usize::MAX),
+ );
+ }
+
+ 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)
+ }
+
+ 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())?;
Review Comment:
The lazy row-ID order is not reserved from the reader memory budget.
`resolve_cache_budgets` can assign all bytes after the steady resident state to
adjacency preload/cache and raw-vector cache, but this check counts only the
steady state, current hot adjacency, the order, and its decode scratch. Once
the order is retained, both shared caches can still grow to their original
limits; if a filtered query loads the order before `optimize_for_search`, the
later adjacency preload also ignores it. I reproduced this with an automatic
ObjectStore budget: `resident + preload + both cache capacities` consumed the
configured budget, and retaining the `4 * vector_count` order pushed the
configured total above it. Could we reserve the order by reducing preload/cache
limits when it is loaded, or include current cache usage and refuse/evict
accordingly?
##########
core/src/index.rs:
##########
@@ -671,36 +1248,118 @@ impl<R: SeekRead> VectorIndexReader<R> {
self.metadata().total_vectors
}
+ pub fn diskann_search_stats(&self) -> Option<DiskAnnSearchStats> {
+ match self {
+ Self::DiskAnn(reader) => Some(reader.last_search_stats()),
+ _ => None,
+ }
+ }
+
+ pub fn effective_storage_profile(&self) -> Option<StorageProfile> {
+ match self {
+ Self::DiskAnn(reader) => Some(reader.effective_storage_profile()),
+ _ => None,
+ }
+ }
+
pub fn optimize_for_search(&mut self) -> io::Result<()> {
match self {
Self::IvfFlat(reader) => reader.ensure_loaded(),
+ Self::IvfSq(reader) => reader.optimize_for_search(),
Self::IvfPq(reader) => reader.optimize_for_search(),
Self::IvfRq(reader) => reader.ensure_loaded(),
- Self::IvfHnswFlat(reader) => reader.ensure_loaded(),
- // IVF_HNSW_SQ warms SQ scan/fallback structures used by filtered
- // searches; normal unfiltered search primarily uses the HNSW
graph.
- Self::IvfHnswSq(reader) => reader.optimize_for_search(),
+ Self::DiskAnn(reader) => reader.optimize_for_search(),
}
}
- pub fn search(
+ /// Warm query-dependent caches with representative queries. DiskANN runs
+ /// the graph and rerank path; other index types perform their normal
+ /// resident optimization because they do not expose a paged query cache.
+ pub fn warmup_queries(
&mut self,
- query: &[f32],
- params: VectorSearchParams,
- ) -> io::Result<(Vec<i64>, Vec<f32>)> {
- validate_query(query, self.dimension())?;
- validate_query_bits_for_index(self.index_type(), params.query_bits)?;
- match self {
- Self::IvfFlat(reader) => reader.search(query, params.top_k,
params.nprobe),
- Self::IvfPq(reader) => search_with_reader(reader, query,
params.top_k, params.nprobe),
- Self::IvfRq(reader) => {
- reader.search_with_query_bits(query, params.top_k,
params.nprobe, params.query_bits)
- }
- Self::IvfHnswFlat(reader) => {
- reader.search(query, params.top_k, params.nprobe,
params.hnsw_ef_search())
+ queries: &[f32],
+ query_count: usize,
+ l_search: usize,
+ ) -> io::Result<()> {
+ let expected_len = query_count
+ .checked_mul(self.dimension())
+ .ok_or_else(|| invalid_input("warmup query count * dimension
overflows usize"))?;
+ if queries.len() != expected_len {
+ return Err(invalid_input(format!(
+ "warmup queries length {} does not match query count *
dimension {}",
+ queries.len(),
+ expected_len
+ )));
+ }
+ validate_finite_values(queries, expected_len, "warmup queries")?;
+ match self {
+ Self::DiskAnn(reader) => reader.warmup_queries(queries, l_search),
+ _ => self.optimize_for_search(),
+ }
+ }
+
+ pub fn calibrate_search_width(
+ &mut self,
+ queries: &[f32],
+ query_count: usize,
+ top_k: usize,
+ ) -> io::Result<usize> {
+ validate_queries(queries, query_count, self.dimension())?;
+ validate_positive(top_k, "top_k")?;
+ match self {
+ Self::DiskAnn(reader) => reader.calibrate_l_search(queries, top_k),
+ _ => Err(invalid_input(
+ "search-width calibration is currently only available for
DiskANN",
+ )),
+ }
+ }
+
+ pub fn search(
+ &mut self,
+ query: &[f32],
+ params: VectorSearchParams,
+ ) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ validate_query(query, self.dimension())?;
Review Comment:
Should the unified reader validate `top_k > 0` before dispatching (and apply
the same check to batch/filtered entry points)? All IVF implementations reject
`k == 0`, while DiskANN returns an empty successful result when an explicit
`l_search` is supplied. I reproduced this through the Python/native API:
IVF-FLAT, IVF-SQ, IVF-PQ, and IVF-RQ all fail with `k must be greater than 0`,
but DiskANN returns two empty arrays. The same public `VectorSearchParams`
therefore has index-dependent validity across Rust, C/C++, JNI, and Python.
--
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]
