use std::sync::Arc;
use arc_swap::ArcSwap;
use parking_lot::RwLock;
use crate::error::{Result, YantrikDbError};
use crate::vector::hnsw::HnswIndex;
#[derive(Clone, Debug)]
pub struct DeltaEntry {
pub rid: String,
pub embedding: Vec<f32>,
pub seq: u64,
pub tombstoned: bool,
}
pub const DEFAULT_DELTA_MAX: usize = 256;
pub const DEFAULT_MAX_DIRTY_AGE: std::time::Duration = std::time::Duration::from_secs(60);
pub struct DeltaIndex {
cold: ArcSwap<HnswIndex>,
delta: RwLock<Vec<DeltaEntry>>,
delta_max: usize,
dim: usize,
oldest_dirty_at: parking_lot::Mutex<Option<std::time::Instant>>,
max_dirty_age: std::time::Duration,
compactor_wake_cv: parking_lot::Condvar,
compactor_wake_mu: parking_lot::Mutex<()>,
}
impl DeltaIndex {
pub fn new(dim: usize) -> Self {
Self::with_capacity(dim, DEFAULT_DELTA_MAX)
}
pub fn with_capacity(dim: usize, delta_max: usize) -> Self {
Self::with_capacity_and_age(dim, delta_max, DEFAULT_MAX_DIRTY_AGE)
}
pub fn with_capacity_and_age(
dim: usize,
delta_max: usize,
max_dirty_age: std::time::Duration,
) -> Self {
Self {
cold: ArcSwap::new(Arc::new(HnswIndex::new(dim))),
delta: RwLock::new(Vec::with_capacity(delta_max.min(4096))),
delta_max,
dim,
oldest_dirty_at: parking_lot::Mutex::new(None),
max_dirty_age,
compactor_wake_cv: parking_lot::Condvar::new(),
compactor_wake_mu: parking_lot::Mutex::new(()),
}
}
pub fn from_cold(cold: HnswIndex, delta_max: usize) -> Self {
Self::from_cold_with_age(cold, delta_max, DEFAULT_MAX_DIRTY_AGE)
}
pub fn from_cold_with_age(
cold: HnswIndex,
delta_max: usize,
max_dirty_age: std::time::Duration,
) -> Self {
let dim = cold.dim();
Self {
cold: ArcSwap::new(Arc::new(cold)),
delta: RwLock::new(Vec::with_capacity(delta_max.min(4096))),
delta_max,
dim,
oldest_dirty_at: parking_lot::Mutex::new(None),
max_dirty_age,
compactor_wake_cv: parking_lot::Condvar::new(),
compactor_wake_mu: parking_lot::Mutex::new(()),
}
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn delta_max(&self) -> usize {
self.delta_max
}
pub fn append(&self, rid: String, embedding: Vec<f32>, seq: u64) -> Result<()> {
if embedding.len() != self.dim {
return Err(YantrikDbError::InvalidInput(format!(
"embedding dimension mismatch: expected {}, got {}",
self.dim,
embedding.len()
)));
}
let mut delta = self.delta.write();
if delta.iter().any(|e| e.rid == rid && e.seq == seq) {
return Ok(());
}
if delta.len() >= self.delta_max {
return Err(YantrikDbError::Backpressure {
pending: delta.len() as i64,
max: self.delta_max as i64,
retry_after_ms: 50,
});
}
let was_empty = delta.is_empty();
delta.push(DeltaEntry {
rid,
embedding,
seq,
tombstoned: false,
});
let new_len = delta.len();
drop(delta);
if was_empty {
*self.oldest_dirty_at.lock() = Some(std::time::Instant::now());
}
if new_len >= self.delta_max * 80 / 100 {
self.compactor_wake_cv.notify_one();
}
Ok(())
}
pub fn tombstone(&self, rid: &str, seq: u64) -> bool {
let mut delta = self.delta.write();
let was_empty = delta.is_empty();
for entry in delta.iter_mut() {
if entry.rid == rid && !entry.tombstoned {
entry.tombstoned = true;
entry.seq = seq;
return true;
}
}
delta.push(DeltaEntry {
rid: rid.to_string(),
embedding: Vec::new(), seq,
tombstoned: true,
});
let new_len = delta.len();
drop(delta);
if was_empty {
*self.oldest_dirty_at.lock() = Some(std::time::Instant::now());
}
if new_len >= self.delta_max * 80 / 100 {
self.compactor_wake_cv.notify_one();
}
false
}
pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<(String, f64)>> {
if query.len() != self.dim {
return Err(YantrikDbError::InvalidInput(format!(
"query dimension mismatch: expected {}, got {}",
self.dim,
query.len()
)));
}
let cold = self.cold.load();
let delta = self.delta.read();
let mut winner_per_rid: std::collections::HashMap<&str, &DeltaEntry> =
std::collections::HashMap::new();
for entry in delta.iter() {
match winner_per_rid.get(entry.rid.as_str()) {
Some(existing) if existing.seq >= entry.seq => {}
_ => {
winner_per_rid.insert(entry.rid.as_str(), entry);
}
}
}
let mut tombstoned: std::collections::HashSet<&str> = std::collections::HashSet::new();
let mut delta_live: Vec<(&DeltaEntry, f64)> = Vec::with_capacity(winner_per_rid.len());
for (rid, entry) in &winner_per_rid {
if entry.tombstoned {
tombstoned.insert(*rid);
} else {
let d = cosine_distance_f64(query, &entry.embedding);
delta_live.push((*entry, d));
}
}
let cold_fetch = k.saturating_mul(2).max(k);
let cold_results = cold.search(query, cold_fetch)?;
let delta_rid_set: std::collections::HashSet<&str> =
delta_live.iter().map(|(e, _)| e.rid.as_str()).collect();
let mut merged: Vec<(String, f64)> =
Vec::with_capacity(cold_results.len() + delta_live.len());
for (rid, dist) in &cold_results {
if tombstoned.contains(rid.as_str()) || delta_rid_set.contains(rid.as_str()) {
continue;
}
merged.push((rid.clone(), *dist));
}
for (entry, dist) in &delta_live {
merged.push((entry.rid.clone(), *dist));
}
merged.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
merged.truncate(k);
Ok(merged)
}
pub fn delta_len(&self) -> usize {
self.delta.read().len()
}
pub fn cold_len(&self) -> usize {
self.cold.load().len()
}
pub fn len(&self) -> usize {
self.cold_len() + self.delta_len()
}
pub fn is_empty(&self) -> bool {
self.cold_len() == 0 && self.delta_len() == 0
}
pub fn snapshot_delta(&self) -> Vec<DeltaEntry> {
self.delta.read().clone()
}
pub fn seal_delta_for_compaction(&self) -> Vec<DeltaEntry> {
let mut delta = self.delta.write();
let sealed = std::mem::replace(&mut *delta, Vec::with_capacity(self.delta_max.min(4096)));
*self.oldest_dirty_at.lock() = None;
sealed
}
pub fn install_cold(&self, new_cold: HnswIndex) {
self.cold.store(Arc::new(new_cold));
}
pub fn compact(&self) -> Result<usize> {
let sealed = self.seal_delta_for_compaction();
if sealed.is_empty() {
return Ok(0);
}
let mut new_cold: HnswIndex = (*self.cold.load_full()).clone();
let mut by_rid: std::collections::HashMap<String, &DeltaEntry> =
std::collections::HashMap::with_capacity(sealed.len());
for entry in &sealed {
match by_rid.get(&entry.rid) {
Some(existing) if existing.seq >= entry.seq => {}
_ => {
by_rid.insert(entry.rid.clone(), entry);
}
}
}
let mut applied = 0usize;
for entry in by_rid.values() {
if entry.tombstoned {
new_cold.remove(&entry.rid);
} else {
new_cold.remove(&entry.rid);
new_cold.insert(&entry.rid, &entry.embedding)?;
}
applied += 1;
}
self.cold.store(Arc::new(new_cold));
Ok(applied)
}
pub fn should_compact(&self) -> bool {
if self.delta_len() >= self.delta_max / 2 {
return true;
}
let stamp = *self.oldest_dirty_at.lock();
match stamp {
Some(t) if self.delta_len() > 0 && t.elapsed() >= self.max_dirty_age => true,
_ => false,
}
}
pub fn wait_for_compaction_signal(&self, timeout: std::time::Duration) -> bool {
let mut guard = self.compactor_wake_mu.lock();
let result = self.compactor_wake_cv.wait_for(&mut guard, timeout);
!result.timed_out()
}
}
fn cosine_distance_f64(a: &[f32], b: &[f32]) -> f64 {
if a.len() != b.len() || a.is_empty() {
return 1.0;
}
let mut dot: f64 = 0.0;
let mut na: f64 = 0.0;
let mut nb: f64 = 0.0;
for (&x, &y) in a.iter().zip(b.iter()) {
let xf = x as f64;
let yf = y as f64;
dot += xf * yf;
na += xf * xf;
nb += yf * yf;
}
let na = na.sqrt();
let nb = nb.sqrt();
if na == 0.0 || nb == 0.0 {
return 1.0;
}
(1.0 - (dot / (na * nb))).clamp(0.0, 2.0)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
fn vec_seed(seed: f32, dim: usize) -> Vec<f32> {
let raw: Vec<f32> = (0..dim).map(|i| (seed + i as f32) * 0.1).collect();
let norm: f32 = raw.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-9);
raw.iter().map(|x| x / norm).collect()
}
#[test]
fn empty_index_search_returns_empty() {
let idx = DeltaIndex::new(64);
let query = vec_seed(1.0, 64);
let r = idx.search(&query, 10).unwrap();
assert!(r.is_empty());
}
#[test]
fn append_then_search_finds_in_delta() {
let idx = DeltaIndex::new(64);
let emb = vec_seed(1.0, 64);
idx.append("rid_1".to_string(), emb.clone(), 1).unwrap();
let query = vec_seed(1.0, 64);
let r = idx.search(&query, 5).unwrap();
assert_eq!(r.len(), 1);
assert_eq!(r[0].0, "rid_1");
assert!(r[0].1 < 0.001, "exact-match distance ~0, got {}", r[0].1);
}
#[test]
fn append_dimension_mismatch_rejected() {
let idx = DeltaIndex::new(64);
let bad = vec![0.0f32; 32];
let err = idx
.append("rid_x".to_string(), bad, 1)
.expect_err("must reject");
assert!(matches!(err, YantrikDbError::InvalidInput(_)));
}
#[test]
fn delta_full_returns_backpressure() {
let idx = DeltaIndex::with_capacity(64, 5);
for i in 0..5 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
let err = idx
.append("rid_overflow".to_string(), vec_seed(99.0, 64), 999)
.expect_err("must backpressure");
match err {
YantrikDbError::Backpressure { pending, max, .. } => {
assert_eq!(pending, 5);
assert_eq!(max, 5);
}
other => panic!("expected Backpressure, got {other:?}"),
}
}
#[test]
fn append_idempotent_on_same_rid_seq() {
let idx = DeltaIndex::new(64);
let emb = vec_seed(1.0, 64);
idx.append("rid_1".to_string(), emb.clone(), 1).unwrap();
idx.append("rid_1".to_string(), emb.clone(), 1).unwrap();
assert_eq!(idx.delta_len(), 1, "second append at same seq is no-op");
}
#[test]
fn tombstone_hides_rid_from_search() {
let idx = DeltaIndex::new(64);
idx.append("rid_keep".to_string(), vec_seed(1.0, 64), 1)
.unwrap();
idx.append("rid_drop".to_string(), vec_seed(2.0, 64), 2)
.unwrap();
let query = vec_seed(2.0, 64);
let r_before = idx.search(&query, 5).unwrap();
assert_eq!(r_before.len(), 2);
idx.tombstone("rid_drop", 3);
let r_after = idx.search(&query, 5).unwrap();
assert_eq!(r_after.len(), 1);
assert_eq!(r_after[0].0, "rid_keep");
}
#[test]
fn tombstone_on_cold_only_rid_appends_marker() {
let idx = DeltaIndex::new(64);
let mut cold = HnswIndex::new(64);
cold.insert("rid_in_cold", &vec_seed(5.0, 64)).unwrap();
idx.install_cold(cold);
let r = idx.search(&vec_seed(5.0, 64), 5).unwrap();
assert_eq!(r.len(), 1);
assert_eq!(r[0].0, "rid_in_cold");
let was_in_delta = idx.tombstone("rid_in_cold", 1);
assert!(!was_in_delta);
assert_eq!(idx.delta_len(), 1, "tombstone marker appended");
let r2 = idx.search(&vec_seed(5.0, 64), 5).unwrap();
assert!(r2.is_empty(), "tombstoned cold rid hidden");
}
#[test]
fn cold_and_delta_merged_in_search() {
let idx = DeltaIndex::new(64);
let mut cold = HnswIndex::new(64);
cold.insert("rid_cold", &vec_seed(1.0, 64)).unwrap();
idx.install_cold(cold);
idx.append("rid_delta".to_string(), vec_seed(2.0, 64), 1)
.unwrap();
let r = idx.search(&vec_seed(1.5, 64), 5).unwrap();
let rids: Vec<&str> = r.iter().map(|(rid, _)| rid.as_str()).collect();
assert!(rids.contains(&"rid_cold"));
assert!(rids.contains(&"rid_delta"));
assert_eq!(rids.len(), 2);
}
#[test]
fn delta_shadows_cold_when_rid_appears_in_both() {
let idx = DeltaIndex::new(64);
let mut cold = HnswIndex::new(64);
cold.insert("rid_dup", &vec_seed(1.0, 64)).unwrap();
idx.install_cold(cold);
idx.append("rid_dup".to_string(), vec_seed(5.0, 64), 1)
.unwrap();
let r = idx.search(&vec_seed(5.0, 64), 5).unwrap();
assert_eq!(r.len(), 1, "no double-count");
assert_eq!(r[0].0, "rid_dup");
assert!(r[0].1 < 0.001);
}
#[test]
fn seal_delta_returns_entries_and_resets() {
let idx = DeltaIndex::new(64);
for i in 0..5 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
let sealed = idx.seal_delta_for_compaction();
assert_eq!(sealed.len(), 5);
assert_eq!(idx.delta_len(), 0, "delta reset to empty");
idx.append("rid_after".to_string(), vec_seed(10.0, 64), 100)
.unwrap();
assert_eq!(idx.delta_len(), 1);
}
#[test]
fn install_cold_atomically_swaps() {
let idx = DeltaIndex::new(64);
assert_eq!(idx.cold_len(), 0);
let mut new_cold = HnswIndex::new(64);
new_cold.insert("rid_a", &vec_seed(1.0, 64)).unwrap();
new_cold.insert("rid_b", &vec_seed(2.0, 64)).unwrap();
idx.install_cold(new_cold);
assert_eq!(idx.cold_len(), 2);
}
#[test]
fn age_trigger_does_not_fire_on_empty_delta() {
let idx = DeltaIndex::with_capacity_and_age(64, 256, Duration::from_millis(20));
std::thread::sleep(Duration::from_millis(40));
assert!(
!idx.should_compact(),
"empty delta never triggers age compaction"
);
}
#[test]
fn age_trigger_fires_after_max_dirty_age_elapses() {
let idx = DeltaIndex::with_capacity_and_age(64, 256, Duration::from_millis(20));
for i in 0..10 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
assert_eq!(idx.delta_len(), 10);
assert!(
!idx.should_compact(),
"below half-cap and within max_dirty_age window must NOT trigger"
);
std::thread::sleep(Duration::from_millis(50));
assert!(
idx.should_compact(),
"10 entries sitting for >max_dirty_age must trigger age path"
);
}
#[test]
fn age_trigger_resets_on_seal() {
let idx = DeltaIndex::with_capacity_and_age(64, 256, Duration::from_millis(20));
idx.append("rid_a".to_string(), vec_seed(1.0, 64), 1)
.unwrap();
std::thread::sleep(Duration::from_millis(30));
assert!(idx.should_compact(), "first window: age trigger fires");
let _ = idx.seal_delta_for_compaction();
assert!(!idx.should_compact(), "seal cleared dirty-age clock");
idx.append("rid_b".to_string(), vec_seed(2.0, 64), 2)
.unwrap();
assert!(!idx.should_compact(), "fresh window after seal");
std::thread::sleep(Duration::from_millis(30));
assert!(
idx.should_compact(),
"second window: age trigger fires again"
);
}
#[test]
fn age_trigger_compacts_low_write_namespace_end_to_end() {
let idx = DeltaIndex::with_capacity_and_age(64, 256, Duration::from_millis(20));
for i in 0..10 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
std::thread::sleep(Duration::from_millis(40));
assert!(idx.should_compact(), "age trigger ready");
let n = idx.compact().unwrap();
assert_eq!(n, 10, "all 10 entries applied to cold");
assert_eq!(idx.delta_len(), 0, "delta drained");
assert_eq!(idx.cold_len(), 10, "cold absorbed all entries");
assert!(!idx.should_compact(), "post-compact: nothing to do");
}
#[test]
fn age_trigger_tombstone_only_delta_also_fires() {
let idx = DeltaIndex::with_capacity_and_age(64, 256, Duration::from_millis(20));
let was_live = idx.tombstone("rid_remote", 1);
assert!(!was_live, "tombstone of unknown rid is appended as marker");
assert_eq!(idx.delta_len(), 1);
std::thread::sleep(Duration::from_millis(30));
assert!(
idx.should_compact(),
"tombstone-only delta also fires by age"
);
}
#[test]
fn concurrent_appends_and_reads_no_corruption() {
let idx = Arc::new(DeltaIndex::with_capacity(64, 1024));
let mut writer_handles = Vec::new();
for w in 0..4 {
let idx_c = Arc::clone(&idx);
writer_handles.push(thread::spawn(move || {
for i in 0..50 {
let rid = format!("w{w}_rid_{i}");
let emb = vec_seed((w * 100 + i) as f32, 64);
idx_c.append(rid, emb, (w * 100 + i) as u64).unwrap();
}
}));
}
let mut reader_handles = Vec::new();
for r in 0..4 {
let idx_c = Arc::clone(&idx);
reader_handles.push(thread::spawn(move || {
for i in 0..100 {
let q = vec_seed((r * 1000 + i) as f32, 64);
let _ = idx_c.search(&q, 10).unwrap();
}
}));
}
for h in writer_handles {
h.join().unwrap();
}
for h in reader_handles {
h.join().unwrap();
}
assert_eq!(idx.delta_len(), 200, "all 4 writers contributed 50 each");
}
#[test]
fn search_returns_top_k_sorted_ascending() {
let idx = DeltaIndex::new(64);
for i in 0..5 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
let r = idx.search(&vec_seed(2.0, 64), 3).unwrap();
assert_eq!(r.len(), 3);
for w in r.windows(2) {
assert!(w[0].1 <= w[1].1, "distances must sort ascending");
}
assert_eq!(r[0].0, "rid_2");
}
#[test]
fn from_cold_preserves_existing_entries() {
let mut cold = HnswIndex::new(64);
cold.insert("rid_a", &vec_seed(1.0, 64)).unwrap();
cold.insert("rid_b", &vec_seed(2.0, 64)).unwrap();
let idx = DeltaIndex::from_cold(cold, 64);
assert_eq!(idx.cold_len(), 2);
assert_eq!(idx.delta_len(), 0);
let r = idx.search(&vec_seed(1.0, 64), 5).unwrap();
assert_eq!(r.len(), 2);
}
#[test]
fn should_compact_at_half_capacity() {
let idx = DeltaIndex::with_capacity(64, 10);
assert!(!idx.should_compact());
for i in 0..4 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
assert!(!idx.should_compact(), "below half cap");
idx.append("rid_5".to_string(), vec_seed(5.0, 64), 5)
.unwrap();
assert!(idx.should_compact(), "at half cap = should compact");
}
#[test]
fn compact_drains_delta_into_cold() {
let idx = DeltaIndex::new(64);
for i in 0..10 {
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
assert_eq!(idx.delta_len(), 10);
assert_eq!(idx.cold_len(), 0);
let n = idx.compact().unwrap();
assert_eq!(n, 10);
assert_eq!(idx.delta_len(), 0, "delta drained");
assert_eq!(idx.cold_len(), 10, "cold has all 10 entries now");
let r = idx.search(&vec_seed(5.0, 64), 3).unwrap();
assert_eq!(r.len(), 3);
assert_eq!(r[0].0, "rid_5", "exact match still found post-compaction");
}
#[test]
fn compact_applies_tombstones_to_cold() {
let idx = DeltaIndex::new(64);
let mut cold = HnswIndex::new(64);
cold.insert("rid_keep", &vec_seed(1.0, 64)).unwrap();
cold.insert("rid_drop", &vec_seed(2.0, 64)).unwrap();
idx.install_cold(cold);
assert_eq!(idx.cold_len(), 2);
idx.tombstone("rid_drop", 1);
let n = idx.compact().unwrap();
assert_eq!(n, 1);
let r = idx.search(&vec_seed(2.0, 64), 5).unwrap();
let rids: Vec<&str> = r.iter().map(|(rid, _)| rid.as_str()).collect();
assert!(!rids.contains(&"rid_drop"), "tombstone applied to cold");
assert!(rids.contains(&"rid_keep"));
}
#[test]
fn compact_applies_archive_then_hydrate_correctly() {
let idx = DeltaIndex::new(64);
idx.append("rid_X".to_string(), vec_seed(5.0, 64), 1)
.unwrap();
idx.tombstone("rid_X", 2);
idx.append("rid_X".to_string(), vec_seed(5.0, 64), 3)
.unwrap();
let n = idx.compact().unwrap();
assert_eq!(n, 1, "highest-seq winner applied once");
let r = idx.search(&vec_seed(5.0, 64), 5).unwrap();
assert_eq!(r.len(), 1);
assert_eq!(r[0].0, "rid_X", "rid alive in cold post-compaction");
}
#[test]
fn compact_idempotent_on_empty_delta() {
let idx = DeltaIndex::new(64);
assert_eq!(idx.compact().unwrap(), 0);
assert_eq!(idx.compact().unwrap(), 0);
}
#[test]
fn compact_preserves_in_flight_writes() {
let idx = DeltaIndex::new(64);
for i in 0..5 {
idx.append(format!("before_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
let sealed = idx.seal_delta_for_compaction();
assert_eq!(sealed.len(), 5);
assert_eq!(idx.delta_len(), 0, "fresh delta after seal");
for i in 0..3 {
idx.append(
format!("after_{i}"),
vec_seed((100 + i) as f32, 64),
(100 + i) as u64,
)
.unwrap();
}
assert_eq!(
idx.delta_len(),
3,
"after-seal writes accumulate in new delta"
);
}
#[test]
fn compact_threshold_drives_compaction() {
let idx = DeltaIndex::with_capacity(64, 10);
for i in 0..50 {
if idx.should_compact() {
idx.compact().unwrap();
}
idx.append(format!("rid_{i}"), vec_seed(i as f32, 64), i as u64)
.unwrap();
}
idx.compact().unwrap(); assert_eq!(idx.cold_len(), 50);
assert_eq!(idx.delta_len(), 0);
}
}