vectordb/examples/run.rs

107 lines
3 KiB
Rust

use std::time::Instant;
use rand::{thread_rng, Rng};
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_many_vectors(&vectors)?;
println!("Inserting vectors - Took {:?}", now.elapsed());
std::thread::sleep(std::time::Duration::from_secs(3));
}
println!();
std::thread::sleep(std::time::Duration::from_secs(10));
let existing_vector = thread_rng().gen::<[f32; 7]>().into();
db.insert_vector(&existing_vector)?;
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(&v, 5)?;
println!("Finding similarities - Took {:?}", now.elapsed());
for id in similarities {
let similar = db.get_vector(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(&v, 5)?;
println!("Finding furthest similarities - Took {:?}", now.elapsed());
for id in similarities {
let similar = db.get_vector(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(&v, 5)?;
println!("Finding dissimilarities - Took {:?}", now.elapsed());
for id in similarities {
let similar = db.get_vector(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(&v, 5)?;
println!("Finding closest dissimilarities - Took {:?}", now.elapsed());
for id in similarities {
let similar = db.get_vector(id)?.expect("Vector exists");
println!(
"dissimilar to {}: {}",
id,
v.squared_euclidean_distance(&similar)
);
}
}
Ok(())
}