diff --git a/crates/ruvector-rabitq/src/index.rs b/crates/ruvector-rabitq/src/index.rs index db99e44aa..b4457864a 100644 --- a/crates/ruvector-rabitq/src/index.rs +++ b/crates/ruvector-rabitq/src/index.rs @@ -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)> { + 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, + ); + } + } + } } diff --git a/crates/ruvector-rulake/src/cache.rs b/crates/ruvector-rulake/src/cache.rs index 83ed58139..afbb91cb3 100644 --- a/crates/ruvector-rulake/src/cache.rs +++ b/crates/ruvector-rulake/src/cache.rs @@ -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, + pos_to_id: Arc>, + ) -> 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, + pos_to_id: Arc>, + ) -> 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` backing `witness` and its + /// `Arc>` 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, Arc>)> { + 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 { let interned = intern_cache_key(key); let inner = self.inner.lock().unwrap(); diff --git a/crates/ruvector-rulake/src/lake.rs b/crates/ruvector-rulake/src/lake.rs index bcced1f79..26a569db9 100644 --- a/crates/ruvector-rulake/src/lake.rs +++ b/crates/ruvector-rulake/src/lake.rs @@ -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` 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) -> Result { + 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 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) -> Result { + 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 = (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( diff --git a/crates/ruvector-rulake/tests/federation_smoke.rs b/crates/ruvector-rulake/tests/federation_smoke.rs index 1dd0f8c6b..95e28fa46 100644 --- a/crates/ruvector-rulake/tests/federation_smoke.rs +++ b/crates/ruvector-rulake/tests/federation_smoke.rs @@ -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 = ref_hits.iter().map(|r| r.id).collect(); + let ref_score_bits: Vec = 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 = warm_hits.iter().map(|r| r.id).collect(); + let warm_score_bits: Vec = 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.