105 lines
3.1 KiB
Rust
105 lines
3.1 KiB
Rust
use std::time::Instant;
|
|
|
|
use rand::{thread_rng, Rng};
|
|
use vectordb::Vector;
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
let db = vectordb::VectorDb::open("./vectordb", 2)?;
|
|
|
|
let total_chunks = 32;
|
|
let chunk_size = 1024 * 8;
|
|
|
|
for _ in 0..total_chunks {
|
|
let vectors = (0..chunk_size)
|
|
.map(|_| thread_rng().gen::<[f32; 7]>().into())
|
|
.collect::<Vec<_>>();
|
|
|
|
println!("Inserting vectors");
|
|
let now = Instant::now();
|
|
db.insert_vectors_blocking(vectors)?;
|
|
println!("Inserting vectors - Took {:?}", now.elapsed());
|
|
}
|
|
|
|
println!();
|
|
|
|
let existing_vector: Vector<7> = thread_rng().gen::<[f32; 7]>().into();
|
|
|
|
db.insert_vector_blocking(existing_vector.clone())?;
|
|
|
|
let test_vecs = [existing_vector]
|
|
.into_iter()
|
|
.chain((0..5).map(|_| thread_rng().gen::<[f32; 7]>().into()))
|
|
.collect::<Vec<_>>();
|
|
|
|
for v in &test_vecs {
|
|
println!("Finding similarities");
|
|
let now = Instant::now();
|
|
let similarities = db.find_similarities_blocking(v.clone(), None, 5)?;
|
|
println!("Finding similarities - Took {:?}", now.elapsed());
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
println!(
|
|
"similar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
}
|
|
|
|
println!();
|
|
|
|
for v in &test_vecs {
|
|
println!("Finding furthest similarities");
|
|
let now = Instant::now();
|
|
let similarities = db.find_furthest_similarities_blocking(v.clone(), None, 5)?;
|
|
println!("Finding furthest similarities - Took {:?}", now.elapsed());
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
println!(
|
|
"similar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
}
|
|
|
|
println!();
|
|
|
|
for v in &test_vecs {
|
|
println!("Finding dissimilarities");
|
|
let now = Instant::now();
|
|
let similarities = db.find_dissimilarities_blocking(v.clone(), None, 5)?;
|
|
println!("Finding dissimilarities - Took {:?}", now.elapsed());
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
println!(
|
|
"dissimilar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
}
|
|
|
|
println!();
|
|
|
|
for v in &test_vecs {
|
|
println!("Finding closest dissimilarities");
|
|
let now = Instant::now();
|
|
let similarities = db.find_closest_dissimilarities_blocking(v.clone(), None, 5)?;
|
|
println!("Finding closest dissimilarities - Took {:?}", now.elapsed());
|
|
|
|
for id in similarities {
|
|
let similar = db.get_vector_blocking(id)?.expect("Vector exists");
|
|
println!(
|
|
"dissimilar to {}: {}",
|
|
id,
|
|
v.squared_euclidean_distance(&similar)
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|