use ndarray::Array1;
use super::container::VectorContainer;
#[derive(Clone,Debug)]
pub struct VectorStore {
pub store: Vec<VectorContainer>,
}
impl VectorStore {
pub fn new() -> Self {
Self { store: Vec::new() }
}
}
#[derive(Clone,Debug)]
pub struct VectorContainerWithSimilarity {
pub vector: Array1<f64>,
pub similarity: f64,
pub name: String,
}
impl VectorStore {
pub fn add(&mut self, container: VectorContainer) {
self.store.push(container);
}
pub fn len(&self) -> usize {
self.store.len()
}
pub fn get(&self, index: usize) -> Option<&VectorContainer> {
self.store.get(index)
}
pub fn find_similers(
&self,
vector: &Array1<f64>,
top_k: usize,
) -> Vec<VectorContainerWithSimilarity> {
let mut top_k_vec = Vec::<VectorContainerWithSimilarity>::new();
for container in &self.store {
top_k_vec.push(VectorContainerWithSimilarity {
vector: container.vector.clone(),
similarity: container.similarity(vector),
name: container.name.clone(),
})
}
top_k_vec.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap());
top_k_vec.truncate(top_k);
top_k_vec
}
}