Skip to main content

abu_rag/vectordb/storage/
memory.rs

1use std::{borrow::Cow, collections::HashMap, convert::Infallible};
2use crate::vectordb::VectorId;
3use super::VectorStorage;
4
5pub struct InMemoryStorage<P> {
6    payloads: HashMap<VectorId, P>
7}
8
9impl<P> InMemoryStorage<P> {
10    pub fn new() -> Self {
11        Self { payloads: HashMap::new() }
12    }
13}
14
15impl<P: Clone + 'static + Send + Sync> VectorStorage for InMemoryStorage<P> {
16    type Payload = P;
17    type Error = Infallible;
18
19    async fn add(&mut self, id: VectorId, payload: P) -> Result<(), Self::Error> {
20        self.payloads.insert(id, payload);
21        Ok(())
22    }
23
24    async fn delete(&mut self, id: VectorId) -> Result<(), Self::Error> {
25        self.payloads.remove(&id);
26        Ok(())
27    }
28
29    async fn get(&self, id: VectorId) -> Result<Option<Cow<P>>, Self::Error> {
30        let value = self.payloads.get(&id);
31        Ok(value.map(|v| Cow::Borrowed(v)))
32    }
33
34    async fn clear(&mut self) -> Result<(), Self::Error> {
35        self.payloads.clear();
36        Ok(())
37    }
38}