use std::sync::Arc;
use uqa_core::{DocId, PostingList};
use super::state::IVFIndex;
use crate::vector_index::VectorIndex;
use crate::StorageBackendResult;
impl VectorIndex for IVFIndex {
fn contains_document(&self, doc_id: DocId) -> StorageBackendResult<bool> {
Ok(self.vectors.lock().contains_key(&(doc_id, 0)))
}
fn dimensions(&self) -> u32 {
self.dimensions
}
fn index_kind(&self) -> &'static str {
"ivf"
}
fn add(&mut self, doc_id: DocId, vector: Vec<f32>) -> StorageBackendResult<()> {
self.replace_document_vectors(doc_id, vec![vector])
}
fn add_many(&mut self, doc_id: DocId, vectors: Vec<Vec<f32>>) -> StorageBackendResult<()> {
self.replace_document_vectors(doc_id, vectors)
}
fn delete(&mut self, doc_id: DocId) -> StorageBackendResult<()> {
self.delete_document(doc_id)
}
fn clear(&mut self) -> StorageBackendResult<()> {
self.clear_index();
Ok(())
}
fn search_knn(&self, query: &[f32], k: usize) -> StorageBackendResult<PostingList> {
self.search_top_k(query, k, None)
}
fn search_threshold(&self, query: &[f32], threshold: f32) -> StorageBackendResult<PostingList> {
self.search_above_threshold(query, threshold, None)
}
fn search_knn_with_control(
&self,
query: &[f32],
k: usize,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<PostingList> {
self.search_top_k(query, k, Some(control))
}
fn search_threshold_with_control(
&self,
query: &[f32],
threshold: f32,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<PostingList> {
self.search_above_threshold(query, threshold, Some(control))
}
fn count(&self) -> StorageBackendResult<usize> {
Ok(self.vectors.lock().len())
}
fn initialize(&mut self) -> StorageBackendResult<()> {
self.train()
}
fn snapshot(&self) -> StorageBackendResult<Arc<dyn VectorIndex>> {
crate::ReadOnlySnapshot::new(Arc::new(self.detached_clone()))
.with_canonical_vectors(None)?
.snapshot()
}
fn snapshot_with_control(
&self,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<Arc<dyn VectorIndex>> {
crate::ReadOnlySnapshot::from_budgeted(self.snapshot_controlled(control)?)?
.with_canonical_vectors(Some(control))?
.with_vector_read_control(control)?
.snapshot()
}
fn vector_read_snapshot(
&self,
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<Option<crate::vector_index::VectorReadSnapshot>> {
self.snapshot_with_control(control)?
.vector_read_snapshot(control)
}
fn matches_document_vectors(
&self,
document: DocId,
vectors: &[Vec<f32>],
control: &crate::read_control::StorageReadControl,
) -> StorageBackendResult<bool> {
crate::vector_index::canonical_vectors_equal(self, document, vectors, control)
}
fn writable_snapshot(&self) -> StorageBackendResult<Box<dyn VectorIndex>> {
Ok(Box::new(self.detached_clone()))
}
}