use super::*;
use crate::{
hnsw_index::{HNSWCanonicalBuilder, HNSWRestoreBuilder},
read_control::StorageReadControl,
};
#[test]
fn graph_larger_than_its_allowance_keeps_topology_scores_and_retained_readers() {
let control = StorageReadControl::with_limit(64 * 1024);
let dimensions = 1024;
let count = 20;
assert!(count * dimensions * size_of::<f32>() > control.memory().limit());
let params = HNSWIndexParams {
m: 2,
ef_construction: 4,
ef_search: 12,
rebuild_threshold: 3,
..HNSWIndexParams::default()
};
let mut resident = HNSWIndex::with_params(dimensions as u32, params).unwrap();
let mut builder = HNSWCanonicalBuilder::new(dimensions as u32, params, &control).unwrap();
for document in 1..=count as u64 {
let value = vector(document, dimensions);
resident.add(document, value.clone()).unwrap();
builder.push(document, 0, &value).unwrap();
}
let (mut spilled, memory) = builder.finish().unwrap().into_parts();
assert!(spilled.nodes.is_spilled());
spilled.validate_invariants().unwrap();
let expected = resident.take_persistence_delta();
let actual = spilled.take_persistence_delta();
assert_eq!(actual.meta, expected.meta);
let mut nodes = actual.nodes();
for node in expected.nodes() {
assert_eq!(*node.unwrap(), *nodes.next().unwrap().unwrap());
}
assert!(nodes.next().is_none());
for query in [3, 15, 90] {
assert_eq!(
spilled.search_knn(&vector(query, dimensions), 7).unwrap(),
resident.search_knn(&vector(query, dimensions), 7).unwrap()
);
}
let retained = spilled.snapshot_with_control(&control).unwrap();
let query = vector(17, dimensions);
let old_results = retained.search_knn(&query, 7).unwrap();
spilled.add(17, vector(250, dimensions)).unwrap();
spilled.delete(8).unwrap();
assert_eq!(retained.search_knn(&query, 7).unwrap(), old_results);
assert_eq!(retained.count().unwrap(), count);
assert_eq!(spilled.count().unwrap(), count - 1);
let changed = spilled.take_persistence_delta();
assert!(!changed.full_rewrite);
let mut changed_nodes = 0;
for node in changed.nodes() {
node.unwrap();
changed_nodes += 1;
}
assert!(changed_nodes > 0);
spilled.delete(9).unwrap();
resident.add(17, vector(250, dimensions)).unwrap();
resident.delete(8).unwrap();
resident.delete(9).unwrap();
let compacted = spilled.take_persistence_delta();
let reference = resident.take_persistence_delta();
assert!(compacted.full_rewrite);
assert_eq!(compacted.meta, reference.meta);
let mut nodes = compacted.nodes();
for expected in reference.nodes() {
assert_eq!(*nodes.next().unwrap().unwrap(), *expected.unwrap());
}
assert!(nodes.next().is_none());
assert_eq!(retained.search_knn(&query, 7).unwrap(), old_results);
assert_eq!(
spilled.search_knn(&query, 7).unwrap(),
resident.search_knn(&query, 7).unwrap()
);
assert!(control.memory().peak() <= control.memory().limit());
drop((retained, changed, actual, compacted, spilled, memory));
assert_eq!(control.memory().used(), 0);
}
#[test]
fn restoration_streams_nodes_and_preserves_bits_with_more_vectors_than_memory() {
let control = StorageReadControl::with_limit(64 * 1024);
let params = HNSWIndexParams {
m: 2,
ef_construction: 4,
..HNSWIndexParams::default()
};
let mut source = HNSWIndex::with_params(1024, params).unwrap();
for document in 1..=20 {
source.add(document, vector(document, 1024)).unwrap();
}
let delta = source.take_persistence_delta();
let mut builder = HNSWRestoreBuilder::new(1024, params, delta.meta, &control).unwrap();
for node in delta.nodes() {
builder.push(node.unwrap().into_parts().0).unwrap();
}
let restored = builder.finish().unwrap();
assert!(restored.nodes.is_spilled());
let query = vector(90, 1024);
assert_eq!(
restored.search_knn(&query, 10).unwrap(),
source.search_knn(&query, 10).unwrap()
);
let actual = restored.delta(&control);
let mut actual = actual;
actual.full_rewrite = true;
let mut nodes = actual.nodes();
for expected in delta.nodes() {
assert_eq!(*nodes.next().unwrap().unwrap(), *expected.unwrap());
}
assert!(nodes.next().is_none());
drop((actual, restored));
assert_eq!(control.memory().used(), 0);
}
#[test]
fn separate_edge_restoration_writes_spilled_vectors_once_per_source() {
let control = StorageReadControl::with_limit(64 * 1024);
let params = HNSWIndexParams {
m: 6,
ef_construction: 12,
..HNSWIndexParams::default()
};
let mut source = HNSWIndex::with_params(1024, params).unwrap();
for document in 1..=24 {
source.add(document, vector(document, 1024)).unwrap();
}
let snapshot = source.persistence_snapshot();
let mut builder = HNSWRestoreBuilder::new(1024, params, snapshot.meta, &control).unwrap();
for node in &snapshot.nodes {
let mut empty = node.clone();
empty.neighbors.iter_mut().for_each(Vec::clear);
builder.push(empty).unwrap();
}
crate::hnsw_index::store::ENCODED_NODES.set(0);
for node in &snapshot.nodes {
for (layer, neighbors) in node.neighbors.iter().enumerate() {
for target in neighbors {
builder.edge(node.node_id, layer, *target).unwrap();
}
}
}
let mut restored = builder.finish().unwrap().into_parts().0;
let encoded = crate::hnsw_index::store::ENCODED_NODES.get();
assert!(restored.nodes.is_spilled());
let mut delta = restored.delta(&control);
delta.full_rewrite = true;
assert_eq!(delta.meta, snapshot.meta);
let mut nodes = delta.nodes();
for expected in &snapshot.nodes {
assert_eq!(&*nodes.next().unwrap().unwrap(), expected);
}
assert!(nodes.next().is_none());
drop(delta);
assert_eq!(
restored.search_knn(&vector(19, 1024), 10).unwrap(),
source.search_knn(&vector(19, 1024), 10).unwrap()
);
assert!(restored.take_persistence_delta().nodes().next().is_none());
assert!(
encoded <= 2 * snapshot.nodes.len(),
"encoded {encoded} vector nodes"
);
assert!(control.memory().peak() <= control.memory().limit());
drop(restored);
assert_eq!(control.memory().used(), 0);
}
#[test]
fn streamed_edge_sources_can_repeat_and_errors_cannot_publish_pending_edges() {
let control = StorageReadControl::with_limit(64 * 1024);
let params = HNSWIndexParams {
m: 2,
ef_construction: 4,
..HNSWIndexParams::default()
};
let mut source = HNSWIndex::with_params(8, params).unwrap();
for document in 1..=12 {
source.add(document, vector(document, 8)).unwrap();
}
let snapshot = source.persistence_snapshot();
let empty_builder = || {
let mut builder = HNSWRestoreBuilder::new(8, params, snapshot.meta, &control).unwrap();
for node in &snapshot.nodes {
let mut node = node.clone();
node.neighbors.iter_mut().for_each(Vec::clear);
builder.push(node).unwrap();
}
builder
};
let edges = snapshot
.nodes
.iter()
.map(|node| {
node.neighbors
.iter()
.enumerate()
.flat_map(|(layer, targets)| {
targets
.iter()
.map(move |target| (node.node_id, layer, *target))
})
.collect::<Vec<_>>()
})
.collect::<Vec<_>>();
let mut builder = empty_builder();
for position in 0..edges.iter().map(Vec::len).max().unwrap() {
for group in &edges {
if let Some(&(source, layer, target)) = group.get(position) {
builder.edge(source, layer, target).unwrap();
}
}
}
let restored = builder.finish().unwrap();
assert_eq!(restored.persistence_snapshot(), snapshot);
drop(restored);
let (source, layer, target) = edges.iter().flatten().next().copied().unwrap();
let mut builder = empty_builder();
builder.edge(source, layer, target).unwrap();
assert!(builder.edge(source, usize::MAX, target).is_err());
assert!(builder.finish().is_err());
let mut builder = empty_builder();
for _ in 0..=params.m * 2 {
if builder.edge(source, 0, target).is_err() {
assert!(builder.finish().is_err());
break;
}
}
let mut builder = empty_builder();
builder.edge(source, layer, target).unwrap();
control.cancellation().cancel();
assert!(builder.finish().is_err());
assert_eq!(control.memory().used(), 0);
}
#[test]
fn adaptive_tensor_search_spills_candidates_without_allocating_one_result_per_vector() {
let control = StorageReadControl::with_limit(32 * 1024);
let params = HNSWIndexParams {
m: 2,
ef_construction: 4,
ef_search: 8,
..HNSWIndexParams::default()
};
let mut source = HNSWIndex::with_params(2, params).unwrap();
source.add_many(7, vec![vec![0.75, 0.25]; 80]).unwrap();
let delta = source.take_persistence_delta();
let mut builder = HNSWRestoreBuilder::new(2, params, delta.meta, &control).unwrap();
for node in delta.nodes() {
builder.push(node.unwrap().into_parts().0).unwrap();
}
let restored = builder.finish().unwrap();
assert!(restored.nodes.is_spilled());
let query = [0.75, 0.25];
assert_eq!(
restored.search_knn(&query, 2).unwrap(),
source.search_knn(&query, 2).unwrap()
);
assert_eq!(
restored.search_threshold(&query, 0.5).unwrap(),
source.search_threshold(&query, 0.5).unwrap()
);
assert!(control.memory().peak() <= control.memory().limit());
drop(restored);
assert_eq!(control.memory().used(), 0);
}
#[test]
fn failed_or_cancelled_stream_builders_cannot_publish_partial_graphs() {
let control = StorageReadControl::with_limit(64 * 1024);
let params = HNSWIndexParams::default();
let mut builder = HNSWCanonicalBuilder::new(2, params, &control).unwrap();
builder.push(1, 0, &[1.0, 0.0]).unwrap();
assert!(builder.push(1, 2, &[1.0, 0.0]).is_err());
assert!(builder.push(1, 1, &[1.0, 0.0]).is_err());
assert!(builder.finish().is_err());
assert_eq!(control.memory().used(), 0);
let mut source = HNSWIndex::new(2);
source.add(1, vec![1.0, 0.0]).unwrap();
let delta = source.take_persistence_delta();
let mut restore = HNSWRestoreBuilder::new(2, params, delta.meta, &control).unwrap();
restore
.push(delta.nodes().next().unwrap().unwrap().into_parts().0)
.unwrap();
assert!(restore.edge(u64::MAX, 0, 1).is_err());
assert!(restore.finish().is_err());
assert_eq!(control.memory().used(), 0);
let builder = HNSWCanonicalBuilder::new(2, params, &control).unwrap();
control.cancellation().cancel();
assert!(builder.finish_delta().is_err());
control.cancellation().reset();
let restore = HNSWRestoreBuilder::new(2, params, delta.meta, &control).unwrap();
control.cancellation().cancel();
assert!(restore.finish().is_err());
assert_eq!(control.memory().used(), 0);
}