use std::{
cmp::Ordering,
fs::{self, File, OpenOptions, metadata},
io::{self, BufReader, BufWriter, Error as IoError, ErrorKind, Read, Seek, SeekFrom, Write},
mem::size_of,
path::{Path, PathBuf},
sync::Arc,
};
use rayon::prelude::*;
use tempfile::{tempdir, tempdir_in};
use crate::superfile::{
BuildError,
format::{
self, FST_SEPARATOR, RESERVED_PREFIX,
checksum::{crc32c, crc32c_append},
vec::{
CELL_DIR_ENTRY_SIZE, CLUSTER_IDX_COUNT_OFFSET, CLUSTER_IDX_ENTRY_BYTES, MAGIC_BYTES,
U32_BYTES, U64_BYTES, cell_dir_entry, sub_hdr,
},
},
vector::{
cell_posting::{MaterializedIvfRow, sq8_residual_norm_sq},
distance::{
Metric, distance, encode_sq16_row, mean_f32_cluster_major, normalize,
sq16_decoded_norm_sq,
},
ivf_merge::MergedIvfSubsection,
kmeans::{assign_to_centroids, kmeans, kmeans_with_assignments},
quant::BitQuantizer,
rerank_codec::{RerankCodec, SQ8_FIXED_OFFSET, SQ8_FIXED_SCALE},
reservoir::{Reservoir, default_kmeans_sample_size, partition_kmeans_sample_size},
rotation::RandomRotation,
spill::{
ChunkedVectorSource, InMemoryVectorSource, MmapVectorSource, SpillWriter,
SpilledCellRows,
},
sq8_simd::{Sq8EncodeConsts, encode_sq8_residual_row, update_min_max},
},
};
const OUTER_HEADER_SIZE: usize = format::vec::OUTER_HEADER_SIZE;
const DIR_ENTRY_SIZE: usize = format::vec::DIR_ENTRY_SIZE;
const SUB_HEADER_SIZE: usize = format::vec::SUB_HEADER_SIZE;
const VECTOR_DIM_MIN: usize = 16;
const VECTOR_DIM_MAX: usize = 4096;
const RESERVOIR_SEED_XOR_MASK: u64 = 0x5a5a_5a5a_5a5a_5a5a;
const KMEANS_ITERS: usize = 5;
const FINE_RUN_SPLIT_BOUND_FACTOR: usize = 2;
const FINE_RUN_SPLIT_MAX_ROUNDS: usize = 4;
const FINE_RUN_SPLIT_SEED_OFFSET: u64 = 101;
const CONSOLIDATED_CELL_ROWS_THRESHOLD: usize = 80_000;
const PASS2_CHUNK_MEM_BUDGET_BYTES: usize = 128 << 20;
const PASS2_CHUNK_ROWS_MIN: usize = 1024;
const PASS2_CHUNK_ROWS_MAX: usize = 65_536;
const MATERIALIZED_BUCKET_CHUNK_BYTES: usize = 16 << 20;
const MATERIALIZED_ASSIGN_BYTES_PER_DIM: usize = 2 + size_of::<f32>();
const N_CENT_LARGE_DOC_THRESHOLD: usize = 5_000_000;
const N_CENT_LARGE: usize = 4096;
const N_CENT_MEDIUM_DOC_THRESHOLD: usize = 100_000;
const N_CENT_MEDIUM: usize = 1024;
const N_CENT_SMALL: usize = 64;
const SQ8_CODE_MAX: f32 = 255.0;
fn n_cent_row_count_cap(n_docs: usize) -> usize {
if n_docs >= N_CENT_LARGE_DOC_THRESHOLD {
N_CENT_LARGE
} else if n_docs >= N_CENT_MEDIUM_DOC_THRESHOLD {
N_CENT_MEDIUM
} else {
N_CENT_SMALL
}
}
fn metric_id(m: Metric) -> u32 {
match m {
Metric::L2Sq => format::vec::METRIC_ID_L2SQ,
Metric::Cosine => format::vec::METRIC_ID_COSINE,
Metric::NegDot => format::vec::METRIC_ID_NEGDOT,
}
}
#[derive(Debug, Clone)]
pub struct VectorConfig {
pub column: String,
pub dim: usize,
pub rot_seed: u64,
pub metric: Metric,
pub rerank_codec: RerankCodec,
pub provided_centroids: Option<std::sync::Arc<[f32]>>,
}
fn default_rerank_codec_for(metric: Metric) -> RerankCodec {
if metric == Metric::Cosine {
RerankCodec::default()
} else {
RerankCodec::Sq8Residual
}
}
impl VectorConfig {
pub fn new(column: String, dim: usize, rot_seed: u64, metric: Metric) -> Self {
Self {
column,
dim,
rot_seed,
metric,
rerank_codec: default_rerank_codec_for(metric),
provided_centroids: None,
}
}
#[must_use]
pub fn with_rerank_codec(mut self, codec: RerankCodec) -> Self {
self.rerank_codec = codec;
self
}
#[must_use]
pub fn with_provided_centroids(mut self, centroids: Option<std::sync::Arc<[f32]>>) -> Self {
self.provided_centroids = centroids;
self
}
}
const DEFAULT_SPILL_THRESHOLD_BYTES: usize = 256 * 1024 * 1024;
struct ColumnState {
config: VectorConfig,
n_docs: u32,
reservoir: Reservoir,
unit_scratch: Vec<f32>,
pre_spill_buffer: Vec<f32>,
spill: Option<SpillWriter>,
spill_threshold_bytes: usize,
materialized_rows: Option<Vec<MaterializedIvfRow>>,
prebuilt_subsection: Option<SubsectionBytes>,
inline_stable_ids: Option<Vec<i128>>,
}
#[derive(Default)]
struct ScratchDir {
parent: Option<PathBuf>,
tempdir: Option<tempfile::TempDir>,
}
impl ScratchDir {
fn in_parent(parent: PathBuf) -> Result<Self, BuildError> {
let meta = metadata(&parent)?;
if !meta.is_dir() {
return Err(BuildError::Io(IoError::new(
ErrorKind::InvalidInput,
format!("VectorBuilder scratch path is not a directory: {parent:?}"),
)));
}
Ok(Self {
parent: Some(parent),
tempdir: None,
})
}
fn path(&mut self) -> Result<&Path, BuildError> {
if self.tempdir.is_none() {
let tmp = if let Some(parent) = &self.parent {
tempfile::TempDir::new_in(parent)?
} else {
tempfile::tempdir()?
};
self.tempdir = Some(tmp);
}
Ok(self
.tempdir
.as_ref()
.expect("scratch tempdir initialized")
.path())
}
}
pub struct VectorBuilder {
columns: Vec<ColumnState>,
scratch_dir: ScratchDir,
spill_threshold_bytes: usize,
}
impl Default for VectorBuilder {
fn default() -> Self {
Self::new()
}
}
impl VectorBuilder {
pub fn new() -> Self {
Self {
columns: Vec::new(),
scratch_dir: ScratchDir::default(),
spill_threshold_bytes: DEFAULT_SPILL_THRESHOLD_BYTES,
}
}
pub fn with_scratch(scratch: PathBuf) -> Result<Self, BuildError> {
Ok(Self {
columns: Vec::new(),
scratch_dir: ScratchDir::in_parent(scratch)?,
spill_threshold_bytes: DEFAULT_SPILL_THRESHOLD_BYTES,
})
}
pub fn set_spill_threshold_bytes(&mut self, threshold: usize) {
self.spill_threshold_bytes = threshold;
}
pub fn register_column(&mut self, config: VectorConfig) -> Result<u32, BuildError> {
if config.column.as_bytes().contains(&FST_SEPARATOR) {
return Err(BuildError::ReservedSeparatorInColumnName(config.column));
}
if config.column.starts_with(RESERVED_PREFIX) {
return Err(BuildError::ReservedPrefixInColumnName(config.column));
}
if !(VECTOR_DIM_MIN..=VECTOR_DIM_MAX).contains(&config.dim) {
return Err(BuildError::VectorDimOutOfRange {
column: config.column.clone(),
dim: config.dim,
});
}
if self
.columns
.iter()
.any(|c| c.config.column == config.column)
{
return Err(BuildError::DuplicateColumnName(config.column));
}
if !config.rerank_codec.is_implemented() {
return Err(BuildError::VectorRerankCodecUnimplemented {
column: config.column.clone(),
codec: config.rerank_codec.name(),
});
}
if !config.rerank_codec.supports_metric(config.metric) {
return Err(BuildError::VectorSchemaMismatch(format!(
"vector index {:?}: codec {} supports cosine metric only",
config.column,
config.rerank_codec.name()
)));
}
let column_id = self.columns.len() as u32;
let sample_size = default_kmeans_sample_size(N_CENT_LARGE);
let reservoir_seed = config.rot_seed ^ RESERVOIR_SEED_XOR_MASK;
let reservoir = Reservoir::new(sample_size, config.dim, reservoir_seed);
let spill_threshold_bytes = self.spill_threshold_bytes;
self.columns.push(ColumnState {
config,
n_docs: 0,
reservoir,
unit_scratch: Vec::new(),
pre_spill_buffer: Vec::new(),
spill: None,
spill_threshold_bytes,
materialized_rows: None,
prebuilt_subsection: None,
inline_stable_ids: None,
});
Ok(column_id)
}
#[allow(dead_code)]
pub(crate) fn load_materialized_rows(
&mut self,
column_id: u32,
rows: Vec<MaterializedIvfRow>,
) -> Result<(), BuildError> {
let idx = column_id as usize;
let col = self
.columns
.get_mut(idx)
.ok_or_else(|| BuildError::FtsColumnTypeInvalid {
column: format!("(unregistered vector column_id {column_id})"),
actual: "n/a".to_string(),
})?;
if !col.config.rerank_codec.is_ivf_mergeable() {
return Err(BuildError::VectorRerankCodecUnimplemented {
column: col.config.column.clone(),
codec: col.config.rerank_codec.name(),
});
}
col.n_docs = rows.len() as u32;
col.materialized_rows = Some(rows);
Ok(())
}
pub(crate) fn set_prebuilt_subsection(
&mut self,
column_id: u32,
subsection: MergedIvfSubsection,
) -> Result<(), BuildError> {
let idx = column_id as usize;
let col = self
.columns
.get_mut(idx)
.ok_or_else(|| BuildError::FtsColumnTypeInvalid {
column: format!("(unregistered vector column_id {column_id})"),
actual: "n/a".to_string(),
})?;
if subsection.rerank_codec != col.config.rerank_codec {
return Err(BuildError::VectorSchemaMismatch(format!(
"prebuilt subsection codec {} does not match destination codec {}",
subsection.rerank_codec.name(),
col.config.rerank_codec.name()
)));
}
col.n_docs = subsection.n_docs;
col.materialized_rows = None;
col.prebuilt_subsection = Some(SubsectionBytes {
bytes: subsection.bytes,
n_cent: subsection.n_cent,
summary_offset_in_sub: subsection.summary_offset_in_sub,
codec_meta_offset_in_sub: subsection.codec_meta_offset_in_sub,
codec_meta_size: subsection.codec_meta_size,
});
Ok(())
}
pub fn set_kmeans_sample_size(
&mut self,
column_id: u32,
sample_size: usize,
) -> Result<(), BuildError> {
let idx = column_id as usize;
if idx >= self.columns.len() {
return Err(BuildError::FtsColumnTypeInvalid {
column: format!("(unregistered vector column_id {column_id})"),
actual: "n/a".to_string(),
});
}
let col = &mut self.columns[idx];
let reservoir_seed = col.config.rot_seed ^ RESERVOIR_SEED_XOR_MASK;
col.reservoir = Reservoir::new(sample_size, col.config.dim, reservoir_seed);
Ok(())
}
pub fn add(&mut self, column_id: u32, vec: &[f32]) -> Result<(), BuildError> {
let idx = column_id as usize;
if idx >= self.columns.len() {
return Err(BuildError::FtsColumnTypeInvalid {
column: format!("(unregistered vector column_id {column_id})"),
actual: "n/a".to_string(),
});
}
{
let col = &mut self.columns[idx];
if vec.len() != col.config.dim {
return Err(BuildError::FtsColumnTypeInvalid {
column: col.config.column.clone(),
actual: format!("vec.len()={} != dim={}", vec.len(), col.config.dim),
});
}
let vec: &[f32] = if col.config.metric == Metric::Cosine {
col.unit_scratch.clear();
col.unit_scratch.extend_from_slice(vec);
normalize(&mut col.unit_scratch);
&col.unit_scratch
} else {
vec
};
col.reservoir.update(vec);
let vec_bytes = vec.len() * 4;
let buf_bytes = col.pre_spill_buffer.len() * 4;
if let Some(spill) = col.spill.as_mut() {
spill.write_vec(vec)?;
col.n_docs += 1;
return Ok(());
}
if buf_bytes + vec_bytes <= col.spill_threshold_bytes {
col.pre_spill_buffer.extend_from_slice(vec);
col.n_docs += 1;
return Ok(());
}
}
let path = self
.scratch_dir
.path()?
.join(format!("infino_input_spill_col{column_id}.bin"));
let col = &mut self.columns[idx];
let mut spill = SpillWriter::create(path)?;
spill.write_all(bytemuck::cast_slice(&col.pre_spill_buffer))?;
spill.write_vec(vec)?;
col.pre_spill_buffer = Vec::new();
col.spill = Some(spill);
col.n_docs += 1;
Ok(())
}
pub fn finish(self) -> Result<Vec<u8>, BuildError> {
let header_dir_hint = OUTER_HEADER_SIZE + (self.columns.len() * DIR_ENTRY_SIZE) + 8;
let mut buf: Vec<u8> = Vec::with_capacity(header_dir_hint);
self.finish_to(&mut buf)?;
Ok(buf)
}
pub fn finish_to<W: Write>(self, mut w: W) -> Result<(), BuildError> {
let VectorBuilder {
columns,
mut scratch_dir,
spill_threshold_bytes: _,
} = self;
let n_columns = columns.len() as u32;
let n_docs: u64 = columns.iter().map(|c| c.n_docs as u64).max().unwrap_or(0);
let column_configs: Vec<(VectorConfig, u32)> = columns
.iter()
.map(|c| (c.config.clone(), c.n_docs))
.collect();
let mut subsections: Vec<SubsectionBytes> = Vec::with_capacity(columns.len());
if !columns.is_empty() {
let scratch_path = scratch_dir.path()?.to_path_buf();
for (col_idx, col) in columns.into_iter().enumerate() {
if let Some(prebuilt) = col.prebuilt_subsection {
subsections.push(prebuilt);
continue;
}
subsections.push(build_subsection_streaming(
col_idx as u32,
col,
&scratch_path,
)?);
}
}
let directory_offset = OUTER_HEADER_SIZE as u64;
let directory_size = (n_columns as usize) * DIR_ENTRY_SIZE;
let mut subsection_start_off =
directory_offset + directory_size as u64 + format::CRC_BYTES as u64;
let mut directory: Vec<u8> = Vec::with_capacity(directory_size);
for (i, sub) in subsections.iter().enumerate() {
let (cfg, _) = &column_configs[i];
let summary_offset_abs = subsection_start_off + sub.summary_offset_in_sub as u64;
directory.extend_from_slice(&(i as u32).to_le_bytes()); directory.extend_from_slice(&(cfg.dim as u32).to_le_bytes()); directory.extend_from_slice(&(sub.n_cent as u32).to_le_bytes()); directory.extend_from_slice(&metric_id(cfg.metric).to_le_bytes()); directory.extend_from_slice(&cfg.rot_seed.to_le_bytes()); directory.extend_from_slice(&subsection_start_off.to_le_bytes()); directory.extend_from_slice(&(sub.bytes.len() as u64).to_le_bytes()); directory.extend_from_slice(&summary_offset_abs.to_le_bytes()); directory.extend_from_slice(&((cfg.dim * 4) as u32).to_le_bytes()); directory.push(cfg.rerank_codec.codec_id()); directory.extend_from_slice(&[0u8; 3]); directory.extend_from_slice(&(sub.codec_meta_offset_in_sub as u32).to_le_bytes());
directory.extend_from_slice(&(sub.codec_meta_size as u32).to_le_bytes());
debug_assert_eq!(directory.len() % DIR_ENTRY_SIZE, 0);
subsection_start_off += sub.bytes.len() as u64;
}
let dir_crc = crc32c(&directory);
let mut outer_header: [u8; OUTER_HEADER_SIZE] = [0; OUTER_HEADER_SIZE];
{
let mut cursor = &mut outer_header[..];
cursor
.write_all(format::vec::OUTER_MAGIC) .map_err(BuildError::Io)?;
cursor
.write_all(&format::vec::VERSION.to_le_bytes()) .map_err(BuildError::Io)?;
cursor
.write_all(&n_columns.to_le_bytes()) .map_err(BuildError::Io)?;
cursor
.write_all(&n_docs.to_le_bytes()) .map_err(BuildError::Io)?;
cursor
.write_all(&directory_offset.to_le_bytes()) .map_err(BuildError::Io)?;
debug_assert!(cursor.is_empty());
}
let mut outer_crc_acc: u32 = 0;
w.write_all(&outer_header).map_err(BuildError::Io)?;
outer_crc_acc = crc32c_append(outer_crc_acc, &outer_header);
w.write_all(&directory).map_err(BuildError::Io)?;
outer_crc_acc = crc32c_append(outer_crc_acc, &directory);
let dir_crc_le = dir_crc.to_le_bytes();
w.write_all(&dir_crc_le).map_err(BuildError::Io)?;
outer_crc_acc = crc32c_append(outer_crc_acc, &dir_crc_le);
drop(directory);
for sub in subsections.drain(..) {
w.write_all(&sub.bytes).map_err(BuildError::Io)?;
outer_crc_acc = crc32c_append(outer_crc_acc, &sub.bytes);
}
let outer_crc_le = outer_crc_acc.to_le_bytes();
w.write_all(&outer_crc_le).map_err(BuildError::Io)?;
drop(scratch_dir);
Ok(())
}
}
pub(crate) trait MultiCellSubsectionSource {
fn cell_id(&self) -> u32;
fn n_docs(&self) -> u32;
fn len(&self) -> u64;
fn rerank_codec(&self) -> RerankCodec;
fn write_to(&self, output: &mut dyn Write) -> Result<(), BuildError>;
}
struct BorrowedMultiCellSubsection<'a> {
cell_id: u32,
subsection: &'a MergedIvfSubsection,
}
impl MultiCellSubsectionSource for BorrowedMultiCellSubsection<'_> {
fn cell_id(&self) -> u32 {
self.cell_id
}
fn n_docs(&self) -> u32 {
self.subsection.n_docs
}
fn len(&self) -> u64 {
self.subsection.bytes.len() as u64
}
fn rerank_codec(&self) -> RerankCodec {
self.subsection.rerank_codec
}
fn write_to(&self, output: &mut dyn Write) -> Result<(), BuildError> {
output
.write_all(&self.subsection.bytes)
.map_err(BuildError::Io)
}
}
struct CrcWriter<'a, W> {
output: &'a mut W,
crc: u32,
}
impl<W: Write> Write for CrcWriter<'_, W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let written = self.output.write(buf)?;
self.crc = crc32c_append(self.crc, &buf[..written]);
Ok(written)
}
fn flush(&mut self) -> io::Result<()> {
self.output.flush()
}
}
pub(crate) fn finish_multi_cell_blob_to<W, S>(cells: &[S], mut output: W) -> Result<(), BuildError>
where
W: Write,
S: MultiCellSubsectionSource,
{
if cells.is_empty() {
return Err(BuildError::VectorSchemaMismatch(
"multi-cell vector blob requires at least one cell IVF".into(),
));
}
let packed_codec = cells[0].rerank_codec();
if cells.iter().any(|cell| cell.rerank_codec() != packed_codec) {
return Err(BuildError::VectorSchemaMismatch(
"multi-cell blob cannot mix rerank codecs".into(),
));
}
for pair in cells.windows(2) {
if pair[0].cell_id() >= pair[1].cell_id() {
return Err(BuildError::VectorSchemaMismatch(
"multi-cell cells must be sorted by unique ascending cell_id".into(),
));
}
}
let n_cells = cells.len() as u32;
let n_docs: u64 = cells.iter().map(|cell| u64::from(cell.n_docs())).sum();
let directory_offset = OUTER_HEADER_SIZE as u64;
let directory_size = cells.len() * CELL_DIR_ENTRY_SIZE;
let mut subsection_start = directory_offset + directory_size as u64 + format::CRC_BYTES as u64;
let mut directory = Vec::with_capacity(directory_size);
for cell in cells {
directory.extend_from_slice(&cell.cell_id().to_le_bytes());
directory.extend_from_slice(&subsection_start.to_le_bytes());
directory.extend_from_slice(&cell.len().to_le_bytes());
directory.extend_from_slice(&u32::from(cell.rerank_codec().codec_id()).to_le_bytes());
debug_assert_eq!(directory.len() % CELL_DIR_ENTRY_SIZE, 0);
let _ = cell_dir_entry::CELL_ID_OFF;
subsection_start += cell.len();
}
let dir_crc = crc32c(&directory);
let mut outer_header = [0u8; OUTER_HEADER_SIZE];
{
let mut cursor = &mut outer_header[..];
cursor
.write_all(format::vec::OUTER_MAGIC)
.map_err(BuildError::Io)?;
cursor
.write_all(&format::vec::VERSION_MULTI_CELL.to_le_bytes())
.map_err(BuildError::Io)?;
cursor
.write_all(&n_cells.to_le_bytes())
.map_err(BuildError::Io)?;
cursor
.write_all(&n_docs.to_le_bytes())
.map_err(BuildError::Io)?;
cursor
.write_all(&directory_offset.to_le_bytes())
.map_err(BuildError::Io)?;
debug_assert!(cursor.is_empty());
}
let outer_crc = {
let mut crc_output = CrcWriter {
output: &mut output,
crc: 0,
};
crc_output
.write_all(&outer_header)
.map_err(BuildError::Io)?;
crc_output.write_all(&directory).map_err(BuildError::Io)?;
crc_output
.write_all(&dir_crc.to_le_bytes())
.map_err(BuildError::Io)?;
for cell in cells {
cell.write_to(&mut crc_output)?;
}
crc_output.flush().map_err(BuildError::Io)?;
crc_output.crc
};
output
.write_all(&outer_crc.to_le_bytes())
.map_err(BuildError::Io)?;
output.flush().map_err(BuildError::Io)?;
Ok(())
}
pub(crate) fn finish_multi_cell_blob(
cells: &[(u32, MergedIvfSubsection)],
) -> Result<Vec<u8>, BuildError> {
let sources: Vec<BorrowedMultiCellSubsection<'_>> = cells
.iter()
.map(|(cell_id, subsection)| BorrowedMultiCellSubsection {
cell_id: *cell_id,
subsection,
})
.collect();
let capacity = OUTER_HEADER_SIZE
+ cells.len() * CELL_DIR_ENTRY_SIZE
+ 2 * format::CRC_BYTES
+ cells
.iter()
.map(|(_, subsection)| subsection.bytes.len())
.sum::<usize>();
let mut output = Vec::with_capacity(capacity);
finish_multi_cell_blob_to(&sources, &mut output)?;
Ok(output)
}
struct SubsectionBytes {
bytes: Vec<u8>,
n_cent: usize,
summary_offset_in_sub: usize,
codec_meta_offset_in_sub: usize,
codec_meta_size: usize,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct StreamedIvfSubsection {
pub(crate) n_docs: u32,
pub(crate) rerank_codec: RerankCodec,
pub(crate) subsection_len: u64,
pub(crate) n_cent: usize,
pub(crate) summary_offset_in_sub: usize,
pub(crate) codec_meta_offset_in_sub: usize,
pub(crate) codec_meta_size: usize,
}
const BUCKET_BUF_SIZE: usize = 64 * 1024;
fn chunk_rows_for_dim(dim: usize) -> usize {
let cap_by_mem = PASS2_CHUNK_MEM_BUDGET_BYTES / (dim.max(1) * 4);
cap_by_mem.clamp(PASS2_CHUNK_ROWS_MIN, PASS2_CHUNK_ROWS_MAX)
}
pub(crate) mod build_phase_timers {
use std::{
sync::{
OnceLock,
atomic::{AtomicU64, Ordering},
},
time::Instant,
};
use crate::config;
pub static TRAIN_US: AtomicU64 = AtomicU64::new(0);
pub static ASSIGN_US: AtomicU64 = AtomicU64::new(0);
pub static CALIB_US: AtomicU64 = AtomicU64::new(0);
pub fn enabled() -> bool {
static ON: OnceLock<bool> = OnceLock::new();
*ON.get_or_init(|| config::global().diagnostics.drain_build_timers)
}
pub fn timed<T>(counter: &AtomicU64, f: impl FnOnce() -> T) -> T {
if !enabled() {
return f();
}
let t = Instant::now();
let out = f();
counter.fetch_add(t.elapsed().as_micros() as u64, Ordering::Relaxed);
out
}
pub fn reset() {
TRAIN_US.store(0, Ordering::Relaxed);
ASSIGN_US.store(0, Ordering::Relaxed);
CALIB_US.store(0, Ordering::Relaxed);
}
pub fn snapshot_ms() -> (f64, f64, f64) {
let ms = |c: &AtomicU64| c.load(Ordering::Relaxed) as f64 / 1000.0;
(ms(&TRAIN_US), ms(&ASSIGN_US), ms(&CALIB_US))
}
}
fn materialized_centroids(
cfg: &VectorConfig,
requested_n_cent: usize,
n_docs: usize,
sample: &[f32],
) -> (usize, Vec<f32>) {
let dim = cfg.dim;
if let Some(global) = cfg.provided_centroids.as_ref() {
debug_assert!(dim > 0 && global.len() % dim == 0);
let n_cent = global.len() / dim.max(1);
return (n_cent, global.to_vec());
}
let consolidated = n_docs > CONSOLIDATED_CELL_ROWS_THRESHOLD;
let requested = if consolidated {
requested_n_cent.max(1).min(n_docs.max(1))
} else {
requested_n_cent
.max(1)
.min(n_cent_row_count_cap(n_docs))
.min(n_docs.max(1))
};
let (mut centroids, assignments) =
kmeans_with_assignments(sample, dim, requested, KMEANS_ITERS, cfg.rot_seed);
let n_cent = split_oversized_fine_runs(
&mut centroids,
sample,
dim,
requested,
cfg.rot_seed,
Some(assignments),
);
order_centroids_geometrically(&mut centroids, dim, n_cent);
(n_cent, centroids)
}
fn split_oversized_fine_runs(
centroids: &mut Vec<f32>,
sample: &[f32],
dim: usize,
requested: usize,
seed: u64,
initial_assignments: Option<Vec<u32>>,
) -> usize {
let mut n_cent = centroids.len() / dim.max(1);
let sample_n = sample.len() / dim.max(1);
if n_cent <= 1 || sample_n == 0 {
return n_cent;
}
let target = sample_n.div_ceil(requested.max(1)).max(1);
let bound = target.saturating_mul(FINE_RUN_SPLIT_BOUND_FACTOR);
let (mut assignments, mut need_assign) = match initial_assignments {
Some(a) if a.len() == sample_n && a.iter().all(|&idx| (idx as usize) < n_cent) => {
(a, false)
}
_ => (vec![0u32; sample_n], true),
};
for round in 0..FINE_RUN_SPLIT_MAX_ROUNDS {
if need_assign {
assign_to_centroids(sample, centroids, dim, n_cent, &mut assignments);
}
need_assign = true;
let mut counts = vec![0usize; n_cent];
for &a in &assignments {
counts[a as usize] += 1;
}
let oversized: Vec<usize> = (0..n_cent).filter(|&c| counts[c] > bound).collect();
if oversized.is_empty() {
break;
}
for &c in &oversized {
let members: Vec<usize> = (0..sample_n)
.filter(|&r| assignments[r] as usize == c)
.collect();
let k = members.len().div_ceil(target).max(2).min(members.len());
let mut rows = Vec::with_capacity(members.len() * dim);
for &r in &members {
rows.extend_from_slice(&sample[r * dim..(r + 1) * dim]);
}
let sub_seed = seed
.wrapping_add(FINE_RUN_SPLIT_SEED_OFFSET)
.wrapping_add(((round as u64) << u32::BITS) | c as u64);
let sub = kmeans(&rows, dim, k, KMEANS_ITERS, sub_seed);
centroids[c * dim..(c + 1) * dim].copy_from_slice(&sub[..dim]);
centroids.extend_from_slice(&sub[dim..]);
}
n_cent = centroids.len() / dim;
}
n_cent
}
fn order_centroids_geometrically(centroids: &mut [f32], dim: usize, n_cent: usize) {
if n_cent <= 2 || centroids.len() != n_cent * dim {
return;
}
let mean = mean_f32_cluster_major(centroids, dim, n_cent);
let dist = |a: &[f32], c: usize| distance(Metric::L2Sq, a, ¢roids[c * dim..(c + 1) * dim]);
let mut visited = vec![false; n_cent];
let mut order = Vec::with_capacity(n_cent);
let mut current = (0..n_cent)
.min_by(|&a, &b| dist(&mean, a).total_cmp(&dist(&mean, b)))
.unwrap_or(0);
visited[current] = true;
order.push(current);
while order.len() < n_cent {
let tail = centroids[current * dim..(current + 1) * dim].to_vec();
let next = (0..n_cent)
.filter(|&c| !visited[c])
.min_by(|&a, &b| dist(&tail, a).total_cmp(&dist(&tail, b)))
.expect("unvisited centroid remains");
visited[next] = true;
order.push(next);
current = next;
}
let mut reordered = vec![0.0f32; centroids.len()];
for (new_idx, &old_idx) in order.iter().enumerate() {
reordered[new_idx * dim..(new_idx + 1) * dim]
.copy_from_slice(¢roids[old_idx * dim..(old_idx + 1) * dim]);
}
centroids.copy_from_slice(&reordered);
}
fn build_subsection_from_materialized(
cfg: VectorConfig,
requested_n_cent: usize,
mut rows: Vec<MaterializedIvfRow>,
) -> Result<SubsectionBytes, BuildError> {
rows.sort_by_key(|r| r.local_doc_id);
let merged =
build_cell_subsection_in_memory(cfg, requested_n_cent, CellPackSource::Rows(rows))?;
Ok(SubsectionBytes {
bytes: merged.bytes,
n_cent: merged.n_cent,
summary_offset_in_sub: merged.summary_offset_in_sub,
codec_meta_offset_in_sub: merged.codec_meta_offset_in_sub,
codec_meta_size: merged.codec_meta_size,
})
}
pub(crate) fn build_merged_subsection_from_materialized(
cfg: VectorConfig,
requested_n_cent: usize,
rows: Vec<MaterializedIvfRow>,
) -> Result<MergedIvfSubsection, BuildError> {
let n_docs = rows.len() as u32;
let rerank_codec = cfg.rerank_codec;
let sub = build_subsection_from_materialized(cfg, requested_n_cent, rows)?;
Ok(MergedIvfSubsection {
bytes: sub.bytes,
n_cent: sub.n_cent,
n_docs,
rerank_codec,
summary_offset_in_sub: sub.summary_offset_in_sub,
codec_meta_offset_in_sub: sub.codec_meta_offset_in_sub,
codec_meta_size: sub.codec_meta_size,
})
}
fn materialized_chunk_rows_for_dim(dim: usize) -> usize {
let row_bytes = dim.max(1).saturating_mul(MATERIALIZED_ASSIGN_BYTES_PER_DIM);
(PASS2_CHUNK_MEM_BUDGET_BYTES / row_bytes).clamp(PASS2_CHUNK_ROWS_MIN, PASS2_CHUNK_ROWS_MAX)
}
#[inline]
fn splitmix64(mut x: u64) -> u64 {
x = x.wrapping_add(0x9e37_79b9_7f4a_7c15);
x = (x ^ (x >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
x = (x ^ (x >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
x ^ (x >> 31)
}
const SAMPLE_JITTER_SEED_XOR: u64 = 0xa5a5_5a5a_c3c3_3c3c;
#[inline]
fn sampled_index(s: usize, sample_size: usize, n_docs: usize, seed: u64) -> usize {
if sample_size >= n_docs {
return s;
}
let base = s * n_docs / sample_size;
if n_docs <= CONSOLIDATED_CELL_ROWS_THRESHOLD {
return base;
}
let next = ((s + 1) * n_docs / sample_size).min(n_docs);
let width = next.saturating_sub(base).max(1) as u64;
let jitter = splitmix64(seed ^ SAMPLE_JITTER_SEED_XOR ^ (s as u64)) % width;
base + jitter as usize
}
enum BucketRecordPayload<'a> {
FixedSq8 {
codes: &'a [u8],
residuals: &'a [u8],
},
FixedPlane(&'a [u8]),
Fp32(&'a [f32]),
}
#[allow(clippy::too_many_arguments)]
fn write_bucket_record(
cluster: usize,
local_doc_id: u32,
stable_id: i128,
rabitq_code: &[u8],
payload: BucketRecordPayload<'_>,
dim: usize,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
stable_ids: &mut BufWriter<File>,
min_max: &mut Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
stable_ids.write_all(&stable_id.to_le_bytes())?;
let writer = &mut bucket_writers[cluster];
writer.write_all(&local_doc_id.to_le_bytes())?;
writer.write_all(rabitq_code)?;
match payload {
BucketRecordPayload::FixedSq8 { codes, residuals } => {
writer.write_all(codes)?;
writer.write_all(residuals)?;
}
BucketRecordPayload::FixedPlane(plane) => {
writer.write_all(plane)?;
}
BucketRecordPayload::Fp32(fp) => {
writer.write_all(bytemuck::cast_slice(fp))?;
if let Some((min, max)) = min_max.as_mut() {
let offset = cluster * dim;
update_min_max(
fp,
&mut min[offset..offset + dim],
&mut max[offset..offset + dim],
);
}
}
}
bucket_counts[cluster] = bucket_counts[cluster].saturating_add(1);
Ok(())
}
fn sample_spilled_materialized_rows(
spill: &SpilledCellRows,
sample_size: usize,
chunk_rows: usize,
seed: u64,
) -> Result<Vec<f32>, BuildError> {
if sample_size == 0 {
return Ok(Vec::new());
}
let n_docs = spill.n_rows();
let dim = spill.dim();
let targets: Vec<usize> = (0..sample_size)
.map(|sample| sampled_index(sample, sample_size, n_docs, seed))
.collect();
let mut sample = vec![0.0f32; sample_size * dim];
let mut reader = spill.reader()?;
let mut row_base = 0usize;
let mut target_idx = 0usize;
while row_base < n_docs {
let rows = reader.next_chunk(chunk_rows)?;
if rows.is_empty() {
break;
}
let row_end = row_base + rows.len();
while target_idx < targets.len() && targets[target_idx] < row_end {
let row = &rows[targets[target_idx] - row_base];
if row.encoded.rerank_codec != spill.rerank_codec() {
return Err(BuildError::VectorSchemaMismatch(
"materialized spill mixes rerank codecs".into(),
));
}
row.encoded
.rerank_codec
.ops()
.expect("materialized spill uses a quantized-rerank codec")
.dequantize_row_into(
&row.encoded.codes,
&row.encoded.residuals,
dim,
&row.encoded.scale,
&row.encoded.offset,
&mut sample[target_idx * dim..(target_idx + 1) * dim],
);
target_idx += 1;
}
row_base = row_end;
}
if target_idx != sample_size {
return Err(BuildError::VectorSchemaMismatch(format!(
"materialized spill yielded {target_idx} of {sample_size} training rows"
)));
}
Ok(sample)
}
#[allow(clippy::too_many_arguments)]
fn bucket_encoded_rows_chunk(
rows: &[MaterializedIvfRow],
base_local: u32,
cfg: &VectorConfig,
centroids: &[f32],
n_cent: usize,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
stable_ids: &mut BufWriter<File>,
min_max: &mut Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
let dim = cfg.dim;
let code_bytes = dim.div_ceil(u8::BITS as usize);
let fixed = cfg.rerank_codec.uses_fixed_quantizer();
let single_plane = cfg.rerank_codec.is_sq16();
let ops = cfg
.rerank_codec
.ops()
.expect("materialized rebuild uses a quantized-rerank codec");
for row in rows {
if row.encoded.rerank_codec != cfg.rerank_codec || row.rabitq_code.len() != code_bytes {
return Err(BuildError::VectorSchemaMismatch(
"materialized row does not match destination vector config".into(),
));
}
}
let mut decoded = vec![0.0f32; rows.len() * dim];
decoded
.par_chunks_mut(dim)
.zip(rows.par_iter())
.for_each(|(out, row)| {
ops.dequantize_row_into(
&row.encoded.codes,
&row.encoded.residuals,
dim,
&row.encoded.scale,
&row.encoded.offset,
out,
);
});
let mut assignments = vec![0u32; rows.len()];
assign_to_centroids(&decoded, centroids, dim, n_cent, &mut assignments);
for (row_idx, (row, &cluster)) in rows.iter().zip(&assignments).enumerate() {
let payload = if single_plane {
BucketRecordPayload::FixedPlane(&row.encoded.codes)
} else if fixed {
BucketRecordPayload::FixedSq8 {
codes: &row.encoded.codes,
residuals: &row.encoded.residuals,
}
} else {
BucketRecordPayload::Fp32(&decoded[row_idx * dim..(row_idx + 1) * dim])
};
write_bucket_record(
cluster as usize,
base_local + row_idx as u32,
row.stable_id,
&row.rabitq_code,
payload,
dim,
bucket_writers,
bucket_counts,
stable_ids,
min_max,
)?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn stream_materialized_rows_to_buckets(
spill: &SpilledCellRows,
cfg: &VectorConfig,
centroids: &[f32],
n_cent: usize,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
stable_ids: &mut BufWriter<File>,
sq8_min_max: Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
let chunk_rows = materialized_chunk_rows_for_dim(cfg.dim);
let mut min_max = sq8_min_max;
let mut reader = spill.reader()?;
let mut next_local = 0u32;
while next_local < spill.n_rows() as u32 {
let rows = reader.next_chunk(chunk_rows)?;
if rows.is_empty() {
break;
}
bucket_encoded_rows_chunk(
&rows,
next_local,
cfg,
centroids,
n_cent,
bucket_writers,
bucket_counts,
stable_ids,
&mut min_max,
)?;
next_local += rows.len() as u32;
}
if next_local as usize != spill.n_rows() {
return Err(BuildError::VectorSchemaMismatch(format!(
"materialized spill streamed {next_local} of {} rows",
spill.n_rows()
)));
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn stream_ram_rows_to_buckets(
rows: &[MaterializedIvfRow],
cfg: &VectorConfig,
centroids: &[f32],
n_cent: usize,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
stable_ids: &mut BufWriter<File>,
sq8_min_max: Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
let chunk_rows = materialized_chunk_rows_for_dim(cfg.dim);
let mut min_max = sq8_min_max;
for (chunk_idx, chunk) in rows.chunks(chunk_rows).enumerate() {
bucket_encoded_rows_chunk(
chunk,
(chunk_idx * chunk_rows) as u32,
cfg,
centroids,
n_cent,
bucket_writers,
bucket_counts,
stable_ids,
&mut min_max,
)?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn stream_fp32_rows_to_buckets(
vectors: &[f32],
stable_ids_in: &[i128],
cfg: &VectorConfig,
centroids: &[f32],
n_cent: usize,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
stable_ids: &mut BufWriter<File>,
sq8_min_max: Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
let dim = cfg.dim;
let n_docs = vectors.len() / dim;
let fixed = cfg.rerank_codec.uses_fixed_quantizer();
let single_plane = cfg.rerank_codec.is_sq16();
let mut min_max = sq8_min_max;
let rotation = RandomRotation::new(dim, cfg.rot_seed);
let quant = BitQuantizer::new(dim);
let code_bytes = quant.code_bytes();
let (fixed_scale, fixed_offset) = fixed_sq8_quantizer(dim);
let residual_encode = (fixed && !single_plane).then(|| {
(
Sq8EncodeConsts::from_scale_offset(&fixed_scale, &fixed_offset),
cfg.rerank_codec
.residual_divisor()
.expect("residual-family codec has divisor"),
)
});
let chunk_rows = materialized_chunk_rows_for_dim(dim);
let mut chunk_rotated = vec![0.0f32; chunk_rows * dim];
let mut chunk_codes = vec![0u8; chunk_rows * code_bytes];
let mut chunk_payload = if fixed {
vec![0u8; chunk_rows * dim * 2]
} else {
Vec::new()
};
let mut row_base = 0usize;
while row_base < n_docs {
let take = (n_docs - row_base).min(chunk_rows);
let chunk = &vectors[row_base * dim..(row_base + take) * dim];
let mut assignments = vec![0u32; take];
assign_to_centroids(chunk, centroids, dim, n_cent, &mut assignments);
chunk_rotated[..take * dim]
.par_chunks_mut(dim)
.zip(chunk.par_chunks(dim))
.for_each(|(dst, src)| rotation.apply(src, dst));
chunk_codes[..take * code_bytes]
.par_chunks_mut(code_bytes)
.zip(chunk_rotated[..take * dim].par_chunks(dim))
.for_each(|(code, rot)| quant.encode_rotated_into(rot, code));
if let Some((encode_consts, divisor)) = &residual_encode {
chunk_payload[..take * dim * 2]
.par_chunks_mut(dim * 2)
.zip(chunk.par_chunks(dim))
.for_each_init(
|| vec![0.0f32; dim],
|recon, (payload, row)| {
let (code_out, residual_out) = payload.split_at_mut(dim);
encode_sq8_residual_row(
row,
encode_consts,
&fixed_scale,
&fixed_offset,
code_out,
residual_out,
recon,
false,
*divisor,
);
},
);
} else if single_plane {
chunk_payload[..take * dim * 2]
.par_chunks_mut(dim * 2)
.zip(chunk.par_chunks(dim))
.for_each(|(payload, row)| encode_sq16_row(row, payload));
}
for i in 0..take {
let payload = if single_plane {
BucketRecordPayload::FixedPlane(&chunk_payload[i * dim * 2..(i + 1) * dim * 2])
} else if fixed {
let (codes, residuals) =
chunk_payload[i * dim * 2..(i + 1) * dim * 2].split_at(dim);
BucketRecordPayload::FixedSq8 { codes, residuals }
} else {
BucketRecordPayload::Fp32(&chunk[i * dim..(i + 1) * dim])
};
write_bucket_record(
assignments[i] as usize,
(row_base + i) as u32,
stable_ids_in[row_base + i],
&chunk_codes[i * code_bytes..(i + 1) * code_bytes],
payload,
dim,
bucket_writers,
bucket_counts,
stable_ids,
&mut min_max,
)?;
}
row_base += take;
}
Ok(())
}
fn write_at(file: &mut File, offset: usize, bytes: &[u8]) -> Result<(), BuildError> {
file.seek(SeekFrom::Start(offset as u64))?;
file.write_all(bytes)?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn stream_bucket_into_subsection(
output: &mut File,
bucket_path: &Path,
block: &ClusterBlock,
code_bytes: usize,
dim: usize,
codec: RerankCodec,
scale: &[f32],
offset: &[f32],
norms_offset: Option<usize>,
) -> Result<(), BuildError> {
let fixed = codec.uses_fixed_quantizer();
let single_plane = codec.is_sq16();
let payload_bytes = if fixed {
dim * 2
} else {
dim * size_of::<f32>()
};
let record_bytes = format::vec::DOC_ID_BYTES + code_bytes + payload_bytes;
let chunk_rows = (MATERIALIZED_BUCKET_CHUNK_BYTES / record_bytes.max(1)).max(1);
let mut reader = BufReader::new(File::open(bucket_path)?);
let mut rows_done = 0usize;
let encode_consts = (!fixed).then(|| Sq8EncodeConsts::from_scale_offset(scale, offset));
let mut recon = vec![0.0f32; dim];
let mut fp_row = vec![0.0f32; dim];
while rows_done < block.count {
let take = (block.count - rows_done).min(chunk_rows);
let mut records = vec![0u8; take * record_bytes];
reader.read_exact(&mut records)?;
let mut ids = vec![0u8; take * format::vec::DOC_ID_BYTES];
let mut codes = vec![0u8; take * code_bytes];
let mut rerank = vec![0u8; take * dim * 2];
let mut norms = norms_offset.map(|_| vec![0u8; take * size_of::<f32>()]);
for row_idx in 0..take {
let record = &records[row_idx * record_bytes..(row_idx + 1) * record_bytes];
let id_end = format::vec::DOC_ID_BYTES;
let code_end = id_end + code_bytes;
ids[row_idx * id_end..(row_idx + 1) * id_end].copy_from_slice(&record[..id_end]);
codes[row_idx * code_bytes..(row_idx + 1) * code_bytes]
.copy_from_slice(&record[id_end..code_end]);
let rerank_row = &mut rerank[row_idx * dim * 2..(row_idx + 1) * dim * 2];
let norm = if single_plane {
rerank_row.copy_from_slice(&record[code_end..code_end + dim * 2]);
norms_offset.map(|_| sq16_decoded_norm_sq(rerank_row, dim))
} else if fixed {
rerank_row.copy_from_slice(&record[code_end..code_end + dim * 2]);
norms_offset.map(|_| {
sq8_residual_norm_sq(
scale,
offset,
&rerank_row[..dim],
&rerank_row[dim..],
codec
.residual_divisor()
.expect("fixed residual codec has divisor"),
)
})
} else {
for (value, bytes) in fp_row
.iter_mut()
.zip(record[code_end..].chunks_exact(size_of::<f32>()))
{
*value = f32::from_le_bytes(bytes.try_into().expect("4-byte f32 bucket value"));
}
let (code_out, residual_out) = rerank_row.split_at_mut(dim);
encode_sq8_residual_row(
&fp_row,
encode_consts
.as_ref()
.expect("non-fixed materialized bucket has encode constants"),
scale,
offset,
code_out,
residual_out,
&mut recon,
norms_offset.is_some(),
codec
.residual_divisor()
.expect("residual-family codec has divisor"),
)
};
if let (Some(norm), Some(norm_bytes)) = (norm, norms.as_mut()) {
let start = row_idx * size_of::<f32>();
norm_bytes[start..start + size_of::<f32>()].copy_from_slice(&norm.to_le_bytes());
}
}
write_at(output, block.codes_base + rows_done * code_bytes, &codes)?;
write_at(
output,
block.ids_base + rows_done * format::vec::DOC_ID_BYTES,
&ids,
)?;
write_at(output, block.rerank_base + rows_done * dim * 2, &rerank)?;
if let (Some(norms_base), Some(norm_bytes)) = (norms_offset, norms) {
write_at(
output,
norms_base + (block.first_row + rows_done) * size_of::<f32>(),
&norm_bytes,
)?;
}
rows_done += take;
}
Ok(())
}
pub(crate) enum CellPackSource<'a> {
Spilled(&'a SpilledCellRows),
Rows(Vec<MaterializedIvfRow>),
Fp32 {
vectors: &'a [f32],
stable_ids: &'a [i128],
},
}
impl CellPackSource<'_> {
fn n_docs(&self, dim: usize) -> usize {
match self {
Self::Spilled(spill) => spill.n_rows(),
Self::Rows(rows) => rows.len(),
Self::Fp32 { vectors, .. } => vectors.len() / dim.max(1),
}
}
}
fn sample_ram_materialized_rows(
rows: &[MaterializedIvfRow],
sample_size: usize,
dim: usize,
seed: u64,
) -> Vec<f32> {
let n_docs = rows.len();
let mut sample = vec![0.0f32; sample_size * dim];
for s in 0..sample_size {
let idx = sampled_index(s, sample_size, n_docs, seed);
let enc = &rows[idx].encoded;
enc.rerank_codec
.ops()
.expect("materialized source uses a quantized-rerank codec")
.dequantize_row_into(
&enc.codes,
&enc.residuals,
dim,
&enc.scale,
&enc.offset,
&mut sample[s * dim..(s + 1) * dim],
);
}
sample
}
fn sample_fp32_rows(vectors: &[f32], sample_size: usize, dim: usize, seed: u64) -> Vec<f32> {
let n_docs = vectors.len() / dim.max(1);
let mut sample = vec![0.0f32; sample_size * dim];
for s in 0..sample_size {
let idx = sampled_index(s, sample_size, n_docs, seed);
sample[s * dim..(s + 1) * dim].copy_from_slice(&vectors[idx * dim..(idx + 1) * dim]);
}
sample
}
pub(crate) fn build_cell_subsection_from_source(
cfg: VectorConfig,
requested_n_cent: usize,
source: CellPackSource<'_>,
subsection_path: &Path,
stable_ids_path: &Path,
scratch: &Path,
) -> Result<StreamedIvfSubsection, BuildError> {
let dim = cfg.dim;
if dim == 0 {
return Err(BuildError::VectorSchemaMismatch(
"cell IVF build requires dim > 0".into(),
));
}
let n_docs = source.n_docs(dim);
if n_docs == 0 {
return Err(BuildError::VectorSchemaMismatch(
"cell IVF build requires at least one row".into(),
));
}
if !cfg.rerank_codec.is_ivf_mergeable() || !cfg.rerank_codec.supports_metric(cfg.metric) {
return Err(BuildError::VectorSchemaMismatch(format!(
"cell IVF build does not support codec {} with metric {:?}",
cfg.rerank_codec.name(),
cfg.metric
)));
}
match &source {
CellPackSource::Spilled(spill) => {
if spill.dim() != dim || cfg.rerank_codec != spill.rerank_codec() {
return Err(BuildError::VectorSchemaMismatch(
"streamed materialized IVF codec or shape mismatch".into(),
));
}
}
CellPackSource::Rows(rows) => {
if rows
.iter()
.any(|row| row.encoded.rerank_codec != cfg.rerank_codec)
{
return Err(BuildError::VectorSchemaMismatch(
"materialized IVF rebuild requires one matching residual-family codec".into(),
));
}
}
CellPackSource::Fp32 {
vectors,
stable_ids,
} => {
if !vectors.len().is_multiple_of(dim) {
return Err(BuildError::VectorSchemaMismatch(format!(
"fp32 corpus length {} is not a multiple of dim {dim}",
vectors.len()
)));
}
if stable_ids.len() != n_docs {
return Err(BuildError::VectorSchemaMismatch(format!(
"fp32 cell IVF stable_ids len {} != n_docs {n_docs}",
stable_ids.len()
)));
}
}
}
let effective_n_cent = if n_docs > CONSOLIDATED_CELL_ROWS_THRESHOLD {
requested_n_cent.max(1).min(n_docs)
} else {
requested_n_cent
.max(1)
.min(n_cent_row_count_cap(n_docs))
.min(n_docs)
};
let sample_size = if cfg.provided_centroids.is_some() {
0
} else {
partition_kmeans_sample_size(effective_n_cent, n_docs).min(n_docs)
};
let chunk_rows = materialized_chunk_rows_for_dim(dim);
let sample = match &source {
CellPackSource::Spilled(spill) => {
sample_spilled_materialized_rows(spill, sample_size, chunk_rows, cfg.rot_seed)?
}
CellPackSource::Rows(rows) => {
sample_ram_materialized_rows(rows, sample_size, dim, cfg.rot_seed)
}
CellPackSource::Fp32 { vectors, .. } => {
sample_fp32_rows(vectors, sample_size, dim, cfg.rot_seed)
}
};
let (n_cent, centroids) = build_phase_timers::timed(&build_phase_timers::TRAIN_US, || {
materialized_centroids(&cfg, requested_n_cent, n_docs, &sample)
});
let summary_centroid = mean_f32_cluster_major(¢roids, dim, n_cent);
let code_bytes = dim.div_ceil(u8::BITS as usize);
let bucket_dir = tempdir_in(scratch)?;
let mut bucket_writers = Vec::with_capacity(n_cent);
for centroid in 0..n_cent {
let path = bucket_dir.path().join(format!("cluster-{centroid}.bin"));
bucket_writers.push(BufWriter::with_capacity(
BUCKET_BUF_SIZE,
File::create(path)?,
));
}
let mut bucket_counts = vec![0u32; n_cent];
let fit_quantizer = matches!(cfg.rerank_codec, RerankCodec::Sq8Residual);
let (mut sq8_min, mut sq8_max) = if fit_quantizer {
(
vec![f32::INFINITY; n_cent * dim],
vec![f32::NEG_INFINITY; n_cent * dim],
)
} else {
(Vec::new(), Vec::new())
};
let mut stable_ids = BufWriter::new(File::create(stable_ids_path)?);
let min_max = fit_quantizer.then_some((sq8_min.as_mut_slice(), sq8_max.as_mut_slice()));
build_phase_timers::timed(&build_phase_timers::ASSIGN_US, || match &source {
CellPackSource::Spilled(spill) => stream_materialized_rows_to_buckets(
spill,
&cfg,
¢roids,
n_cent,
&mut bucket_writers,
&mut bucket_counts,
&mut stable_ids,
min_max,
),
CellPackSource::Rows(rows) => stream_ram_rows_to_buckets(
rows,
&cfg,
¢roids,
n_cent,
&mut bucket_writers,
&mut bucket_counts,
&mut stable_ids,
min_max,
),
CellPackSource::Fp32 {
vectors,
stable_ids: inline_ids,
} => stream_fp32_rows_to_buckets(
vectors,
inline_ids,
&cfg,
¢roids,
n_cent,
&mut bucket_writers,
&mut bucket_counts,
&mut stable_ids,
min_max,
),
})?;
stable_ids.flush()?;
stable_ids.get_ref().sync_all()?;
for writer in bucket_writers {
writer
.into_inner()
.map_err(|error| BuildError::Io(error.into_error()))?
.sync_all()?;
}
let quantizers: Vec<(Vec<f32>, Vec<f32>)> = if fit_quantizer {
(0..n_cent)
.map(|centroid| {
let start = centroid * dim;
derive_sq8_quantizer_from_min_max(
&sq8_min[start..start + dim],
&sq8_max[start..start + dim],
)
})
.collect()
} else {
(0..n_cent).map(|_| fixed_sq8_quantizer(dim)).collect()
};
let codec = cfg.rerank_codec;
let codec_meta_size = codec.codec_meta_bytes(dim, n_docs, n_cent, cfg.metric);
let per_vec_bytes = codec.per_vector_bytes(dim);
let cluster_stride = code_bytes + format::vec::DOC_ID_BYTES + per_vec_bytes;
let stable_ids_region_bytes = n_docs * format::vec::STABLE_ID_BYTES;
let layout = IvfSubsectionLayout::compute(
dim,
n_cent,
n_docs,
cluster_stride,
codec_meta_size,
stable_ids_region_bytes,
);
let cluster_order = centroid_storage_order(¢roids, n_cent, dim);
let planned = plan_ivf_cluster_blocks(
&layout,
&cluster_order,
&bucket_counts,
code_bytes,
per_vec_bytes,
);
let open_region_len = layout
.stable_ids_off
.unwrap_or(layout.per_cluster_blocks_off);
let mut open_region = vec![0u8; open_region_len];
write_ivf_subsection_header(
&mut open_region,
&layout,
codec_meta_size,
&summary_centroid,
¢roids,
);
for planned_block in &planned {
let idx = planned_block.cluster_idx_offset;
open_region[idx..idx + CLUSTER_IDX_COUNT_OFFSET]
.copy_from_slice(&(planned_block.doc_offset as u32).to_le_bytes());
open_region[idx + CLUSTER_IDX_COUNT_OFFSET..idx + CLUSTER_IDX_ENTRY_BYTES]
.copy_from_slice(&(planned_block.count as u32).to_le_bytes());
}
let meta_layout = codec
.ops()
.expect("build-from-source uses a quantized-rerank codec")
.codec_meta_layout(layout.codec_meta_off, n_cent, dim, cfg.metric);
if let (Some(scale_offset), Some(offset_offset)) =
(meta_layout.scale_off, meta_layout.offset_off)
{
for (centroid, (scale, offset)) in quantizers.iter().enumerate() {
let start = centroid * dim * size_of::<f32>();
open_region[scale_offset + start..scale_offset + start + dim * size_of::<f32>()]
.copy_from_slice(bytemuck::cast_slice(scale));
open_region[offset_offset + start..offset_offset + start + dim * size_of::<f32>()]
.copy_from_slice(bytemuck::cast_slice(offset));
}
}
let norms_offset = meta_layout.norms_off;
let mut output = OpenOptions::new()
.create(true)
.truncate(true)
.read(true)
.write(true)
.open(subsection_path)?;
output.set_len((layout.total_size_before_crc + format::CRC_BYTES) as u64)?;
write_at(&mut output, 0, &open_region)?;
let stable_ids_offset = layout
.stable_ids_off
.expect("streamed materialized IVF always has stable ids");
output.seek(SeekFrom::Start(stable_ids_offset as u64))?;
let copied = io::copy(
&mut BufReader::new(File::open(stable_ids_path)?),
&mut output,
)?;
if copied != stable_ids_region_bytes as u64 {
return Err(BuildError::VectorSchemaMismatch(format!(
"streamed stable-id bytes {copied} != expected {stable_ids_region_bytes}"
)));
}
for planned_block in &planned {
let Some(block) = planned_block.block.as_ref() else {
continue;
};
let path = bucket_dir
.path()
.join(format!("cluster-{}.bin", planned_block.centroid_id));
let (scale, offset) = &quantizers[planned_block.centroid_id];
stream_bucket_into_subsection(
&mut output,
&path,
block,
code_bytes,
dim,
codec,
scale,
offset,
norms_offset,
)?;
}
output.flush()?;
output.seek(SeekFrom::Start(0))?;
let mut remaining = layout.total_size_before_crc;
let mut crc = 0u32;
let mut crc_buffer = vec![0u8; MATERIALIZED_BUCKET_CHUNK_BYTES];
while remaining > 0 {
let take = remaining.min(crc_buffer.len());
output.read_exact(&mut crc_buffer[..take])?;
crc = crc32c_append(crc, &crc_buffer[..take]);
remaining -= take;
}
output.seek(SeekFrom::Start(layout.total_size_before_crc as u64))?;
output.write_all(&crc.to_le_bytes())?;
output.flush()?;
output.sync_all()?;
Ok(StreamedIvfSubsection {
n_docs: n_docs as u32,
rerank_codec: codec,
subsection_len: (layout.total_size_before_crc + format::CRC_BYTES) as u64,
n_cent,
summary_offset_in_sub: layout.summary_off,
codec_meta_offset_in_sub: if codec_meta_size == 0 {
0
} else {
layout.codec_meta_off
},
codec_meta_size,
})
}
pub(crate) fn build_merged_subsection_from_spilled_materialized(
cfg: VectorConfig,
requested_n_cent: usize,
spill: &SpilledCellRows,
subsection_path: &Path,
stable_ids_path: &Path,
scratch: &Path,
) -> Result<StreamedIvfSubsection, BuildError> {
build_cell_subsection_from_source(
cfg,
requested_n_cent,
CellPackSource::Spilled(spill),
subsection_path,
stable_ids_path,
scratch,
)
}
fn build_cell_subsection_in_memory(
cfg: VectorConfig,
requested_n_cent: usize,
source: CellPackSource<'_>,
) -> Result<MergedIvfSubsection, BuildError> {
let scratch = tempdir()?;
let subsection_path = scratch.path().join("cell.ivf");
let stable_ids_path = scratch.path().join("cell.ids");
let built = build_cell_subsection_from_source(
cfg,
requested_n_cent,
source,
&subsection_path,
&stable_ids_path,
scratch.path(),
)?;
let bytes = fs::read(&subsection_path)?;
if bytes.len() as u64 != built.subsection_len {
return Err(BuildError::VectorSchemaMismatch(format!(
"cell subsection read-back {} bytes != built {}",
bytes.len(),
built.subsection_len
)));
}
Ok(MergedIvfSubsection {
bytes,
n_cent: built.n_cent,
n_docs: built.n_docs,
rerank_codec: built.rerank_codec,
summary_offset_in_sub: built.summary_offset_in_sub,
codec_meta_offset_in_sub: built.codec_meta_offset_in_sub,
codec_meta_size: built.codec_meta_size,
})
}
pub(crate) fn build_merged_subsection_from_fp32(
cfg: VectorConfig,
requested_n_cent: usize,
vectors: Arc<Vec<f32>>,
stable_ids: &[i128],
) -> Result<MergedIvfSubsection, BuildError> {
build_cell_subsection_in_memory(
cfg,
requested_n_cent,
CellPackSource::Fp32 {
vectors: &vectors,
stable_ids,
},
)
}
fn build_subsection_streaming(
column_id: u32,
col: ColumnState,
scratch: &Path,
) -> Result<SubsectionBytes, BuildError> {
let ColumnState {
config: cfg,
n_docs: n_docs_u32,
reservoir,
unit_scratch: _,
pre_spill_buffer,
spill,
spill_threshold_bytes: _,
materialized_rows,
prebuilt_subsection: _,
inline_stable_ids,
} = col;
if let Some(rows) = materialized_rows {
drop(reservoir);
drop(inline_stable_ids);
let rows_len = rows.len();
let requested_n_cent = n_cent_row_count_cap(rows_len).min(rows_len).max(1);
return build_subsection_from_materialized(cfg, requested_n_cent, rows);
}
let dim = cfg.dim;
let n_docs = n_docs_u32 as usize;
let sample_rows = reservoir.n_rows();
let (n_cent, centroids) = if let Some(global) = cfg.provided_centroids.as_ref() {
debug_assert!(dim > 0 && global.len() % dim == 0);
let nc = (global.len() / dim.max(1)).max(1);
drop(reservoir);
(nc, global.to_vec())
} else {
let n_cent = n_cent_row_count_cap(n_docs)
.min(n_docs.max(1))
.min(sample_rows.max(1))
.max(1);
let centroids = if sample_rows == 0 || n_docs == 0 {
vec![0.0f32; n_cent * dim]
} else {
kmeans(reservoir.sample(), dim, n_cent, KMEANS_ITERS, cfg.rot_seed)
};
drop(reservoir);
(n_cent, centroids)
};
let summary_centroid = mean_f32_cluster_major(¢roids, dim, n_cent);
let rotation = RandomRotation::new(dim, cfg.rot_seed);
let quant = BitQuantizer::new(dim);
let code_bytes = quant.code_bytes();
let mut bucket_writers: Vec<BufWriter<File>> = Vec::with_capacity(n_cent);
for c in 0..n_cent {
let path = scratch.join(format!("infino_bucket_col{column_id}_c{c}.bin"));
let file = File::create(&path)?;
bucket_writers.push(BufWriter::with_capacity(BUCKET_BUF_SIZE, file));
}
let mut bucket_counts = vec![0u32; n_cent];
let chunk_rows = chunk_rows_for_dim(dim);
let codec = cfg.rerank_codec;
if !codec.supports_metric(cfg.metric) {
return Err(BuildError::VectorSchemaMismatch(format!(
"vector index {:?}: codec {} supports cosine metric only",
cfg.column,
codec.name()
)));
}
let sq8_family = codec.is_sq8_residual_family();
let fit_sq8_quantizer = matches!(codec, RerankCodec::Sq8Residual);
let (mut sq8_min_arr, mut sq8_max_arr): (Vec<f32>, Vec<f32>) = if fit_sq8_quantizer {
(
vec![f32::INFINITY; n_cent * dim],
vec![f32::NEG_INFINITY; n_cent * dim],
)
} else {
(Vec::new(), Vec::new())
};
if n_docs > 0 {
let mut source: Box<dyn ChunkedVectorSource> = if let Some(spill) = spill {
debug_assert!(
pre_spill_buffer.is_empty(),
"spill active but pre_spill_buffer still has {} f32s",
pre_spill_buffer.len()
);
let path = spill.finish()?;
Box::new(MmapVectorSource::open(&path, dim, chunk_rows)?)
} else {
Box::new(InMemoryVectorSource::new(
Arc::new(pre_spill_buffer),
dim,
chunk_rows,
))
};
let sq8_acc: Option<(&mut [f32], &mut [f32])> = if fit_sq8_quantizer {
Some((&mut sq8_min_arr, &mut sq8_max_arr))
} else {
None
};
run_pass2(
source.as_mut(),
dim,
n_cent,
code_bytes,
¢roids,
&rotation,
&quant,
&mut bucket_writers,
&mut bucket_counts,
codec,
sq8_acc,
)?;
}
let sq8_quantizers: Vec<(Vec<f32>, Vec<f32>)> = if fit_sq8_quantizer {
(0..n_cent)
.map(|c| {
let off = c * dim;
derive_sq8_quantizer_from_min_max(
&sq8_min_arr[off..off + dim],
&sq8_max_arr[off..off + dim],
)
})
.collect()
} else if codec.uses_fixed_quantizer() && codec.is_sq8_residual_family() {
(0..n_cent).map(|_| fixed_sq8_quantizer(dim)).collect()
} else {
Vec::new()
};
drop(sq8_min_arr);
drop(sq8_max_arr);
let mut bucket_files: Vec<File> = Vec::with_capacity(n_cent);
for w in bucket_writers {
let mut inner = w.into_inner().map_err(|e| BuildError::Io(e.into_error()))?;
inner.flush()?;
bucket_files.push(inner);
}
drop(bucket_files);
let cluster_order = centroid_storage_order(¢roids, n_cent, dim);
let codec_meta_size = codec.codec_meta_bytes(dim, n_docs, n_cent, cfg.metric);
let per_vec_bytes = codec.per_vector_bytes(dim);
let cluster_stride = code_bytes + format::vec::DOC_ID_BYTES + per_vec_bytes;
let stable_ids_region_bytes = match &inline_stable_ids {
Some(ids) if ids.len() == n_docs => n_docs * format::vec::STABLE_ID_BYTES,
Some(ids) => {
return Err(BuildError::VectorSchemaMismatch(format!(
"streaming inline_stable_ids len {} != n_docs {n_docs}",
ids.len()
)));
}
None => 0,
};
let layout = IvfSubsectionLayout::compute(
dim,
n_cent,
n_docs,
cluster_stride,
codec_meta_size,
stable_ids_region_bytes,
);
let total_size_before_crc = layout.total_size_before_crc;
let mut bytes =
alloc_ivf_subsection_with_header(&layout, codec_meta_size, &summary_centroid, ¢roids);
let sq8_scale_block_off = layout.codec_meta_off;
let sq8_offset_block_off = sq8_scale_block_off + n_cent * dim * 4;
let sq8_norms_block_off = if sq8_family && matches!(cfg.metric, Metric::L2Sq | Metric::Cosine) {
Some(sq8_offset_block_off + n_cent * dim * 4)
} else {
None
};
let sq16_norms_block_off =
if codec.is_sq16() && matches!(cfg.metric, Metric::L2Sq | Metric::Cosine) {
Some(layout.codec_meta_off)
} else {
None
};
if sq8_family {
for (cid, (scale_c, offset_c)) in sq8_quantizers.iter().enumerate().take(n_cent) {
let sc_off = sq8_scale_block_off + cid * dim * 4;
bytes[sc_off..sc_off + dim * 4].copy_from_slice(bytemuck::cast_slice(scale_c));
let oc_off = sq8_offset_block_off + cid * dim * 4;
bytes[oc_off..oc_off + dim * 4].copy_from_slice(bytemuck::cast_slice(offset_c));
}
}
let full_row_bytes_in_bucket = if codec.writes_full() { dim * 4 } else { 0 };
let mut id_block: Vec<u8> = Vec::new();
let mut code_block: Vec<u8> = Vec::new();
let mut full_block: Vec<u8> = Vec::new();
write_ivf_cluster_blocks(
&mut bytes,
&layout,
&cluster_order,
&bucket_counts,
code_bytes,
per_vec_bytes,
|bytes, centroid_id, blk| {
let path = scratch.join(format!("infino_bucket_col{column_id}_c{centroid_id}.bin"));
let mut reader = BufReader::with_capacity(BUCKET_BUF_SIZE, File::open(&path)?);
id_block.resize(blk.count * format::vec::DOC_ID_BYTES, 0);
code_block.resize(blk.count * code_bytes, 0);
if full_row_bytes_in_bucket > 0 {
full_block.resize(blk.count * full_row_bytes_in_bucket, 0);
}
for i in 0..blk.count {
reader.read_exact(&mut id_block[i * 4..(i + 1) * 4])?;
reader.read_exact(&mut code_block[i * code_bytes..(i + 1) * code_bytes])?;
if full_row_bytes_in_bucket > 0 {
let off = i * full_row_bytes_in_bucket;
reader.read_exact(&mut full_block[off..off + full_row_bytes_in_bucket])?;
}
}
bytes[blk.codes_base..blk.codes_base + blk.count * code_bytes]
.copy_from_slice(&code_block);
bytes[blk.ids_base..blk.ids_base + blk.count * format::vec::DOC_ID_BYTES]
.copy_from_slice(&id_block);
match codec {
RerankCodec::RabitqOnly => {}
RerankCodec::Fp32 => {
bytes[blk.rerank_base..blk.rerank_base + blk.count * dim * 4]
.copy_from_slice(&full_block);
}
RerankCodec::Sq16 => {
let cluster_rows: &[f32] = bytemuck::cast_slice(&full_block);
for i in 0..blk.count {
let src = &cluster_rows[i * dim..(i + 1) * dim];
let row_off = blk.rerank_base + i * per_vec_bytes;
encode_sq16_row(src, &mut bytes[row_off..row_off + per_vec_bytes]);
if let Some(norms_off) = sq16_norms_block_off {
let n_sq =
sq16_decoded_norm_sq(&bytes[row_off..row_off + per_vec_bytes], dim);
let pos = blk.first_row + i;
let n_off = norms_off + pos * 4;
bytes[n_off..n_off + 4].copy_from_slice(&n_sq.to_le_bytes());
}
}
}
RerankCodec::Sq8Residual | RerankCodec::Sq8FixedResidual => {
let cluster_rows: &[f32] = bytemuck::cast_slice(&full_block);
let (scale_c, offset_c) = &sq8_quantizers[centroid_id];
encode_sq8_residual_cluster_simd(
cluster_rows,
dim,
blk.count,
blk.first_row,
blk.rerank_base,
sq8_norms_block_off,
scale_c,
offset_c,
bytes,
codec
.residual_divisor()
.expect("residual-family codec has divisor"),
);
}
}
Ok(())
},
)?;
if let (Some(stable_ids_off), Some(ids)) = (layout.stable_ids_off, inline_stable_ids.as_ref()) {
for (local, &stable_id) in ids.iter().enumerate() {
let off = stable_ids_off + local * format::vec::STABLE_ID_BYTES;
bytes[off..off + format::vec::STABLE_ID_BYTES]
.copy_from_slice(&stable_id.to_le_bytes());
}
}
debug_assert_eq!(bytes.len(), total_size_before_crc);
let crc = crc32c(&bytes);
let mut out = bytes;
out.extend_from_slice(&crc.to_le_bytes());
Ok(SubsectionBytes {
bytes: out,
n_cent,
summary_offset_in_sub: layout.summary_off,
codec_meta_offset_in_sub: if codec_meta_size == 0 {
0
} else {
layout.codec_meta_off
},
codec_meta_size,
})
}
#[allow(clippy::too_many_arguments)]
fn encode_sq8_residual_cluster_simd(
cluster_rows: &[f32],
dim: usize,
cluster_count: usize,
cluster_doc_off: usize,
full_chunk_base: usize,
sq8_norms_block_off: Option<usize>,
scale_c: &[f32],
offset_c: &[f32],
bytes: &mut [u8],
residual_divisor: f32,
) {
debug_assert_eq!(cluster_rows.len(), cluster_count * dim);
let row_bytes = dim * 2;
let consts = Sq8EncodeConsts::from_scale_offset(scale_c, offset_c);
let store_norm = sq8_norms_block_off.is_some();
let mut recon = vec![0f32; dim];
for i in 0..cluster_count {
let src = &cluster_rows[i * dim..(i + 1) * dim];
let pos = cluster_doc_off + i;
let row_off = full_chunk_base + i * row_bytes;
let row_bytes_mut = &mut bytes[row_off..row_off + row_bytes];
let (code_slice, res_slice) = row_bytes_mut.split_at_mut(dim);
let norm = encode_sq8_residual_row(
src,
&consts,
scale_c,
offset_c,
code_slice,
res_slice,
&mut recon,
store_norm,
residual_divisor,
);
if let (Some(norms_off), Some(n_sq)) = (sq8_norms_block_off, norm) {
let n_off = norms_off + pos * 4;
bytes[n_off..n_off + 4].copy_from_slice(&n_sq.to_le_bytes());
}
}
}
#[inline]
pub(crate) fn derive_sq8_quantizer_from_min_max(min: &[f32], max: &[f32]) -> (Vec<f32>, Vec<f32>) {
debug_assert_eq!(min.len(), max.len());
let dim = min.len();
let mut scale = vec![0.0f32; dim];
let mut offset = vec![0.0f32; dim];
for d in 0..dim {
let span = max[d] - min[d];
if span > 0.0 && span.is_finite() {
offset[d] = min[d];
scale[d] = span / SQ8_CODE_MAX;
} else {
offset[d] = if min[d].is_finite() { min[d] } else { 0.0 };
scale[d] = 1.0;
}
}
(scale, offset)
}
pub(crate) fn fixed_sq8_quantizer(dim: usize) -> (Vec<f32>, Vec<f32>) {
(vec![SQ8_FIXED_SCALE; dim], vec![SQ8_FIXED_OFFSET; dim])
}
pub(crate) fn centroid_storage_order(centroids: &[f32], n_cent: usize, dim: usize) -> Vec<usize> {
let mut order: Vec<usize> = (0..n_cent).collect();
order_centroids_recursive(&mut order, centroids, dim);
order
}
fn order_centroids_recursive(order: &mut [usize], centroids: &[f32], dim: usize) {
if order.len() <= 1 || dim == 0 {
return;
}
let mut best_dim = 0usize;
let mut best_span = 0.0f32;
for d in 0..dim {
let mut lo = f32::INFINITY;
let mut hi = f32::NEG_INFINITY;
for &c in order.iter() {
let v = centroids[c * dim + d];
lo = lo.min(v);
hi = hi.max(v);
}
let span = hi - lo;
if span > best_span {
best_span = span;
best_dim = d;
}
}
order.sort_unstable_by(|&a, &b| {
centroids[a * dim + best_dim]
.partial_cmp(¢roids[b * dim + best_dim])
.unwrap_or(Ordering::Equal)
.then_with(|| a.cmp(&b))
});
let mid = order.len() / 2;
let (left, right) = order.split_at_mut(mid);
order_centroids_recursive(left, centroids, dim);
order_centroids_recursive(right, centroids, dim);
}
pub(crate) struct IvfSubsectionLayout {
pub summary_off: usize,
pub centroids_off: usize,
pub cluster_idx_off: usize,
pub codec_meta_off: usize,
pub per_cluster_blocks_off: usize,
pub stable_ids_off: Option<usize>,
pub total_size_before_crc: usize,
}
impl IvfSubsectionLayout {
pub(crate) fn compute(
dim: usize,
n_cent: usize,
n_docs: usize,
per_cluster_stride: usize,
codec_meta_size: usize,
stable_ids_region_bytes: usize,
) -> Self {
let summary_off = SUB_HEADER_SIZE;
let centroids_off = summary_off + dim * 4;
let cluster_idx_off = centroids_off + n_cent * dim * 4;
let codec_meta_off = cluster_idx_off + n_cent * CLUSTER_IDX_ENTRY_BYTES;
let codec_meta_end = codec_meta_off + codec_meta_size;
let stable_ids_off = (stable_ids_region_bytes > 0).then_some(codec_meta_end);
let per_cluster_blocks_off = codec_meta_end + stable_ids_region_bytes;
let total_size_before_crc = per_cluster_blocks_off + n_docs * per_cluster_stride;
Self {
summary_off,
centroids_off,
cluster_idx_off,
codec_meta_off,
per_cluster_blocks_off,
stable_ids_off,
total_size_before_crc,
}
}
}
pub(crate) fn alloc_ivf_subsection_with_header(
layout: &IvfSubsectionLayout,
codec_meta_size: usize,
summary_centroid: &[f32],
centroids: &[f32],
) -> Vec<u8> {
let mut bytes = vec![0u8; layout.total_size_before_crc];
write_ivf_subsection_header(
&mut bytes,
layout,
codec_meta_size,
summary_centroid,
centroids,
);
bytes
}
fn write_ivf_subsection_header(
bytes: &mut [u8],
layout: &IvfSubsectionLayout,
codec_meta_size: usize,
summary_centroid: &[f32],
centroids: &[f32],
) {
debug_assert!(bytes.len() >= layout.codec_meta_off);
bytes[0..MAGIC_BYTES].copy_from_slice(format::vec::SUB_MAGIC);
bytes[sub_hdr::VERSION_OFF..sub_hdr::VERSION_OFF + U32_BYTES]
.copy_from_slice(&format::vec::SUBSECTION_VERSION.to_le_bytes());
bytes[sub_hdr::CODEC_META_SIZE_OFF..sub_hdr::CODEC_META_SIZE_OFF + U32_BYTES]
.copy_from_slice(&(codec_meta_size as u32).to_le_bytes());
bytes[sub_hdr::SUMMARY_OFF_OFF..sub_hdr::SUMMARY_OFF_OFF + U64_BYTES]
.copy_from_slice(&(layout.summary_off as u64).to_le_bytes());
bytes[sub_hdr::CENTROIDS_OFF_OFF..sub_hdr::CENTROIDS_OFF_OFF + U64_BYTES]
.copy_from_slice(&(layout.centroids_off as u64).to_le_bytes());
bytes[sub_hdr::CLUSTER_IDX_OFF_OFF..sub_hdr::CLUSTER_IDX_OFF_OFF + U64_BYTES]
.copy_from_slice(&(layout.cluster_idx_off as u64).to_le_bytes());
bytes[sub_hdr::PER_CLUSTER_BLOCKS_OFF_OFF..sub_hdr::PER_CLUSTER_BLOCKS_OFF_OFF + U64_BYTES]
.copy_from_slice(&(layout.per_cluster_blocks_off as u64).to_le_bytes());
bytes[layout.summary_off..layout.summary_off + summary_centroid.len() * 4]
.copy_from_slice(bytemuck::cast_slice(summary_centroid));
bytes[layout.centroids_off..layout.centroids_off + centroids.len() * 4]
.copy_from_slice(bytemuck::cast_slice(centroids));
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct ClusterBlock {
pub codes_base: usize,
pub ids_base: usize,
pub rerank_base: usize,
pub first_row: usize,
pub count: usize,
}
#[derive(Debug, Clone, Copy)]
struct PlannedClusterBlock {
centroid_id: usize,
cluster_idx_offset: usize,
doc_offset: usize,
count: usize,
block: Option<ClusterBlock>,
}
fn plan_ivf_cluster_blocks(
layout: &IvfSubsectionLayout,
cluster_order: &[usize],
cluster_counts: &[u32],
code_bytes: usize,
per_vec_bytes: usize,
) -> Vec<PlannedClusterBlock> {
let cluster_stride = code_bytes + format::vec::DOC_ID_BYTES + per_vec_bytes;
let mut block_cursor = 0usize;
let mut doc_offset = 0usize;
let mut planned = Vec::with_capacity(cluster_order.len());
for ¢roid_id in cluster_order {
let count = cluster_counts[centroid_id] as usize;
let cluster_idx_offset = layout.cluster_idx_off + centroid_id * CLUSTER_IDX_ENTRY_BYTES;
let block = (count > 0).then(|| {
let block_base = layout.per_cluster_blocks_off + block_cursor;
let codes_len = count * code_bytes;
let ids_len = count * format::vec::DOC_ID_BYTES;
ClusterBlock {
codes_base: block_base,
ids_base: block_base + codes_len,
rerank_base: block_base + codes_len + ids_len,
first_row: doc_offset,
count,
}
});
planned.push(PlannedClusterBlock {
centroid_id,
cluster_idx_offset,
doc_offset,
count,
block,
});
block_cursor += count * cluster_stride;
doc_offset += count;
}
planned
}
pub(crate) fn write_ivf_cluster_blocks<F>(
bytes: &mut [u8],
layout: &IvfSubsectionLayout,
cluster_order: &[usize],
cluster_counts: &[u32],
code_bytes: usize,
per_vec_bytes: usize,
mut write_cluster: F,
) -> Result<(), BuildError>
where
F: FnMut(&mut [u8], usize, &ClusterBlock) -> Result<(), BuildError>,
{
for planned in plan_ivf_cluster_blocks(
layout,
cluster_order,
cluster_counts,
code_bytes,
per_vec_bytes,
) {
let idx_base = planned.cluster_idx_offset;
bytes[idx_base..idx_base + CLUSTER_IDX_COUNT_OFFSET]
.copy_from_slice(&(planned.doc_offset as u32).to_le_bytes());
bytes[idx_base + CLUSTER_IDX_COUNT_OFFSET..idx_base + CLUSTER_IDX_ENTRY_BYTES]
.copy_from_slice(&(planned.count as u32).to_le_bytes());
if let Some(block) = planned.block {
write_cluster(bytes, planned.centroid_id, &block)?;
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn run_pass2(
source: &mut dyn ChunkedVectorSource,
dim: usize,
n_cent: usize,
code_bytes: usize,
centroids: &[f32],
rotation: &RandomRotation,
quant: &BitQuantizer,
bucket_writers: &mut [BufWriter<File>],
bucket_counts: &mut [u32],
codec: RerankCodec,
mut sq8_min_max: Option<(&mut [f32], &mut [f32])>,
) -> Result<(), BuildError> {
let chunk_rows_cap = source.chunk_rows();
let mut chunk_rotated = vec![0f32; chunk_rows_cap * dim];
let mut chunk_assignments = vec![0u32; chunk_rows_cap];
let mut chunk_codes = vec![0u8; chunk_rows_cap * code_bytes];
let mut global_doc_id: u32 = 0;
while let Some(chunk) = source.next_chunk() {
let actual_rows = chunk.len() / dim;
debug_assert!(actual_rows <= chunk_rows_cap);
let asgn = &mut chunk_assignments[..actual_rows];
assign_to_centroids(&chunk[..actual_rows * dim], centroids, dim, n_cent, asgn);
chunk_rotated[..actual_rows * dim]
.par_chunks_mut(dim)
.zip(chunk[..actual_rows * dim].par_chunks(dim))
.for_each(|(dst, src)| rotation.apply(src, dst));
chunk_codes[..actual_rows * code_bytes]
.par_chunks_mut(code_bytes)
.enumerate()
.for_each(|(r, code_out)| {
let rot_row = &chunk_rotated[r * dim..(r + 1) * dim];
quant.encode_rotated_into(rot_row, code_out);
});
let write_full = codec.writes_full();
let mut sq8_acc = sq8_min_max.as_mut();
for r in 0..actual_rows {
let cid = asgn[r] as usize;
let local_doc_id = global_doc_id + r as u32;
let writer = &mut bucket_writers[cid];
writer.write_all(&local_doc_id.to_le_bytes())?;
writer.write_all(&chunk_codes[r * code_bytes..(r + 1) * code_bytes])?;
if write_full {
writer.write_all(bytemuck::cast_slice(&chunk[r * dim..(r + 1) * dim]))?;
}
if let Some((mn, mx)) = sq8_acc.as_deref_mut() {
let row = &chunk[r * dim..(r + 1) * dim];
let off = cid * dim;
update_min_max(row, &mut mn[off..off + dim], &mut mx[off..off + dim]);
}
bucket_counts[cid] += 1;
}
global_doc_id += actual_rows as u32;
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::fs::{read, write};
use bytes::Bytes;
use tempfile::tempdir;
use super::*;
use crate::superfile::vector::{
cell_posting::EncodedCellRow, reader::VectorReader, spill::MaterializedRowSpillWriter,
};
fn block_on<F: std::future::Future>(f: F) -> F::Output {
tokio::runtime::Builder::new_current_thread()
.build()
.expect("build current-thread runtime")
.block_on(f)
}
fn cfg(name: &str, dim: usize) -> VectorConfig {
VectorConfig {
column: name.to_string(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Fp32,
provided_centroids: None,
}
}
#[test]
fn n_cent_row_count_cap_steps_at_thresholds() {
assert_eq!(n_cent_row_count_cap(0), N_CENT_SMALL);
assert_eq!(
n_cent_row_count_cap(N_CENT_MEDIUM_DOC_THRESHOLD - 1),
N_CENT_SMALL
);
assert_eq!(
n_cent_row_count_cap(N_CENT_MEDIUM_DOC_THRESHOLD),
N_CENT_MEDIUM
);
assert_eq!(
n_cent_row_count_cap(N_CENT_LARGE_DOC_THRESHOLD - 1),
N_CENT_MEDIUM
);
assert_eq!(
n_cent_row_count_cap(N_CENT_LARGE_DOC_THRESHOLD),
N_CENT_LARGE
);
}
#[test]
fn register_column_returns_sequential_ids() {
let mut b = VectorBuilder::new();
assert_eq!(b.register_column(cfg("a", 16)).expect("register column"), 0);
assert_eq!(b.register_column(cfg("b", 32)).expect("register column"), 1);
}
#[test]
fn register_column_rejects_separator_in_name() {
let mut b = VectorBuilder::new();
let bad = cfg("a\x1Fb", 16);
let err = b.register_column(bad).expect_err("expected error");
assert!(matches!(err, BuildError::ReservedSeparatorInColumnName(_)));
}
#[test]
fn register_column_rejects_inf_prefix() {
let mut b = VectorBuilder::new();
let bad = cfg("inf.embedding", 16);
let err = b.register_column(bad).expect_err("expected error");
assert!(matches!(err, BuildError::ReservedPrefixInColumnName(_)));
}
#[test]
fn register_column_rejects_dim_too_small() {
let mut b = VectorBuilder::new();
let err = b.register_column(cfg("a", 8)).expect_err("expected error");
assert!(matches!(err, BuildError::VectorDimOutOfRange { .. }));
}
#[test]
fn register_column_rejects_dim_too_large() {
let mut b = VectorBuilder::new();
let err = b
.register_column(cfg("a", 5000))
.expect_err("expected error");
assert!(matches!(err, BuildError::VectorDimOutOfRange { .. }));
}
#[test]
fn register_column_rejects_duplicate() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
let err = b.register_column(cfg("a", 32)).expect_err("expected error");
assert!(matches!(err, BuildError::DuplicateColumnName(_)));
}
#[test]
fn add_rejects_unknown_column_id() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
let err = b.add(99, &[0.0; 16]).expect_err("expected error");
assert!(matches!(err, BuildError::FtsColumnTypeInvalid { .. }));
}
#[test]
fn add_rejects_wrong_dim() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
let err = b.add(0, &[0.0; 8]).expect_err("expected error");
assert!(matches!(err, BuildError::FtsColumnTypeInvalid { .. }));
}
#[test]
fn finish_emits_valid_outer_header() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
for i in 0..32 {
let v: Vec<f32> = (0..16).map(|j| (i + j) as f32).collect();
b.add(0, &v).expect("add to vector builder");
}
let blob = b.finish().expect("finish");
assert_eq!(&blob[0..8], format::vec::OUTER_MAGIC);
let version = u32::from_le_bytes([blob[8], blob[9], blob[10], blob[11]]);
assert_eq!(version, format::vec::VERSION);
let n_cols = u32::from_le_bytes([blob[12], blob[13], blob[14], blob[15]]);
assert_eq!(n_cols, 1);
}
#[test]
fn finish_with_no_docs_produces_valid_blob() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
let blob = b.finish().expect("finish");
assert_eq!(&blob[0..8], format::vec::OUTER_MAGIC);
let mut buf = [0u8; 8];
buf.copy_from_slice(&blob[16..24]);
assert_eq!(u64::from_le_bytes(buf), 0);
}
#[test]
fn materialized_build_round_trips_inline_stable_ids() {
use bytes::Bytes;
use crate::superfile::vector::reader::VectorReader;
let dim = 16;
let n = 24usize;
let json =
format!(r#"[{{"column":"v","dim":{dim},"n_cent":4,"rot_seed":7,"metric":"cosine"}}]"#);
let cfg = || VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let mut b = VectorBuilder::new();
b.register_column(cfg()).expect("register");
for i in 0..n {
let mut v = vec![0.0f32; dim];
v[i % dim] = 1.0 + (i as f32);
v[(i * 7) % dim] += 0.5;
b.add(0, &v).expect("add");
}
let stream_blob = b.finish().expect("finish streaming");
let stream_reader =
VectorReader::open(Bytes::from(stream_blob), &json).expect("open streaming");
let mut stream_rows =
block_on(stream_reader.materialized_index_rows_async("v")).expect("streaming rows");
assert!(
stream_rows.iter().all(|r| r.stable_id == 0),
"streaming build must not carry inline stable_ids"
);
let want = |local: u32| -> i128 { 1_700_000_000_000i128 + local as i128 };
for r in &mut stream_rows {
r.stable_id = want(r.local_doc_id);
r.encoded.stable_id = r.stable_id;
}
let mut mb = VectorBuilder::new();
mb.register_column(cfg()).expect("register mat");
mb.load_materialized_rows(0, stream_rows)
.expect("load materialized");
let mat_blob = mb.finish().expect("finish materialized");
let mat_reader =
VectorReader::open(Bytes::from(mat_blob), &json).expect("open materialized");
assert_eq!(mat_reader.n_docs(), n as u64);
let mat_rows =
block_on(mat_reader.materialized_index_rows_async("v")).expect("materialized rows");
assert_eq!(mat_rows.len(), n);
for r in &mat_rows {
assert_eq!(
r.stable_id,
want(r.local_doc_id),
"inline stable_id must round-trip for local {}",
r.local_doc_id
);
assert_eq!(
r.encoded.stable_id, r.stable_id,
"EncodedCellRow.stable_id must match"
);
}
}
#[test]
fn sq8_merge_carries_inline_stable_ids_through_compaction() {
use bytes::Bytes;
use crate::superfile::vector::{
ivf_merge::merge_sq8_ivf_subsections, reader::VectorReader,
};
let dim = 16;
let json =
format!(r#"[{{"column":"v","dim":{dim},"n_cent":4,"rot_seed":7,"metric":"cosine"}}]"#);
let cfg = || VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let build_cell = |n: usize, id_base: i128| -> Bytes {
let mut b = VectorBuilder::new();
b.register_column(cfg()).expect("register");
for i in 0..n {
let mut v = vec![0.0f32; dim];
v[i % dim] = 1.0 + (i as f32);
v[(i * 5) % dim] += 0.25;
b.add(0, &v).expect("add");
}
let stream = b.finish().expect("finish streaming");
let r = VectorReader::open(Bytes::from(stream), &json).expect("open streaming");
let mut rows = block_on(r.materialized_index_rows_async("v")).expect("rows");
for row in &mut rows {
row.stable_id = id_base + row.local_doc_id as i128;
row.encoded.stable_id = row.stable_id;
}
let mut mb = VectorBuilder::new();
mb.register_column(cfg()).expect("register mat");
mb.load_materialized_rows(0, rows)
.expect("load materialized");
Bytes::from(mb.finish().expect("finish materialized"))
};
let (na, nb) = (10usize, 10usize);
let blob_a = build_cell(na, 5_000);
let blob_b = build_cell(nb, 9_000);
let reader_a = VectorReader::open(blob_a, &json).expect("open A");
let reader_b = VectorReader::open(blob_b, &json).expect("open B");
let merged = merge_sq8_ivf_subsections(&[(&reader_a, "v", 0), (&reader_b, "v", na as u32)])
.expect("merge");
assert_eq!(merged.n_docs as usize, na + nb);
let mut wb = VectorBuilder::new();
wb.register_column(cfg()).expect("register merged");
wb.set_prebuilt_subsection(0, merged).expect("set prebuilt");
let merged_blob = wb.finish().expect("finish merged");
let reader_m = VectorReader::open(Bytes::from(merged_blob), &json).expect("open merged");
let rows = block_on(reader_m.materialized_index_rows_async("v")).expect("merged rows");
assert_eq!(rows.len(), na + nb);
for r in &rows {
let want = if (r.local_doc_id as usize) < na {
5_000 + r.local_doc_id as i128
} else {
9_000 + (r.local_doc_id as i128 - na as i128)
};
assert_eq!(
r.stable_id, want,
"merged inline stable_id wrong for local {}",
r.local_doc_id
);
}
}
#[test]
fn sq8_tiny_shard_writes_physical_n_cent_to_directory() {
use bytes::Bytes;
use crate::superfile::vector::reader::VectorReader;
let dim = 16;
let configured_n_cent = 4;
let mut b = VectorBuilder::new();
b.register_column(VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
})
.expect("register sq8 column");
b.add(0, &[1.0; 16]).expect("add single row");
let blob = b.finish().expect("finish tiny sq8 shard");
let dir_off = OUTER_HEADER_SIZE;
let physical_n_cent = u32::from_le_bytes(
blob[dir_off + 8..dir_off + 12]
.try_into()
.expect("n_cent bytes"),
);
assert_eq!(
physical_n_cent, 1,
"directory must describe physical IVF layout, not configured n_cent"
);
let json = format!(
r#"[{{"column":"v","dim":{dim},"n_cent":{configured_n_cent},"rot_seed":7,"metric":"cosine"}}]"#
);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open tiny sq8 shard");
assert_eq!(reader.n_docs(), 1);
}
#[test]
fn order_centroids_geometrically_chains_neighbors() {
const DIM: usize = 4;
let positions = [33.0f32, 10.0, 0.0, 21.0];
let mut centroids = Vec::with_capacity(positions.len() * DIM);
for p in positions {
centroids.extend(std::iter::repeat_n(p, DIM));
}
order_centroids_geometrically(&mut centroids, DIM, positions.len());
let ordered: Vec<f32> = (0..positions.len()).map(|c| centroids[c * DIM]).collect();
assert_eq!(ordered, vec![21.0, 10.0, 0.0, 33.0]);
let positions_b = [0.0f32, 33.0, 21.0, 10.0];
let mut centroids_b = Vec::with_capacity(positions_b.len() * DIM);
for p in positions_b {
centroids_b.extend(std::iter::repeat_n(p, DIM));
}
order_centroids_geometrically(&mut centroids_b, DIM, positions_b.len());
let ordered_b: Vec<f32> = (0..positions_b.len())
.map(|c| centroids_b[c * DIM])
.collect();
assert_eq!(ordered_b, ordered, "chain order is input-order invariant");
}
const SPLIT_DIM: usize = 8;
const SPLIT_REQUESTED: usize = 10;
fn two_blob_sample() -> Vec<f32> {
let mut rows = Vec::with_capacity(600 * SPLIT_DIM);
let mut state = 0x2545_f491_4f6c_dd1du64;
let mut jitter = || {
state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1);
((state >> 33) % 1000) as f32 / 1000.0 - 0.5
};
for r in 0..600 {
let base = if r < 500 { 0.0f32 } else { 10.0 };
for _ in 0..SPLIT_DIM {
rows.push(base + jitter());
}
}
rows
}
fn run_counts(sample: &[f32], centroids: &[f32], n_cent: usize) -> Vec<usize> {
let mut assignments = vec![0u32; sample.len() / SPLIT_DIM];
assign_to_centroids(sample, centroids, SPLIT_DIM, n_cent, &mut assignments);
let mut counts = vec![0usize; n_cent];
for &a in &assignments {
counts[a as usize] += 1;
}
counts
}
#[test]
fn split_oversized_fine_runs_bounds_every_run() {
let sample = two_blob_sample();
let sample_n = sample.len() / SPLIT_DIM;
let target = sample_n.div_ceil(SPLIT_REQUESTED);
let bound = target * FINE_RUN_SPLIT_BOUND_FACTOR;
let mut centroids = vec![0.0f32; 5 * SPLIT_DIM];
for c in 1..5 {
for d in 0..SPLIT_DIM {
centroids[c * SPLIT_DIM + d] = 1_000.0 + c as f32;
}
}
let before = run_counts(&sample, ¢roids, 5);
assert!(
before.iter().any(|&n| n > bound),
"fixture must start oversized (max run {} ≤ bound {bound})",
before.iter().max().expect("nonempty")
);
let n_cent =
split_oversized_fine_runs(&mut centroids, &sample, SPLIT_DIM, SPLIT_REQUESTED, 7, None);
assert_eq!(centroids.len(), n_cent * SPLIT_DIM);
assert!(n_cent > 5, "split must add sub-centroids");
let after = run_counts(&sample, ¢roids, n_cent);
assert!(
after.iter().all(|&n| n <= bound),
"every run must fit the bound {bound}: {after:?}"
);
}
#[test]
fn sampled_index_breaks_periodic_aliasing_on_consolidated_cells() {
let n_docs = CONSOLIDATED_CELL_ROWS_THRESHOLD * 2;
let sample_size = n_docs / 4;
let seed = 7u64;
let indices: Vec<usize> = (0..sample_size)
.map(|s| sampled_index(s, sample_size, n_docs, seed))
.collect();
let mut class_seen = [false; 4];
for (s, &idx) in indices.iter().enumerate() {
assert!(idx < n_docs, "index {idx} out of bounds");
if s > 0 {
assert!(
idx > indices[s - 1],
"indices must be strictly increasing: {} then {idx}",
indices[s - 1]
);
}
class_seen[idx % 4] = true;
}
assert_eq!(
class_seen, [true; 4],
"every periodic class must be sampled (plain stride 4 sees one)"
);
let again: Vec<usize> = (0..sample_size)
.map(|s| sampled_index(s, sample_size, n_docs, seed))
.collect();
assert_eq!(indices, again, "same seed must select the same rows");
let small = CONSOLIDATED_CELL_ROWS_THRESHOLD;
for s in [0usize, 7, 1000] {
assert_eq!(
sampled_index(s, small / 4, small, seed),
s * 4,
"sub-threshold sampling must stay the plain stride"
);
}
assert_eq!(sampled_index(3, 8, 8, seed), 3);
}
#[test]
fn split_oversized_fine_runs_leaves_single_run_untouched() {
let sample = two_blob_sample();
let mut centroids = vec![0.5f32; SPLIT_DIM];
let original = centroids.clone();
let n_cent = split_oversized_fine_runs(&mut centroids, &sample, SPLIT_DIM, 1, 7, None);
assert_eq!(n_cent, 1);
assert_eq!(centroids, original);
}
#[test]
fn materialized_centroids_caps_only_below_consolidated_threshold() {
let dim = 8;
let requested = 92;
let sample: Vec<f32> = (0..256 * dim).map(|i| (i % 251) as f32 * 0.01).collect();
let mk = |n_docs: usize| {
let cfg = VectorConfig {
column: "emb".into(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
materialized_centroids(&cfg, requested, n_docs, &sample).0
};
let below = mk(CONSOLIDATED_CELL_ROWS_THRESHOLD);
let above = mk(CONSOLIDATED_CELL_ROWS_THRESHOLD + 1);
assert!(
below >= N_CENT_SMALL && below < requested,
"sub-threshold cell stays near the legacy cap (got {below}, cap {N_CENT_SMALL}, byte-target {requested})"
);
assert!(
above >= requested,
"consolidated cell takes the byte-target count uncapped"
);
}
#[test]
fn materialized_centroids_splits_oversized_runs_below_threshold() {
let sample = two_blob_sample();
let sample_n = sample.len() / SPLIT_DIM;
let target = sample_n.div_ceil(SPLIT_REQUESTED);
let bound = target * FINE_RUN_SPLIT_BOUND_FACTOR;
let cfg = VectorConfig {
column: "emb".into(),
dim: SPLIT_DIM,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let (n_cent, centroids) = materialized_centroids(
&cfg,
SPLIT_REQUESTED,
CONSOLIDATED_CELL_ROWS_THRESHOLD,
&sample,
);
assert!(
n_cent >= SPLIT_REQUESTED,
"split may grow past the request, never shrink it"
);
let after = run_counts(&sample, ¢roids, n_cent);
assert!(
after.iter().all(|&n| n <= bound),
"sub-threshold cell must still bound every fine run at {bound}: {after:?}"
);
}
#[test]
fn split_oversized_fine_runs_is_deterministic() {
let sample = two_blob_sample();
let make = || {
let mut centroids = vec![0.0f32; 5 * SPLIT_DIM];
for c in 1..5 {
for d in 0..SPLIT_DIM {
centroids[c * SPLIT_DIM + d] = 1_000.0 + c as f32;
}
}
let n = split_oversized_fine_runs(
&mut centroids,
&sample,
SPLIT_DIM,
SPLIT_REQUESTED,
7,
None,
);
(n, centroids)
};
assert_eq!(make(), make());
}
#[test]
fn build_merged_subsection_from_fp32_stays_in_memory() {
use bytes::Bytes;
use crate::superfile::vector::reader::VectorReader;
let dim = 16;
let n = 5;
let mut corpus = Vec::with_capacity(n * dim);
for r in 0..n {
for c in 0..dim {
corpus.push((r as f32) * 0.01 + (c as f32) * 0.001);
}
}
let cfg = VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let ids: Vec<i128> = (0..n as i128).map(|i| 9_000 + i).collect();
let sub = build_merged_subsection_from_fp32(cfg.clone(), 64, Arc::new(corpus), &ids)
.expect("fp32 build");
assert_eq!(sub.n_docs, n as u32);
assert!(sub.n_cent >= 1);
assert!(!sub.bytes.is_empty());
let blob = finish_multi_cell_blob(&[(0, sub)]).expect("multi-cell blob");
let json =
format!(r#"[{{"column":"v","dim":{dim},"n_cent":64,"rot_seed":7,"metric":"l2sq"}}]"#);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open fp32 cell pack");
let resolved = reader
.inline_stable_ids_for_locals(&[0, 1, 2])
.expect("inline stable ids");
assert_eq!(resolved, ids[0..3]);
}
#[test]
fn finish_two_columns_at_different_dims() {
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
b.register_column(cfg("b", 32)).expect("register column");
for _ in 0..16 {
b.add(0, &[1.0; 16]).expect("add to vector builder");
b.add(1, &[1.0; 32]).expect("add to vector builder");
}
let blob = b.finish().expect("finish");
let n_cols = u32::from_le_bytes([blob[12], blob[13], blob[14], blob[15]]);
assert_eq!(n_cols, 2);
let dir_off = OUTER_HEADER_SIZE;
let entry_a_dim = u32::from_le_bytes([
blob[dir_off + 4],
blob[dir_off + 5],
blob[dir_off + 6],
blob[dir_off + 7],
]);
let entry_b_dim = u32::from_le_bytes([
blob[dir_off + DIR_ENTRY_SIZE + 4],
blob[dir_off + DIR_ENTRY_SIZE + 5],
blob[dir_off + DIR_ENTRY_SIZE + 6],
blob[dir_off + DIR_ENTRY_SIZE + 7],
]);
assert_eq!(entry_a_dim, 16);
assert_eq!(entry_b_dim, 32);
}
#[test]
fn build_via_forced_spill_path_round_trips() {
let dim = 16;
let n_docs = 64usize;
let mut b = VectorBuilder::new();
b.set_spill_threshold_bytes(0);
b.register_column(VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Fp32,
provided_centroids: None,
})
.expect("register column");
let mut corpus = Vec::with_capacity(n_docs * dim);
for d in 0..n_docs {
let mut row = vec![0.0f32; dim];
row[0] = d as f32;
row[1] = (d as f32) * 0.5;
row[2] = -(d as f32);
corpus.extend_from_slice(&row);
b.add(0, &row).expect("add via forced-spill path");
}
let blob = b.finish().expect("finish via forced-spill path");
assert_eq!(&blob[0..8], format::vec::OUTER_MAGIC);
let n_cols = u32::from_le_bytes([blob[12], blob[13], blob[14], blob[15]]);
assert_eq!(n_cols, 1);
let n_docs_hdr = u64::from_le_bytes(blob[16..24].try_into().expect("8 bytes"));
assert_eq!(n_docs_hdr, n_docs as u64);
}
#[tokio::test]
async fn forced_spill_path_matches_in_ram_path_on_self_nn() {
use bytes::Bytes;
use crate::superfile::vector::reader::VectorReader;
let dim = 16;
let n_docs = 50;
let n_cent = 4;
let mut corpus = Vec::with_capacity(n_docs * dim);
for d in 0..n_docs {
let mut row = vec![0.0f32; dim];
for (j, slot) in row.iter_mut().enumerate() {
*slot = ((d as f32) * 0.07 + (j as f32) * 0.13).sin();
}
corpus.extend_from_slice(&row);
}
let build = |force_spill: bool| -> Vec<u8> {
let mut b = VectorBuilder::new();
if force_spill {
b.set_spill_threshold_bytes(0);
}
b.register_column(VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Fp32,
provided_centroids: None,
})
.expect("register column");
for d in 0..n_docs {
b.add(0, &corpus[d * dim..(d + 1) * dim])
.expect("add to vector builder");
}
b.finish().expect("finish")
};
let blob_ram = build(false);
let blob_spill = build(true);
let json = format!(
r#"[{{"column":"v","dim":{dim},"n_cent":{n_cent},"rot_seed":7,"metric":"l2sq"}}]"#
);
let r_ram = VectorReader::open(Bytes::from(blob_ram), &json).expect("open ram");
let r_spill = VectorReader::open(Bytes::from(blob_spill), &json).expect("open spill");
let nprobe = n_cent;
let rerank_mult = n_docs + 1;
for q in 0..n_docs {
let query = &corpus[q * dim..(q + 1) * dim];
let top_ram = r_ram
.search("v", query, 1, nprobe, rerank_mult)
.await
.expect("search ram");
let top_spill = r_spill
.search("v", query, 1, nprobe, rerank_mult)
.await
.expect("search spill");
assert_eq!(
top_ram[0].0 as usize, q,
"in-RAM path missed self-NN at q={q}"
);
assert_eq!(
top_spill[0].0 as usize, q,
"spill path missed self-NN at q={q}"
);
}
}
#[test]
fn finish_to_matches_finish_byte_for_byte() {
let build = || -> VectorBuilder {
let mut b = VectorBuilder::new();
b.register_column(cfg("v", 16)).expect("register column");
for i in 0..32 {
let v: Vec<f32> = (0..16).map(|j| ((i + j) as f32) * 0.1).collect();
b.add(0, &v).expect("add to vector builder");
}
b
};
let blob_finish = build().finish().expect("finish");
let mut blob_finish_to: Vec<u8> = Vec::new();
build()
.finish_to(&mut blob_finish_to)
.expect("finish_to Vec<u8>");
assert_eq!(
blob_finish, blob_finish_to,
"finish_to must produce identical bytes to finish"
);
}
#[test]
fn finish_to_cursor_round_trips_outer_crc() {
use std::io::Cursor;
let mut b = VectorBuilder::new();
b.register_column(cfg("v", 16)).expect("register column");
for i in 0..32 {
let v: Vec<f32> = (0..16).map(|j| ((i + j) as f32) * 0.1).collect();
b.add(0, &v).expect("add to vector builder");
}
let mut buf: Vec<u8> = Vec::new();
{
let cursor = Cursor::new(&mut buf);
b.finish_to(cursor).expect("finish_to Cursor");
}
assert_eq!(
&buf[0..8],
format::vec::OUTER_MAGIC,
"outer magic preserved"
);
assert!(
buf.len() >= OUTER_HEADER_SIZE + DIR_ENTRY_SIZE + 4 + 4,
"blob too short: {} bytes",
buf.len()
);
let body_len = buf.len() - 4;
let trailing_crc = u32::from_le_bytes([
buf[body_len],
buf[body_len + 1],
buf[body_len + 2],
buf[body_len + 3],
]);
let recomputed = crc32c(&buf[..body_len]);
assert_eq!(
trailing_crc, recomputed,
"trailing outer CRC32C must match recomputed body CRC"
);
}
#[tokio::test]
async fn finish_to_temp_file_round_trips_through_reader() {
use std::io::BufWriter;
use bytes::Bytes;
use crate::superfile::vector::reader::VectorReader;
let dim = 16usize;
let n_docs = 32usize;
let n_cent = 4usize;
let mut b = VectorBuilder::new();
b.register_column(VectorConfig {
column: "v".into(),
dim,
rot_seed: 7,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Fp32,
provided_centroids: None,
})
.expect("register column");
for d in 0..n_docs {
let row: Vec<f32> = (0..dim)
.map(|j| ((d as f32) * 0.07 + (j as f32) * 0.13).sin())
.collect();
b.add(0, &row).expect("add to vector builder");
}
let tmp = tempfile::tempdir().expect("tempdir");
let path = tmp.path().join("vector_blob.bin");
{
let file = File::create(&path).expect("create blob file");
let writer = BufWriter::new(file);
b.finish_to(writer).expect("finish_to BufWriter<File>");
}
let blob = read(&path).expect("read blob file");
let json = format!(
r#"[{{"column":"v","dim":{dim},"n_cent":{n_cent},"rot_seed":7,"metric":"l2sq"}}]"#
);
let reader = VectorReader::open(Bytes::from(blob), &json)
.expect("open VectorReader from streamed blob");
let query: Vec<f32> = (0..dim).map(|j| ((j as f32) * 0.13).sin()).collect();
let hits = reader
.search("v", &query, 5, n_cent, n_docs + 1)
.await
.expect("kNN search");
assert!(!hits.is_empty(), "search returned no hits");
}
#[test]
fn vector_config_new_and_with_rerank_codec() {
let dim = 16usize;
let rot_seed = 7u64;
let base = VectorConfig::new("v".into(), dim, rot_seed, Metric::Cosine);
assert_eq!(base.column, "v");
assert_eq!(base.dim, dim);
assert_eq!(base.rot_seed, rot_seed);
assert_eq!(base.metric, Metric::Cosine);
assert_eq!(base.rerank_codec, RerankCodec::default());
let overridden = base.with_rerank_codec(RerankCodec::Fp32);
assert_eq!(overridden.rerank_codec, RerankCodec::Fp32);
assert_eq!(overridden.column, "v");
}
#[test]
fn vector_builder_default_matches_new() {
let mut b = VectorBuilder::default();
assert_eq!(b.register_column(cfg("a", 16)).expect("register column"), 0);
}
#[test]
fn set_kmeans_sample_size_ok_and_unregistered() {
const SAMPLE_SIZE: usize = 1024;
let mut b = VectorBuilder::new();
b.register_column(cfg("a", 16)).expect("register column");
b.set_kmeans_sample_size(0, SAMPLE_SIZE)
.expect("resize sample for registered column");
let err = b
.set_kmeans_sample_size(9, SAMPLE_SIZE)
.expect_err("unregistered column id");
assert!(matches!(err, BuildError::FtsColumnTypeInvalid { .. }));
}
#[test]
fn with_scratch_accepts_dir_and_rejects_file() {
let dir = tempfile::tempdir().expect("tempdir");
let mut b = VectorBuilder::with_scratch(dir.path().to_path_buf())
.expect("scratch under existing dir");
assert_eq!(b.register_column(cfg("a", 16)).expect("register column"), 0);
let file_path = dir.path().join("not-a-dir");
write(&file_path, b"x").expect("write file");
match VectorBuilder::with_scratch(file_path) {
Ok(_) => panic!("scratch path is a file, expected rejection"),
Err(err) => assert!(matches!(err, BuildError::Io(_))),
}
}
#[test]
fn multi_cell_blob_round_trips_cell_directory_and_centroids() {
use std::sync::Arc;
use bytes::Bytes;
use crate::superfile::vector::{
builder::{build_merged_subsection_from_materialized, finish_multi_cell_blob},
cell_posting::EncodedCellRow,
reader::VectorReader,
};
let dim = 16;
let make_rows = |cell: u32, n: usize| -> Vec<MaterializedIvfRow> {
let scale: Arc<[f32]> = Arc::from(vec![1.0f32; dim]);
let offset: Arc<[f32]> = Arc::from(vec![0.0f32; dim]);
(0..n)
.map(|i| {
let local = i as u32;
let stable_id = (cell as i128) * 1_000 + local as i128;
let mut codes = vec![0u8; dim];
codes[0] = (cell as u8).wrapping_add(i as u8);
let encoded = EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8Residual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes,
residuals: vec![0u8; dim],
norm_sq: Some(1.0),
};
MaterializedIvfRow {
local_doc_id: local,
stable_id,
cluster: 0,
rabitq_code: vec![0u8; dim.div_ceil(8)],
encoded,
}
})
.collect()
};
let cfg = || VectorConfig {
column: "emb".into(),
dim,
rot_seed: 1,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let sub0 = build_merged_subsection_from_materialized(cfg(), 2, make_rows(0, 4))
.expect("cell 0 subsection");
let sub1 = build_merged_subsection_from_materialized(cfg(), 2, make_rows(1, 3))
.expect("cell 1 subsection");
let cells = vec![(0, sub0), (1, sub1)];
let blob = finish_multi_cell_blob(&cells).expect("pack");
let streamed_sources: Vec<BorrowedMultiCellSubsection<'_>> = cells
.iter()
.map(|(cell_id, subsection)| BorrowedMultiCellSubsection {
cell_id: *cell_id,
subsection,
})
.collect();
let mut streamed = Vec::new();
finish_multi_cell_blob_to(&streamed_sources, &mut streamed).expect("stream pack");
assert_eq!(
streamed, blob,
"streamed and in-memory multi-cell assembly must be byte-identical"
);
let json =
format!(r#"[{{"column":"emb","dim":{dim},"n_cent":2,"rot_seed":1,"metric":"l2sq"}}]"#);
let reader = VectorReader::open(Bytes::from(blob.clone()), &json).expect("open multi-cell");
assert!(reader.is_multi_cell());
assert_eq!(reader.packed_cell_ids(), &[0, 1]);
assert_eq!(reader.n_docs(), 7);
let (n_cent, got_dim, _centroids, counts) = reader
.cluster_centroids("emb")
.expect("concatenated centroids");
assert_eq!(got_dim, dim as u32);
assert_eq!(n_cent, 4, "flat n_cent must sum packed cells, got {n_cent}");
assert_eq!(counts.len(), n_cent as usize);
assert!(counts.iter().any(|&c| c > 0));
assert_eq!(reader.resolve_flat_cluster(0), Some((0, 0)));
assert_eq!(reader.resolve_flat_cluster(2), Some((1, 0)));
assert_eq!(reader.resolve_flat_cluster(3), Some((1, 1)));
let mut zero_codec = blob;
let directory_start = OUTER_HEADER_SIZE;
let directory_size = 2 * CELL_DIR_ENTRY_SIZE;
for entry in 0..2 {
let codec_off =
directory_start + entry * CELL_DIR_ENTRY_SIZE + cell_dir_entry::CODEC_ID_OFF;
zero_codec[codec_off..codec_off + U32_BYTES].copy_from_slice(&0u32.to_le_bytes());
}
let directory_crc = crc32c(&zero_codec[directory_start..directory_start + directory_size]);
let crc_off = directory_start + directory_size;
zero_codec[crc_off..crc_off + format::CRC_BYTES]
.copy_from_slice(&directory_crc.to_le_bytes());
assert!(
VectorReader::open(Bytes::from(zero_codec), &json).is_err(),
"zero codec id has no v2 compatibility fallback"
);
}
#[test]
fn sq16_merge_round_trips_norms() {
use std::sync::Arc;
use bytes::Bytes;
use crate::superfile::vector::{
cell_posting::EncodedCellRow,
distance::{Metric, sq16_decoded_norm_sq},
reader::VectorReader,
rerank_codec::RerankCodec,
};
let dim = 16usize;
let n = 6usize;
let vecgen = |i: usize| -> Vec<f32> {
(0..dim)
.map(|j| (((i as f32 + 1.0) * 0.31 + (j as f32) * 0.11).sin()) * 0.6)
.collect::<Vec<f32>>()
};
let rows: Vec<MaterializedIvfRow> = (0..n)
.map(|i| {
let v = vecgen(i);
let mut codes = vec![0u8; dim * 2];
encode_sq16_row(&v, &mut codes);
let norm_sq = sq16_decoded_norm_sq(&codes, dim);
MaterializedIvfRow {
local_doc_id: i as u32,
stable_id: i as i128,
cluster: 0,
rabitq_code: vec![0u8; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id: i as i128,
rerank_codec: RerankCodec::Sq16,
scale: Arc::from(Vec::<f32>::new()),
offset: Arc::from(Vec::<f32>::new()),
codes,
residuals: Vec::new(),
norm_sq: Some(norm_sq),
},
}
})
.collect();
let mut expected: Vec<f32> = rows.iter().filter_map(|r| r.encoded.norm_sq).collect();
let cfg = VectorConfig {
column: "emb".into(),
dim,
rot_seed: 1,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq16,
provided_centroids: None,
};
let sub = build_merged_subsection_from_materialized(cfg, 2, rows).expect("Sq16 merge");
let cells = vec![(0u32, sub)];
let blob = finish_multi_cell_blob(&cells).expect("pack Sq16");
let json = format!(
r#"[{{"column":"emb","dim":{dim},"n_cent":2,"rot_seed":1,"metric":"cosine"}}]"#
);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open merged Sq16");
assert_eq!(reader.n_docs(), n as u64);
let norms = reader
.per_doc_norms_for_test("emb")
.expect("Sq16 cosine carries a per-doc norm table");
assert_eq!(norms.len(), n, "all Sq16 rows survived the merge");
let mut have: Vec<f32> = norms.to_vec();
expected.sort_by(|a, b| a.total_cmp(b));
have.sort_by(|a, b| a.total_cmp(b));
for (e, h) in expected.iter().zip(have.iter()) {
assert!((e - h).abs() < 1e-5, "Sq16 merged norm {h} != stored {e}");
}
let mat = block_on(reader.materialized_cells_rows_async(None))
.expect("Sq16 drain-read path (materialized_cells_rows_async) must yield rows");
let total: usize = mat.iter().map(|(_, rows)| rows.len()).sum();
assert_eq!(total, n, "drain materialize must return every Sq16 row");
for (_, rows) in &mat {
for r in rows {
assert_eq!(r.encoded.rerank_codec, RerankCodec::Sq16);
assert!(
r.encoded.norm_sq.is_some(),
"Sq16 cosine row carries a norm"
);
assert!(
r.encoded.scale.is_empty() && r.encoded.offset.is_empty(),
"Sq16 materialized row carries no per-cluster scale/offset"
);
}
}
}
#[test]
fn fixed_residual_multi_cell_rebuild_preserves_payload_bytes() {
let dim = 16;
let make_rows = |cell: u32| -> Vec<MaterializedIvfRow> {
let scale: Arc<[f32]> = Arc::from(vec![SQ8_FIXED_SCALE; dim]);
let offset: Arc<[f32]> = Arc::from(vec![SQ8_FIXED_OFFSET; dim]);
(0..4)
.map(|i| {
let stable_id = i128::from(cell) * 100 + i;
MaterializedIvfRow {
local_doc_id: i as u32,
stable_id,
cluster: 0,
rabitq_code: vec![0; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8FixedResidual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes: vec![64 + i as u8; dim],
residuals: vec![i as i8 as u8; dim],
norm_sq: None,
},
}
})
.collect()
};
let config = VectorConfig {
column: "emb".into(),
dim,
rot_seed: 7,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq8FixedResidual,
provided_centroids: None,
};
let source_rows: Vec<MaterializedIvfRow> = [make_rows(0), make_rows(1)].concat();
let sub0 = build_merged_subsection_from_materialized(config.clone(), 2, make_rows(0))
.expect("fixed cell 0");
let sub1 = build_merged_subsection_from_materialized(config, 2, make_rows(1))
.expect("fixed cell 1");
let blob = finish_multi_cell_blob(&[(0, sub0), (1, sub1)]).expect("pack fixed cells");
let json = r#"[{"column":"emb","dim":16,"n_cent":2,"rot_seed":7,"metric":"cosine"}]"#;
let reader = VectorReader::open(Bytes::from(blob), json).expect("open fixed multi-cell");
assert!(
reader
.vector_columns_config()
.all(|column| column.rerank_codec == RerankCodec::Sq8FixedResidual)
);
let mut rebuilt =
block_on(reader.materialized_index_rows_async("emb")).expect("materialize fixed rows");
rebuilt.sort_by_key(|row| row.stable_id);
let mut expected = source_rows;
expected.sort_by_key(|row| row.stable_id);
for (before, after) in expected.iter().zip(&rebuilt) {
assert_eq!(before.stable_id, after.stable_id);
assert_eq!(before.encoded.codes, after.encoded.codes);
assert_eq!(before.encoded.residuals, after.encoded.residuals);
assert_eq!(after.encoded.rerank_codec, RerankCodec::Sq8FixedResidual);
}
}
#[test]
fn streamed_materialized_cell_matches_in_memory_fixed_residual() {
let dim = 16;
let scale: Arc<[f32]> = Arc::from(vec![SQ8_FIXED_SCALE; dim]);
let offset: Arc<[f32]> = Arc::from(vec![SQ8_FIXED_OFFSET; dim]);
let rows: Vec<MaterializedIvfRow> = (0..32)
.map(|row| {
let stable_id = i128::from(row) * 17 + 3;
MaterializedIvfRow {
local_doc_id: row,
stable_id,
cluster: 0,
rabitq_code: vec![row as u8; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8FixedResidual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes: vec![32 + row as u8; dim],
residuals: vec![(row as i8 - 16) as u8; dim],
norm_sq: None,
},
}
})
.collect();
let config = VectorConfig {
column: "emb".into(),
dim,
rot_seed: 7,
metric: Metric::Cosine,
rerank_codec: RerankCodec::Sq8FixedResidual,
provided_centroids: None,
};
let expected = build_merged_subsection_from_materialized(config.clone(), 4, rows.clone())
.expect("in-memory materialized build");
let directory = tempdir().expect("tempdir");
let mut spill_writer =
MaterializedRowSpillWriter::create(directory.path(), 11, dim, dim.div_ceil(8))
.expect("spill writer");
for row in &rows {
spill_writer.append(row).expect("spill row");
}
let spill = spill_writer.finish().expect("finish spill");
let subsection_path = directory.path().join("streamed.ivf");
let stable_ids_path = directory.path().join("streamed.ids");
let built = build_merged_subsection_from_spilled_materialized(
config,
4,
&spill,
&subsection_path,
&stable_ids_path,
directory.path(),
)
.expect("streamed materialized build");
assert_eq!(built.n_docs, rows.len() as u32);
assert_eq!(
read(&subsection_path).expect("read streamed subsection"),
expected.bytes,
"streamed and in-memory materialized builders must be byte-identical"
);
let expected_ids: Vec<u8> = rows
.iter()
.flat_map(|row| row.stable_id.to_le_bytes())
.collect();
assert_eq!(
read(&stable_ids_path).expect("read stable ids"),
expected_ids
);
}
#[tokio::test]
async fn multi_cell_search_returns_ascending_distance() {
use std::sync::Arc;
use bytes::Bytes;
use crate::superfile::vector::{
builder::{build_merged_subsection_from_materialized, finish_multi_cell_blob},
cell_posting::EncodedCellRow,
reader::VectorReader,
};
let dim = 16;
let make_rows = |cell: u32, n: usize| -> Vec<MaterializedIvfRow> {
let scale: Arc<[f32]> = Arc::from(vec![1.0f32; dim]);
let offset: Arc<[f32]> = Arc::from(vec![0.0f32; dim]);
(0..n)
.map(|i| {
let local = i as u32;
let stable_id = (cell as i128) * 1_000 + local as i128;
let mut codes = vec![0u8; dim];
codes[0] = (cell as u8).wrapping_add(i as u8);
MaterializedIvfRow {
local_doc_id: local,
stable_id,
cluster: 0,
rabitq_code: vec![0u8; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8Residual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes,
residuals: vec![0u8; dim],
norm_sq: Some(1.0),
},
}
})
.collect()
};
let cfg = || VectorConfig {
column: "emb".into(),
dim,
rot_seed: 1,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let sub0 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(0, 4)).expect("cell 0");
let sub1 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(1, 3)).expect("cell 1");
let blob = finish_multi_cell_blob(&[(0, sub0), (1, sub1)]).expect("pack");
let json =
format!(r#"[{{"column":"emb","dim":{dim},"n_cent":2,"rot_seed":1,"metric":"l2sq"}}]"#);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open");
let q = vec![0.0f32; dim];
let hits = reader
.search_clusters_async("emb", &q, 3, &[0, 1, 2, 3], 8, None, None, None, None)
.await
.expect("search");
assert_eq!(hits.len(), 3);
assert!(
hits[0].1 <= hits[1].1 && hits[1].1 <= hits[2].1,
"multi-cell hits must be ascending distance, got {hits:?}"
);
}
#[tokio::test]
async fn multi_cell_search_remaps_file_local_allow() {
use std::sync::Arc;
use bytes::Bytes;
use roaring::RoaringBitmap;
use crate::superfile::vector::{
builder::{build_merged_subsection_from_materialized, finish_multi_cell_blob},
cell_posting::EncodedCellRow,
reader::VectorReader,
};
let dim = 16;
let make_rows = |cell: u32, n: usize| -> Vec<MaterializedIvfRow> {
let scale: Arc<[f32]> = Arc::from(vec![1.0f32; dim]);
let offset: Arc<[f32]> = Arc::from(vec![0.0f32; dim]);
(0..n)
.map(|i| {
let local = i as u32;
let stable_id = (cell as i128) * 1_000 + local as i128;
let mut codes = vec![0u8; dim];
codes[0] = (cell as u8).wrapping_add(i as u8);
MaterializedIvfRow {
local_doc_id: local,
stable_id,
cluster: 0,
rabitq_code: vec![0u8; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8Residual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes,
residuals: vec![0u8; dim],
norm_sq: Some(1.0),
},
}
})
.collect()
};
let cfg = || VectorConfig {
column: "emb".into(),
dim,
rot_seed: 1,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let sub0 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(0, 4)).expect("cell 0");
let sub1 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(1, 3)).expect("cell 1");
let blob = finish_multi_cell_blob(&[(0, sub0), (1, sub1)]).expect("pack");
let json =
format!(r#"[{{"column":"emb","dim":{dim},"n_cent":2,"rot_seed":1,"metric":"l2sq"}}]"#);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open");
let mut allow = RoaringBitmap::new();
allow.insert(4);
allow.insert(5);
allow.insert(6);
let q = vec![0.0f32; dim];
let hits = reader
.search_clusters_async(
"emb",
&q,
3,
&[0, 1, 2, 3],
8,
Some(Arc::new(allow)),
None,
None,
None,
)
.await
.expect("filtered multi-cell search");
assert_eq!(
hits.len(),
3,
"expected all three allowed cell1 rows, got {hits:?}"
);
for (file_local, _) in &hits {
assert!(
*file_local >= 4,
"allow was cell1-only (file-local 4..6); got hit {file_local}"
);
}
}
#[tokio::test]
async fn multi_cell_materialize_filters_by_cell_directory() {
use std::sync::Arc;
use bytes::Bytes;
use crate::superfile::vector::{
builder::{build_merged_subsection_from_materialized, finish_multi_cell_blob},
cell_posting::EncodedCellRow,
reader::VectorReader,
};
let dim = 16;
let make_rows = |cell: u32, n: usize| -> Vec<MaterializedIvfRow> {
let scale: Arc<[f32]> = Arc::from(vec![1.0f32; dim]);
let offset: Arc<[f32]> = Arc::from(vec![0.0f32; dim]);
(0..n)
.map(|i| {
let local = i as u32;
let stable_id = (cell as i128) * 1_000 + local as i128;
MaterializedIvfRow {
local_doc_id: local,
stable_id,
cluster: 0,
rabitq_code: vec![0u8; dim.div_ceil(8)],
encoded: EncodedCellRow {
stable_id,
rerank_codec: RerankCodec::Sq8Residual,
scale: Arc::clone(&scale),
offset: Arc::clone(&offset),
codes: vec![cell as u8; dim],
residuals: vec![0u8; dim],
norm_sq: Some(1.0),
},
}
})
.collect()
};
let cfg = || VectorConfig {
column: "emb".into(),
dim,
rot_seed: 1,
metric: Metric::L2Sq,
rerank_codec: RerankCodec::Sq8Residual,
provided_centroids: None,
};
let sub0 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(7, 3)).expect("cell 7");
let sub1 =
build_merged_subsection_from_materialized(cfg(), 2, make_rows(15, 2)).expect("cell 15");
let blob = finish_multi_cell_blob(&[(7, sub0), (15, sub1)]).expect("pack");
let json =
format!(r#"[{{"column":"emb","dim":{dim},"n_cent":2,"rot_seed":1,"metric":"l2sq"}}]"#);
let reader = VectorReader::open(Bytes::from(blob), &json).expect("open");
let only15 = reader
.materialized_cells_rows_async(Some(&[15]))
.await
.expect("materialize cell 15");
assert_eq!(only15.len(), 1);
assert_eq!(only15[0].0, 15);
assert_eq!(only15[0].1.len(), 2);
assert!(only15[0].1.iter().all(|r| r.stable_id / 1000 == 15));
let all = reader
.materialized_index_rows_async("emb")
.await
.expect("all cells");
assert_eq!(all.len(), 5);
let locals: Vec<u32> = all.iter().map(|r| r.local_doc_id).collect();
assert_eq!(locals, vec![0, 1, 2, 3, 4]);
assert_eq!(reader.packed_cell_n_docs(7), Some(3));
assert_eq!(reader.packed_cell_n_docs(15), Some(2));
assert_eq!(reader.packed_cell_n_docs(99), None);
let counts: Vec<(u32, u32)> = reader
.packed_cell_ids()
.iter()
.filter_map(|&c| reader.packed_cell_n_docs(c).map(|n| (c, n)))
.collect();
let overflow = counts
.iter()
.copied()
.max_by_key(|(_, n)| *n)
.expect("counts");
assert_eq!(overflow, (7, 3));
let packed_summary = reader.summary("emb").expect("packed summary");
assert_eq!(packed_summary.len(), dim);
let cell7_only = {
let sub0 = build_merged_subsection_from_materialized(cfg(), 2, make_rows(7, 3))
.expect("cell 7 alone");
let blob0 = finish_multi_cell_blob(&[(7, sub0)]).expect("pack one");
VectorReader::open(Bytes::from(blob0), &json)
.expect("open")
.summary("emb")
.expect("cell7 summary")
};
let cell15_only = {
let sub1 = build_merged_subsection_from_materialized(cfg(), 2, make_rows(15, 2))
.expect("cell 15 alone");
let blob1 = finish_multi_cell_blob(&[(15, sub1)]).expect("pack one");
VectorReader::open(Bytes::from(blob1), &json)
.expect("open")
.summary("emb")
.expect("cell15 summary")
};
for d in 0..dim {
let expected = (cell7_only[d] * 3.0 + cell15_only[d] * 2.0) / 5.0;
assert!(
(packed_summary[d] - expected).abs() < 1e-5,
"dim {d}: packed={} expected={expected}",
packed_summary[d]
);
}
let file_locals: Vec<u32> = (0..5).collect();
let resolved = reader
.inline_stable_ids_for_locals(&file_locals)
.expect("multi-cell file-local stable ids");
assert_eq!(
resolved,
vec![7000, 7001, 7002, 15000, 15001],
"file-local → stable_id must span both packed cells"
);
}
}