9728194f16
Resolves every finding in docs/reviews/M0-M10-code-review-2026-06-08-pass2.md across the engine, network, server, and CLI crates: session restore, replication/CRDT, WAL format and recovery, storage indexes, query/ranking executors, cohort/community governance, and scatter-gather routing. Adds regression tests: - review_pass2_creator_search_filter - review_pass2_d_replication - review_pass2_query_for_session - review_pass2_storage_indexes_bitmap_cache - review_pass2_zone_a_sessions Verified: cargo clippy -D warnings and full test suite green across all crates.
1072 lines
44 KiB
Rust
1072 lines
44 KiB
Rust
//! Embedding slot registry: maps `(EntityKind, slot_name)` to HNSW index state.
|
||
//!
|
||
//! The [`EmbeddingSlotRegistry`] is the central authority for embedding slot
|
||
//! configuration. It is constructed from the schema at `TidalDB::open()` time
|
||
//! and provides the lookup path that the embedding lifecycle operations use
|
||
//! to find the correct index for a given entity kind and slot name.
|
||
//!
|
||
//! Each entity type can define up to 4 embedding slots (per Entity Model
|
||
//! Specification). Each slot has its own HNSW index, dimensionality,
|
||
//! quantization level, and source (external vs. database-managed).
|
||
|
||
use std::collections::HashMap;
|
||
|
||
use super::{QuantizationLevel, VectorError, VectorIndex, VectorIndexConfig};
|
||
use crate::schema::EntityKind;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Vector backend selection (production HNSW vs. brute-force)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Minimum live-vector count for a slot before the production HNSW engine
|
||
/// (`UsearchIndex`) is preferred over the exact `BruteForceIndex`.
|
||
///
|
||
/// `BruteForceIndex` computes an L2 distance against every stored vector under
|
||
/// an `RwLock` read on each search — exact, but O(n) per query and serialized
|
||
/// against writes. It is the right choice for tiny slots (no graph-build cost,
|
||
/// exact results) but blows past the latency budget at scale. `UsearchIndex`
|
||
/// (HNSW) is the mandated production engine (`CODING_GUIDELINES` §4: "126K+ QPS",
|
||
/// §8: "ANN retrieval at 1M vectors <10ms p99"). Gating on slot size needs NO
|
||
/// external config plumbing: at process restart, `rebuild_from_store` already
|
||
/// knows each slot's vector count, so a slot that crossed this threshold is
|
||
/// rebuilt as HNSW automatically, while a small slot stays exact. The threshold
|
||
/// is deliberately conservative — well inside the range where brute force is
|
||
/// still fast — so correctness-sensitive small slots keep exact results.
|
||
pub(crate) const USEARCH_MIN_VECTORS: usize = 10_000;
|
||
|
||
/// Build the vector index backend for a slot, gating on the expected vector
|
||
/// `count`: `UsearchIndex` (production HNSW) at or above [`USEARCH_MIN_VECTORS`],
|
||
/// `BruteForceIndex` (exact) below it.
|
||
///
|
||
/// When the `USearch` backend is selected, capacity is reserved up front
|
||
/// (`UsearchIndex::add` requires a prior `reserve`), so the subsequent inserts
|
||
/// cannot fail for lack of capacity. If `UsearchIndex` construction OR the
|
||
/// reservation fails for any
|
||
/// reason, this DEGRADES to `BruteForceIndex` with a `tracing::warn!` rather than
|
||
/// propagating — losing the latency win is acceptable; losing the slot entirely
|
||
/// (and therefore all ANN results for the kind) is not. This is the single place
|
||
/// that picks the engine; the write-path auto-register and the rebuild path both
|
||
/// route through it so the choice cannot drift.
|
||
pub(crate) fn build_slot_index(dimensions: usize, count: usize) -> Box<dyn VectorIndex> {
|
||
let config = VectorIndexConfig {
|
||
dimensions,
|
||
..VectorIndexConfig::default()
|
||
};
|
||
if count < USEARCH_MIN_VECTORS {
|
||
return Box::new(super::BruteForceIndex::new(config));
|
||
}
|
||
match super::UsearchIndex::new(config.clone()) {
|
||
Ok(index) => {
|
||
// USearch's `add` requires capacity reserved up front. Reserve the
|
||
// expected count so the rebuild's inserts all fit; the write path
|
||
// grows it further as needed.
|
||
if let Err(e) = index.reserve(count) {
|
||
tracing::warn!(
|
||
dimensions,
|
||
count,
|
||
error = %e,
|
||
"USearch reserve failed; falling back to brute-force for this slot"
|
||
);
|
||
return Box::new(super::BruteForceIndex::new(config));
|
||
}
|
||
tracing::info!(
|
||
dimensions,
|
||
count,
|
||
"using USearch (HNSW) backend for embedding slot"
|
||
);
|
||
Box::new(index)
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
dimensions,
|
||
count,
|
||
error = %e,
|
||
"USearch index construction failed; falling back to brute-force for this slot"
|
||
);
|
||
Box::new(super::BruteForceIndex::new(config))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Number of bytes per vector component at a given quantization level.
|
||
///
|
||
/// Used by [`EmbeddingSlotRegistry::index_stats`] to estimate index footprint
|
||
/// from the slot's *registered* precision rather than assuming f16.
|
||
const fn bytes_per_component(quantization: QuantizationLevel) -> u64 {
|
||
match quantization {
|
||
QuantizationLevel::F32 => 4,
|
||
QuantizationLevel::F16 => 2,
|
||
QuantizationLevel::Int8 => 1,
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// HNSW default parameters (single source of truth)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Default HNSW connectivity (M parameter): maximum bi-directional links per node.
|
||
///
|
||
/// M=16 is the sweet spot for quality vs. memory (Malkov & Yashunin, 2018).
|
||
/// This is the single authoritative source for the default; both [`HnswParams`]
|
||
/// and the vector index config derive their default `connectivity` from it.
|
||
pub const DEFAULT_CONNECTIVITY: usize = 16;
|
||
|
||
/// Default HNSW construction-time beam width (`ef_construction`).
|
||
///
|
||
/// `ef_construction=400` raises recall@10 to >0.99 at 100K vectors (128D) with
|
||
/// ~2× build overhead vs. ef=200 — acceptable for a write-once index. Grid search
|
||
/// at 100K vectors / 128D confirmed recall@10 ≈ 0.993 vs. 0.978 for ef=200
|
||
/// (see `docs/profiling/usearch-tuning.md`). Single authoritative source.
|
||
pub const DEFAULT_EF_CONSTRUCTION: usize = 400;
|
||
|
||
/// Default HNSW search-time beam width (`ef_search`). Can be overridden per-query.
|
||
/// Single authoritative source for the default.
|
||
pub const DEFAULT_EF_SEARCH: usize = 200;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// HNSW parameters
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// HNSW graph parameters for an embedding slot.
|
||
///
|
||
/// These are set at slot registration time and used to construct the HNSW
|
||
/// index. They cannot be changed without rebuilding the index.
|
||
#[derive(Debug, Clone)]
|
||
pub struct HnswParams {
|
||
/// Maximum number of bi-directional links per node (M parameter).
|
||
/// Higher = better recall, more memory. Default: [`DEFAULT_CONNECTIVITY`].
|
||
pub connectivity: usize,
|
||
/// Beam width during index construction. Higher = better graph quality,
|
||
/// slower builds. Default: [`DEFAULT_EF_CONSTRUCTION`].
|
||
pub ef_construction: usize,
|
||
/// Default beam width during search. Can be overridden per-query.
|
||
/// Default: [`DEFAULT_EF_SEARCH`].
|
||
pub ef_search: usize,
|
||
}
|
||
|
||
impl Default for HnswParams {
|
||
fn default() -> Self {
|
||
// All defaults and their tuning rationale live on the DEFAULT_* consts
|
||
// above — the single source of truth. Do not inline literals here.
|
||
Self {
|
||
connectivity: DEFAULT_CONNECTIVITY,
|
||
ef_construction: DEFAULT_EF_CONSTRUCTION,
|
||
ef_search: DEFAULT_EF_SEARCH,
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Embedding source
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Source of an embedding slot's vectors.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum EmbeddingSource {
|
||
/// Provided by the application via `write_item()` or `write_user()`.
|
||
External,
|
||
/// Computed and maintained by tidalDB (e.g., user preference vector,
|
||
/// creator catalog embedding).
|
||
DatabaseManaged,
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Slot state
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// State for a single embedding slot.
|
||
///
|
||
/// Combines the HNSW index with its configuration metadata. The index is
|
||
/// boxed behind the [`VectorIndex`] trait for swappability between
|
||
/// `UsearchIndex` (production) and `BruteForceIndex` (testing/small data).
|
||
pub struct EmbeddingSlotState {
|
||
/// The HNSW index for this slot.
|
||
pub index: Box<dyn VectorIndex>,
|
||
/// Number of dimensions for this slot.
|
||
pub dimensions: usize,
|
||
/// Quantization level used in the HNSW index.
|
||
pub quantization: QuantizationLevel,
|
||
/// Whether this embedding is externally provided or database-managed.
|
||
pub source: EmbeddingSource,
|
||
/// HNSW graph parameters.
|
||
pub params: HnswParams,
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Registry
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Registry of all embedding slots across all entity types.
|
||
///
|
||
/// Constructed from the schema at `TidalDB::open()` time. Each entity type
|
||
/// can define up to 4 embedding slots (per Entity Model Specification).
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```text
|
||
/// Item "content" -> 1536d, f16, External, M=16
|
||
/// Item "visual" -> 512d, f16, External, M=16
|
||
/// User "preference" -> 1536d, f16, DatabaseManaged, M=16
|
||
/// ```
|
||
pub struct EmbeddingSlotRegistry {
|
||
/// Two-level map: `EntityKind -> (slot_name -> state)`. The nested form lets
|
||
/// `get`/`get_mut` probe the inner map with a borrowed `&str`, avoiding the
|
||
/// per-call `String` heap allocation that a flat `(EntityKind, String)` key
|
||
/// would force on the per-query / per-write hot path. Slot counts are tiny
|
||
/// (<=4 per kind), so the extra indirection is negligible.
|
||
slots: HashMap<EntityKind, HashMap<String, EmbeddingSlotState>>,
|
||
}
|
||
|
||
impl EmbeddingSlotRegistry {
|
||
/// Create an empty registry.
|
||
#[must_use]
|
||
pub fn new() -> Self {
|
||
Self {
|
||
slots: HashMap::new(),
|
||
}
|
||
}
|
||
|
||
/// Register an embedding slot.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns [`VectorError::Backend`] if a slot with the same
|
||
/// `(entity_kind, slot_name)` is already registered.
|
||
pub fn register(
|
||
&mut self,
|
||
entity_kind: EntityKind,
|
||
slot_name: String,
|
||
state: EmbeddingSlotState,
|
||
) -> Result<(), VectorError> {
|
||
let inner = self.slots.entry(entity_kind).or_default();
|
||
if inner.contains_key(&slot_name) {
|
||
return Err(VectorError::Backend(format!(
|
||
"slot already registered: ({entity_kind}, \"{slot_name}\")"
|
||
)));
|
||
}
|
||
inner.insert(slot_name, state);
|
||
Ok(())
|
||
}
|
||
|
||
/// Look up an embedding slot by entity kind and slot name.
|
||
///
|
||
/// Returns `None` if the slot is not registered. Probes the inner map with a
|
||
/// borrowed `&str` — no allocation on this hot path.
|
||
#[must_use]
|
||
pub fn get(&self, entity_kind: EntityKind, slot_name: &str) -> Option<&EmbeddingSlotState> {
|
||
self.slots.get(&entity_kind)?.get(slot_name)
|
||
}
|
||
|
||
/// Look up an embedding slot mutably. Zero-alloc inner probe (see [`Self::get`]).
|
||
pub fn get_mut(
|
||
&mut self,
|
||
entity_kind: EntityKind,
|
||
slot_name: &str,
|
||
) -> Option<&mut EmbeddingSlotState> {
|
||
self.slots.get_mut(&entity_kind)?.get_mut(slot_name)
|
||
}
|
||
|
||
/// List all slot names for a given entity kind.
|
||
///
|
||
/// The returned names are in arbitrary order (`HashMap` iteration order).
|
||
#[must_use]
|
||
pub fn slots_for(&self, entity_kind: EntityKind) -> Vec<&str> {
|
||
self.slots
|
||
.get(&entity_kind)
|
||
.map(|inner| inner.keys().map(String::as_str).collect())
|
||
.unwrap_or_default()
|
||
}
|
||
|
||
/// Total number of registered slots across all entity kinds.
|
||
#[must_use]
|
||
pub fn slot_count(&self) -> usize {
|
||
self.slots.values().map(HashMap::len).sum()
|
||
}
|
||
|
||
/// Return the total vector count and estimated byte size across all slots.
|
||
///
|
||
/// Byte size estimate: `count * dimensions * bytes_per_component`, where
|
||
/// `bytes_per_component` is derived from each slot's *registered* quantization
|
||
/// (F32 = 4, F16 = 2, Int8 = 1). The previous implementation hardcoded 2
|
||
/// (f16), which undercounted by 2x for the F32 slots that production actually
|
||
/// registers (see `db::items::write_item_embedding`). This is still a lower
|
||
/// bound that excludes HNSW graph overhead.
|
||
#[must_use]
|
||
pub fn index_stats(&self) -> (u64, u64) {
|
||
let mut total_vectors: u64 = 0;
|
||
let mut total_bytes: u64 = 0;
|
||
for inner in self.slots.values() {
|
||
for slot in inner.values() {
|
||
let count = slot.index.len() as u64;
|
||
let dim = slot.dimensions as u64;
|
||
let bytes_per_component = bytes_per_component(slot.quantization);
|
||
total_vectors = total_vectors.saturating_add(count);
|
||
// Saturating arithmetic: this is a lower-bound footprint estimate,
|
||
// so saturating to u64::MAX at absurd scale is the correct semantic
|
||
// — it can never panic in debug nor silently wrap in release.
|
||
total_bytes = total_bytes.saturating_add(
|
||
count
|
||
.saturating_mul(dim)
|
||
.saturating_mul(bytes_per_component),
|
||
);
|
||
}
|
||
}
|
||
(total_vectors, total_bytes)
|
||
}
|
||
|
||
/// Rebuild every embedding slot index for `entity_kind` from durable storage.
|
||
///
|
||
/// The entity store is the source of truth for full-precision f32 embeddings:
|
||
/// each one is persisted under `Tag::Meta` with the `EMB:{slot}` suffix (see
|
||
/// [`crate::storage::vector::embedding_store_key`]). The in-memory ANN index,
|
||
/// by contrast, is *derived* state that does not survive a process restart.
|
||
/// Without this rebuild, ANN search silently returns nothing after every
|
||
/// reopen until each entity's embedding is written again -- a documented
|
||
/// requirement that previously had no implementation.
|
||
///
|
||
/// This method scans `storage` for `EMB:` keys, deserializes each embedding,
|
||
/// lazily registers the slot the first time it is seen (deriving the
|
||
/// dimensionality from the stored vector and matching the precision/source of
|
||
/// the write path -- F32, [`EmbeddingSource::External`]), and inserts the
|
||
/// vector into the slot's index. It is idempotent: re-running it over the same
|
||
/// store re-inserts the same `(id, vector)` pairs.
|
||
///
|
||
/// `schema` is accepted so future precision/source declarations can override
|
||
/// the write-path defaults; today [`crate::schema::EmbeddingSlotDef`] carries
|
||
/// only `name`/`entity_kind`/`dimensions`, so it is used purely to scope which
|
||
/// slot names are expected for `entity_kind`.
|
||
///
|
||
/// Returns the number of vectors inserted across all slots for this kind.
|
||
///
|
||
/// # Backend selection
|
||
///
|
||
/// Each slot's backend is chosen by [`build_slot_index`] from the slot's live
|
||
/// vector count: a slot at or above [`USEARCH_MIN_VECTORS`] is rebuilt as the
|
||
/// production HNSW engine (`UsearchIndex`); smaller slots stay exact
|
||
/// (`BruteForceIndex`). This is the live path that makes `UsearchIndex`
|
||
/// reachable in production after a restart, with no config plumbing.
|
||
///
|
||
/// # Fault isolation
|
||
///
|
||
/// A single un-indexable embedding (un-deserializable, zero-dimension, or a
|
||
/// dimensionality that mismatches its slot) is SKIPPED with a `tracing::warn!`
|
||
/// and counted, not propagated — one poison row must never abort the rebuild
|
||
/// for an entire entity kind (which would leave ANN search empty for every
|
||
/// good embedding too). The skipped count is logged at the end.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// - [`VectorError::Io`] if a storage scan entry fails to read (a whole-scan
|
||
/// failure is genuinely unrecoverable, so it still propagates).
|
||
/// - [`VectorError::Backend`] if slot registration fails.
|
||
#[allow(clippy::too_many_lines)]
|
||
pub(crate) fn rebuild_from_store(
|
||
&mut self,
|
||
entity_kind: EntityKind,
|
||
storage: &dyn crate::storage::StorageEngine,
|
||
schema: &crate::schema::Schema,
|
||
) -> Result<usize, VectorError> {
|
||
use std::collections::HashSet;
|
||
|
||
use super::lifecycle::ops::ARCHIVED_PREFIX;
|
||
use crate::storage::{Tag, keys::parse_key};
|
||
|
||
// Suffix marker that distinguishes embedding keys from other Tag::Meta keys.
|
||
const EMB_PREFIX: &[u8] = b"EMB:";
|
||
|
||
// Schema-declared dimensions for this kind, keyed by slot name. Used only
|
||
// as a cross-check against the authoritative stored header.
|
||
let declared_dims: HashMap<&str, usize> = schema
|
||
.embedding_slots()
|
||
.iter()
|
||
.filter(|s| s.entity_kind == entity_kind)
|
||
.map(|s| (s.name.as_str(), s.dimensions))
|
||
.collect();
|
||
|
||
// Single scan: tombstones (ARCHIVED:) and embeddings (EMB:) are interleaved
|
||
// under Tag::Meta in lexicographic order, and a tombstone may appear AFTER
|
||
// its embedding, so we cannot decide an embedding's fate the moment we see
|
||
// it. Instead, one pass partitions the keyspace into the archived set and
|
||
// the embedding candidates (already deserialized), then we filter and index
|
||
// afterwards. This replaces the previous two full O(N) store scans, and the
|
||
// candidate owns its slot `String` so the archived-set lookup borrows it
|
||
// rather than allocating a throwaway `String` per embedding key.
|
||
let mut archived: HashSet<(u64, String)> = HashSet::new();
|
||
let mut candidates: Vec<(crate::schema::EntityId, String, Vec<f32>)> = Vec::new();
|
||
// Count of embeddings skipped due to a per-row fault (un-deserializable,
|
||
// zero-dim, or un-indexable). Reported at the end so a degraded rebuild is
|
||
// observable rather than silent.
|
||
let mut skipped = 0usize;
|
||
for entry in storage.scan_prefix(&[]) {
|
||
let (key, value) = entry.map_err(|e| {
|
||
VectorError::Io(std::io::Error::other(format!(
|
||
"embedding rebuild scan failed: {e}"
|
||
)))
|
||
})?;
|
||
|
||
// Only Tag::Meta keys carry embeddings or archive tombstones.
|
||
let Some((entity_id, Tag::Meta, suffix)) = parse_key(&key) else {
|
||
continue;
|
||
};
|
||
|
||
if let Some(slot_bytes) = suffix.strip_prefix(ARCHIVED_PREFIX) {
|
||
// A non-UTF-8 slot name cannot have been produced by the writer; skip.
|
||
if let Ok(slot_name) = std::str::from_utf8(slot_bytes) {
|
||
archived.insert((entity_id.as_u64(), slot_name.to_string()));
|
||
}
|
||
continue;
|
||
}
|
||
|
||
let Some(slot_bytes) = suffix.strip_prefix(EMB_PREFIX) else {
|
||
// Plain metadata (title, etc.) — not an embedding or tombstone.
|
||
continue;
|
||
};
|
||
// A non-UTF-8 slot name cannot have been produced by
|
||
// embedding_store_key; skip it rather than abort the rebuild.
|
||
let Ok(slot_name) = std::str::from_utf8(slot_bytes) else {
|
||
tracing::warn!(
|
||
entity_id = entity_id.as_u64(),
|
||
"skipping embedding key with non-UTF-8 slot name during rebuild"
|
||
);
|
||
continue;
|
||
};
|
||
|
||
// Deserialize the source-of-truth vector. A single un-deserializable
|
||
// embedding (corrupt header, hand-edited store) must NOT abort the
|
||
// rebuild for the whole entity kind — that would leave ANN search empty
|
||
// for every good embedding too. Isolate the fault: warn, count it, and
|
||
// skip this one row. The durable value is untouched; the next write
|
||
// re-inserts it. (W: one poison row must not brick the kind.)
|
||
let vector = match super::deserialize_embedding(&value) {
|
||
Ok(v) => v,
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
entity_id = entity_id.as_u64(),
|
||
slot = slot_name,
|
||
error = %e,
|
||
"skipping un-deserializable embedding during rebuild"
|
||
);
|
||
skipped += 1;
|
||
continue;
|
||
}
|
||
};
|
||
candidates.push((entity_id, slot_name.to_string(), vector));
|
||
}
|
||
|
||
// Group the surviving (non-archived, non-degenerate) embeddings by slot so
|
||
// each slot's vector count is known before its backend is built. The count
|
||
// drives the backend selection (USearch HNSW for large slots, brute force
|
||
// for small) AND the capacity reservation, with NO external config plumbing.
|
||
let mut by_slot: HashMap<String, Vec<(crate::schema::EntityId, Vec<f32>)>> = HashMap::new();
|
||
for (entity_id, slot_name, vector) in candidates {
|
||
// Skip embeddings that were soft-deleted (archived). The source-of-truth
|
||
// value is intentionally retained on disk, but it must not re-enter the
|
||
// ANN index — that is the whole point of the durable tombstone.
|
||
if archived.contains(&(entity_id.as_u64(), slot_name.clone())) {
|
||
continue;
|
||
}
|
||
|
||
// Reject a zero-dimension vector. The write path can never produce one
|
||
// (`l2_normalize` rejects an empty vector before anything is stored),
|
||
// but `serialize_embedding(&[])` is a valid 4-byte payload that
|
||
// `deserialize_embedding` accepts, so a corrupt / hand-crafted store
|
||
// could carry one. A 0-dim slot makes `l2_distance_sq` return 0.0 for
|
||
// every query, so the entity would tie for nearest on every search and
|
||
// pollute results. Skip it loudly. (W: corrupt 0-dim slot.)
|
||
if vector.is_empty() {
|
||
tracing::warn!(
|
||
entity_id = entity_id.as_u64(),
|
||
slot = slot_name.as_str(),
|
||
"skipping zero-dimension embedding during rebuild (corrupt store)"
|
||
);
|
||
skipped += 1;
|
||
continue;
|
||
}
|
||
|
||
by_slot
|
||
.entry(slot_name)
|
||
.or_default()
|
||
.push((entity_id, vector));
|
||
}
|
||
|
||
let mut inserted = 0usize;
|
||
|
||
for (slot_name, entries) in by_slot {
|
||
// The slot's dimensionality is the FIRST stored vector's length; any
|
||
// later vector of a different length is fault-isolated below. Cross-
|
||
// check against the schema declaration (advisory only — the stored
|
||
// header is authoritative for what was actually written).
|
||
let Some((_, first_vec)) = entries.first() else {
|
||
continue;
|
||
};
|
||
let dimensions = first_vec.len();
|
||
|
||
if let Some(&declared) = declared_dims.get(slot_name.as_str())
|
||
&& declared != dimensions
|
||
{
|
||
tracing::warn!(
|
||
slot = slot_name.as_str(),
|
||
declared,
|
||
stored = dimensions,
|
||
"embedding dimension mismatch between schema and stored vector during rebuild; \
|
||
using stored dimensionality"
|
||
);
|
||
}
|
||
|
||
// Lazily register the slot on first sight, routing the backend through
|
||
// the size-gated factory: a slot with >= USEARCH_MIN_VECTORS live
|
||
// vectors is rebuilt as the production HNSW (USearch) engine; smaller
|
||
// slots stay exact (brute force). This is the live path that makes
|
||
// USearch reachable in production restarts — no config flag required.
|
||
if self.get(entity_kind, &slot_name).is_none() {
|
||
let state = EmbeddingSlotState {
|
||
index: build_slot_index(dimensions, entries.len()),
|
||
dimensions,
|
||
quantization: QuantizationLevel::F32,
|
||
source: EmbeddingSource::External,
|
||
params: HnswParams::default(),
|
||
};
|
||
self.register(entity_kind, slot_name.clone(), state)?;
|
||
}
|
||
|
||
let slot = self.get(entity_kind, &slot_name).ok_or_else(|| {
|
||
VectorError::Backend(format!(
|
||
"slot ({entity_kind}, \"{slot_name}\") missing immediately after registration"
|
||
))
|
||
})?;
|
||
|
||
for (entity_id, vector) in entries {
|
||
// Fault-isolate the per-embedding insert: a single row whose
|
||
// dimensionality differs from the slot's (e.g. a schema-revision
|
||
// mismatch that still self-deserializes) returns DimensionMismatch.
|
||
// Skip it loudly rather than abort the kind's rebuild — the
|
||
// hundreds of thousands of good embeddings for the same slot must
|
||
// still be indexed. (W: one poison row must not brick the kind.)
|
||
if let Err(e) = slot.index.insert(entity_id.as_u64(), &vector) {
|
||
tracing::warn!(
|
||
entity_id = entity_id.as_u64(),
|
||
slot = slot_name.as_str(),
|
||
error = %e,
|
||
"skipping un-indexable embedding during rebuild"
|
||
);
|
||
skipped += 1;
|
||
continue;
|
||
}
|
||
inserted += 1;
|
||
}
|
||
}
|
||
|
||
if inserted > 0 || skipped > 0 {
|
||
tracing::info!(
|
||
entity_kind = %entity_kind,
|
||
vectors = inserted,
|
||
skipped,
|
||
"embedding indexes rebuilt from durable storage"
|
||
);
|
||
}
|
||
|
||
Ok(inserted)
|
||
}
|
||
}
|
||
|
||
impl Default for EmbeddingSlotRegistry {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[cfg(test)]
|
||
#[allow(clippy::unwrap_used)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::storage::engine::StorageEngine;
|
||
use crate::{
|
||
schema::EntityKind,
|
||
storage::vector::{BruteForceIndex, QuantizationLevel, VectorIndexConfig},
|
||
};
|
||
|
||
/// Helper: create a slot state with the given dimensions.
|
||
fn make_slot(dimensions: usize, source: EmbeddingSource) -> EmbeddingSlotState {
|
||
let config = VectorIndexConfig {
|
||
dimensions,
|
||
..VectorIndexConfig::default()
|
||
};
|
||
EmbeddingSlotState {
|
||
index: Box::new(BruteForceIndex::new(config)),
|
||
dimensions,
|
||
quantization: QuantizationLevel::F16,
|
||
source,
|
||
params: HnswParams::default(),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn registry_register_and_lookup() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let state = make_slot(1536, EmbeddingSource::External);
|
||
|
||
registry
|
||
.register(EntityKind::Item, "content".into(), state)
|
||
.unwrap();
|
||
|
||
let slot = registry.get(EntityKind::Item, "content");
|
||
assert!(slot.is_some());
|
||
let slot = slot.unwrap();
|
||
assert_eq!(slot.dimensions, 1536);
|
||
assert_eq!(slot.source, EmbeddingSource::External);
|
||
}
|
||
|
||
#[test]
|
||
fn registry_duplicate_slot_fails() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
|
||
let state1 = make_slot(1536, EmbeddingSource::External);
|
||
let state2 = make_slot(1536, EmbeddingSource::External);
|
||
|
||
registry
|
||
.register(EntityKind::Item, "content".into(), state1)
|
||
.unwrap();
|
||
let result = registry.register(EntityKind::Item, "content".into(), state2);
|
||
assert!(result.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn registry_different_entity_kinds_same_name() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
|
||
let state_item = make_slot(1536, EmbeddingSource::External);
|
||
let state_user = make_slot(1536, EmbeddingSource::DatabaseManaged);
|
||
|
||
registry
|
||
.register(EntityKind::Item, "content".into(), state_item)
|
||
.unwrap();
|
||
registry
|
||
.register(EntityKind::User, "content".into(), state_user)
|
||
.unwrap();
|
||
|
||
let item_slot = registry.get(EntityKind::Item, "content").unwrap();
|
||
let user_slot = registry.get(EntityKind::User, "content").unwrap();
|
||
assert_eq!(item_slot.source, EmbeddingSource::External);
|
||
assert_eq!(user_slot.source, EmbeddingSource::DatabaseManaged);
|
||
}
|
||
|
||
#[test]
|
||
fn registry_slots_for_entity_kind() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
|
||
for name in &["content", "visual", "audio"] {
|
||
let state = make_slot(128, EmbeddingSource::External);
|
||
registry
|
||
.register(EntityKind::Item, (*name).to_string(), state)
|
||
.unwrap();
|
||
}
|
||
|
||
let mut slots = registry.slots_for(EntityKind::Item);
|
||
slots.sort_unstable(); // HashMap order is nondeterministic
|
||
assert_eq!(slots.len(), 3);
|
||
assert!(slots.contains(&"content"));
|
||
assert!(slots.contains(&"visual"));
|
||
assert!(slots.contains(&"audio"));
|
||
|
||
// No user slots
|
||
let user_slots = registry.slots_for(EntityKind::User);
|
||
assert!(user_slots.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn registry_nonexistent_slot_returns_none() {
|
||
let registry = EmbeddingSlotRegistry::new();
|
||
assert!(registry.get(EntityKind::Item, "content").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn registry_slot_count() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
assert_eq!(registry.slot_count(), 0);
|
||
|
||
registry
|
||
.register(
|
||
EntityKind::Item,
|
||
"content".into(),
|
||
make_slot(1536, EmbeddingSource::External),
|
||
)
|
||
.unwrap();
|
||
assert_eq!(registry.slot_count(), 1);
|
||
|
||
registry
|
||
.register(
|
||
EntityKind::User,
|
||
"preference".into(),
|
||
make_slot(1536, EmbeddingSource::DatabaseManaged),
|
||
)
|
||
.unwrap();
|
||
assert_eq!(registry.slot_count(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn registry_get_mut() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
registry
|
||
.register(
|
||
EntityKind::Item,
|
||
"content".into(),
|
||
make_slot(1536, EmbeddingSource::External),
|
||
)
|
||
.unwrap();
|
||
|
||
let slot = registry.get_mut(EntityKind::Item, "content").unwrap();
|
||
// Mutate dimensions to prove we have mutable access
|
||
slot.dimensions = 512;
|
||
|
||
let slot = registry.get(EntityKind::Item, "content").unwrap();
|
||
assert_eq!(slot.dimensions, 512);
|
||
}
|
||
|
||
#[test]
|
||
fn index_stats_empty_registry() {
|
||
let registry = EmbeddingSlotRegistry::new();
|
||
let (vectors, bytes) = registry.index_stats();
|
||
assert_eq!(vectors, 0);
|
||
assert_eq!(bytes, 0);
|
||
}
|
||
|
||
#[test]
|
||
fn index_stats_with_vectors() {
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let dim = 128;
|
||
let state = make_slot(dim, EmbeddingSource::External);
|
||
registry
|
||
.register(EntityKind::Item, "content".into(), state)
|
||
.unwrap();
|
||
|
||
// Insert a vector.
|
||
let slot = registry.get(EntityKind::Item, "content").unwrap();
|
||
let embedding = vec![0.1_f32; dim];
|
||
slot.index.insert(1, &embedding).unwrap();
|
||
|
||
let (vectors, bytes) = registry.index_stats();
|
||
assert_eq!(vectors, 1);
|
||
// 1 vector * 128 dims * 2 bytes (f16) = 256 bytes
|
||
assert_eq!(bytes, 256);
|
||
}
|
||
|
||
#[test]
|
||
fn registry_default_is_empty() {
|
||
let registry = EmbeddingSlotRegistry::default();
|
||
assert_eq!(registry.slot_count(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn hnsw_params_default() {
|
||
let params = HnswParams::default();
|
||
assert_eq!(params.connectivity, 16);
|
||
assert_eq!(params.ef_construction, 400);
|
||
assert_eq!(params.ef_search, 200);
|
||
}
|
||
|
||
#[test]
|
||
fn index_stats_uses_registered_precision_for_f32() {
|
||
// Regression: index_stats previously hardcoded f16 (*2), undercounting
|
||
// the F32 slots that production actually registers by 2x.
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let dim = 128;
|
||
let config = VectorIndexConfig {
|
||
dimensions: dim,
|
||
..VectorIndexConfig::default()
|
||
};
|
||
let state = EmbeddingSlotState {
|
||
index: Box::new(BruteForceIndex::new(config)),
|
||
dimensions: dim,
|
||
quantization: QuantizationLevel::F32,
|
||
source: EmbeddingSource::External,
|
||
params: HnswParams::default(),
|
||
};
|
||
registry
|
||
.register(EntityKind::Item, "content".into(), state)
|
||
.unwrap();
|
||
|
||
let slot = registry.get(EntityKind::Item, "content").unwrap();
|
||
slot.index.insert(1, &vec![0.1_f32; dim]).unwrap();
|
||
|
||
let (vectors, bytes) = registry.index_stats();
|
||
assert_eq!(vectors, 1);
|
||
// 1 vector * 128 dims * 4 bytes (f32) = 512 bytes.
|
||
assert_eq!(bytes, 512);
|
||
}
|
||
|
||
// -------------------------------------------------------------------
|
||
// rebuild_from_store
|
||
// -------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn rebuild_from_store_repopulates_index_and_is_searchable() {
|
||
use crate::{
|
||
schema::{DecaySpec, EntityId, SchemaBuilder},
|
||
storage::{
|
||
memory::InMemoryBackend,
|
||
vector::{embedding_store_key, serialize_embedding},
|
||
},
|
||
};
|
||
|
||
// Schema declaring a single Item embedding slot (dimensions cross-checked
|
||
// against the stored header during rebuild).
|
||
let dim = 4;
|
||
let mut builder = SchemaBuilder::new();
|
||
let _ = builder
|
||
.signal("view", EntityKind::Item, DecaySpec::Permanent)
|
||
.add();
|
||
builder.embedding_slot("content", EntityKind::Item, dim);
|
||
let schema = builder.build().expect("valid schema");
|
||
|
||
// Pre-populate a store with durable EMB:content embeddings, exactly as
|
||
// the write path would. Vectors are NOT normalized here on purpose:
|
||
// rebuild stores whatever is durable, the source of truth.
|
||
let storage = InMemoryBackend::new();
|
||
let vectors: [(u64, [f32; 4]); 3] = [
|
||
(1, [1.0, 0.0, 0.0, 0.0]),
|
||
(2, [0.0, 1.0, 0.0, 0.0]),
|
||
(3, [0.0, 0.0, 1.0, 0.0]),
|
||
];
|
||
for (id, v) in &vectors {
|
||
let key = embedding_store_key(EntityId::new(*id), "content");
|
||
storage.put(&key, &serialize_embedding(v)).unwrap();
|
||
}
|
||
// A non-embedding Tag::Meta key must be ignored by the scan.
|
||
let meta_key =
|
||
crate::storage::encode_key(EntityId::new(1), crate::storage::Tag::Meta, b"title");
|
||
storage.put(&meta_key, b"hello").unwrap();
|
||
|
||
// Rebuild an empty registry from the store.
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let inserted = registry
|
||
.rebuild_from_store(EntityKind::Item, &storage, &schema)
|
||
.expect("rebuild succeeds");
|
||
|
||
assert_eq!(inserted, 3, "all three embeddings should be re-indexed");
|
||
|
||
let slot = registry
|
||
.get(EntityKind::Item, "content")
|
||
.expect("slot registered during rebuild");
|
||
assert_eq!(slot.dimensions, dim);
|
||
assert_eq!(slot.quantization, QuantizationLevel::F32);
|
||
assert_eq!(slot.source, EmbeddingSource::External);
|
||
assert_eq!(slot.index.len(), 3);
|
||
|
||
// ANN search must now return the rebuilt vectors (the whole point of the
|
||
// fix: ANN search returned nothing after reopen before this existed).
|
||
let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 1, 0).unwrap();
|
||
assert_eq!(results.len(), 1);
|
||
assert_eq!(results[0].id, 1, "nearest neighbor of e0 is vector id 1");
|
||
}
|
||
|
||
#[test]
|
||
fn rebuild_skips_soft_deleted_embedding() {
|
||
// Finding-1 regression: a soft delete must survive a restart. We write
|
||
// an embedding through the real lifecycle path, soft-delete it (writing a
|
||
// durable tombstone), then rebuild a *fresh* registry from the same store
|
||
// — simulating a process restart — and assert the vector is absent from
|
||
// both the index and ANN search results.
|
||
use crate::{
|
||
schema::{DecaySpec, EntityId, SchemaBuilder},
|
||
storage::{
|
||
memory::InMemoryBackend,
|
||
vector::{BruteForceIndex, VectorIndexConfig, delete_embedding, insert_embedding},
|
||
},
|
||
};
|
||
|
||
let dim = 4;
|
||
let mut builder = SchemaBuilder::new();
|
||
let _ = builder
|
||
.signal("view", EntityKind::Item, DecaySpec::Permanent)
|
||
.add();
|
||
builder.embedding_slot("content", EntityKind::Item, dim);
|
||
let schema = builder.build().expect("valid schema");
|
||
|
||
let storage = InMemoryBackend::new();
|
||
// A scratch index used only to drive the live write/delete path; it is
|
||
// discarded before the rebuild (the rebuild constructs its own).
|
||
let scratch = BruteForceIndex::new(VectorIndexConfig {
|
||
dimensions: dim,
|
||
..VectorIndexConfig::default()
|
||
});
|
||
|
||
// Two embeddings; soft-delete the first, keep the second.
|
||
insert_embedding(
|
||
EntityId::new(1),
|
||
"content",
|
||
&[1.0, 0.0, 0.0, 0.0],
|
||
dim,
|
||
&scratch,
|
||
&storage,
|
||
)
|
||
.unwrap();
|
||
insert_embedding(
|
||
EntityId::new(2),
|
||
"content",
|
||
&[0.0, 1.0, 0.0, 0.0],
|
||
dim,
|
||
&scratch,
|
||
&storage,
|
||
)
|
||
.unwrap();
|
||
delete_embedding(EntityId::new(1), "content", &scratch, &storage, false).unwrap();
|
||
drop(scratch);
|
||
|
||
// Simulate restart: fresh registry rebuilt from durable storage.
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let inserted = registry
|
||
.rebuild_from_store(EntityKind::Item, &storage, &schema)
|
||
.expect("rebuild succeeds");
|
||
|
||
// Only the non-archived embedding (id 2) is re-indexed.
|
||
assert_eq!(inserted, 1, "soft-deleted embedding must not be re-indexed");
|
||
let slot = registry
|
||
.get(EntityKind::Item, "content")
|
||
.expect("slot registered during rebuild");
|
||
assert_eq!(slot.index.len(), 1);
|
||
|
||
// ANN search must not surface the soft-deleted vector even when queried
|
||
// with its own embedding.
|
||
let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 10, 0).unwrap();
|
||
assert!(
|
||
results.iter().all(|r| r.id != 1),
|
||
"soft-deleted vector 1 resurfaced after rebuild"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn rebuild_from_store_empty_store_inserts_nothing() {
|
||
use crate::{
|
||
schema::{DecaySpec, SchemaBuilder},
|
||
storage::memory::InMemoryBackend,
|
||
};
|
||
|
||
let mut builder = SchemaBuilder::new();
|
||
let _ = builder
|
||
.signal("view", EntityKind::Item, DecaySpec::Permanent)
|
||
.add();
|
||
builder.embedding_slot("content", EntityKind::Item, 4);
|
||
let schema = builder.build().expect("valid schema");
|
||
|
||
let storage = InMemoryBackend::new();
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let inserted = registry
|
||
.rebuild_from_store(EntityKind::Item, &storage, &schema)
|
||
.expect("rebuild succeeds on empty store");
|
||
|
||
assert_eq!(inserted, 0);
|
||
// No slot is registered when there is nothing durable to index.
|
||
assert!(registry.get(EntityKind::Item, "content").is_none());
|
||
}
|
||
|
||
// -------------------------------------------------------------------
|
||
// build_slot_index backend selection (USearch wiring) — finding 5
|
||
// -------------------------------------------------------------------
|
||
|
||
/// CRITICAL regression: the size-gated factory must build the production
|
||
/// HNSW (`UsearchIndex`) at/above the threshold and the exact `BruteForceIndex`
|
||
/// below it, and BOTH must be functional. Before the fix `USearch` was wired
|
||
/// into nothing and every slot was brute force.
|
||
#[test]
|
||
fn build_slot_index_selects_backend_by_count() {
|
||
// Below threshold: exact brute force, functional.
|
||
let small = build_slot_index(4, 0);
|
||
small.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
|
||
small.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();
|
||
let r = small.search(&[1.0, 0.0, 0.0, 0.0], 1, 0).unwrap();
|
||
assert_eq!(r[0].id, 1, "brute-force backend must return the nearest");
|
||
|
||
// At threshold: USearch (HNSW), with capacity reserved by the factory so
|
||
// inserts succeed without a manual `reserve`. Searchable end-to-end —
|
||
// proving the production engine is actually reachable, not dead code.
|
||
let big = build_slot_index(4, USEARCH_MIN_VECTORS);
|
||
big.insert(10, &[1.0, 0.0, 0.0, 0.0]).unwrap();
|
||
big.insert(20, &[0.0, 1.0, 0.0, 0.0]).unwrap();
|
||
let r = big.search(&[1.0, 0.0, 0.0, 0.0], 1, big.len()).unwrap();
|
||
assert_eq!(
|
||
r[0].id, 10,
|
||
"USearch backend must be constructed AND searchable at the threshold"
|
||
);
|
||
}
|
||
|
||
// -------------------------------------------------------------------
|
||
// rebuild fault isolation — findings 8 and 13
|
||
// -------------------------------------------------------------------
|
||
|
||
/// W regression: a single un-indexable embedding (a different dimension under
|
||
/// the same slot) must NOT abort the rebuild — the good embeddings for the
|
||
/// kind are still indexed and searchable; only the poison row is skipped.
|
||
#[test]
|
||
fn rebuild_isolates_dimension_mismatch_and_zero_dim() {
|
||
use crate::{
|
||
schema::{DecaySpec, EntityId, SchemaBuilder},
|
||
storage::{
|
||
memory::InMemoryBackend,
|
||
vector::{embedding_store_key, serialize_embedding},
|
||
},
|
||
};
|
||
|
||
let dim = 4;
|
||
let mut builder = SchemaBuilder::new();
|
||
let _ = builder
|
||
.signal("view", EntityKind::Item, DecaySpec::Permanent)
|
||
.add();
|
||
builder.embedding_slot("content", EntityKind::Item, dim);
|
||
let schema = builder.build().expect("valid schema");
|
||
|
||
let storage = InMemoryBackend::new();
|
||
// Two valid 4-D embeddings.
|
||
for (id, v) in &[
|
||
(1u64, [1.0f32, 0.0, 0.0, 0.0]),
|
||
(2u64, [0.0f32, 1.0, 0.0, 0.0]),
|
||
] {
|
||
let key = embedding_store_key(EntityId::new(*id), "content");
|
||
storage.put(&key, &serialize_embedding(v)).unwrap();
|
||
}
|
||
// One poison row: a 2-D vector under the SAME slot. Its stored header is
|
||
// self-consistent (deserializes fine) but its dimension differs from the
|
||
// slot's (4) → DimensionMismatch on insert. It must be skipped, not abort.
|
||
let bad_key = embedding_store_key(EntityId::new(3), "content");
|
||
storage
|
||
.put(&bad_key, &serialize_embedding(&[1.0, 1.0]))
|
||
.unwrap();
|
||
// One zero-dimension row (corrupt store): a valid 4-byte payload that
|
||
// deserializes to an empty vector. Must be skipped (finding 13).
|
||
let empty_key = embedding_store_key(EntityId::new(4), "content");
|
||
storage.put(&empty_key, &serialize_embedding(&[])).unwrap();
|
||
|
||
let mut registry = EmbeddingSlotRegistry::new();
|
||
let inserted = registry
|
||
.rebuild_from_store(EntityKind::Item, &storage, &schema)
|
||
.expect("rebuild must not abort on a poison row");
|
||
|
||
assert_eq!(
|
||
inserted, 2,
|
||
"the two valid embeddings are indexed; the mismatch and zero-dim rows are skipped"
|
||
);
|
||
let slot = registry
|
||
.get(EntityKind::Item, "content")
|
||
.expect("slot registered from the valid embeddings");
|
||
assert_eq!(
|
||
slot.dimensions, dim,
|
||
"the slot must take the valid (4-D) dimensionality, not the corrupt 0-dim row"
|
||
);
|
||
assert_eq!(slot.index.len(), 2);
|
||
|
||
// The two valid embeddings are searchable; the bad ids never resurface.
|
||
let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 10, 0).unwrap();
|
||
assert!(
|
||
results.iter().all(|r| r.id != 3 && r.id != 4),
|
||
"poison rows (ids 3, 4) must not be indexed"
|
||
);
|
||
assert!(
|
||
results.iter().any(|r| r.id == 1),
|
||
"valid embedding 1 must be searchable"
|
||
);
|
||
}
|
||
}
|