#[cfg(feature = "native")]
use std::cmp::Reverse;
#[cfg(feature = "native")]
use std::collections::BinaryHeap;
use std::io;
#[cfg(feature = "native")]
use std::io::Write;
use std::ops::Range;
#[cfg(feature = "native")]
use byteorder::WriteBytesExt;
use byteorder::{LittleEndian, ReadBytesExt};
use crate::directories::OwnedBytes;
use crate::dsl::IvfRoutingMode;
#[cfg(feature = "native")]
use crate::structures::BinaryIvfIndex;
use crate::structures::vector::index::{BoundedAnnCollector, BoundedUniqueAnnCollector};
type CombinedBinaryCandidates = (Vec<AnnDocumentCandidate>, Vec<(u32, u16, f32)>);
const ANN_HEADER_MAGIC: u32 = 0x3152_4e41; const ANN_FOOTER_MAGIC: u32 = 0x3146_4e41; const ANN_DISK_VERSION: u16 = 1;
const ANN_HEADER_SIZE: usize = 56;
const ANN_RUN_SIZE: usize = 48;
const ANN_FOOTER_SIZE: usize = 24;
#[cfg(feature = "native")]
const COPY_CHUNK: usize = 8 * 1024 * 1024;
#[cfg(feature = "native")]
const PREFETCH_COALESCE_GAP: usize = 4 * 1024;
const BINARY_SCORE_BATCH: usize = 8_192;
const TQ_PRUNE_ESTIMATE_BOUND: f32 = 1.3;
#[cfg(feature = "native")]
const TQ_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536;
#[cfg(feature = "native")]
const TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum AnnKind {
BinaryIvf = 2,
TqFlat = 3,
IvfTq = 4,
}
impl AnnKind {
fn from_u8(value: u8) -> io::Result<Self> {
match value {
1 => Err(invalid_data(
"ANN kind 1 (IVF-PQ) is no longer supported; recreate the index \
with `ivf_tq` and reindex",
)),
2 => Ok(Self::BinaryIvf),
3 => Ok(Self::TqFlat),
4 => Ok(Self::IvfTq),
_ => Err(invalid_data(format!("unknown ANN kind {value}"))),
}
}
}
fn expected_codes_column_len(kind: AnnKind, count: usize, code_size: usize) -> io::Result<usize> {
match kind {
AnnKind::BinaryIvf => count
.checked_mul(code_size)
.ok_or_else(|| invalid_data("ANN code column size overflows usize")),
AnnKind::TqFlat => {
crate::structures::vector::quantization::tq_codes_column_len_checked(count, code_size)
.ok_or_else(|| invalid_data("TQ code column size overflows usize"))
}
AnnKind::IvfTq => crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
count, code_size,
)
.ok_or_else(|| invalid_data("IVF-TQ code column size overflows usize")),
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AnnDiskHeader {
pub kind: AnnKind,
pub routing: IvfRoutingMode,
pub dim: usize,
pub code_size: usize,
pub num_clusters: u32,
pub quantizer_version: u64,
pub codebook_version: u64,
pub vector_count: usize,
}
#[derive(Debug)]
struct AnnRun {
cluster_id: u32,
doc_base: u32,
max_doc_id: u32,
count: usize,
doc_ids: Range<usize>,
ordinals: Range<usize>,
codes: Range<usize>,
}
pub(crate) struct AnnDiskIndex {
#[cfg(feature = "native")]
heap_pins: crate::segment::pin::HeapPinSet,
raw: OwnedBytes,
header: AnnDiskHeader,
runs: Vec<AnnRun>,
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AnnHealth {
pub vectors: u64,
pub clusters_nonempty: u32,
pub clusters_total: u32,
pub runs: u32,
pub largest_cluster: u32,
pub largest_cluster_vectors: u64,
pub imbalance: f64,
pub payload_bytes: u64,
}
impl AnnHealth {
pub fn fragmentation(&self) -> f64 {
if self.clusters_nonempty == 0 {
return 0.0;
}
f64::from(self.runs) / f64::from(self.clusters_nonempty)
}
pub fn largest_cluster_share(&self) -> f64 {
if self.vectors == 0 {
return 0.0;
}
self.largest_cluster_vectors as f64 / self.vectors as f64
}
}
const ANN_SKEW_WARN_SHARE: f64 = 0.05;
const ANN_SKEW_WARN_MIN_VECTORS: u64 = 100_000;
const ANN_FRAGMENTATION_WARN: f64 = 8.0;
#[derive(Debug, Clone, Copy, PartialEq)]
pub(crate) struct AnnDocumentCandidate {
pub(crate) doc_id: u32,
pub(crate) score: f32,
}
impl AnnDiskIndex {
pub(crate) fn open(
raw: OwnedBytes,
expected_kind: AnnKind,
total_docs: u32,
) -> io::Result<Self> {
if raw.len() < ANN_HEADER_SIZE + ANN_FOOTER_SIZE {
return Err(invalid_data("ANN payload is shorter than header + footer"));
}
let bytes = raw.as_slice();
let mut header_cursor = std::io::Cursor::new(&bytes[..ANN_HEADER_SIZE]);
if header_cursor.read_u32::<LittleEndian>()? != ANN_HEADER_MAGIC {
return Err(invalid_data("ANN payload has unsupported header magic"));
}
let kind = AnnKind::from_u8(header_cursor.read_u8()?)?;
if kind != expected_kind {
return Err(invalid_data(format!(
"ANN payload kind {kind:?} does not match expected {expected_kind:?}"
)));
}
let routing = routing_from_u8(header_cursor.read_u8()?)?;
if header_cursor.read_u16::<LittleEndian>()? != ANN_DISK_VERSION {
return Err(invalid_data("ANN payload has unsupported format version"));
}
let dim = header_cursor.read_u32::<LittleEndian>()? as usize;
let code_size = header_cursor.read_u32::<LittleEndian>()? as usize;
let num_clusters = header_cursor.read_u32::<LittleEndian>()?;
if header_cursor.read_u32::<LittleEndian>()? != 0 {
return Err(invalid_data("ANN header reserved field is non-zero"));
}
let quantizer_version = header_cursor.read_u64::<LittleEndian>()?;
let codebook_version = header_cursor.read_u64::<LittleEndian>()?;
let vector_count = usize::try_from(header_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN vector count exceeds usize"))?;
if header_cursor.read_u64::<LittleEndian>()? != 0 {
return Err(invalid_data("ANN header tail is non-zero"));
}
let header = AnnDiskHeader {
kind,
routing,
dim,
code_size,
num_clusters,
quantizer_version,
codebook_version,
vector_count,
};
validate_header(&header)?;
let footer_start = bytes.len() - ANN_FOOTER_SIZE;
let mut footer_cursor = std::io::Cursor::new(&bytes[footer_start..]);
let directory_offset = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN directory offset exceeds usize"))?;
let num_runs = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN run count exceeds usize"))?;
if footer_cursor.read_u32::<LittleEndian>()? != ANN_FOOTER_MAGIC
|| footer_cursor.read_u32::<LittleEndian>()? != u32::from(ANN_DISK_VERSION)
{
return Err(invalid_data("ANN payload has unsupported footer"));
}
if num_runs == 0 {
return Err(invalid_data("ANN payload has no cluster runs"));
}
let directory_len = num_runs
.checked_mul(ANN_RUN_SIZE)
.ok_or_else(|| invalid_data("ANN directory size overflows usize"))?;
if directory_offset < ANN_HEADER_SIZE
|| directory_offset.checked_add(directory_len) != Some(footer_start)
{
return Err(invalid_data("ANN directory does not end at the footer"));
}
let mut runs = Vec::with_capacity(num_runs);
let mut directory_cursor = std::io::Cursor::new(&bytes[directory_offset..footer_start]);
let mut previous_cluster = None;
let mut counted_vectors = 0usize;
for _ in 0..num_runs {
let cluster_id = directory_cursor.read_u32::<LittleEndian>()?;
let doc_base = directory_cursor.read_u32::<LittleEndian>()?;
let count = directory_cursor.read_u32::<LittleEndian>()? as usize;
let max_doc_id = directory_cursor.read_u32::<LittleEndian>()?;
let doc_ids_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN doc-ID offset exceeds usize"))?;
let ordinals_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN ordinal offset exceeds usize"))?;
let codes_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN code offset exceeds usize"))?;
let codes_len = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
.map_err(|_| invalid_data("ANN code length exceeds usize"))?;
if count == 0
|| cluster_id >= num_clusters
|| previous_cluster.is_some_and(|previous| previous > cluster_id)
|| doc_base
.checked_add(max_doc_id)
.is_none_or(|doc_id| doc_id >= total_docs)
{
return Err(invalid_data("ANN run metadata is invalid"));
}
previous_cluster = Some(cluster_id);
let doc_ids_len = count
.checked_mul(std::mem::size_of::<u32>())
.ok_or_else(|| invalid_data("ANN doc-ID column size overflows usize"))?;
let ordinals_len = count
.checked_mul(std::mem::size_of::<u16>())
.ok_or_else(|| invalid_data("ANN ordinal column size overflows usize"))?;
let expected_codes_len = expected_codes_column_len(kind, count, code_size)?;
let doc_ids_end = doc_ids_offset
.checked_add(doc_ids_len)
.ok_or_else(|| invalid_data("ANN doc-ID range overflows usize"))?;
let ordinals_end = ordinals_offset
.checked_add(ordinals_len)
.ok_or_else(|| invalid_data("ANN ordinal range overflows usize"))?;
let codes_end = codes_offset
.checked_add(codes_len)
.ok_or_else(|| invalid_data("ANN code range overflows usize"))?;
if doc_ids_offset < ANN_HEADER_SIZE
|| ordinals_offset != doc_ids_end
|| codes_offset != ordinals_end
|| codes_len != expected_codes_len
|| codes_end > directory_offset
{
return Err(invalid_data("ANN run columns are not contiguous/in bounds"));
}
runs.push(AnnRun {
cluster_id,
doc_base,
max_doc_id,
count,
doc_ids: doc_ids_offset..doc_ids_end,
ordinals: ordinals_offset..ordinals_end,
codes: codes_offset..codes_end,
});
counted_vectors = counted_vectors
.checked_add(count)
.ok_or_else(|| invalid_data("ANN run vector count overflows usize"))?;
}
let mut payload_order: Vec<usize> = (0..runs.len()).collect();
payload_order.sort_unstable_by_key(|&index| runs[index].doc_ids.start);
let mut expected_payload_offset = ANN_HEADER_SIZE;
for index in payload_order {
let run = &runs[index];
if run.doc_ids.start != expected_payload_offset {
return Err(invalid_data("ANN payload runs overlap or contain gaps"));
}
expected_payload_offset = run.codes.end;
}
if expected_payload_offset != directory_offset || counted_vectors != vector_count {
return Err(invalid_data(
"ANN payload coverage/vector count is inconsistent",
));
}
#[cfg(feature = "native")]
raw.madvise_range(
ANN_HEADER_SIZE..directory_offset,
if kind == AnnKind::TqFlat {
libc::MADV_SEQUENTIAL
} else {
libc::MADV_RANDOM
},
);
Ok(Self {
#[cfg(feature = "native")]
heap_pins: Default::default(),
raw,
header,
runs,
})
}
pub(crate) fn health(&self) -> AnnHealth {
let mut vectors = 0u64;
let mut clusters_nonempty = 0u32;
let mut payload_bytes = 0u64;
let mut largest = (0u32, 0u64);
let mut sum_squares = 0f64;
let mut index = 0usize;
while index < self.runs.len() {
let cluster_id = self.runs[index].cluster_id;
let mut cluster_vectors = 0u64;
while index < self.runs.len() && self.runs[index].cluster_id == cluster_id {
let run = &self.runs[index];
cluster_vectors += run.count as u64;
payload_bytes += (run.codes.end - run.codes.start) as u64;
index += 1;
}
vectors += cluster_vectors;
clusters_nonempty += 1;
sum_squares += (cluster_vectors as f64) * (cluster_vectors as f64);
if cluster_vectors > largest.1 {
largest = (cluster_id, cluster_vectors);
}
}
let imbalance = if vectors == 0 || clusters_nonempty == 0 {
0.0
} else {
f64::from(clusters_nonempty) * sum_squares / ((vectors as f64) * (vectors as f64))
};
AnnHealth {
vectors,
clusters_nonempty,
clusters_total: self.header.num_clusters,
runs: self.runs.len() as u32,
largest_cluster: largest.0,
largest_cluster_vectors: largest.1,
imbalance,
payload_bytes,
}
}
pub(crate) fn report_health(&self, index_label: &str, field_id: u32, segment_id: u128) {
let health = self.health();
let share = health.largest_cluster_share();
let fragmentation = health.fragmentation();
log::info!(
"[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
vectors={} clusters={}/{} runs={} fragmentation={fragmentation:.2} \
imbalance={:.2} largest_leaf={:.2}% payload={}",
health.vectors,
health.clusters_nonempty,
health.clusters_total,
health.runs,
health.imbalance,
100.0 * share,
crate::format_bytes(health.payload_bytes),
);
crate::observe::ann_health(
index_label,
field_id,
health.imbalance,
fragmentation,
share,
);
if share >= ANN_SKEW_WARN_SHARE && health.vectors >= ANN_SKEW_WARN_MIN_VECTORS {
log::warn!(
"[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
leaf {} holds {:.1}% of {} vectors — every query probing it scans that leaf \
in full; degenerate embeddings collapse into one leaf exactly like this",
health.largest_cluster,
100.0 * share,
health.vectors,
);
}
if fragmentation >= ANN_FRAGMENTATION_WARN {
log::warn!(
"[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
{fragmentation:.1} extents per probed cluster ({} runs / {} clusters) — \
cold probes pay that many seeks; the next merge or vector-generation rewrite \
compacts to 1.0",
health.runs,
health.clusters_nonempty,
);
}
}
pub(crate) fn header(&self) -> &AnnDiskHeader {
&self.header
}
pub(crate) fn estimated_heap_bytes(&self) -> usize {
std::mem::size_of::<Self>() + self.runs.capacity() * std::mem::size_of::<AnnRun>()
}
#[cfg(feature = "native")]
pub(crate) fn pin_lookup_directory(
&mut self,
mode: crate::segment::pin::PinMode,
remaining: &mut u64,
report: &mut crate::segment::pin::PinReport,
) {
let before = self.heap_pins.report();
self.heap_pins
.pin_slice(&self.runs, "ANN cluster-run directory", mode, remaining);
let after = self.heap_pins.report();
report.intended_bytes += after.intended_bytes - before.intended_bytes;
report.pinned_bytes += after.pinned_bytes - before.pinned_bytes;
report.skipped_budget_bytes += after.skipped_budget_bytes - before.skipped_budget_bytes;
report.failed_bytes += after.failed_bytes - before.failed_bytes;
report.heap_copy_bytes += after.heap_copy_bytes - before.heap_copy_bytes;
}
fn cluster_runs(&self, cluster_id: u32) -> &[AnnRun] {
let start = self.runs.partition_point(|run| run.cluster_id < cluster_id);
let end = self
.runs
.partition_point(|run| run.cluster_id <= cluster_id);
&self.runs[start..end]
}
#[cfg(feature = "native")]
fn prefetch_cluster_runs(&self, cluster_ids: &[u32]) {
if cluster_ids.is_empty() || !self.raw.is_mmap() {
return;
}
let mut ranges = Vec::with_capacity(cluster_ids.len());
for &cluster_id in cluster_ids {
ranges.extend(
self.cluster_runs(cluster_id)
.iter()
.map(|run| run.doc_ids.start..run.codes.end),
);
}
coalesce_prefetch_ranges(&mut ranges);
for range in ranges {
self.raw.madvise_range(range, libc::MADV_WILLNEED);
}
}
pub(crate) fn search_tq_combined_documents(
&self,
k: usize,
plan: &crate::structures::TqQueryPlan,
combiner: crate::query::MultiValueCombiner,
) -> io::Result<Vec<AnnDocumentCandidate>> {
use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
combiner
.validate()
.map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?;
if plan.padded_dim() != self.header.code_size * 2 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"TQ query plan does not match the payload dimension",
));
}
if k == 0 {
return Ok(Vec::new());
}
let block_bytes = tq_block_bytes(self.header.code_size);
let bytes = self.raw.as_slice();
let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
let mut ordinal_scores = Vec::new();
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for run in &self.runs {
let mut current_doc = None;
let codes = &bytes[run.codes.clone()];
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let index = lane_base + lane;
let doc_id = run_doc_id(bytes, run, index)?;
if current_doc.is_some_and(|previous| doc_id < previous) {
return Err(invalid_data("flat TQ run is not grouped by document ID"));
}
if current_doc.is_some_and(|previous| doc_id != previous) {
retain_combined_document(
&mut top_documents,
current_doc.expect("current document is present"),
&ordinal_scores,
combiner,
);
ordinal_scores.clear();
}
current_doc = Some(doc_id);
if score.is_finite() {
ordinal_scores.push((
u32::from(read_u16(bytes, run.ordinals.start + index * 2)),
score,
));
}
}
}
if let Some(doc_id) = current_doc {
retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
ordinal_scores.clear();
}
}
Ok(top_documents
.into_sorted_results()
.into_iter()
.map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
.collect())
}
pub(crate) fn search_ivf_tq_combined_documents(
&self,
k: usize,
plan: &crate::structures::TqIvfQueryPlan,
combiner: crate::query::MultiValueCombiner,
) -> io::Result<Vec<AnnDocumentCandidate>> {
use crate::structures::vector::quantization::{
TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
};
validate_combined_search(combiner)?;
let tq_plan = plan.tq_plan();
self.validate_ivf_tq_query_plan(plan)?;
if k == 0 {
return Ok(Vec::new());
}
#[cfg(feature = "native")]
self.prefetch_cluster_runs(&plan.cluster_ids);
let block_bytes = tq_ivf_block_bytes(self.header.code_size);
let bytes = self.raw.as_slice();
let mut ordinal_scores = Vec::new();
ordinal_scores
.try_reserve_exact(probed_posting_count(self, &plan.cluster_ids)?)
.map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?;
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for (cluster_id, cluster_dot) in plan.cluster_dots() {
for run in self.cluster_runs(cluster_id) {
let codes = &bytes[run.codes.clone()];
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
if !score.is_finite() {
continue;
}
let index = lane_base + lane;
ordinal_scores.push((
run_doc_id(bytes, run, index)?,
read_u16(bytes, run.ordinals.start + index * 2),
score,
));
}
}
}
}
Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
}
pub(crate) fn search_binary_combined_documents(
&self,
k: usize,
query: &[u8],
cluster_ids: &[u32],
combiner: crate::query::MultiValueCombiner,
) -> io::Result<CombinedBinaryCandidates> {
validate_combined_search(combiner)?;
if query.len() != self.header.code_size {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"binary ANN query has the wrong byte length",
));
}
if k == 0 {
return Ok((Vec::new(), Vec::new()));
}
#[cfg(feature = "native")]
self.prefetch_cluster_runs(cluster_ids);
let bytes = self.raw.as_slice();
let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
let mut ordinal_scores = Vec::new();
ordinal_scores
.try_reserve_exact(probed_posting_count(self, cluster_ids)?)
.map_err(|_| invalid_data("binary combined score buffer allocation failed"))?;
score_binary_cluster_runs(
self,
bytes,
query,
cluster_ids,
&mut score_batch,
&mut ordinal_scores,
)?;
Ok(combine_scored_ordinals_retaining(
ordinal_scores,
k,
combiner,
))
}
pub(crate) fn search_tq_distinct(
&self,
k: usize,
plan: &crate::structures::TqQueryPlan,
) -> io::Result<Vec<(u32, u16, f32)>> {
use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
if plan.padded_dim() != self.header.code_size * 2 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"TQ query plan does not match the payload dimension",
));
}
let block_bytes = tq_block_bytes(self.header.code_size);
let bytes = self.raw.as_slice();
#[cfg(feature = "native")]
if self.header.vector_count >= TQ_PARALLEL_SCAN_MIN_VECTORS {
use rayon::prelude::*;
let collector = self
.runs
.par_iter()
.flat_map(|run| {
let codes = &bytes[run.codes.clone()];
let blocks = codes.len() / block_bytes;
(0..blocks.div_ceil(TQ_PARALLEL_SCAN_CHUNK_BLOCKS))
.into_par_iter()
.map(move |chunk| (run, chunk * TQ_PARALLEL_SCAN_CHUNK_BLOCKS, blocks))
})
.try_fold(
|| BoundedAnnCollector::<true, true>::new(k),
|mut collector, (run, first_block, total_blocks)| {
let codes = &bytes[run.codes.clone()];
let last_block =
(first_block + TQ_PARALLEL_SCAN_CHUNK_BLOCKS).min(total_blocks);
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for block_index in first_block..last_block {
let block = &codes[block_index * block_bytes..][..block_bytes];
crate::structures::vector::quantization::tq_score_block(
plan,
block,
&mut scores,
);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let index = lane_base + lane;
collector.insert(
run_doc_id(bytes, run, index)?,
read_u16(bytes, run.ordinals.start + index * 2),
score,
);
}
}
Ok::<_, io::Error>(collector)
},
)
.try_reduce(
|| BoundedAnnCollector::<true, true>::new(k),
|mut collector, partial| {
collector.merge_from(partial);
Ok(collector)
},
)?;
return Ok(collector.into_sorted_results());
}
let mut collector = BoundedAnnCollector::<true, true>::new(k);
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for run in &self.runs {
let codes = &bytes[run.codes.clone()];
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let index = lane_base + lane;
collector.insert(
run_doc_id(bytes, run, index)?,
read_u16(bytes, run.ordinals.start + index * 2),
score,
);
}
}
}
Ok(collector.into_sorted_results())
}
pub(crate) fn search_ivf_tq_distinct(
&self,
k: usize,
plan: &crate::structures::TqIvfQueryPlan,
) -> io::Result<Vec<(u32, u16, f32)>> {
use crate::structures::vector::quantization::{
TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
};
let tq_plan = plan.tq_plan();
self.validate_ivf_tq_query_plan(plan)?;
#[cfg(feature = "native")]
self.prefetch_cluster_runs(&plan.cluster_ids);
let block_bytes = tq_ivf_block_bytes(self.header.code_size);
let bytes = self.raw.as_slice();
let mut collector = BoundedAnnCollector::<true, true>::new(k);
let mut scores = [0.0f32; TQ_BLOCK_LANES];
let mut pruned_blocks = 0usize;
let mut scored_blocks = 0usize;
for (cluster_id, cluster_dot) in plan.cluster_dots() {
for run in self.cluster_runs(cluster_id) {
let codes = &bytes[run.codes.clone()];
let block_count = codes.len() / block_bytes;
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
if let Some(threshold) = collector.pruning_threshold() {
let max_scale = tq_ivf_block_max_scale(block);
if cluster_dot + max_scale * TQ_PRUNE_ESTIMATE_BOUND <= threshold {
pruned_blocks += block_count - block_index;
break;
}
}
scored_blocks += 1;
tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let index = lane_base + lane;
collector.insert(
run_doc_id(bytes, run, index)?,
read_u16(bytes, run.ordinals.start + index * 2),
score,
);
}
}
}
}
if pruned_blocks > 0 {
log::debug!(
"[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds",
pruned_blocks + scored_blocks,
);
}
Ok(collector.into_sorted_results())
}
fn validate_ivf_tq_query_plan(
&self,
plan: &crate::structures::TqIvfQueryPlan,
) -> io::Result<()> {
if self.header.kind != AnnKind::IvfTq
|| !crate::structures::is_ivf_tq_cosine_generation(self.header.quantizer_version)
{
return Err(invalid_data(
"legacy raw IVF-TQ payloads cannot be searched; rebuild the index",
));
}
if plan.tq_plan().padded_dim() != self.header.code_size * 2
|| plan.quantizer_version != self.header.quantizer_version
|| plan.fingerprint != self.header.codebook_version
{
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"IVF-TQ query plan does not match the payload generation",
));
}
Ok(())
}
pub(crate) fn search_binary_clusters<const BY_DOCUMENT: bool>(
&self,
query: &[u8],
k: usize,
cluster_ids: &[u32],
) -> io::Result<Vec<(u32, u16, f32)>> {
if query.len() != self.header.code_size {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"binary ANN query has the wrong byte length",
));
}
#[cfg(feature = "native")]
self.prefetch_cluster_runs(cluster_ids);
let bytes = self.raw.as_slice();
let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
if BY_DOCUMENT {
let mut collector = BoundedAnnCollector::<true, true>::new(k);
score_binary_cluster_runs(
self,
bytes,
query,
cluster_ids,
&mut scores,
&mut collector,
)?;
return Ok(collector.into_sorted_results());
}
debug_assert!(
{
let mut seen = rustc_hash::FxHashSet::default();
cluster_ids
.iter()
.all(|cluster_id| seen.insert(*cluster_id))
},
"an IVF probe plan must not repeat a cluster",
);
let mut collector = BoundedUniqueAnnCollector::<true>::new(k);
score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?;
Ok(collector.into_sorted_results())
}
}
#[cfg(feature = "native")]
fn coalesce_prefetch_ranges(ranges: &mut Vec<Range<usize>>) {
if ranges.len() < 2 {
return;
}
ranges.sort_unstable_by_key(|range| range.start);
let mut output_len = 1usize;
for input_index in 1..ranges.len() {
let next_start = ranges[input_index].start;
let next_end = ranges[input_index].end;
let previous = &mut ranges[output_len - 1];
if next_start <= previous.end.saturating_add(PREFETCH_COALESCE_GAP) {
previous.end = previous.end.max(next_end);
} else {
ranges[output_len] = next_start..next_end;
output_len += 1;
}
}
ranges.truncate(output_len);
}
trait AnnScoreSink {
fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32);
}
impl<const BY_DOCUMENT: bool> AnnScoreSink for BoundedAnnCollector<BY_DOCUMENT, true> {
#[inline]
fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
self.insert(doc_id, ordinal, score);
}
}
impl AnnScoreSink for BoundedUniqueAnnCollector<true> {
#[inline]
fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
self.insert(doc_id, ordinal, score);
}
}
impl AnnScoreSink for Vec<(u32, u16, f32)> {
#[inline]
fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
if score.is_finite() {
self.push((doc_id, ordinal, score));
}
}
}
fn probed_posting_count(index: &AnnDiskIndex, cluster_ids: &[u32]) -> io::Result<usize> {
cluster_ids.iter().try_fold(0usize, |count, &cluster_id| {
index
.cluster_runs(cluster_id)
.iter()
.try_fold(count, |count, run| {
count
.checked_add(run.count)
.ok_or_else(|| invalid_data("ANN probed posting count overflows usize"))
})
})
}
fn score_binary_cluster_runs(
index: &AnnDiskIndex,
bytes: &[u8],
query: &[u8],
cluster_ids: &[u32],
scores: &mut [f32],
collector: &mut impl AnnScoreSink,
) -> io::Result<()> {
for &cluster_id in cluster_ids {
for run in index.cluster_runs(cluster_id) {
score_binary_run(
bytes,
run,
query,
index.header.dim,
index.header.code_size,
scores,
collector,
)?;
}
}
Ok(())
}
fn score_binary_run(
bytes: &[u8],
run: &AnnRun,
query: &[u8],
dim_bits: usize,
code_size: usize,
scores: &mut [f32],
collector: &mut impl AnnScoreSink,
) -> io::Result<()> {
for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) {
let batch_count = BINARY_SCORE_BATCH.min(run.count - batch_start);
let code_start = run.codes.start + batch_start * code_size;
let code_end = code_start + batch_count * code_size;
crate::structures::simd::batch_hamming_scores(
query,
&bytes[code_start..code_end],
code_size,
dim_bits,
&mut scores[..batch_count],
);
for (batch_index, &score) in scores.iter().enumerate().take(batch_count) {
let index = batch_start + batch_index;
collector.insert_score(
run_doc_id(bytes, run, index)?,
read_u16(bytes, run.ordinals.start + index * 2),
score,
);
}
}
Ok(())
}
#[inline]
fn retain_combined_document(
collector: &mut BoundedAnnCollector<true, true>,
doc_id: u32,
ordinal_scores: &[(u32, f32)],
combiner: crate::query::MultiValueCombiner,
) {
if !ordinal_scores.is_empty() {
collector.insert(doc_id, 0, combiner.combine(ordinal_scores));
}
}
fn validate_combined_search(combiner: crate::query::MultiValueCombiner) -> io::Result<()> {
combiner
.validate()
.map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))
}
fn combine_scored_ordinals(
scores: Vec<(u32, u16, f32)>,
k: usize,
combiner: crate::query::MultiValueCombiner,
) -> Vec<AnnDocumentCandidate> {
combine_scored_ordinals_retaining(scores, k, combiner).0
}
fn combine_scored_ordinals_retaining(
mut scores: Vec<(u32, u16, f32)>,
k: usize,
combiner: crate::query::MultiValueCombiner,
) -> CombinedBinaryCandidates {
if k == 0 || scores.is_empty() {
return (Vec::new(), Vec::new());
}
scores.retain(|entry| entry.2.is_finite());
scores.sort_unstable_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)));
let mut unique_len = 0usize;
for read_index in 0..scores.len() {
let candidate = scores[read_index];
if unique_len > 0
&& scores[unique_len - 1].0 == candidate.0
&& scores[unique_len - 1].1 == candidate.1
{
if candidate.2.total_cmp(&scores[unique_len - 1].2).is_gt() {
scores[unique_len - 1].2 = candidate.2;
}
} else {
scores[unique_len] = candidate;
unique_len += 1;
}
}
scores.truncate(unique_len);
let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
let mut current_doc = None;
let mut ordinal_scores = Vec::new();
for &(doc_id, ordinal, score) in &scores {
if current_doc.is_some_and(|current| current != doc_id) {
retain_combined_document(
&mut top_documents,
current_doc.expect("current document is present"),
&ordinal_scores,
combiner,
);
ordinal_scores.clear();
}
current_doc = Some(doc_id);
ordinal_scores.push((u32::from(ordinal), score));
}
if let Some(doc_id) = current_doc {
retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
}
let candidates: Vec<AnnDocumentCandidate> = top_documents
.into_sorted_results()
.into_iter()
.map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
.collect();
let mut retained_ids: Vec<u32> = candidates
.iter()
.map(|candidate| candidate.doc_id)
.collect();
retained_ids.sort_unstable();
let mut retained = Vec::with_capacity(scores.len().min(retained_ids.len().saturating_mul(4)));
let mut cursor = 0usize;
for entry in scores {
while cursor < retained_ids.len() && retained_ids[cursor] < entry.0 {
cursor += 1;
}
if retained_ids.get(cursor) == Some(&entry.0) {
retained.push(entry);
}
}
(candidates, retained)
}
#[cfg(feature = "native")]
pub(crate) fn write_built_binary_ivf(
index: &BinaryIvfIndex,
routing: IvfRoutingMode,
writer: &mut (impl Write + ?Sized),
) -> io::Result<u64> {
let runs: Vec<_> = index
.clusters
.iter()
.map(|(cluster_id, cluster)| BuildRun {
cluster_id: *cluster_id,
doc_ids: &cluster.doc_ids,
ordinals: &cluster.ordinals,
codes: &cluster.codes,
})
.collect();
write_built_runs(
AnnDiskHeader {
kind: AnnKind::BinaryIvf,
routing,
dim: index.dim_bits,
code_size: index.dim_bits.div_ceil(8),
num_clusters: index.num_clusters,
quantizer_version: index.quantizer_version,
codebook_version: 0,
vector_count: index.len(),
},
&runs,
writer,
)
}
#[cfg(feature = "native")]
pub(crate) fn write_built_ivf_tq(
index: &crate::structures::IvfTqIndex,
num_clusters: u32,
writer: &mut (impl Write + ?Sized),
) -> io::Result<u64> {
use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_pack_ivf_block};
if !crate::structures::is_ivf_tq_cosine_generation(index.centroids_version) {
return Err(invalid_data(
"legacy raw IVF-TQ generations cannot be serialized; rebuild the index",
));
}
let codec = index.codec();
let padded_dim = codec.padded_dim();
let mut clusters: Vec<_> = index.clusters.iter().collect();
clusters.sort_unstable_by_key(|(cluster_id, _)| **cluster_id);
struct PackedCluster {
cluster_id: u32,
doc_ids: Vec<u32>,
ordinals: Vec<u16>,
codes: Vec<u8>,
}
let packed: Vec<PackedCluster> = clusters
.iter()
.map(|&(&cluster_id, cluster)| {
let count = cluster.doc_ids.len();
let mut order: Vec<usize> = (0..count).collect();
order.sort_by(|&a, &b| {
cluster.scales[b]
.total_cmp(&cluster.scales[a])
.then_with(|| a.cmp(&b))
});
let doc_ids: Vec<u32> = order.iter().map(|&i| cluster.doc_ids[i]).collect();
let ordinals: Vec<u16> = order.iter().map(|&i| cluster.ordinals[i]).collect();
let scales: Vec<f32> = order.iter().map(|&i| cluster.scales[i]).collect();
let gammas: Vec<f32> = order.iter().map(|&i| cluster.gammas[i]).collect();
let mut codes = Vec::with_capacity(
crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
count,
codec.code_size(),
)
.unwrap_or_default(),
);
for block_start in (0..count).step_by(TQ_BLOCK_LANES) {
let lanes = TQ_BLOCK_LANES.min(count - block_start);
let rows: Vec<&[u8]> = order[block_start..block_start + lanes]
.iter()
.map(|&row| &cluster.rows[row * padded_dim..(row + 1) * padded_dim])
.collect();
tq_pack_ivf_block(
&rows,
&scales[block_start..block_start + lanes],
&gammas[block_start..block_start + lanes],
padded_dim,
&mut codes,
);
}
PackedCluster {
cluster_id,
doc_ids,
ordinals,
codes,
}
})
.collect();
let runs: Vec<_> = packed
.iter()
.map(|cluster| BuildRun {
cluster_id: cluster.cluster_id,
doc_ids: &cluster.doc_ids,
ordinals: &cluster.ordinals,
codes: &cluster.codes,
})
.collect();
write_built_runs(
AnnDiskHeader {
kind: AnnKind::IvfTq,
routing: index.routing,
dim: index.dim,
code_size: codec.code_size(),
num_clusters,
quantizer_version: index.centroids_version,
codebook_version: codec.fingerprint(),
vector_count: index.len(),
},
&runs,
writer,
)
}
#[cfg(feature = "native")]
pub(crate) fn tq_builder_extra_run(builder: &crate::structures::TqFlatBuilder) -> BuildRun<'_> {
BuildRun {
cluster_id: 0,
doc_ids: &builder.doc_ids,
ordinals: &builder.ordinals,
codes: &builder.codes,
}
}
#[cfg(feature = "native")]
pub(crate) fn write_built_tq_flat(
builder: &crate::structures::TqFlatBuilder,
writer: &mut (impl Write + ?Sized),
) -> io::Result<u64> {
let codec = builder.codec();
let runs = [BuildRun {
cluster_id: 0,
doc_ids: &builder.doc_ids,
ordinals: &builder.ordinals,
codes: &builder.codes,
}];
write_built_runs(
AnnDiskHeader {
kind: AnnKind::TqFlat,
routing: IvfRoutingMode::Flat,
dim: codec.dim(),
code_size: codec.code_size(),
num_clusters: 1,
quantizer_version: codec.fingerprint(),
codebook_version: 0,
vector_count: builder.len(),
},
&runs,
writer,
)
}
#[cfg(feature = "native")]
pub(crate) struct BuildRun<'a> {
cluster_id: u32,
doc_ids: &'a [u32],
ordinals: &'a [u16],
codes: &'a [u8],
}
#[cfg(feature = "native")]
struct RunRecord {
cluster_id: u32,
doc_base: u32,
count: u32,
max_doc_id: u32,
doc_ids_offset: u64,
ordinals_offset: u64,
codes_offset: u64,
codes_len: u64,
}
#[cfg(feature = "native")]
fn write_built_runs(
header: AnnDiskHeader,
runs: &[BuildRun<'_>],
writer: &mut (impl Write + ?Sized),
) -> io::Result<u64> {
if runs.is_empty() || header.vector_count == 0 {
return Err(invalid_data("cannot write an empty ANN payload"));
}
validate_header(&header)?;
write_header(writer, &header)?;
let mut offset = ANN_HEADER_SIZE as u64;
let mut records = Vec::with_capacity(runs.len());
let mut counted = 0usize;
let mut scratch = Vec::new();
let mut previous_cluster = None;
for run in runs {
let count = run.doc_ids.len();
if count == 0
|| run.cluster_id >= header.num_clusters
|| previous_cluster.is_some_and(|cluster| cluster >= run.cluster_id)
|| run.ordinals.len() != count
|| run.codes.len() != expected_codes_column_len(header.kind, count, header.code_size)?
{
return Err(invalid_data("ANN build run columns are inconsistent"));
}
previous_cluster = Some(run.cluster_id);
let count_u32 = u32::try_from(count)
.map_err(|_| invalid_data("ANN cluster run exceeds u32 vectors"))?;
let max_doc_id = run.doc_ids.iter().copied().max().unwrap_or(0);
let doc_ids_offset = offset;
write_u32_column(writer, run.doc_ids, &mut scratch)?;
offset = offset
.checked_add(
u64::try_from(count)
.ok()
.and_then(|count| count.checked_mul(4))
.ok_or_else(|| invalid_data("ANN doc-ID output size overflows u64"))?,
)
.ok_or_else(|| invalid_data("ANN output offset overflow"))?;
let ordinals_offset = offset;
write_u16_column(writer, run.ordinals, &mut scratch)?;
offset = offset
.checked_add(
u64::try_from(count)
.ok()
.and_then(|count| count.checked_mul(2))
.ok_or_else(|| invalid_data("ANN ordinal output size overflows u64"))?,
)
.ok_or_else(|| invalid_data("ANN output offset overflow"))?;
let codes_offset = offset;
writer.write_all(run.codes)?;
offset = offset
.checked_add(
u64::try_from(run.codes.len())
.map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
)
.ok_or_else(|| invalid_data("ANN output offset overflow"))?;
records.push(RunRecord {
cluster_id: run.cluster_id,
doc_base: 0,
count: count_u32,
max_doc_id,
doc_ids_offset,
ordinals_offset,
codes_offset,
codes_len: u64::try_from(run.codes.len())
.map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
});
counted = counted
.checked_add(count)
.ok_or_else(|| invalid_data("ANN vector count overflow"))?;
}
if counted != header.vector_count {
return Err(invalid_data("ANN header/build vector counts disagree"));
}
finish_layout(writer, offset, &records)
}
#[cfg(all(feature = "native", test))]
pub(crate) fn write_merged_ann(
sources: &[(&AnnDiskIndex, u32)],
writer: &mut (impl Write + ?Sized),
) -> io::Result<u64> {
write_merged_ann_impl(sources, &[], writer, None)
}
#[cfg(feature = "native")]
pub(crate) fn write_merged_ann_cancellable(
sources: &[(&AnnDiskIndex, u32)],
writer: &mut (impl Write + ?Sized),
cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> io::Result<u64> {
write_merged_ann_impl(sources, &[], writer, cancellation)
}
#[cfg(feature = "native")]
pub(crate) fn predicted_merge_fragmentation(sources: &[(&AnnDiskIndex, u32)]) -> f64 {
let mut total_runs = 0usize;
let mut cursors: Vec<std::iter::Peekable<std::slice::Iter<'_, AnnRun>>> = sources
.iter()
.map(|(source, _)| {
total_runs += source.runs.len();
source.runs.iter().peekable()
})
.collect();
let mut distinct = 0usize;
while let Some(cluster) = cursors
.iter_mut()
.filter_map(|cursor| cursor.peek().map(|run| run.cluster_id))
.min()
{
distinct += 1;
for cursor in &mut cursors {
while cursor.peek().is_some_and(|run| run.cluster_id == cluster) {
cursor.next();
}
}
}
if distinct == 0 {
0.0
} else {
total_runs as f64 / distinct as f64
}
}
#[cfg(feature = "native")]
pub(crate) const ANN_COMPACTION_FRAGMENTATION_THRESHOLD: f64 = 4.0;
#[cfg(feature = "native")]
const DOC_ID_REWRITE_CHUNK: usize = 64 * 1024;
#[cfg(feature = "native")]
pub(crate) fn write_compacted_ann_cancellable(
sources: &[(&AnnDiskIndex, u32)],
writer: &mut (impl Write + ?Sized),
cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> io::Result<u64> {
let Some((first, _)) = sources.first() else {
return Err(invalid_data("cannot compact an empty ANN source list"));
};
if first.header.kind != AnnKind::BinaryIvf {
return Err(invalid_data(
"ANN run compaction is only defined for binary IVF payloads",
));
}
let mut header = first.header.clone();
header.vector_count = 0;
for &(source, _) in sources {
if !headers_compatible(&first.header, &source.header) {
return Err(invalid_data(
"ANN compaction sources use incompatible generations",
));
}
header.vector_count = header
.vector_count
.checked_add(source.header.vector_count)
.ok_or_else(|| invalid_data("compacted ANN vector count overflows usize"))?;
}
validate_header(&header)?;
write_header(writer, &header)?;
let code_size = header.code_size;
let mut offset = ANN_HEADER_SIZE as u64;
let mut records: Vec<RunRecord> = Vec::new();
let mut scratch = Vec::new();
let mut cursors: Vec<usize> = vec![0; sources.len()];
loop {
if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed)) {
return Err(io::Error::new(
io::ErrorKind::Interrupted,
"ANN compaction cancelled",
));
}
let Some(cluster_id) = sources
.iter()
.zip(&cursors)
.filter_map(|((source, _), &cursor)| source.runs.get(cursor).map(|run| run.cluster_id))
.min()
else {
break;
};
let mut count = 0usize;
let mut max_doc_id = 0u32;
let doc_ids_offset = offset;
for (source_index, &(source, segment_base)) in sources.iter().enumerate() {
let mut cursor = cursors[source_index];
while let Some(run) = source
.runs
.get(cursor)
.filter(|run| run.cluster_id == cluster_id)
{
let base = run
.doc_base
.checked_add(segment_base)
.ok_or_else(|| invalid_data("compacted ANN document base overflows u32"))?;
let bytes = source.raw.as_slice();
for chunk_start in (0..run.count).step_by(DOC_ID_REWRITE_CHUNK) {
let chunk_end = (chunk_start + DOC_ID_REWRITE_CHUNK).min(run.count);
scratch.clear();
scratch.reserve((chunk_end - chunk_start) * 4);
for index in chunk_start..chunk_end {
let doc_id = run_doc_id_with_base(bytes, run, index, base)?;
max_doc_id = max_doc_id.max(doc_id);
scratch.extend_from_slice(&doc_id.to_le_bytes());
}
writer.write_all(&scratch)?;
if cancellation
.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
{
return Err(io::Error::new(
io::ErrorKind::Interrupted,
"ANN compaction cancelled",
));
}
}
offset = checked_advance(offset, run.count * 4)?;
count = count
.checked_add(run.count)
.ok_or_else(|| invalid_data("compacted ANN run count overflows usize"))?;
cursor += 1;
}
}
let ordinals_offset = offset;
for (source_index, &(source, _)) in sources.iter().enumerate() {
let mut cursor = cursors[source_index];
while let Some(run) = source
.runs
.get(cursor)
.filter(|run| run.cluster_id == cluster_id)
{
copy_range(writer, &source.raw, run.ordinals.clone(), cancellation)?;
offset = checked_advance(offset, run.ordinals.len())?;
cursor += 1;
}
}
let codes_offset = offset;
for (source_index, &(source, _)) in sources.iter().enumerate() {
let mut cursor = cursors[source_index];
while let Some(run) = source
.runs
.get(cursor)
.filter(|run| run.cluster_id == cluster_id)
{
copy_range(writer, &source.raw, run.codes.clone(), cancellation)?;
offset = checked_advance(offset, run.codes.len())?;
cursor += 1;
}
}
for (source_index, &(source, _)) in sources.iter().enumerate() {
while source
.runs
.get(cursors[source_index])
.is_some_and(|run| run.cluster_id == cluster_id)
{
cursors[source_index] += 1;
}
}
records.push(RunRecord {
cluster_id,
doc_base: 0,
count: u32::try_from(count)
.map_err(|_| invalid_data("compacted ANN run exceeds u32 vectors"))?,
max_doc_id,
doc_ids_offset,
ordinals_offset,
codes_offset,
codes_len: u64::try_from(expected_codes_column_len(
AnnKind::BinaryIvf,
count,
code_size,
)?)
.map_err(|_| invalid_data("compacted ANN code length exceeds u64"))?,
});
}
if records.is_empty() {
return Err(invalid_data("cannot compact an ANN payload with no runs"));
}
finish_layout(writer, offset, &records)
}
#[cfg(feature = "native")]
fn run_doc_id_with_base(bytes: &[u8], run: &AnnRun, index: usize, base: u32) -> io::Result<u32> {
let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
if local_doc_id > run.max_doc_id {
return Err(invalid_data(
"ANN run contains a document above its declared maximum",
));
}
base.checked_add(local_doc_id)
.ok_or_else(|| invalid_data("compacted ANN document ID overflows u32"))
}
#[cfg(feature = "native")]
pub(crate) fn write_merged_ann_with_extra(
sources: &[(&AnnDiskIndex, u32)],
extra: &[BuildRun<'_>],
writer: &mut (impl Write + ?Sized),
cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> io::Result<u64> {
write_merged_ann_impl(sources, extra, writer, cancellation)
}
#[cfg(feature = "native")]
fn write_merged_ann_impl(
sources: &[(&AnnDiskIndex, u32)],
extra: &[BuildRun<'_>],
writer: &mut (impl Write + ?Sized),
cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> io::Result<u64> {
let Some((first, _)) = sources.first() else {
return Err(invalid_data("cannot merge an empty ANN source list"));
};
if first.header.kind == AnnKind::IvfTq
&& !crate::structures::is_ivf_tq_cosine_generation(first.header.quantizer_version)
{
return Err(invalid_data(
"legacy raw IVF-TQ generations cannot be merged; rebuild the index",
));
}
let mut header = first.header.clone();
header.vector_count = 0;
for &(source, _) in sources {
if !headers_compatible(&first.header, &source.header) {
return Err(invalid_data(
"ANN merge sources use incompatible generations",
));
}
header.vector_count = header
.vector_count
.checked_add(source.header.vector_count)
.ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
}
for run in extra {
if run.doc_ids.is_empty()
|| run.cluster_id >= header.num_clusters
|| run.ordinals.len() != run.doc_ids.len()
|| run.codes.len()
!= expected_codes_column_len(header.kind, run.doc_ids.len(), header.code_size)?
{
return Err(invalid_data("extra ANN merge run columns are inconsistent"));
}
header.vector_count = header
.vector_count
.checked_add(run.doc_ids.len())
.ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
}
validate_header(&header)?;
write_header(writer, &header)?;
let mut offset = ANN_HEADER_SIZE as u64;
let run_capacity = sources.iter().try_fold(extra.len(), |count, (source, _)| {
count
.checked_add(source.runs.len())
.ok_or_else(|| invalid_data("merged ANN run count overflows usize"))
})?;
let mut output_payload_starts = Vec::with_capacity(sources.len());
for &(source, _) in sources {
let payload_end = source
.runs
.iter()
.map(|run| run.codes.end)
.max()
.ok_or_else(|| invalid_data("ANN source has no payload runs"))?;
let output_payload_start = offset;
output_payload_starts.push(output_payload_start);
copy_range(
writer,
&source.raw,
ANN_HEADER_SIZE..payload_end,
cancellation,
)?;
offset = checked_advance(offset, payload_end - ANN_HEADER_SIZE)?;
}
let mut extra_records = Vec::with_capacity(extra.len());
let mut scratch = Vec::new();
for run in extra {
let count = run.doc_ids.len();
let doc_ids_offset = offset;
write_u32_column(writer, run.doc_ids, &mut scratch)?;
offset = checked_advance(offset, count * 4)?;
let ordinals_offset = offset;
write_u16_column(writer, run.ordinals, &mut scratch)?;
offset = checked_advance(offset, count * 2)?;
let codes_offset = offset;
writer.write_all(run.codes)?;
offset = checked_advance(offset, run.codes.len())?;
extra_records.push(RunRecord {
cluster_id: run.cluster_id,
doc_base: 0,
count: u32::try_from(count)
.map_err(|_| invalid_data("extra ANN run exceeds u32 vectors"))?,
max_doc_id: run.doc_ids.iter().copied().max().unwrap_or(0),
doc_ids_offset,
ordinals_offset,
codes_offset,
codes_len: u64::try_from(run.codes.len())
.map_err(|_| invalid_data("extra ANN code length exceeds u64"))?,
});
}
let directory_offset = offset;
let mut pending = BinaryHeap::with_capacity(sources.len() + 1);
for (source_index, (source, _)) in sources.iter().enumerate() {
pending.push(Reverse((source.runs[0].cluster_id, source_index, 0usize)));
}
if let Some(first_extra) = extra_records.first() {
pending.push(Reverse((first_extra.cluster_id, sources.len(), 0usize)));
}
let mut written_runs = 0usize;
while let Some(Reverse((_, source_index, run_index))) = pending.pop() {
if source_index == sources.len() {
write_run_record(writer, &extra_records[run_index])?;
written_runs += 1;
if let Some(next) = extra_records.get(run_index + 1) {
pending.push(Reverse((next.cluster_id, source_index, run_index + 1)));
}
continue;
}
let (source, segment_base) = sources[source_index];
let run = &source.runs[run_index];
write_run_record(
writer,
&RunRecord {
cluster_id: run.cluster_id,
doc_base: run
.doc_base
.checked_add(segment_base)
.ok_or_else(|| invalid_data("merged ANN document base overflows u32"))?,
count: u32::try_from(run.count)
.map_err(|_| invalid_data("ANN source run exceeds u32 vectors"))?,
max_doc_id: run.max_doc_id,
doc_ids_offset: relocate_payload_offset(
output_payload_starts[source_index],
run.doc_ids.start,
)?,
ordinals_offset: relocate_payload_offset(
output_payload_starts[source_index],
run.ordinals.start,
)?,
codes_offset: relocate_payload_offset(
output_payload_starts[source_index],
run.codes.start,
)?,
codes_len: u64::try_from(run.codes.len())
.map_err(|_| invalid_data("ANN source code length exceeds u64"))?,
},
)?;
written_runs = written_runs
.checked_add(1)
.ok_or_else(|| invalid_data("merged ANN run count overflows usize"))?;
let next_run_index = run_index + 1;
if let Some(next_run) = source.runs.get(next_run_index) {
pending.push(Reverse((next_run.cluster_id, source_index, next_run_index)));
}
}
if written_runs != run_capacity {
return Err(invalid_data("merged ANN directory lost source runs"));
}
finish_footer(writer, directory_offset, written_runs)
}
#[cfg(feature = "native")]
fn relocate_payload_offset(output_payload_start: u64, source_offset: usize) -> io::Result<u64> {
let relative = source_offset
.checked_sub(ANN_HEADER_SIZE)
.ok_or_else(|| invalid_data("ANN source offset precedes its payload"))?;
output_payload_start
.checked_add(
u64::try_from(relative)
.map_err(|_| invalid_data("ANN source relative offset exceeds u64"))?,
)
.ok_or_else(|| invalid_data("merged ANN payload offset overflows u64"))
}
#[cfg(feature = "native")]
fn headers_compatible(left: &AnnDiskHeader, right: &AnnDiskHeader) -> bool {
left.kind == right.kind
&& left.routing == right.routing
&& left.dim == right.dim
&& left.code_size == right.code_size
&& left.num_clusters == right.num_clusters
&& left.quantizer_version == right.quantizer_version
&& left.codebook_version == right.codebook_version
}
#[cfg(feature = "native")]
fn finish_layout(
writer: &mut (impl Write + ?Sized),
directory_offset: u64,
records: &[RunRecord],
) -> io::Result<u64> {
for record in records {
write_run_record(writer, record)?;
}
finish_footer(writer, directory_offset, records.len())
}
#[cfg(feature = "native")]
fn write_run_record(writer: &mut (impl Write + ?Sized), record: &RunRecord) -> io::Result<()> {
writer.write_u32::<LittleEndian>(record.cluster_id)?;
writer.write_u32::<LittleEndian>(record.doc_base)?;
writer.write_u32::<LittleEndian>(record.count)?;
writer.write_u32::<LittleEndian>(record.max_doc_id)?;
writer.write_u64::<LittleEndian>(record.doc_ids_offset)?;
writer.write_u64::<LittleEndian>(record.ordinals_offset)?;
writer.write_u64::<LittleEndian>(record.codes_offset)?;
writer.write_u64::<LittleEndian>(record.codes_len)?;
Ok(())
}
#[cfg(feature = "native")]
fn finish_footer(
writer: &mut (impl Write + ?Sized),
directory_offset: u64,
num_records: usize,
) -> io::Result<u64> {
writer.write_u64::<LittleEndian>(directory_offset)?;
writer.write_u64::<LittleEndian>(
u64::try_from(num_records).map_err(|_| invalid_data("ANN run count exceeds u64"))?,
)?;
writer.write_u32::<LittleEndian>(ANN_FOOTER_MAGIC)?;
writer.write_u32::<LittleEndian>(u32::from(ANN_DISK_VERSION))?;
let tail_size = num_records
.checked_mul(ANN_RUN_SIZE)
.and_then(|size| size.checked_add(ANN_FOOTER_SIZE))
.and_then(|size| u64::try_from(size).ok())
.ok_or_else(|| invalid_data("ANN final tail size overflows u64"))?;
directory_offset
.checked_add(tail_size)
.ok_or_else(|| invalid_data("ANN final size overflows u64"))
}
#[cfg(feature = "native")]
fn write_header(writer: &mut (impl Write + ?Sized), header: &AnnDiskHeader) -> io::Result<()> {
writer.write_u32::<LittleEndian>(ANN_HEADER_MAGIC)?;
writer.write_u8(header.kind as u8)?;
writer.write_u8(routing_to_u8(header.routing))?;
writer.write_u16::<LittleEndian>(ANN_DISK_VERSION)?;
writer.write_u32::<LittleEndian>(
u32::try_from(header.dim).map_err(|_| invalid_data("ANN dimension exceeds u32"))?,
)?;
writer.write_u32::<LittleEndian>(
u32::try_from(header.code_size).map_err(|_| invalid_data("ANN code size exceeds u32"))?,
)?;
writer.write_u32::<LittleEndian>(header.num_clusters)?;
writer.write_u32::<LittleEndian>(0)?;
writer.write_u64::<LittleEndian>(header.quantizer_version)?;
writer.write_u64::<LittleEndian>(header.codebook_version)?;
writer.write_u64::<LittleEndian>(
u64::try_from(header.vector_count)
.map_err(|_| invalid_data("ANN vector count exceeds u64"))?,
)?;
writer.write_u64::<LittleEndian>(0)?;
Ok(())
}
#[cfg(feature = "native")]
fn write_u32_column(
writer: &mut (impl Write + ?Sized),
values: &[u32],
scratch: &mut Vec<u8>,
) -> io::Result<()> {
for chunk in values.chunks(64 * 1024) {
scratch.clear();
scratch.reserve(chunk.len() * 4);
for value in chunk {
scratch.extend_from_slice(&value.to_le_bytes());
}
writer.write_all(scratch)?;
}
Ok(())
}
#[cfg(feature = "native")]
fn write_u16_column(
writer: &mut (impl Write + ?Sized),
values: &[u16],
scratch: &mut Vec<u8>,
) -> io::Result<()> {
for chunk in values.chunks(64 * 1024) {
scratch.clear();
scratch.reserve(chunk.len() * 2);
for value in chunk {
scratch.extend_from_slice(&value.to_le_bytes());
}
writer.write_all(scratch)?;
}
Ok(())
}
#[cfg(feature = "native")]
fn copy_range(
writer: &mut (impl Write + ?Sized),
bytes: &OwnedBytes,
range: Range<usize>,
cancellation: Option<&std::sync::atomic::AtomicBool>,
) -> io::Result<()> {
if range.is_empty() {
return Ok(());
}
let range_end = range.end;
let mut chunk_start = range.start;
let first_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
bytes.madvise_range(chunk_start..first_end, libc::MADV_WILLNEED);
while chunk_start < range_end {
if cancellation
.is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
{
return Err(io::Error::new(
io::ErrorKind::Interrupted,
"ANN merge copy cancelled",
));
}
let chunk_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
let next_end = chunk_end.saturating_add(COPY_CHUNK).min(range_end);
if chunk_end < next_end {
bytes.madvise_range(chunk_end..next_end, libc::MADV_WILLNEED);
}
writer.write_all(&bytes.as_slice()[chunk_start..chunk_end])?;
chunk_start = chunk_end;
}
Ok(())
}
#[cfg(feature = "native")]
fn checked_advance(offset: u64, length: usize) -> io::Result<u64> {
offset
.checked_add(
u64::try_from(length).map_err(|_| invalid_data("ANN copy length exceeds u64"))?,
)
.ok_or_else(|| invalid_data("ANN output offset overflows u64"))
}
fn validate_header(header: &AnnDiskHeader) -> io::Result<()> {
if header.kind == AnnKind::IvfTq
&& !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
{
return Err(invalid_data(
"IVF-TQ payload uses an unsupported legacy generation; rebuild the index",
));
}
if header.dim == 0
|| header.code_size == 0
|| header.num_clusters == 0
|| header.quantizer_version == 0
|| header.vector_count == 0
|| (header.kind == AnnKind::BinaryIvf
&& (header.codebook_version != 0
|| !header.dim.is_multiple_of(8)
|| header.code_size != header.dim.div_ceil(8)))
|| (header.kind == AnnKind::TqFlat
&& (header.codebook_version != 0
|| header.num_clusters != 1
|| header.routing != IvfRoutingMode::Flat
|| header.code_size * 2
!= crate::structures::vector::quantization::tq_padded_dim(header.dim)))
|| (header.kind == AnnKind::IvfTq
&& (header.codebook_version == 0
|| header.code_size * 2
!= crate::structures::vector::quantization::tq_padded_dim(header.dim)))
{
return Err(invalid_data("ANN header contains invalid metadata"));
}
Ok(())
}
fn read_u32(bytes: &[u8], offset: usize) -> u32 {
u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
}
fn read_u16(bytes: &[u8], offset: usize) -> u16 {
u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
}
#[inline]
fn tq_ivf_block_max_scale(block: &[u8]) -> f32 {
f32::from_le_bytes(
block[..size_of::<f32>()]
.try_into()
.expect("scale is one f32"),
)
}
fn run_doc_id(bytes: &[u8], run: &AnnRun, index: usize) -> io::Result<u32> {
let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
if local_doc_id > run.max_doc_id {
return Err(invalid_data(
"ANN run contains a document above its declared maximum",
));
}
run.doc_base
.checked_add(local_doc_id)
.ok_or_else(|| invalid_data("ANN run document ID overflows u32"))
}
#[cfg(feature = "native")]
fn routing_to_u8(routing: IvfRoutingMode) -> u8 {
match routing {
IvfRoutingMode::Auto => 0,
IvfRoutingMode::Flat => 1,
IvfRoutingMode::TwoLevel => 2,
IvfRoutingMode::Hnsw => 3,
}
}
fn routing_from_u8(value: u8) -> io::Result<IvfRoutingMode> {
match value {
0 => Ok(IvfRoutingMode::Auto),
1 => Ok(IvfRoutingMode::Flat),
2 => Ok(IvfRoutingMode::TwoLevel),
3 => Ok(IvfRoutingMode::Hnsw),
_ => Err(invalid_data(format!("unknown ANN routing mode {value}"))),
}
}
fn invalid_data(message: impl Into<String>) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message.into())
}
#[cfg(all(test, feature = "native"))]
mod tests {
use super::*;
#[test]
fn compacted_merge_matches_byte_copy_and_resets_fragmentation() {
let a0_docs = [0u32, 1, 2];
let a0_ords = [0u16, 0, 1];
let a0_codes = [0x11u8, 0x22, 0x33];
let a5_docs = [3u32, 4];
let a5_ords = [0u16; 2];
let a5_codes = [0x44u8, 0x55];
let a_runs = [
BuildRun {
cluster_id: 0,
doc_ids: &a0_docs,
ordinals: &a0_ords,
codes: &a0_codes,
},
BuildRun {
cluster_id: 5,
doc_ids: &a5_docs,
ordinals: &a5_ords,
codes: &a5_codes,
},
];
let mut header = binary_header(5);
header.num_clusters = 8;
let mut a_bytes = Vec::new();
write_built_runs(header.clone(), &a_runs, &mut a_bytes).unwrap();
let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 5).unwrap();
let b0_docs = [0u32, 2];
let b0_ords = [1u16, 0];
let b0_codes = [0x66u8, 0x77];
let b2_docs = [1u32];
let b2_ords = [0u16];
let b2_codes = [0x88u8];
let b_runs = [
BuildRun {
cluster_id: 0,
doc_ids: &b0_docs,
ordinals: &b0_ords,
codes: &b0_codes,
},
BuildRun {
cluster_id: 2,
doc_ids: &b2_docs,
ordinals: &b2_ords,
codes: &b2_codes,
},
];
let mut header_b = binary_header(3);
header_b.num_clusters = 8;
let mut b_bytes = Vec::new();
write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 3).unwrap();
let mut gen1_bytes = Vec::new();
write_merged_ann(&[(&a, 0), (&b, 5)], &mut gen1_bytes).unwrap();
let gen1 = AnnDiskIndex::open(OwnedBytes::new(gen1_bytes), AnnKind::BinaryIvf, 8).unwrap();
let sources: [(&AnnDiskIndex, u32); 2] = [(&gen1, 0), (&a, 8)];
let predicted = predicted_merge_fragmentation(&sources);
assert!((predicted - 2.0).abs() < 1e-9, "{predicted}");
let mut copied_bytes = Vec::new();
write_merged_ann(&sources, &mut copied_bytes).unwrap();
let copied =
AnnDiskIndex::open(OwnedBytes::new(copied_bytes), AnnKind::BinaryIvf, 13).unwrap();
let mut compacted_bytes = Vec::new();
write_compacted_ann_cancellable(&sources, &mut compacted_bytes, None).unwrap();
let compacted =
AnnDiskIndex::open(OwnedBytes::new(compacted_bytes), AnnKind::BinaryIvf, 13).unwrap();
let copied_health = copied.health();
let compacted_health = compacted.health();
assert!((copied_health.fragmentation() - 2.0).abs() < 1e-9);
assert!((compacted_health.fragmentation() - 1.0).abs() < 1e-9);
assert_eq!(compacted_health.runs, 3, "one run per non-empty cluster");
assert_eq!(copied_health.vectors, compacted_health.vectors);
assert_eq!(copied_health.payload_bytes, compacted_health.payload_bytes);
assert_eq!(
copied_health.largest_cluster_vectors,
compacted_health.largest_cluster_vectors
);
for cluster in 0..8u32 {
let query = [0x5Au8];
let from_copy = copied
.search_binary_clusters::<false>(&query, 16, &[cluster])
.unwrap();
let from_compact = compacted
.search_binary_clusters::<false>(&query, 16, &[cluster])
.unwrap();
assert_eq!(from_copy, from_compact, "cluster {cluster} diverged");
}
assert!(compacted.runs.iter().all(|run| run.doc_base == 0));
let mut generation3 = Vec::new();
write_merged_ann(&[(&compacted, 0), (&b, 13)], &mut generation3).unwrap();
let generation3 =
AnnDiskIndex::open(OwnedBytes::new(generation3), AnnKind::BinaryIvf, 16).unwrap();
assert_eq!(generation3.health().vectors, 16);
let all: Vec<(u32, u16, f32)> = compacted
.search_binary_clusters::<false>(&[0x5A], 32, &[0, 2, 5])
.unwrap();
let mut docs: Vec<u32> = all.iter().map(|&(doc, _, _)| doc).collect();
docs.sort_unstable();
assert_eq!(docs, (0..=12).collect::<Vec<u32>>());
}
#[test]
#[ignore]
fn ann_merge_throughput_byte_copy_vs_compaction() {
let code_size = 320usize;
let clusters = 4_096u32;
let vectors_per_source = 262_144usize;
let sources_count = 4usize;
let mut sources_bytes = Vec::new();
for source_index in 0..sources_count {
let mut per_cluster: Vec<(Vec<u32>, Vec<u16>, Vec<u8>)> = Vec::new();
let vectors_per_cluster = vectors_per_source / clusters as usize;
let mut doc = 0u32;
for cluster in 0..clusters {
let mut docs = Vec::with_capacity(vectors_per_cluster);
let mut ords = Vec::with_capacity(vectors_per_cluster);
let mut codes = Vec::with_capacity(vectors_per_cluster * code_size);
for _ in 0..vectors_per_cluster {
docs.push(doc);
ords.push(0u16);
codes.extend(std::iter::repeat_n(
(doc ^ cluster ^ source_index as u32) as u8,
code_size,
));
doc += 1;
}
per_cluster.push((docs, ords, codes));
}
let runs: Vec<BuildRun<'_>> = per_cluster
.iter()
.enumerate()
.map(|(cluster, (docs, ords, codes))| BuildRun {
cluster_id: cluster as u32,
doc_ids: docs,
ordinals: ords,
codes,
})
.collect();
let header = AnnDiskHeader {
kind: AnnKind::BinaryIvf,
routing: IvfRoutingMode::Hnsw,
dim: code_size * 8,
code_size,
num_clusters: clusters,
quantizer_version: 42,
codebook_version: 0,
vector_count: vectors_per_source,
};
let mut bytes = Vec::new();
write_built_runs(header, &runs, &mut bytes).unwrap();
sources_bytes.push(bytes);
}
let sources_open: Vec<AnnDiskIndex> = sources_bytes
.iter()
.map(|bytes| {
AnnDiskIndex::open(
OwnedBytes::new(bytes.clone()),
AnnKind::BinaryIvf,
(vectors_per_source * sources_count) as u32,
)
.unwrap()
})
.collect();
let sources: Vec<(&AnnDiskIndex, u32)> = sources_open
.iter()
.enumerate()
.map(|(index, source)| (source, (index * vectors_per_source) as u32))
.collect();
let payload_bytes = sources_bytes.iter().map(Vec::len).sum::<usize>();
let mut out = Vec::with_capacity(payload_bytes + (1 << 20));
let start = std::time::Instant::now();
write_merged_ann(&sources, &mut out).unwrap();
let copy_secs = start.elapsed().as_secs_f64();
out.clear();
let start = std::time::Instant::now();
write_compacted_ann_cancellable(&sources, &mut out, None).unwrap();
let compact_secs = start.elapsed().as_secs_f64();
let compacted = AnnDiskIndex::open(
OwnedBytes::new(out),
AnnKind::BinaryIvf,
(vectors_per_source * sources_count) as u32,
)
.unwrap();
assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
let gib = payload_bytes as f64 / (1u64 << 30) as f64;
println!(
"ann merge {:.2} GiB: byte-copy {:.3}s ({:.2} GiB/s), compaction {:.3}s \
({:.2} GiB/s), overhead {:.1}%",
gib,
copy_secs,
gib / copy_secs,
compact_secs,
gib / compact_secs,
100.0 * (compact_secs - copy_secs) / copy_secs,
);
}
#[test]
fn compaction_refuses_non_binary_payloads() {
let codec = std::sync::Arc::new(crate::structures::TqCodec::new(8));
let mut builder = crate::structures::TqFlatBuilder::new(codec);
builder
.add_batch(
&[(0, 0), (1, 0)],
&[
1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
],
)
.unwrap();
builder.finish();
let mut bytes = Vec::new();
write_built_tq_flat(&builder, &mut bytes).unwrap();
let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
let error = write_compacted_ann_cancellable(&[(&disk, 0)], &mut Vec::new(), None)
.expect_err("TQ payloads must not compact");
assert!(error.to_string().contains("binary"), "{error}");
}
#[test]
fn ann_health_measures_skew_and_fragmentation() {
let a0_docs = [0u32, 1, 2, 3, 4, 5];
let a0_ords = [0u16; 6];
let a0_codes = [0xAAu8; 6];
let a3_docs = [6u32, 7];
let a3_ords = [0u16; 2];
let a3_codes = [0x0Fu8; 2];
let a_runs = [
BuildRun {
cluster_id: 0,
doc_ids: &a0_docs,
ordinals: &a0_ords,
codes: &a0_codes,
},
BuildRun {
cluster_id: 3,
doc_ids: &a3_docs,
ordinals: &a3_ords,
codes: &a3_codes,
},
];
let mut header = binary_header(8);
header.num_clusters = 8;
let mut a_bytes = Vec::new();
write_built_runs(header, &a_runs, &mut a_bytes).unwrap();
let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 8).unwrap();
let b0_docs = [0u32, 1];
let b0_ords = [0u16; 2];
let b0_codes = [0xBBu8; 2];
let b_runs = [BuildRun {
cluster_id: 0,
doc_ids: &b0_docs,
ordinals: &b0_ords,
codes: &b0_codes,
}];
let mut header_b = binary_header(2);
header_b.num_clusters = 8;
let mut b_bytes = Vec::new();
write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 2).unwrap();
let mut merged_bytes = Vec::new();
write_merged_ann(&[(&a, 0), (&b, 8)], &mut merged_bytes).unwrap();
let disk =
AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::BinaryIvf, 10).unwrap();
let health = disk.health();
assert_eq!(health.vectors, 10);
assert_eq!(health.clusters_nonempty, 2);
assert_eq!(health.clusters_total, 8);
assert_eq!(health.runs, 3);
assert_eq!(health.largest_cluster, 0);
assert_eq!(health.largest_cluster_vectors, 8);
assert!((health.largest_cluster_share() - 0.8).abs() < 1e-9);
assert!((health.fragmentation() - 1.5).abs() < 1e-9);
assert!(
(health.imbalance - 1.36).abs() < 1e-9,
"{}",
health.imbalance
);
assert_eq!(health.payload_bytes, 10);
}
#[test]
fn ann_health_is_balanced_at_one() {
let docs: Vec<Vec<u32>> = (0..4).map(|c| vec![c * 2, c * 2 + 1]).collect();
let ords = [0u16; 2];
let codes = [0x55u8; 2];
let runs: Vec<BuildRun<'_>> = docs
.iter()
.enumerate()
.map(|(cluster, doc_ids)| BuildRun {
cluster_id: cluster as u32,
doc_ids,
ordinals: &ords,
codes: &codes,
})
.collect();
let mut header = binary_header(8);
header.num_clusters = 4;
let mut bytes = Vec::new();
write_built_runs(header, &runs, &mut bytes).unwrap();
let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 8).unwrap();
let health = disk.health();
assert!((health.imbalance - 1.0).abs() < 1e-9);
assert!((health.fragmentation() - 1.0).abs() < 1e-9);
assert!((health.largest_cluster_share() - 0.25).abs() < 1e-9);
}
fn binary_header(vector_count: usize) -> AnnDiskHeader {
AnnDiskHeader {
kind: AnnKind::BinaryIvf,
routing: IvfRoutingMode::Hnsw,
dim: 8,
code_size: 1,
num_clusters: 2,
quantizer_version: 42,
codebook_version: 0,
vector_count,
}
}
fn payload_end(index: &AnnDiskIndex) -> usize {
index.runs.iter().map(|run| run.codes.end).max().unwrap()
}
#[test]
fn ann_prefetch_ranges_are_sorted_and_only_merge_page_near_extents() {
let mut ranges = vec![
15_000..16_000,
0..1_000,
9_000..10_000,
1_000..2_000,
7_000..8_000,
];
coalesce_prefetch_ranges(&mut ranges);
assert_eq!(ranges, [0..2_000, 7_000..10_000, 15_000..16_000]);
}
#[test]
fn normal_merge_copies_ann_payload_columns_byte_for_byte() {
let first_doc_0 = [0u32];
let first_doc_1 = [1u32];
let first_ord_0 = [0u16];
let first_ord_1 = [2u16];
let first_code_0 = [0x00u8];
let first_code_1 = [0xffu8];
let first_runs = [
BuildRun {
cluster_id: 0,
doc_ids: &first_doc_0,
ordinals: &first_ord_0,
codes: &first_code_0,
},
BuildRun {
cluster_id: 1,
doc_ids: &first_doc_1,
ordinals: &first_ord_1,
codes: &first_code_1,
},
];
let mut first_bytes = Vec::new();
write_built_runs(binary_header(2), &first_runs, &mut first_bytes).unwrap();
let first = AnnDiskIndex::open(OwnedBytes::new(first_bytes.clone()), AnnKind::BinaryIvf, 2)
.unwrap();
let second_docs = [0u32, 1u32];
let second_ords = [1u16, 0u16];
let second_codes = [0x0fu8, 0xf0u8];
let second_runs = [BuildRun {
cluster_id: 0,
doc_ids: &second_docs,
ordinals: &second_ords,
codes: &second_codes,
}];
let mut second_bytes = Vec::new();
write_built_runs(binary_header(2), &second_runs, &mut second_bytes).unwrap();
let second =
AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::BinaryIvf, 2)
.unwrap();
let mut merged_bytes = Vec::new();
write_merged_ann(&[(&first, 0), (&second, 2)], &mut merged_bytes).unwrap();
let merged =
AnnDiskIndex::open(OwnedBytes::new(merged_bytes.clone()), AnnKind::BinaryIvf, 4)
.unwrap();
let mut expected_payload = first_bytes[ANN_HEADER_SIZE..payload_end(&first)].to_vec();
expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
assert_eq!(
&merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
expected_payload.as_slice(),
"normal merge must not decode or rewrite any corpus-sized ANN column",
);
let mut docs: Vec<u32> = merged
.search_binary_clusters::<false>(&[0], 4, &[0, 1])
.unwrap()
.into_iter()
.map(|result| result.0)
.collect();
docs.sort_unstable();
assert_eq!(docs, [0, 1, 2, 3]);
let mut second_merge_bytes = Vec::new();
write_merged_ann(&[(&merged, 0), (&first, 4)], &mut second_merge_bytes).unwrap();
let second_merge = AnnDiskIndex::open(
OwnedBytes::new(second_merge_bytes.clone()),
AnnKind::BinaryIvf,
6,
)
.unwrap();
let mut expected_second_payload =
merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)].to_vec();
expected_second_payload
.extend_from_slice(&first_bytes[ANN_HEADER_SIZE..payload_end(&first)]);
assert_eq!(
&second_merge_bytes[ANN_HEADER_SIZE..payload_end(&second_merge)],
expected_second_payload.as_slice(),
);
let mut docs: Vec<u32> = second_merge
.search_binary_clusters::<false>(&[0], 6, &[0, 1])
.unwrap()
.into_iter()
.map(|result| result.0)
.collect();
docs.sort_unstable();
assert_eq!(docs, [0, 1, 2, 3, 4, 5]);
}
#[test]
fn legacy_ivf_tq_payload_is_rejected_while_opening() {
let dim = 8;
let marked_version = crate::structures::mark_ivf_tq_cosine_generation(7);
let centroids = crate::structures::CoarseCentroids {
num_clusters: 1,
dim,
centroids: vec![0.0; dim],
version: marked_version,
soar_config: None,
routing_index: None,
};
let mut bytes = crate::segment::ann_build::build_ivf_tq(
dim,
IvfRoutingMode::Flat,
¢roids,
&[(0, 0)],
&[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
)
.unwrap();
bytes[24..32].copy_from_slice(&7u64.to_le_bytes());
let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
.err()
.expect("legacy IVF-TQ payload must fail while opening")
.to_string();
assert!(error.contains("unsupported legacy generation"), "{error}");
}
#[test]
fn binary_combined_search_deduplicates_soar_and_bounds_document_results() {
let cluster_0_docs = [0u32, 0, 1];
let cluster_0_ordinals = [0u16, 1, 0];
let cluster_0_codes = [0x00u8, 0xff, 0x03];
let cluster_1_docs = [0u32, 1, 2];
let cluster_1_ordinals = [0u16, 1, 0];
let cluster_1_codes = [0x00u8, 0x0c, 0xf0];
let runs = [
BuildRun {
cluster_id: 0,
doc_ids: &cluster_0_docs,
ordinals: &cluster_0_ordinals,
codes: &cluster_0_codes,
},
BuildRun {
cluster_id: 1,
doc_ids: &cluster_1_docs,
ordinals: &cluster_1_ordinals,
codes: &cluster_1_codes,
},
];
let mut bytes = Vec::new();
write_built_runs(binary_header(6), &runs, &mut bytes).unwrap();
let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap();
for combiner in [
crate::query::MultiValueCombiner::Sum,
crate::query::MultiValueCombiner::default(),
] {
let (result, probed) = disk
.search_binary_combined_documents(1, &[0], &[0, 1], combiner)
.unwrap();
assert_eq!(result.len(), 1, "combined search must honor k");
assert_eq!(
result[0].doc_id, 1,
"SOAR duplicate changed {combiner:?} ranking: {result:?}",
);
assert_eq!(
probed
.iter()
.map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
.collect::<Vec<_>>(),
vec![(1, 0), (1, 1)],
"{combiner:?}",
);
}
let (top_two, probed) = disk
.search_binary_combined_documents(
2,
&[0],
&[0, 1],
crate::query::MultiValueCombiner::Sum,
)
.unwrap();
assert_eq!(
top_two
.iter()
.map(|candidate| candidate.doc_id)
.collect::<Vec<_>>(),
vec![1, 0],
);
assert_eq!(top_two.len(), 2, "full probing must still return at most k");
assert_eq!(
probed
.iter()
.map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
.collect::<Vec<_>>(),
vec![(0, 0), (0, 1), (1, 0), (1, 1)],
);
}
#[test]
fn combined_ordinal_reduction_handles_out_of_order_runs_for_every_combiner() {
let out_of_order_with_duplicate = vec![
(7, 1, 0.4),
(3, 0, 0.8),
(7, 0, 0.6),
(3, 1, 0.2),
(7, 1, 0.5), ];
for combiner in [
crate::query::MultiValueCombiner::Max,
crate::query::MultiValueCombiner::Sum,
crate::query::MultiValueCombiner::Avg,
crate::query::MultiValueCombiner::default(),
crate::query::MultiValueCombiner::WeightedTopK { k: 2, decay: 0.7 },
] {
let actual = combine_scored_ordinals(out_of_order_with_duplicate.clone(), 2, combiner);
let mut expected = vec![
AnnDocumentCandidate {
doc_id: 3,
score: combiner.combine(&[(0, 0.8), (1, 0.2)]),
},
AnnDocumentCandidate {
doc_id: 7,
score: combiner.combine(&[(0, 0.6), (1, 0.5)]),
},
];
expected.sort_unstable_by(|left, right| {
right
.score
.total_cmp(&left.score)
.then_with(|| left.doc_id.cmp(&right.doc_id))
});
assert_eq!(actual, expected, "combiner {combiner:?}");
}
}
fn build_tq_payload(dim: usize, count: usize, seed: u64) -> (Vec<u8>, Vec<Vec<f32>>) {
let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
let mut state = seed;
let mut vectors = Vec::new();
let mut flat = Vec::new();
for _ in 0..count {
let vector: Vec<f32> = (0..dim)
.map(|_| {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
})
.collect();
flat.extend_from_slice(&vector);
vectors.push(vector);
}
let labels: Vec<(u32, u16)> = (0..count).map(|index| (index as u32, 0)).collect();
builder.add_batch(&labels, &flat).unwrap();
builder.finish();
let mut bytes = Vec::new();
write_built_tq_flat(&builder, &mut bytes).unwrap();
(bytes, vectors)
}
#[test]
fn tq_combined_scan_ranks_complete_documents_instead_of_individual_values() {
let dim = 8;
let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
let labels = [(0u32, 0u16), (1, 0), (1, 1)];
let vectors = [
1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.8, 0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.8, -0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
];
builder.add_batch(&labels, &vectors).unwrap();
builder.finish();
let mut bytes = Vec::new();
write_built_tq_flat(&builder, &mut bytes).unwrap();
let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
let query = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
let plan = crate::structures::TqQueryPlan::build(&codec, &query);
let max = disk.search_tq_distinct(1, &plan).unwrap();
let sum = disk
.search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
.unwrap();
let default_combiner = disk
.search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::default())
.unwrap();
assert_eq!(max[0].0, 0, "fixture must favor doc 0 by Max: {max:?}");
assert_eq!(
sum[0].doc_id, 1,
"two complete values must make doc 1 win by Sum: {sum:?}"
);
assert_eq!(
default_combiner[0].doc_id, 1,
"default LogSumExp must aggregate complete documents: {default_combiner:?}"
);
assert!(sum[0].score > max[0].2);
let (single_bytes, _) = build_tq_payload(dim, 1, 91);
let single = AnnDiskIndex::open(OwnedBytes::new(single_bytes), AnnKind::TqFlat, 1).unwrap();
let mut merged_bytes = Vec::new();
write_merged_ann(&[(&disk, 1), (&single, 0)], &mut merged_bytes).unwrap();
let merged = AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::TqFlat, 3).unwrap();
let merged_sum = merged
.search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
.unwrap();
assert_eq!(merged_sum[0].doc_id, 2);
}
#[test]
fn tq_payload_roundtrip_search_and_pure_copy_merge() {
let dim = 20; let count = 21;
let (bytes, vectors) = build_tq_payload(dim, count, 42);
let index = AnnDiskIndex::open(
OwnedBytes::new(bytes.clone()),
AnnKind::TqFlat,
count as u32,
)
.unwrap();
assert_eq!(index.header().vector_count, count);
let codec = crate::structures::TqCodec::new(dim);
for target in [0usize, 7, 20] {
let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[target]);
let results = index.search_tq_distinct(3, &plan).unwrap();
assert_eq!(
results[0].0, target as u32,
"query duplicating vector {target} must rank it first: {results:?}"
);
}
let (second_bytes, _) = build_tq_payload(dim, 5, 77);
let second =
AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::TqFlat, 5).unwrap();
let mut merged_bytes = Vec::new();
write_merged_ann(&[(&index, 0), (&second, count as u32)], &mut merged_bytes).unwrap();
let merged = AnnDiskIndex::open(
OwnedBytes::new(merged_bytes.clone()),
AnnKind::TqFlat,
count as u32 + 5,
)
.unwrap();
let mut expected_payload = bytes[ANN_HEADER_SIZE..payload_end(&index)].to_vec();
expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
assert_eq!(
&merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
expected_payload.as_slice(),
"TQ merge must be a pure byte copy of the source columns",
);
let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[7]);
let results = merged.search_tq_distinct(1, &plan).unwrap();
assert_eq!(results[0].0, 7, "merged payload must keep doc bases");
}
#[test]
fn tq_parallel_fold_matches_a_sequential_scan() {
use crate::structures::vector::quantization::{
TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
};
let dim = 8;
let count = TQ_PARALLEL_SCAN_MIN_VECTORS;
let (bytes, vectors) = build_tq_payload(dim, count, 87);
let disk =
AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, count as u32).unwrap();
let codec = crate::structures::TqCodec::new(dim);
let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[count / 3]);
let k = 31;
let block_bytes = tq_block_bytes(disk.header().code_size);
assert_eq!(
disk.runs.len(),
1,
"the regression must parallelize chunks inside one run"
);
assert!(
disk.runs[0].codes.len() / block_bytes > TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
"the single run must span multiple parallel chunks"
);
let raw = disk.raw.as_slice();
let mut reference = BoundedAnnCollector::<true, true>::new(k);
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for run in &disk.runs {
let codes = &raw[run.codes.clone()];
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
tq_score_block(&plan, block, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let index = lane_base + lane;
reference.insert(
run_doc_id(raw, run, index).unwrap(),
read_u16(raw, run.ordinals.start + index * 2),
score,
);
}
}
}
assert_eq!(
disk.search_tq_distinct(k, &plan).unwrap(),
reference.into_sorted_results(),
);
}
#[test]
fn ivf_tq_scale_pruning_matches_the_unpruned_scan() {
use crate::structures::vector::ivf::{CoarseCentroids, CoarseConfig};
use crate::structures::vector::quantization::{
TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
};
use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
let dim = 32;
let count = 400usize;
let codec = std::sync::Arc::new(TqCodec::new(dim));
let mut state = 5u64;
let mut next = move || {
state = state
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
};
let vectors: Vec<Vec<f32>> = (0..count)
.map(|_| {
let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
v.iter_mut().for_each(|x| *x /= norm);
v
})
.collect();
let mut centroids = CoarseCentroids::train(&CoarseConfig::new(dim, 8), &vectors, "test");
centroids.version = crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
let mut index = IvfTqIndex::new(
dim,
crate::dsl::IvfRoutingMode::Flat,
centroids.version,
std::sync::Arc::clone(&codec),
);
let mut scratch = TqIvfEncodeScratch::default();
for (i, vector) in vectors.iter().enumerate() {
index.add_vector(
¢roids,
(i / 2) as u32,
(i % 2) as u16,
vector,
&mut scratch,
);
}
let mut bytes = Vec::new();
write_built_ivf_tq(&index, centroids.num_clusters, &mut bytes).unwrap();
let disk =
AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, count as u32).unwrap();
let block_bytes = tq_ivf_block_bytes(disk.header().code_size);
let raw = disk.raw.as_slice();
for run in &disk.runs {
let codes = &raw[run.codes.clone()];
let mut previous_scale = f32::INFINITY;
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
let mut block_scales = block[..TQ_BLOCK_LANES * size_of::<f32>()]
.chunks_exact(size_of::<f32>())
.take(lanes)
.map(|lane| f32::from_le_bytes(lane.try_into().unwrap()));
let first_scale = block_scales.next().unwrap();
assert_eq!(tq_ivf_block_max_scale(block), first_scale);
assert!(first_scale <= previous_scale);
previous_scale = first_scale;
for scale in block_scales {
assert!(scale <= previous_scale);
previous_scale = scale;
}
}
}
let k = 10;
for query_seed in [1u64, 9, 42] {
let mut qstate = query_seed;
let mut qnext = move || {
qstate = qstate
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((qstate >> 33) as f32 / (1u64 << 31) as f32) - 0.5
};
let query: Vec<f32> = (0..dim).map(|_| qnext()).collect();
let plan = TqIvfQueryPlan::build(
¢roids,
&codec,
&query,
8,
crate::dsl::IvfRoutingMode::Flat,
);
let mut reference = BoundedAnnCollector::<true, true>::new(k);
let mut unpruned_scores = Vec::new();
let mut scores = [0.0f32; TQ_BLOCK_LANES];
for (cluster_id, cluster_dot) in plan.cluster_dots() {
for run in disk.cluster_runs(cluster_id) {
let codes = &raw[run.codes.clone()];
for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
tq_score_ivf_block(plan.tq_plan(), block, cluster_dot, &mut scores);
let lane_base = block_index * TQ_BLOCK_LANES;
let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
for (lane, &score) in scores.iter().enumerate().take(lanes) {
let idx = lane_base + lane;
reference.insert(
run_doc_id(raw, run, idx).unwrap(),
read_u16(raw, run.ordinals.start + idx * 2),
score,
);
unpruned_scores.push((
run_doc_id(raw, run, idx).unwrap(),
read_u16(raw, run.ordinals.start + idx * 2),
score,
));
}
}
}
}
let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap();
let reference = reference.into_sorted_results();
assert_eq!(
pruned, reference,
"scale-bound pruning must not change the estimated top-k (seed {query_seed})"
);
for combiner in [
crate::query::MultiValueCombiner::Max,
crate::query::MultiValueCombiner::Sum,
crate::query::MultiValueCombiner::Avg,
crate::query::MultiValueCombiner::default(),
crate::query::MultiValueCombiner::WeightedTopK { k: 3, decay: 0.7 },
] {
let expected = combine_scored_ordinals(unpruned_scores.clone(), k, combiner);
let combined = disk
.search_ivf_tq_combined_documents(k, &plan, combiner)
.unwrap();
assert_eq!(
combined, expected,
"combined IVF-TQ scan diverged from the unpruned reference \
for {combiner:?} (seed {query_seed})",
);
assert!(combined.len() <= k);
}
}
}
#[test]
fn open_rejects_tq_payload_with_inconsistent_geometry() {
let (bytes, _) = build_tq_payload(20, 4, 9);
let mut corrupted = bytes.clone();
corrupted[12..16].copy_from_slice(&15u32.to_le_bytes());
assert!(
AnnDiskIndex::open(OwnedBytes::new(corrupted), AnnKind::TqFlat, 4).is_err(),
"TQ header with code_size != padded_dim/2 must be refused"
);
let (short_bytes, _) = build_tq_payload(20, 4, 9);
let mut wrong_kind = short_bytes.clone();
wrong_kind[4] = AnnKind::BinaryIvf as u8;
assert!(
AnnDiskIndex::open(OwnedBytes::new(wrong_kind), AnnKind::BinaryIvf, 4).is_err(),
"TQ block-padded columns must not validate under another kind"
);
let (legacy_bytes, _) = build_tq_payload(20, 4, 9);
let mut legacy_kind = legacy_bytes.clone();
legacy_kind[4] = 1;
let Err(error) = AnnDiskIndex::open(OwnedBytes::new(legacy_kind), AnnKind::TqFlat, 4)
else {
panic!("retired IVF-PQ payloads must not open");
};
assert!(
error.to_string().contains("IVF-PQ"),
"error must name the retired format: {error}"
);
assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 4).is_ok());
}
#[test]
fn open_rejects_old_or_out_of_range_ann_payloads() {
let mut legacy = vec![0u8; ANN_HEADER_SIZE + ANN_FOOTER_SIZE];
legacy[..4].copy_from_slice(b"old!");
assert!(AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::BinaryIvf, 1).is_err());
let docs = [0u32];
let ordinals = [0u16];
let codes = [0u8];
let runs = [BuildRun {
cluster_id: 0,
doc_ids: &docs,
ordinals: &ordinals,
codes: &codes,
}];
let mut bytes = Vec::new();
write_built_runs(binary_header(1), &runs, &mut bytes).unwrap();
let footer = bytes.len() - ANN_FOOTER_SIZE;
let directory = usize::try_from(u64::from_le_bytes(
bytes[footer..footer + 8].try_into().unwrap(),
))
.unwrap();
bytes[directory + 12..directory + 16].copy_from_slice(&10u32.to_le_bytes());
assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 1).is_err());
}
}