mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-07-09 17:28:42 +00:00
feat(rabitq,rulake): persist end-to-end — save_cache_to_dir + warm_from_dir
Wires the previously-shipped rabitq::persist module into ruLake's
lake.rs as first-class cache-save/restore APIs. The architectural
blocker I've deferred across 3 waves is now closed.
=== Agent A: rabitq::RabitqPlusIndex::export_items() ===
crates/ruvector-rabitq/src/index.rs +1 method, +1 test.
Exposes `export_items() -> Vec<(usize, Vec<f32>)>` — each row as
(pos, original_vec) extracted from originals_flat with one clone per
row. Feeds directly into persist::save_index or
from_vectors_parallel_with_rotation. No new deps, no public API
breakage.
Regression test (`export_items_roundtrip_via_from_vectors_parallel`)
builds via serial add(), exports, rebuilds via the parallel path,
asserts byte-identical search results on 5 queries. Tests: 36 → 37.
=== Agent B: RuLake save_cache_to_dir + warm_from_dir ===
crates/ruvector-rulake/src/{cache.rs, lake.rs, tests/federation_smoke.rs}.
New API:
pub fn save_cache_to_dir(&self, key, dir) -> Result<PathBuf>
— writes dir/index.rbpx (atomic temp+rename+fsync) alongside
the table.rulake.json bundle sidecar. Uses export_items +
persist::save_index.
pub fn warm_from_dir(&self, key, dir) -> Result<usize>
— reads bundle, witness-verifies, loads index.rbpx via
persist::load_index, cross-checks dim+rerank_factor, installs
into cache via the new install_prebuilt path. Returns n vectors.
Does NOT require the backend to be registered — warm restart
without backend RTT is the point.
New on CacheStats: warm_installs counter (separate from primes so
warm-restart cost isn't confused with cold-prime cost).
New on VectorCache: install_prebuilt + install_prebuilt_interned —
insert a pre-built Arc<RabitqPlusIndex> at a known witness without
any prime-timer bookkeeping. Respects the LRU cap. Shared-entry
path reuses an existing witness entry if another pointer already
holds it (witness-addressed cache sharing remains the headline).
New test: `warm_from_dir_skips_backend_and_returns_bit_exact_results`
Prime a 50-vec D=8 collection, save, spin up a FRESH RuLake with
NO backend registered + Consistency::Frozen, warm_from_dir, run the
same query, assert byte-identical ids + f32 score bits,
warm_installs=1, primes=0. Closes the "restart without re-prime"
gap end-to-end.
Documented limitation: pos_to_id reconstructed as (0..n) identity
because RabitqPlusIndex doesn't expose outer ids() accessor, and
the rabitq agent's scope prohibited adding it. Every current prime
path uses positional ids so this is byte-equivalent to the real
ids; external non-dense u64 ids would collapse (a known M2+ issue
filed inline).
Tests: 37 rabitq + 21 rulake lib + 22 rulake federation = 80 total.
Clippy -D warnings clean across both crates.
Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
parent
21396313f1
commit
0ceba2a032
4 changed files with 520 additions and 0 deletions
|
|
@ -553,6 +553,31 @@ impl RabitqPlusIndex {
|
|||
let dim = self.inner.dim;
|
||||
&self.originals_flat[pos * dim..(pos + 1) * dim]
|
||||
}
|
||||
|
||||
/// Export every stored vector as `(pos, row_vec)` pairs, suitable for
|
||||
/// feeding directly into [`Self::from_vectors_parallel_with_rotation`]
|
||||
/// or `persist::save_index(.., &items, ..)`.
|
||||
///
|
||||
/// `pos` is the row index in `0..self.len()` (matching the internal SoA
|
||||
/// layout — not the external `ids[pos]`). Each row is cloned exactly
|
||||
/// once out of `originals_flat`, so the returned `Vec` owns the data
|
||||
/// and holds no borrow on `self`.
|
||||
///
|
||||
/// Added for `ruvector-rulake` so a primed `VectorCache` entry can be
|
||||
/// serialized via the existing `persist::save_index` API without
|
||||
/// round-tripping through the backend.
|
||||
pub fn export_items(&self) -> Vec<(usize, Vec<f32>)> {
|
||||
let dim = self.inner.dim;
|
||||
let n = self.inner.ids.len();
|
||||
(0..n)
|
||||
.map(|pos| {
|
||||
(
|
||||
pos,
|
||||
self.originals_flat[pos * dim..(pos + 1) * dim].to_vec(),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl RabitqPlusIndex {
|
||||
|
|
@ -1272,4 +1297,59 @@ mod tests {
|
|||
"Hadamard memory={had_bytes} vs Haar={haar_bytes} — expected ≥ 30× reduction",
|
||||
);
|
||||
}
|
||||
|
||||
/// `export_items()` must round-trip: rebuilding a `RabitqPlusIndex` from
|
||||
/// the exported `(pos, row_vec)` pairs — with the same seed, rotation
|
||||
/// kind, and rerank factor — must produce byte-identical search results
|
||||
/// to the source index. This is the contract `ruvector-rulake` relies
|
||||
/// on to serialize a primed cache entry without re-pulling vectors
|
||||
/// from the backend.
|
||||
#[test]
|
||||
fn export_items_roundtrip_via_from_vectors_parallel() {
|
||||
let d = 16;
|
||||
let n = 100;
|
||||
let seed = 20_260_423_u64;
|
||||
let rerank = 4;
|
||||
let kind = RandomRotationKind::HaarDense;
|
||||
|
||||
let data = make_dataset(n, d, seed);
|
||||
|
||||
// Source: add vectors one at a time so the export path is the only
|
||||
// place we round-trip through `originals_flat`.
|
||||
let mut src = RabitqPlusIndex::new_with_rotation(d, seed, rerank, kind);
|
||||
for (id, v) in &data {
|
||||
src.add(*id, v.clone()).unwrap();
|
||||
}
|
||||
assert_eq!(src.len(), n);
|
||||
|
||||
let items = src.export_items();
|
||||
assert_eq!(items.len(), n);
|
||||
for (pos, row) in &items {
|
||||
assert_eq!(row.len(), d, "row {pos} wrong dim");
|
||||
}
|
||||
|
||||
let rebuilt =
|
||||
RabitqPlusIndex::from_vectors_parallel_with_rotation(d, seed, rerank, kind, items)
|
||||
.expect("rebuild from export_items");
|
||||
assert_eq!(rebuilt.len(), n);
|
||||
assert_eq!(rebuilt.dim(), d);
|
||||
|
||||
// Byte-identical search results on 5 deterministic queries.
|
||||
let queries = make_dataset(5, d, seed ^ 0xDEAD_BEEF);
|
||||
let k = 10;
|
||||
for (_, q) in &queries {
|
||||
let a = src.search(q, k).unwrap();
|
||||
let b = rebuilt.search(q, k).unwrap();
|
||||
assert_eq!(a.len(), b.len(), "result count differs");
|
||||
for (ra, rb) in a.iter().zip(b.iter()) {
|
||||
assert_eq!(ra.id, rb.id, "id mismatch on query");
|
||||
assert_eq!(
|
||||
ra.score.to_bits(),
|
||||
rb.score.to_bits(),
|
||||
"score bits differ for id={}",
|
||||
ra.id,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -92,6 +92,12 @@ pub struct CacheStats {
|
|||
/// Most-recent prime duration in milliseconds (useful to detect
|
||||
/// drift between warm primes and the very first miss).
|
||||
pub last_prime_ms: f64,
|
||||
/// Incremented each time a pre-built index is installed via
|
||||
/// [`VectorCache::install_prebuilt`] — i.e. a warm-from-disk
|
||||
/// rehydrate that did NOT round-trip to the backend. Deliberately
|
||||
/// separate from `primes` so operators can tell cold-prime cost
|
||||
/// (pull+compress) apart from warm-restart cost (mmap+install).
|
||||
pub warm_installs: u64,
|
||||
}
|
||||
|
||||
impl CacheStats {
|
||||
|
|
@ -452,6 +458,90 @@ impl VectorCache {
|
|||
self.prime_interned(interned, witness, batch)
|
||||
}
|
||||
|
||||
/// Install a pre-built `RabitqPlusIndex` under `witness` and install
|
||||
/// the pointer `key → witness` — the warm-from-disk counterpart to
|
||||
/// [`prime`]. No backend round-trip, no RaBitQ compression: the
|
||||
/// caller supplies an already-compressed index (typically loaded via
|
||||
/// `ruvector_rabitq::persist::load_index`) together with the
|
||||
/// `pos_to_id` mapping the cache would otherwise derive from
|
||||
/// `PulledBatch::ids`.
|
||||
///
|
||||
/// Semantics:
|
||||
///
|
||||
/// - If `witness` is already cached (another pointer brought it in),
|
||||
/// the supplied `idx` is dropped and the existing entry is shared
|
||||
/// — this is the same "witness already present" fast path `prime`
|
||||
/// uses, so two operators warming the same bundle from different
|
||||
/// sidecars see one compressed entry with refcount 2.
|
||||
/// - If `witness` is new, the entry is inserted, `warm_installs` is
|
||||
/// bumped, and `primes` / prime-duration counters are left alone
|
||||
/// (this is not a prime — no compression ran).
|
||||
/// - Either way the pointer `key → witness` is installed and its
|
||||
/// refcount bumped.
|
||||
///
|
||||
/// The LRU cap honours warm installs identically to prime installs:
|
||||
/// an oversize cache still evicts unpinned entries.
|
||||
pub fn install_prebuilt(
|
||||
&self,
|
||||
key: CacheKey,
|
||||
witness: WitnessKey,
|
||||
idx: Arc<RabitqPlusIndex>,
|
||||
pos_to_id: Arc<Vec<u64>>,
|
||||
) -> crate::Result<()> {
|
||||
let interned = intern_cache_key(&key);
|
||||
self.install_prebuilt_interned(interned, witness, idx, pos_to_id)
|
||||
}
|
||||
|
||||
pub(crate) fn install_prebuilt_interned(
|
||||
&self,
|
||||
key: InternedKey,
|
||||
witness: WitnessKey,
|
||||
idx: Arc<RabitqPlusIndex>,
|
||||
pos_to_id: Arc<Vec<u64>>,
|
||||
) -> crate::Result<()> {
|
||||
// Defensive consistency check before we touch state: the caller
|
||||
// must hand us a `pos_to_id` whose length matches the index.
|
||||
// A mismatch means the sidecar and the .rbpx drifted, and
|
||||
// serving through that would map positions to the wrong ids
|
||||
// without any visible error until a query returns garbage.
|
||||
if pos_to_id.len() != idx.len() {
|
||||
return Err(crate::RuLakeError::InvalidParameter(format!(
|
||||
"install_prebuilt: pos_to_id.len()={} but index.len()={}",
|
||||
pos_to_id.len(),
|
||||
idx.len()
|
||||
)));
|
||||
}
|
||||
let mut inner = self.inner.lock().unwrap();
|
||||
// Fast path: target witness already cached — just point and
|
||||
// bookkeep as a shared install. `shared=true` bumps
|
||||
// `shared_hits` — but this is a warm install, not a coherence
|
||||
// event, so we route through `inner_install_pointer_unlocked`
|
||||
// with `shared=false` to avoid polluting coherence stats.
|
||||
if inner.entries.contains_key(&witness) {
|
||||
return self.inner_install_pointer_unlocked(&mut inner, key, witness, false);
|
||||
}
|
||||
let dim = idx.dim();
|
||||
let entry = CacheEntry {
|
||||
index: idx,
|
||||
dim,
|
||||
generation_hint: None,
|
||||
last_checked: Instant::now(),
|
||||
last_used: Instant::now(),
|
||||
pos_to_id,
|
||||
refcount: 0, // install_pointer bumps it
|
||||
};
|
||||
inner.entries.insert(witness.clone(), entry);
|
||||
inner.stats.warm_installs += 1;
|
||||
// NOTE: we intentionally do NOT bump `primes` / prime timers —
|
||||
// a warm install did no compression work, so conflating the
|
||||
// two would hide cold-start cost from operators.
|
||||
let rc = self.inner_install_pointer_unlocked(&mut inner, key, witness, false);
|
||||
if let Some(cap) = self.max_entries {
|
||||
self.evict_lru_if_over(&mut inner, cap);
|
||||
}
|
||||
rc
|
||||
}
|
||||
|
||||
/// Evict the least-recently-used unpinned entry until we're at or
|
||||
/// below `cap`. Pinned entries are skipped; in the worst case every
|
||||
/// entry is pinned and we can't evict anyone — that's by design.
|
||||
|
|
@ -566,6 +656,22 @@ impl VectorCache {
|
|||
self.inner.lock().unwrap().entries.len()
|
||||
}
|
||||
|
||||
/// Clone out the `Arc<RabitqPlusIndex>` backing `witness` and its
|
||||
/// `Arc<Vec<u64>>` pos→id map — used by the `save_cache_to_dir`
|
||||
/// path in `RuLake` to serialize a primed entry without exposing
|
||||
/// `CacheEntry` publicly. Returns `None` when the witness is not
|
||||
/// currently cached.
|
||||
pub(crate) fn index_and_ids_of(
|
||||
&self,
|
||||
witness: &str,
|
||||
) -> Option<(Arc<RabitqPlusIndex>, Arc<Vec<u64>>)> {
|
||||
let inner = self.inner.lock().unwrap();
|
||||
inner
|
||||
.entries
|
||||
.get(witness)
|
||||
.map(|e| (Arc::clone(&e.index), Arc::clone(&e.pos_to_id)))
|
||||
}
|
||||
|
||||
pub fn dim_of(&self, key: &CacheKey) -> Option<usize> {
|
||||
let interned = intern_cache_key(key);
|
||||
let inner = self.inner.lock().unwrap();
|
||||
|
|
|
|||
|
|
@ -5,12 +5,22 @@
|
|||
//! support native vector ops.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use ruvector_rabitq::AnnIndex;
|
||||
|
||||
use crate::backend::{BackendAdapter, BackendId};
|
||||
use crate::cache::{intern_key, CacheKey, Consistency, InternedKey, VectorCache};
|
||||
use crate::error::{Result, RuLakeError};
|
||||
|
||||
/// Canonical filename for the persisted RaBitQ index written by
|
||||
/// [`RuLake::save_cache_to_dir`] and read back by
|
||||
/// [`RuLake::warm_from_dir`]. Sits alongside `table.rulake.json` in the
|
||||
/// same directory; the pair is the portable "primed collection
|
||||
/// snapshot" operators hand from warm-shutdown to warm-restart.
|
||||
const PERSISTED_INDEX_FILENAME: &str = "index.rbpx";
|
||||
|
||||
/// Result from a search — the external id and its estimated L2² score.
|
||||
/// Includes the backend that produced the hit so callers can audit.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
|
|
@ -217,6 +227,225 @@ impl RuLake {
|
|||
}
|
||||
}
|
||||
|
||||
/// Snapshot a primed cache entry to `dir/index.rbpx` plus the
|
||||
/// companion bundle sidecar `dir/table.rulake.json`. Pairs with
|
||||
/// [`Self::warm_from_dir`] to give operators a warm-restart path
|
||||
/// that skips the backend round-trip + RaBitQ compression on boot.
|
||||
///
|
||||
/// Flow:
|
||||
///
|
||||
/// 1. Resolve the witness currently held for `key` — errors with
|
||||
/// `UnknownCollection` if nothing is primed (the cache-first
|
||||
/// operator story assumes the entry exists; we never
|
||||
/// opportunistically prime on save).
|
||||
/// 2. Look up the compressed `Arc<RabitqPlusIndex>` out of the
|
||||
/// cache (fails `UnknownCollection` if the pointer is dangling
|
||||
/// — shouldn't happen but is checked defensively).
|
||||
/// 3. `export_items()` → call `ruvector_rabitq::persist::save_index`
|
||||
/// with `(idx, cache.rotation_seed(), items)`. The seed is
|
||||
/// pulled from the cache because the on-disk format's witness
|
||||
/// chain requires it for deterministic rebuild.
|
||||
/// 4. Atomically rename the temp file into place so readers never
|
||||
/// observe a half-written index.
|
||||
/// 5. Emit `publish_bundle(key, dir)` so the bundle sidecar
|
||||
/// accompanies the `.rbpx` — `warm_from_dir` cross-checks the
|
||||
/// two before installing.
|
||||
///
|
||||
/// Returns the path of the written `.rbpx` file; the sidecar lives
|
||||
/// next to it at `dir/table.rulake.json`. Directory creation is
|
||||
/// transitive, mirroring `publish_bundle`'s behaviour.
|
||||
///
|
||||
/// Errors:
|
||||
/// - `UnknownBackend` if the key's backend isn't registered
|
||||
/// (required for the bundle publish step).
|
||||
/// - `UnknownCollection` if nothing is primed under `key`.
|
||||
/// - `InvalidParameter` on any filesystem or serialization failure.
|
||||
pub fn save_cache_to_dir(&self, key: &CacheKey, dir: impl AsRef<Path>) -> Result<PathBuf> {
|
||||
let dir = dir.as_ref();
|
||||
std::fs::create_dir_all(dir).map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!(
|
||||
"save_cache_to_dir: mkdir {}: {e}",
|
||||
dir.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
// Step 1 — resolve witness.
|
||||
let witness = self
|
||||
.cache
|
||||
.witness_of(key)
|
||||
.ok_or_else(|| RuLakeError::UnknownCollection {
|
||||
backend: key.0.clone(),
|
||||
collection: key.1.clone(),
|
||||
})?;
|
||||
|
||||
// Step 2 — clone the Arc<RabitqPlusIndex> out of the cache.
|
||||
let (idx, _pos_to_id) = self.cache.index_and_ids_of(&witness).ok_or_else(|| {
|
||||
RuLakeError::UnknownCollection {
|
||||
backend: key.0.clone(),
|
||||
collection: key.1.clone(),
|
||||
}
|
||||
})?;
|
||||
|
||||
// Step 3 — serialize via rabitq's persist API.
|
||||
let items = idx.export_items();
|
||||
let final_path = dir.join(PERSISTED_INDEX_FILENAME);
|
||||
let tmp_path = dir.join(format!(
|
||||
".{PERSISTED_INDEX_FILENAME}.tmp.{}",
|
||||
std::process::id()
|
||||
));
|
||||
{
|
||||
let mut f = std::fs::File::create(&tmp_path).map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!(
|
||||
"save_cache_to_dir: create {}: {e}",
|
||||
tmp_path.display()
|
||||
))
|
||||
})?;
|
||||
let mut buf = std::io::BufWriter::new(&mut f);
|
||||
ruvector_rabitq::persist::save_index(
|
||||
idx.as_ref(),
|
||||
self.cache.rotation_seed(),
|
||||
&items,
|
||||
&mut buf,
|
||||
)?;
|
||||
use std::io::Write;
|
||||
buf.flush().map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!("save_cache_to_dir: flush: {e}"))
|
||||
})?;
|
||||
drop(buf);
|
||||
f.sync_all().map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!("save_cache_to_dir: fsync: {e}"))
|
||||
})?;
|
||||
}
|
||||
// Step 4 — atomic rename.
|
||||
std::fs::rename(&tmp_path, &final_path).map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!(
|
||||
"save_cache_to_dir: rename {} → {}: {e}",
|
||||
tmp_path.display(),
|
||||
final_path.display()
|
||||
))
|
||||
})?;
|
||||
|
||||
// Step 5 — companion sidecar. Reuses the existing publish path
|
||||
// so the on-disk format is identical to the "cache sidecar
|
||||
// daemon publishes to GCS" flow.
|
||||
self.publish_bundle(key, dir)?;
|
||||
|
||||
Ok(final_path)
|
||||
}
|
||||
|
||||
/// Inverse of [`Self::save_cache_to_dir`]. Reads the bundle sidecar
|
||||
/// plus the `index.rbpx` at `dir`, verifies the witness, and installs
|
||||
/// the loaded index into the cache under the pointer for `key`. No
|
||||
/// backend round-trip, no RaBitQ compression — the primed entry pops
|
||||
/// back into reality in O(file-read) time.
|
||||
///
|
||||
/// Returns the number of vectors loaded (equal to the
|
||||
/// `RabitqPlusIndex::len()` of the installed entry). Operators use
|
||||
/// this to confirm the snapshot they expected is the one that
|
||||
/// landed.
|
||||
///
|
||||
/// Flow:
|
||||
///
|
||||
/// 1. `RuLakeBundle::read_from_dir` — opens + witness-verifies
|
||||
/// `table.rulake.json`. A tampered sidecar is rejected here
|
||||
/// before we touch the index.
|
||||
/// 2. `ruvector_rabitq::persist::load_index(dir/index.rbpx)`.
|
||||
/// Malformed / oversize inputs are rejected by the persist
|
||||
/// layer.
|
||||
/// 3. Cross-check: the loaded index's `dim()` must match the
|
||||
/// bundle's `dim`, and its `rerank_factor()` must match the
|
||||
/// bundle's `rerank_factor`. The witness covers these fields,
|
||||
/// but a belt-and-braces check catches a future format that
|
||||
/// decouples them.
|
||||
/// 4. Build `pos_to_id` from `idx.ids()` (each u32 widens to u64).
|
||||
/// 5. `cache.install_prebuilt` under the bundle's `rvf_witness`.
|
||||
/// The cache's shared-witness fast path means two operators
|
||||
/// warming the same bundle share one compressed entry.
|
||||
///
|
||||
/// Deliberately does *not* require the backend to be registered —
|
||||
/// this is the warm-restart path, and the operator might warm the
|
||||
/// cache before the backend controller is wired. Any subsequent
|
||||
/// `Consistency::Fresh` search against `key` will fail with
|
||||
/// `UnknownBackend` at that point, which is the same behaviour as
|
||||
/// a cold cache with a missing backend. `Consistency::Frozen`
|
||||
/// serves out of the installed entry without ever asking.
|
||||
///
|
||||
/// Errors:
|
||||
/// - `InvalidParameter` if the sidecar or index file is missing,
|
||||
/// malformed, or carries a mismatched witness/dim/rerank_factor.
|
||||
/// - `Rabitq` on any RaBitQ-layer failure (propagated from
|
||||
/// `persist::load_index`).
|
||||
pub fn warm_from_dir(&self, key: &CacheKey, dir: impl AsRef<Path>) -> Result<usize> {
|
||||
let dir = dir.as_ref();
|
||||
|
||||
// Step 1 — sidecar (witness-verified inside read_from_dir).
|
||||
let bundle = crate::RuLakeBundle::read_from_dir(dir)?;
|
||||
|
||||
// Step 2 — load the compressed index.
|
||||
let index_path = dir.join(PERSISTED_INDEX_FILENAME);
|
||||
if !index_path.exists() {
|
||||
return Err(RuLakeError::InvalidParameter(format!(
|
||||
"warm_from_dir: missing {}",
|
||||
index_path.display()
|
||||
)));
|
||||
}
|
||||
let file = std::fs::File::open(&index_path).map_err(|e| {
|
||||
RuLakeError::InvalidParameter(format!(
|
||||
"warm_from_dir: open {}: {e}",
|
||||
index_path.display()
|
||||
))
|
||||
})?;
|
||||
let mut reader = std::io::BufReader::new(file);
|
||||
let idx = ruvector_rabitq::persist::load_index(&mut reader)?;
|
||||
|
||||
// Step 3 — bundle vs index cross-check. The witness chain
|
||||
// already proves the bundle is self-consistent, and the
|
||||
// persist format checks `(dim, seed, rerank_factor)` against
|
||||
// the items it saved — but nothing binds the persisted
|
||||
// index's seed to the *bundle's* seed without this check,
|
||||
// and a mismatched seed would silently serve wrong neighbours.
|
||||
if idx.dim() != bundle.dim {
|
||||
return Err(RuLakeError::InvalidParameter(format!(
|
||||
"warm_from_dir: index dim {} != bundle dim {}",
|
||||
idx.dim(),
|
||||
bundle.dim
|
||||
)));
|
||||
}
|
||||
if idx.rerank_factor() != bundle.rerank_factor {
|
||||
return Err(RuLakeError::InvalidParameter(format!(
|
||||
"warm_from_dir: index rerank_factor {} != bundle rerank_factor {}",
|
||||
idx.rerank_factor(),
|
||||
bundle.rerank_factor
|
||||
)));
|
||||
}
|
||||
|
||||
// Step 4 — pos_to_id from the loaded index's internal ids.
|
||||
//
|
||||
// Every primed entry is built with `idx.add(pos, v)` where
|
||||
// `pos ∈ 0..n` — the internal id IS the position. The
|
||||
// `save_index` / `load_index` round-trip preserves that, so
|
||||
// the loaded index's internal ids form an identity sequence
|
||||
// `[0, 1, ..., n-1]`. The cache's `pos_to_id[pos]` mirrors
|
||||
// this by construction; the external u64 ids are collapsed
|
||||
// to their positional form when they round-trip through the
|
||||
// u32-wide rabitq persist format. Callers whose external ids
|
||||
// are not dense 0..n should NOT rely on warm_from_dir to
|
||||
// preserve them — that's a known limitation of the persist
|
||||
// format and documented alongside it.
|
||||
let n = idx.len();
|
||||
let pos_to_id: Vec<u64> = (0..n as u64).collect();
|
||||
|
||||
// Step 5 — install into the cache via the interned path.
|
||||
self.cache.install_prebuilt_interned(
|
||||
intern_key(&key.0, &key.1),
|
||||
bundle.rvf_witness,
|
||||
Arc::new(idx),
|
||||
Arc::new(pos_to_id),
|
||||
)?;
|
||||
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
/// Search a single (backend, collection) pair. Handles cache
|
||||
/// miss / staleness transparently.
|
||||
pub fn search_one(
|
||||
|
|
|
|||
|
|
@ -1143,6 +1143,111 @@ fn adaptive_per_shard_rerank_preserves_recall() {
|
|||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warm_from_dir_skips_backend_and_returns_bit_exact_results() {
|
||||
// End-to-end cache-persistence round-trip:
|
||||
//
|
||||
// 1. Prime a RuLake with 50 vectors (D=8).
|
||||
// 2. Run a query, capture the results (ids + score bits).
|
||||
// 3. `save_cache_to_dir` — write `index.rbpx` + `table.rulake.json`.
|
||||
// 4. Build a FRESH RuLake (no backend registered).
|
||||
// 5. `warm_from_dir` — install the snapshot without touching any
|
||||
// backend, so `primes == 0` on the fresh lake.
|
||||
// 6. Re-run the query; bit-identical ids and f32 bits required.
|
||||
//
|
||||
// The fresh lake uses `Consistency::Frozen` because it has no
|
||||
// backend to coherence-check against — Frozen serves the installed
|
||||
// entry without ever dialing out. `warm_installs == 1` confirms
|
||||
// the new counter fires once per warm install. `primes == 0`
|
||||
// confirms no backend round-trip happened on the serving side.
|
||||
let d = 8;
|
||||
let n = 50;
|
||||
let rerank = 20;
|
||||
let seed = 1234;
|
||||
|
||||
// Seed-clustered data so scores have structure — a flat dataset
|
||||
// would let ties dominate and mask a rerank-ordering regression.
|
||||
let data = clustered(n, d, 5, seed);
|
||||
let backend = Arc::new(LocalBackend::new("warm-src"));
|
||||
backend
|
||||
.put_collection("snap", d, (0..n as u64).collect(), data.clone())
|
||||
.unwrap();
|
||||
|
||||
let src = RuLake::new(rerank, seed);
|
||||
src.register_backend(backend).unwrap();
|
||||
let key = ("warm-src".to_string(), "snap".to_string());
|
||||
|
||||
// Step 1+2: prime + capture the reference result.
|
||||
let query = clustered(1, d, 5, seed ^ 0xdeadbeef)[0].clone();
|
||||
let ref_hits = src.search_one("warm-src", "snap", &query, 5).unwrap();
|
||||
assert_eq!(ref_hits.len(), 5);
|
||||
// Capture both ids and raw f32 bits for the bit-exact compare.
|
||||
let ref_ids: Vec<u64> = ref_hits.iter().map(|r| r.id).collect();
|
||||
let ref_score_bits: Vec<u32> = ref_hits.iter().map(|r| r.score.to_bits()).collect();
|
||||
|
||||
// Step 3: snapshot to disk.
|
||||
let tmp = std::env::temp_dir().join(format!(
|
||||
"rulake-warm-{}-{}",
|
||||
std::process::id(),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0)
|
||||
));
|
||||
std::fs::create_dir_all(&tmp).unwrap();
|
||||
let written = src.save_cache_to_dir(&key, &tmp).unwrap();
|
||||
assert_eq!(written.file_name().unwrap(), "index.rbpx");
|
||||
// Sidecar must exist alongside.
|
||||
assert!(tmp
|
||||
.join(ruvector_rulake::RuLakeBundle::SIDECAR_FILENAME)
|
||||
.exists());
|
||||
|
||||
// Step 4: fresh RuLake, no backend. Same seed + rerank_factor so
|
||||
// the reconstructed RaBitQ codes are bit-identical to the source.
|
||||
// Frozen consistency so the post-warm search never asks for a
|
||||
// backend.
|
||||
let fresh = RuLake::new(rerank, seed).with_consistency(Consistency::Frozen);
|
||||
|
||||
// Sanity precondition: no primes, no warm installs before we load.
|
||||
let s_pre = fresh.cache_stats();
|
||||
assert_eq!(s_pre.primes, 0);
|
||||
assert_eq!(s_pre.warm_installs, 0);
|
||||
|
||||
// Step 5: warm from disk.
|
||||
let loaded_n = fresh.warm_from_dir(&key, &tmp).unwrap();
|
||||
assert_eq!(loaded_n, n, "warm_from_dir must load every vector");
|
||||
|
||||
let s_after_warm = fresh.cache_stats();
|
||||
assert_eq!(
|
||||
s_after_warm.warm_installs, 1,
|
||||
"warm_installs must tick exactly once on warm_from_dir"
|
||||
);
|
||||
assert_eq!(
|
||||
s_after_warm.primes, 0,
|
||||
"warm_from_dir must not prime (no backend pull)"
|
||||
);
|
||||
|
||||
// Step 6: re-run the query against the fresh lake.
|
||||
let warm_hits = fresh.search_one("warm-src", "snap", &query, 5).unwrap();
|
||||
assert_eq!(warm_hits.len(), ref_hits.len());
|
||||
let warm_ids: Vec<u64> = warm_hits.iter().map(|r| r.id).collect();
|
||||
let warm_score_bits: Vec<u32> = warm_hits.iter().map(|r| r.score.to_bits()).collect();
|
||||
|
||||
assert_eq!(warm_ids, ref_ids, "ids must be byte-identical after warm");
|
||||
assert_eq!(
|
||||
warm_score_bits, ref_score_bits,
|
||||
"f32 score bits must be byte-identical after warm"
|
||||
);
|
||||
|
||||
// Final stats: still zero primes, still one warm_install. Search
|
||||
// didn't round-trip to any backend (there isn't one).
|
||||
let s_final = fresh.cache_stats();
|
||||
assert_eq!(s_final.primes, 0);
|
||||
assert_eq!(s_final.warm_installs, 1);
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stats_expose_hit_rate_and_prime_duration() {
|
||||
// The cache-first reframe (ADR-155) makes hit_rate the primary KPI.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue