#![allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::float_cmp
)]
#![allow(clippy::useless_vec)]
use super::native_index::*;
use crate::distance::DistanceMetric;
use crate::index::VectorIndex;
use tempfile::tempdir;
#[test]
fn test_native_index_new() {
let index = NativeHnswIndex::new(64, DistanceMetric::Euclidean).expect("test");
assert_eq!(index.dimension(), 64);
assert_eq!(index.metric(), DistanceMetric::Euclidean);
assert!(index.is_empty());
}
#[test]
fn test_native_index_insert_search() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
for i in 0..50 {
let vec: Vec<f32> = (0..32).map(|j| (i * 32 + j) as f32 * 0.01).collect();
index.insert(i, &vec).expect("test");
}
assert_eq!(index.len(), 50);
let query: Vec<f32> = (0..32).map(|j| j as f32 * 0.01).collect();
let results = index.search(&query, 5);
assert!(!results.is_empty());
assert!(results.len() <= 5);
assert_eq!(results[0].id, 0);
}
#[test]
fn test_native_index_batch_insert() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
let items: Vec<(u64, Vec<f32>)> = (0..50).map(|i| (i, vec![i as f32 * 0.01; 32])).collect();
index.insert_batch(&items).expect("test");
assert_eq!(index.len(), 50);
}
#[test]
fn test_native_index_persistence() {
let dir = tempdir().unwrap();
let index = NativeHnswIndex::new(32, DistanceMetric::Cosine).expect("test");
for i in 0..30 {
index.insert(i, &vec![i as f32 * 0.1; 32]).expect("test");
}
index.save(dir.path()).unwrap();
let loaded = NativeHnswIndex::load(dir.path(), 32, DistanceMetric::Cosine).unwrap();
assert_eq!(loaded.dimension(), 32);
assert_eq!(loaded.metric(), DistanceMetric::Cosine);
assert_eq!(loaded.len(), 30);
let results = loaded.search(&vec![0.0; 32], 5);
assert!(!results.is_empty());
let brute_force = loaded.brute_force_search_parallel(&vec![0.0; 32], 5);
assert_eq!(brute_force.len(), 5);
}
#[test]
fn test_native_index_fast_insert_save_does_not_persist_vectors() {
let dir = tempdir().unwrap();
let index = NativeHnswIndex::new_fast_insert(16, DistanceMetric::Cosine).expect("test");
for i in 0..10 {
index.insert(i, &vec![i as f32 * 0.1; 16]).expect("test");
}
index.save(dir.path()).unwrap();
assert!(!dir.path().join("native_vectors.bin").exists());
let loaded = NativeHnswIndex::load(dir.path(), 16, DistanceMetric::Cosine).unwrap();
assert!(!loaded.has_vector_storage());
assert!(loaded
.brute_force_search_parallel(&vec![0.0; 16], 5)
.is_empty());
}
#[test]
fn test_native_index_fast_insert_save_removes_stale_vectors_file() {
let dir = tempdir().unwrap();
let regular = NativeHnswIndex::new(16, DistanceMetric::Cosine).expect("test");
regular.insert(1, &vec![0.1; 16]).expect("test");
regular.save(dir.path()).unwrap();
assert!(dir.path().join("native_vectors.bin").exists());
let fast = NativeHnswIndex::new_fast_insert(16, DistanceMetric::Cosine).expect("test");
fast.insert(2, &vec![0.2; 16]).expect("test");
fast.save(dir.path()).unwrap();
assert!(!dir.path().join("native_vectors.bin").exists());
}
#[test]
fn test_native_index_delete() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
index.insert(1, &vec![0.1; 32]).expect("test");
index.insert(2, &vec![0.2; 32]).expect("test");
assert!(index.remove(1));
assert!(!index.remove(999));
}
#[test]
fn test_native_index_vector_index_trait() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
<NativeHnswIndex as VectorIndex>::insert(&index, 1, &vec![0.1; 32]);
assert_eq!(<NativeHnswIndex as VectorIndex>::len(&index), 1);
let results = <NativeHnswIndex as VectorIndex>::search(&index, &vec![0.1; 32], 1);
assert_eq!(results.len(), 1);
assert_eq!(results[0].id, 1);
}
#[test]
fn test_native_index_brute_force_search() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
for i in 0..20u64 {
let vec: Vec<f32> = (0..32u64).map(|j| (i * 32 + j) as f32 * 0.001).collect();
index.insert(i, &vec).expect("test");
}
let query: Vec<f32> = (0..32).map(|j| j as f32 * 0.001).collect();
let results = index.brute_force_search_parallel(&query, 5);
assert_eq!(results.len(), 5);
assert_eq!(results[0].id, 0);
for i in 1..results.len() {
assert!(
results[i].score >= results[i - 1].score,
"Results not sorted"
);
}
}
#[test]
fn test_native_index_brute_force_empty() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
let query = vec![0.0; 32];
let results = index.brute_force_search_parallel(&query, 5);
assert!(results.is_empty());
}
#[test]
fn test_native_index_brute_force_k_larger_than_size() {
let index = NativeHnswIndex::new(32, DistanceMetric::Euclidean).expect("test");
index.insert(1, &vec![0.1; 32]).expect("test");
index.insert(2, &vec![0.2; 32]).expect("test");
let results = index.brute_force_search_parallel(&vec![0.0; 32], 10);
assert_eq!(results.len(), 2);
}
#[test]
fn test_native_insert_same_id_updates_vector() {
let index = NativeHnswIndex::new(4, DistanceMetric::Cosine).expect("test");
let vector_a = [1.0, 0.0, 0.0, 0.0];
index.insert(1, &vector_a).expect("test");
let vector_b = [0.0, 1.0, 0.0, 0.0];
index.insert(1, &vector_b).expect("test");
assert_eq!(index.len(), 1, "Upsert must not create duplicate entries");
let results = index.search(&vector_b, 1);
assert_eq!(results.len(), 1, "Should find exactly one result");
assert_eq!(results[0].id, 1, "Result must be id=1");
assert!(
results[0].score > 0.9,
"Similarity to updated vector B should be > 0.9, got {}",
results[0].score,
);
}
#[test]
fn test_native_batch_upsert_updates_existing() {
let index = NativeHnswIndex::new(4, DistanceMetric::Cosine).expect("test");
let initial: Vec<(u64, Vec<f32>)> = (0..10).map(|i| (i, vec![1.0, 0.0, 0.0, 0.0])).collect();
index.insert_batch(&initial).expect("test");
assert_eq!(
index.len(),
10,
"Should have 10 vectors after initial batch"
);
let updates: Vec<(u64, Vec<f32>)> = (0..5).map(|i| (i, vec![0.0, 1.0, 0.0, 0.0])).collect();
index.insert_batch(&updates).expect("test");
assert_eq!(index.len(), 10, "Upsert batch must not inflate count");
let query = [0.0, 1.0, 0.0, 0.0];
let results = index.search(&query, 5);
assert!(!results.is_empty(), "Search must return results");
assert!(
results[0].id < 5,
"Top result should be an updated vector (id < 5), got id={}",
results[0].id,
);
assert!(
results[0].score > 0.9,
"Updated vector similarity should be > 0.9, got {}",
results[0].score,
);
}
#[test]
fn test_native_remove_cleans_up_vector_storage() {
let index = NativeHnswIndex::new(4, DistanceMetric::Cosine).expect("test");
index.insert(1, &[1.0, 0.0, 0.0, 0.0]).expect("test");
let before = index.brute_force_search_parallel(&[1.0, 0.0, 0.0, 0.0], 1);
assert_eq!(before.len(), 1, "Should find vector before removal");
assert!(index.remove(1), "Remove should return true for existing ID");
let after = index.brute_force_search_parallel(&[1.0, 0.0, 0.0, 0.0], 1);
assert!(
after.is_empty(),
"Brute-force should find nothing after removal, got {} results",
after.len(),
);
}