Skip to main content

storage/
traits.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use common::{NamespaceId, Result, Vector, VectorId};
5
6/// Index types that can be persisted
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
8pub enum IndexType {
9    /// HNSW graph index
10    Hnsw,
11    /// Product Quantization index
12    Pq,
13    /// IVF (Inverted File) index
14    Ivf,
15    /// SPFresh index
16    SpFresh,
17    /// Full-text inverted index
18    FullText,
19}
20
21impl IndexType {
22    pub fn as_str(&self) -> &'static str {
23        match self {
24            IndexType::Hnsw => "hnsw",
25            IndexType::Pq => "pq",
26            IndexType::Ivf => "ivf",
27            IndexType::SpFresh => "spfresh",
28            IndexType::FullText => "fulltext",
29        }
30    }
31}
32
33/// Storage trait for persisting index data
34#[async_trait]
35pub trait IndexStorage: Send + Sync {
36    /// Save index data for a namespace
37    async fn save_index(
38        &self,
39        namespace: &NamespaceId,
40        index_type: IndexType,
41        data: Vec<u8>,
42    ) -> Result<()>;
43
44    /// Load index data for a namespace
45    async fn load_index(
46        &self,
47        namespace: &NamespaceId,
48        index_type: IndexType,
49    ) -> Result<Option<Vec<u8>>>;
50
51    /// Delete index data for a namespace
52    async fn delete_index(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool>;
53
54    /// Check if index exists for a namespace
55    async fn index_exists(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool>;
56
57    /// List all indexes for a namespace
58    async fn list_indexes(&self, namespace: &NamespaceId) -> Result<Vec<IndexType>>;
59}
60
61/// Blanket implementation for Arc<T>
62#[async_trait]
63impl<T: IndexStorage + ?Sized> IndexStorage for Arc<T> {
64    async fn save_index(
65        &self,
66        namespace: &NamespaceId,
67        index_type: IndexType,
68        data: Vec<u8>,
69    ) -> Result<()> {
70        (**self).save_index(namespace, index_type, data).await
71    }
72
73    async fn load_index(
74        &self,
75        namespace: &NamespaceId,
76        index_type: IndexType,
77    ) -> Result<Option<Vec<u8>>> {
78        (**self).load_index(namespace, index_type).await
79    }
80
81    async fn delete_index(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool> {
82        (**self).delete_index(namespace, index_type).await
83    }
84
85    async fn index_exists(&self, namespace: &NamespaceId, index_type: IndexType) -> Result<bool> {
86        (**self).index_exists(namespace, index_type).await
87    }
88
89    async fn list_indexes(&self, namespace: &NamespaceId) -> Result<Vec<IndexType>> {
90        (**self).list_indexes(namespace).await
91    }
92}
93
94/// Core storage abstraction - implementations can be in-memory, S3, etc.
95#[async_trait]
96pub trait VectorStorage: Send + Sync {
97    /// Store or update vectors in a namespace
98    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize>;
99
100    /// Get vectors by IDs
101    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>>;
102
103    /// Get all vectors in a namespace (for brute-force search)
104    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>>;
105
106    /// Get all vectors in a namespace WITHOUT their embedding values (DAK-7387).
107    ///
108    /// Metadata-only projection for scan paths (e.g. batch recall) that filter on
109    /// metadata and never read `values` — skipping the per-vector embedding clone.
110    /// Returned vectors are identical to [`VectorStorage::get_all`] except `values`
111    /// is empty. The default implementation falls back to `get_all` and drops the
112    /// values, so wrapper/remote backends stay correct; in-process backends should
113    /// override to avoid cloning embeddings at all.
114    async fn get_all_meta(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
115        let mut all = self.get_all(namespace).await?;
116        for v in &mut all {
117            v.values = Vec::new();
118        }
119        Ok(all)
120    }
121
122    /// Delete vectors by IDs
123    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize>;
124
125    /// Check if namespace exists
126    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool>;
127
128    /// Create namespace if it doesn't exist
129    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()>;
130
131    /// Get vector count in namespace
132    async fn count(&self, namespace: &NamespaceId) -> Result<usize>;
133
134    /// Get vector dimension for namespace (None if empty)
135    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>>;
136
137    /// List all namespaces
138    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>>;
139
140    /// Delete a namespace and all its vectors
141    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool>;
142
143    /// Clean up expired vectors in a namespace
144    /// Returns the number of vectors removed
145    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize>;
146
147    /// Clean up expired vectors in all namespaces
148    /// Returns total number of vectors removed
149    async fn cleanup_all_expired(&self) -> Result<usize>;
150
151    /// Retrieve a page of vectors from a namespace for memory-bounded iteration (DAK-5716).
152    ///
153    /// Default implementation loads the entire namespace and slices — override in
154    /// storage backends that support native cursor-based pagination for O(page) memory.
155    async fn get_page(
156        &self,
157        namespace: &NamespaceId,
158        offset: usize,
159        limit: usize,
160    ) -> Result<Vec<Vector>> {
161        let all = self.get_all(namespace).await?;
162        Ok(all.into_iter().skip(offset).take(limit).collect())
163    }
164
165    /// Best-effort release of derived cache layers held by this storage stack
166    /// (e.g. the Moka L1 vector cache), called under memory pressure before new
167    /// work is shed (DAK-7337). Durable data MUST NOT be dropped — only layers
168    /// that can be transparently repopulated from the source of truth. Default:
169    /// no-op for backends with no derived layers.
170    async fn reclaim_derived_caches(&self) {}
171}
172
173/// Blanket implementation for Arc<T> to enable dynamic dispatch
174#[async_trait]
175impl<T: VectorStorage + ?Sized> VectorStorage for Arc<T> {
176    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
177        (**self).upsert(namespace, vectors).await
178    }
179
180    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
181        (**self).get(namespace, ids).await
182    }
183
184    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
185        (**self).get_all(namespace).await
186    }
187
188    async fn get_all_meta(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
189        (**self).get_all_meta(namespace).await
190    }
191
192    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
193        (**self).delete(namespace, ids).await
194    }
195
196    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
197        (**self).namespace_exists(namespace).await
198    }
199
200    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
201        (**self).ensure_namespace(namespace).await
202    }
203
204    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
205        (**self).count(namespace).await
206    }
207
208    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
209        (**self).dimension(namespace).await
210    }
211
212    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
213        (**self).list_namespaces().await
214    }
215
216    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
217        (**self).delete_namespace(namespace).await
218    }
219
220    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
221        (**self).cleanup_expired(namespace).await
222    }
223
224    async fn cleanup_all_expired(&self) -> Result<usize> {
225        (**self).cleanup_all_expired().await
226    }
227
228    async fn get_page(
229        &self,
230        namespace: &NamespaceId,
231        offset: usize,
232        limit: usize,
233    ) -> Result<Vec<Vector>> {
234        (**self).get_page(namespace, offset, limit).await
235    }
236
237    async fn reclaim_derived_caches(&self) {
238        (**self).reclaim_derived_caches().await
239    }
240}