use super::*;
fn open(base: &std::path::Path) -> StorageHandle {
open_store(StorageConfig {
buffer_pool_pages: 8,
wal_dir: base.join("wal"),
wal_segment_max_bytes: 1 << 20,
manifest_path: base.join("ir.manifest"),
sstable_dir: base.join("sst"),
})
.unwrap()
}
#[test]
fn ann_space_for_query_routes_by_dimension() {
let base = temp_dir("ann_route_by_dim");
let mut handle = open(&base);
let space2d = 10_u32;
let space3d = 11_u32;
for i in 1..=5_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space2d, VectorMetric::Cosine, &[i as f32, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
for i in 6..=10_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space3d, VectorMetric::Cosine, &[i as f32, 0.0, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
let routed = ann_space_for_query(&handle, VectorMetric::Cosine, Some(2));
assert_eq!(
routed,
Some(space2d),
"dim=2 query must route to space {space2d}"
);
let routed = ann_space_for_query(&handle, VectorMetric::Cosine, Some(3));
assert_eq!(
routed,
Some(space3d),
"dim=3 query must route to space {space3d}"
);
}
#[test]
fn ann_space_for_query_selects_largest_graph_when_ambiguous() {
let base = temp_dir("ann_select_largest");
let mut handle = open(&base);
let space_small = 20_u32;
let space_large = 21_u32;
for i in 1..=2_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space_small, VectorMetric::Cosine, &[i as f32, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
for i in 3..=12_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space_large, VectorMetric::Cosine, &[i as f32, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
let routed = ann_space_for_query(&handle, VectorMetric::Cosine, None);
assert_eq!(
routed,
Some(space_large),
"must select the larger graph when dimension is ambiguous"
);
}
#[test]
fn ann_space_for_query_skips_empty_sibling_graph() {
let base = temp_dir("ann_skip_empty_sibling");
let mut handle = open(&base);
let space_live = 30_u32;
let space_drained = 31_u32;
for i in 1..=3_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space_live, VectorMetric::Cosine, &[i as f32, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
for i in 10..=11_u64 {
put_full_node(&mut handle, i, 1, &[]).unwrap();
let payload =
encode_vector_payload_f32(space_drained, VectorMetric::Cosine, &[i as f32, 0.0], false);
put_vector_delta(&mut handle, &encode_delta(i, 2, &payload)).unwrap();
}
put_tombstone(&mut handle, 10, 3).unwrap();
put_tombstone(&mut handle, 11, 3).unwrap();
rebuild_vector_space(&mut handle, space_drained).unwrap();
let routed = ann_space_for_query(&handle, VectorMetric::Cosine, Some(2));
assert_eq!(
routed,
Some(space_live),
"must skip empty sibling graph and route to space with live vectors"
);
}