use serde::{Deserialize, Serialize};
use uqa_core::{memory::Budgeted, DocId};
use super::codec::{
blob_to_vector, decode_u64_value, decode_value, encode_value, other_error, read_u64,
usize_to_u64, vector_to_blob,
};
use super::index_keys::{
ivf_assignment_doc_prefix, ivf_assignment_key, ivf_assignment_prefix, ivf_centroid_key,
ivf_centroid_prefix, ivf_metadata_key,
};
use super::{KeyValueBatch, KeyValueRead, KeyValueVectorIndex};
use crate::ivf_index::{
IVFMetadataSnapshot, IVFPreparedMetadata, IVFReadIndex, IVFRestoreBuilder, IVFState,
};
use crate::vector_index::IVFIndexParams;
use crate::StorageBackendResult;
pub(super) const IVF_FORMAT_VERSION: u32 = 1;
#[cfg(test)]
thread_local! {
pub(super) static RESTORED_INDEXES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
#[derive(Debug, Serialize, Deserialize)]
pub(super) struct PersistedIVFMetadata {
pub(super) format_version: u32,
pub(super) dimensions: u32,
pub(super) nlist: u64,
pub(super) nprobe: u64,
pub(super) train_threshold: u64,
pub(super) state: PersistedIVFState,
pub(super) trained_size: u64,
pub(super) deletes_since_train: u64,
pub(super) vector_count: u64,
pub(super) revision: u64,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub(super) enum PersistedIVFState {
Untrained,
Trained,
Stale,
}
pub(super) fn restore_state(
store: &dyn KeyValueRead,
raw: &KeyValueVectorIndex,
table: &str,
field: &str,
dimensions: u32,
params: IVFIndexParams,
) -> StorageBackendResult<(Budgeted<IVFReadIndex>, u64)> {
#[cfg(test)]
RESTORED_INDEXES.set(RESTORED_INDEXES.get() + 1);
let metadata = load_metadata(store, table, field)?.ok_or_else(|| {
other_error(format!(
"missing persisted IVF metadata for {table}.{field}"
))
})?;
validate_metadata(&metadata, table, field, dimensions, params)?;
let trained_size = checked_usize(metadata.trained_size, "IVF trained_size")?;
let deletes_since_train =
checked_usize(metadata.deletes_since_train, "IVF deletes_since_train")?;
let vector_count = checked_usize(metadata.vector_count, "IVF vector_count")?;
let mut builder = IVFRestoreBuilder::new(
dimensions,
params,
IVFMetadataSnapshot {
state: metadata.state.into(),
centroids: Vec::new(),
assignments: Vec::new(),
trained_size,
deletes_since_train,
vector_count,
},
store.control(),
)?;
load_centroids(store, table, field, &mut builder)?;
load_assignments(store, table, field, &mut builder)?;
raw.visit_canonical_from(store, |document, ordinal, vector| {
builder.vector(document, ordinal, vector)
})?;
let index = IVFReadIndex::new(builder.finish()?)?;
Ok((index, metadata.revision))
}
pub(super) fn load_revision(
store: &dyn KeyValueRead,
table: &str,
field: &str,
) -> StorageBackendResult<Option<u64>> {
Ok(load_metadata(store, table, field)?.map(|metadata| metadata.revision))
}
#[expect(
clippy::too_many_arguments,
reason = "keeps persisted write inputs aligned"
)]
pub(super) fn stage_snapshot(
batch: &mut dyn KeyValueBatch,
table: &str,
field: &str,
dimensions: u32,
params: IVFIndexParams,
candidate: &IVFPreparedMetadata,
revision: u64,
full_rewrite: bool,
changed_doc: Option<DocId>,
preview: bool,
) -> StorageBackendResult<()> {
let snapshot = candidate.header();
if !preview {
batch.fence_ivf_prefix(&ivf_metadata_key(table, field)?)?;
}
put_record(
batch,
preview,
&ivf_metadata_key(table, field)?,
&encode_value(&metadata_from_snapshot(
dimensions, params, snapshot, revision,
)?)?,
)?;
if full_rewrite {
delete_prefix(batch, preview, &ivf_centroid_prefix(table, field)?)?;
delete_prefix(batch, preview, &ivf_assignment_prefix(table, field)?)?;
for (centroid, vector) in snapshot.centroids.iter().enumerate() {
put_record(
batch,
preview,
&ivf_centroid_key(table, field, centroid)?,
&vector_to_blob(vector)?,
)?;
}
for assignment in candidate.assignments() {
let (doc_id, ordinal, centroid) = assignment?;
put_assignment(batch, table, field, doc_id, ordinal, centroid, preview)?;
}
} else if let Some(doc_id) = changed_doc {
delete_prefix(
batch,
preview,
&ivf_assignment_doc_prefix(table, field, doc_id)?,
)?;
for assignment in candidate.document_assignments(doc_id) {
let (ordinal, centroid) = assignment?;
put_assignment(batch, table, field, doc_id, ordinal, centroid, preview)?;
}
}
Ok(())
}
fn put_assignment(
batch: &mut dyn KeyValueBatch,
table: &str,
field: &str,
doc_id: DocId,
ordinal: u32,
centroid: usize,
preview: bool,
) -> StorageBackendResult<()> {
put_record(
batch,
preview,
&ivf_assignment_key(table, field, doc_id, ordinal)?,
&usize_to_u64(centroid, "IVF centroid assignment")?.to_be_bytes(),
)
}
fn load_metadata(
store: &dyn KeyValueRead,
table: &str,
field: &str,
) -> StorageBackendResult<Option<PersistedIVFMetadata>> {
store
.get(&ivf_metadata_key(table, field)?)?
.map(|bytes| decode_value(&bytes))
.transpose()
}
fn load_centroids(
store: &dyn KeyValueRead,
table: &str,
field: &str,
builder: &mut IVFRestoreBuilder,
) -> StorageBackendResult<()> {
let prefix = ivf_centroid_prefix(table, field)?;
store.visit_prefix(&prefix, &mut |key, value| {
let mut offset = prefix.len();
let found = checked_usize(read_u64(key, &mut offset)?, "IVF centroid id")?;
if offset != key.len() {
return Err(other_error("corrupt IVF centroid key sequence"));
}
let _memory = store.control().memory().reserve(value.len())?;
builder.centroid(found, &blob_to_vector(value)?)
})
}
fn load_assignments(
store: &dyn KeyValueRead,
table: &str,
field: &str,
builder: &mut IVFRestoreBuilder,
) -> StorageBackendResult<()> {
let prefix = ivf_assignment_prefix(table, field)?;
store.visit_prefix(&prefix, &mut |key, value| {
let mut offset = prefix.len();
let doc_id = read_u64(key, &mut offset)?;
let ordinal = u32::try_from(read_u64(key, &mut offset)?)
.map_err(|_| other_error("persisted IVF ordinal exceeds u32"))?;
if offset != key.len() {
return Err(other_error(
"persisted IVF assignment key has trailing bytes",
));
}
let centroid = checked_usize(decode_u64_value(value)?, "IVF centroid assignment")?;
builder.assignment(doc_id, ordinal, centroid)
})
}
pub(super) fn metadata_from_snapshot(
dimensions: u32,
params: IVFIndexParams,
snapshot: &IVFMetadataSnapshot,
revision: u64,
) -> StorageBackendResult<PersistedIVFMetadata> {
Ok(PersistedIVFMetadata {
format_version: IVF_FORMAT_VERSION,
dimensions,
nlist: usize_to_u64(params.nlist, "IVF nlist")?,
nprobe: usize_to_u64(params.nprobe, "IVF nprobe")?,
train_threshold: usize_to_u64(params.train_threshold, "IVF train_threshold")?,
state: snapshot.state.into(),
trained_size: usize_to_u64(snapshot.trained_size, "IVF trained_size")?,
deletes_since_train: usize_to_u64(snapshot.deletes_since_train, "IVF deletes_since_train")?,
vector_count: usize_to_u64(snapshot.vector_count, "IVF vector_count")?,
revision,
})
}
fn validate_metadata(
metadata: &PersistedIVFMetadata,
table: &str,
field: &str,
dimensions: u32,
params: IVFIndexParams,
) -> StorageBackendResult<()> {
let persisted_params = IVFIndexParams {
nlist: checked_usize(metadata.nlist, "IVF nlist")?,
nprobe: checked_usize(metadata.nprobe, "IVF nprobe")?,
train_threshold: checked_usize(metadata.train_threshold, "IVF train_threshold")?,
};
if metadata.format_version != IVF_FORMAT_VERSION
|| metadata.dimensions != dimensions
|| persisted_params != params
{
return Err(other_error(format!(
"persisted IVF metadata does not match the catalog for {table}.{field}"
)));
}
Ok(())
}
fn checked_usize(value: u64, field: &str) -> StorageBackendResult<usize> {
usize::try_from(value).map_err(|_| other_error(format!("{field} exceeds usize")))
}
impl From<IVFState> for PersistedIVFState {
fn from(value: IVFState) -> Self {
match value {
IVFState::Untrained => Self::Untrained,
IVFState::Trained => Self::Trained,
IVFState::Stale => Self::Stale,
}
}
}
impl From<PersistedIVFState> for IVFState {
fn from(value: PersistedIVFState) -> Self {
match value {
PersistedIVFState::Untrained => Self::Untrained,
PersistedIVFState::Trained => Self::Trained,
PersistedIVFState::Stale => Self::Stale,
}
}
}
fn put_record(
batch: &mut dyn KeyValueBatch,
preview: bool,
key: &[u8],
value: &[u8],
) -> StorageBackendResult<()> {
if preview {
batch.preview_ivf_record(key, Some(value))
} else {
batch.put(key, value)
}
}
fn delete_prefix(
batch: &mut dyn KeyValueBatch,
preview: bool,
prefix: &[u8],
) -> StorageBackendResult<()> {
if preview {
batch.preview_ivf_prefix(prefix)
} else {
batch.delete_prefix(prefix)
}
}