vectordb/examples/run.rs

66 lines
1.8 KiB
Rust

use std::time::Instant;
use itertools::Itertools;
use rand::{thread_rng, Rng};
fn main() {
let db = vectordb::VectorDb::open("./db.redb", String::from("moments")).expect("Launched");
let total = 1024 * 1024;
println!("Inserting");
let now = Instant::now();
let chunks = (0..total)
.map(|_| thread_rng().gen::<[f32; 7]>().into())
.chunks(1024);
let iter = chunks.into_iter().map(|chunk| {
let vectors = chunk.collect::<Vec<_>>();
db.insert_many_vectors(&vectors)
});
for res in iter {
res.expect("Insert works");
}
println!("Inserting - Took {:?}", now.elapsed());
let existing_vector = thread_rng().gen::<[f32; 7]>().into();
db.insert_vector(&existing_vector).expect("Insert works");
for _ in 0..4 {
println!("Creating Hypertree");
let now = Instant::now();
db.create_hypertree().expect("Tree created");
println!("Creating Hypertree - Took {:?}", now.elapsed());
}
let test_vecs = [existing_vector]
.into_iter()
.chain((0..50).map(|_| thread_rng().gen::<[f32; 7]>().into()));
for v in test_vecs {
let compare_id = db.insert_vector(&v).expect("Inserted");
println!("Finding similarities");
let now = Instant::now();
let similarities = db
.find_similarities(compare_id, 5)
.expect("Similar vectors exist");
println!("Finding similarities - Took {:?}", now.elapsed());
for id in similarities {
let similar = db
.get_vector(id)
.expect("Got vector")
.expect("Vector exists");
println!(
"{} is similar to {}: {}",
compare_id,
id,
v.squared_euclidean_distance(&similar)
);
}
}
}