a3s_memory/vector/
in_memory.rs

1use super::search::search_snapshot;
2use super::{
3    VectorBudgetResource, VectorIndex, VectorIndexDescriptor, VectorIndexError, VectorIndexStatus,
4    VectorNormalization, VectorRecord, VectorResult, VectorRevision, VectorSearchRequest,
5    VectorSearchResult,
6};
7use std::collections::{BTreeMap, BTreeSet};
8use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
9
10/// Exact, session-ephemeral vector index backed by immutable partition blocks.
11#[derive(Clone)]
12pub struct InMemoryVectorIndex {
13    inner: Arc<IndexInner>,
14}
15
16impl std::fmt::Debug for InMemoryVectorIndex {
17    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        formatter
19            .debug_struct("InMemoryVectorIndex")
20            .field("descriptor", &self.inner.descriptor)
21            .field("status", &self.status())
22            .finish()
23    }
24}
25
26struct IndexInner {
27    descriptor: VectorIndexDescriptor,
28    snapshot: RwLock<Arc<IndexSnapshot>>,
29}
30
31#[derive(Default)]
32pub(super) struct IndexSnapshot {
33    revision: VectorRevision,
34    pub(super) partitions: BTreeMap<String, Arc<PartitionBlock>>,
35    record_count: usize,
36    byte_count: usize,
37}
38
39pub(super) struct PartitionBlock {
40    pub(super) name: String,
41    pub(super) ids: Vec<String>,
42    pub(super) labels: Vec<BTreeMap<String, String>>,
43    pub(super) vectors: Vec<f32>,
44    byte_count: usize,
45}
46
47impl PartitionBlock {
48    pub(super) fn record_count(&self) -> usize {
49        self.ids.len()
50    }
51}
52
53impl IndexSnapshot {
54    pub(super) fn status(&self) -> VectorIndexStatus {
55        VectorIndexStatus {
56            revision: self.revision,
57            partition_count: self.partitions.len(),
58            record_count: self.record_count,
59            byte_count: self.byte_count,
60        }
61    }
62}
63
64impl InMemoryVectorIndex {
65    pub fn new(descriptor: VectorIndexDescriptor) -> VectorResult<Self> {
66        descriptor.validate()?;
67        Ok(Self {
68            inner: Arc::new(IndexInner {
69                descriptor,
70                snapshot: RwLock::new(Arc::new(IndexSnapshot::default())),
71            }),
72        })
73    }
74
75    fn snapshot(&self) -> Arc<IndexSnapshot> {
76        read_unpoisoned(&self.inner.snapshot).clone()
77    }
78}
79
80#[async_trait::async_trait]
81impl VectorIndex for InMemoryVectorIndex {
82    fn descriptor(&self) -> &VectorIndexDescriptor {
83        &self.inner.descriptor
84    }
85
86    fn status(&self) -> VectorIndexStatus {
87        self.snapshot().status()
88    }
89
90    async fn replace_partition(
91        &self,
92        partition: &str,
93        records: Vec<VectorRecord>,
94    ) -> VectorResult<VectorIndexStatus> {
95        let partition = validate_partition(partition)?.to_string();
96        let inner = Arc::clone(&self.inner);
97        run_blocking(move || {
98            let block = build_partition(&inner.descriptor, partition, records)?;
99            publish_partition(&inner, block)
100        })
101        .await
102    }
103
104    async fn remove_partition(&self, partition: &str) -> VectorResult<VectorIndexStatus> {
105        let partition = validate_partition(partition)?.to_string();
106        let inner = Arc::clone(&self.inner);
107        run_blocking(move || remove_partition(&inner, &partition)).await
108    }
109
110    async fn search(&self, mut request: VectorSearchRequest) -> VectorResult<VectorSearchResult> {
111        validate_request_filters(&request)?;
112        if request.limit == 0 {
113            return Err(VectorIndexError::InvalidRequest(
114                "limit must be greater than zero".to_string(),
115            ));
116        }
117        let query = prepare_vector(
118            std::mem::take(&mut request.embedding),
119            &self.inner.descriptor,
120            "query".to_string(),
121        )?;
122        let descriptor = self.inner.descriptor.clone();
123        let snapshot = self.snapshot();
124        run_blocking(move || search_snapshot(snapshot, &descriptor, query, request)).await
125    }
126
127    async fn clear(&self) -> VectorResult<VectorIndexStatus> {
128        let inner = Arc::clone(&self.inner);
129        run_blocking(move || clear_index(&inner)).await
130    }
131}
132
133async fn run_blocking<T, F>(operation: F) -> VectorResult<T>
134where
135    T: Send + 'static,
136    F: FnOnce() -> VectorResult<T> + Send + 'static,
137{
138    tokio::task::spawn_blocking(operation)
139        .await
140        .map_err(|error| VectorIndexError::WorkerFailed(error.to_string()))?
141}
142
143fn validate_partition(partition: &str) -> VectorResult<&str> {
144    let partition = partition.trim();
145    if partition.is_empty() {
146        Err(VectorIndexError::InvalidPartition)
147    } else {
148        Ok(partition)
149    }
150}
151
152fn validate_request_filters(request: &VectorSearchRequest) -> VectorResult<()> {
153    if request
154        .partitions
155        .iter()
156        .any(|partition| partition.trim().is_empty())
157    {
158        return Err(VectorIndexError::InvalidPartition);
159    }
160    if request.labels.keys().any(|key| key.trim().is_empty()) {
161        return Err(VectorIndexError::InvalidLabel {
162            context: "query filter".to_string(),
163        });
164    }
165    Ok(())
166}
167
168fn build_partition(
169    descriptor: &VectorIndexDescriptor,
170    name: String,
171    records: Vec<VectorRecord>,
172) -> VectorResult<Arc<PartitionBlock>> {
173    if records.len() > descriptor.max_records {
174        return Err(VectorIndexError::BudgetExceeded {
175            resource: VectorBudgetResource::Records,
176            limit: descriptor.max_records,
177            required: records.len(),
178        });
179    }
180    let minimum_vector_bytes = records
181        .len()
182        .checked_mul(descriptor.dimension)
183        .and_then(|elements| elements.checked_mul(std::mem::size_of::<f32>()))
184        .ok_or(VectorIndexError::SizeOverflow)?;
185    if minimum_vector_bytes > descriptor.max_bytes {
186        return Err(VectorIndexError::BudgetExceeded {
187            resource: VectorBudgetResource::Bytes,
188            limit: descriptor.max_bytes,
189            required: minimum_vector_bytes,
190        });
191    }
192    let mut seen = BTreeSet::new();
193    let mut byte_count = std::mem::size_of::<PartitionBlock>()
194        .checked_add(name.len())
195        .ok_or(VectorIndexError::SizeOverflow)?;
196
197    for (record_index, record) in records.iter().enumerate() {
198        if record.id.trim().is_empty() {
199            return Err(VectorIndexError::InvalidRecordId {
200                partition: name.clone(),
201                record_index,
202            });
203        }
204        if !seen.insert(record.id.clone()) {
205            return Err(VectorIndexError::DuplicateRecordId {
206                partition: name.clone(),
207                id: record.id.clone(),
208            });
209        }
210        if record.labels.keys().any(|key| key.trim().is_empty()) {
211            return Err(VectorIndexError::InvalidLabel {
212                context: format!("record '{}' in partition '{name}'", record.id),
213            });
214        }
215        let context = format!("record '{}' in partition '{name}'", record.id);
216        validate_vector(&record.embedding, descriptor, context)?;
217        byte_count = accounted_record_bytes(byte_count, &record.id, &record.labels, descriptor)?;
218        if byte_count > descriptor.max_bytes {
219            return Err(VectorIndexError::BudgetExceeded {
220                resource: VectorBudgetResource::Bytes,
221                limit: descriptor.max_bytes,
222                required: byte_count,
223            });
224        }
225    }
226
227    let vector_capacity = records
228        .len()
229        .checked_mul(descriptor.dimension)
230        .ok_or(VectorIndexError::SizeOverflow)?;
231    let mut ids = Vec::with_capacity(records.len());
232    let mut labels = Vec::with_capacity(records.len());
233    let mut vectors = Vec::with_capacity(vector_capacity);
234    for record in records {
235        let context = format!("record '{}' in partition '{name}'", record.id);
236        let embedding = prepare_vector(record.embedding, descriptor, context)?;
237        ids.push(record.id);
238        labels.push(record.labels);
239        vectors.extend(embedding);
240    }
241
242    Ok(Arc::new(PartitionBlock {
243        name,
244        ids,
245        labels,
246        vectors,
247        byte_count,
248    }))
249}
250
251fn accounted_record_bytes(
252    current: usize,
253    id: &str,
254    labels: &BTreeMap<String, String>,
255    descriptor: &VectorIndexDescriptor,
256) -> VectorResult<usize> {
257    let label_bytes = labels.iter().try_fold(0usize, |total, (key, value)| {
258        total
259            .checked_add(key.len())
260            .and_then(|total| total.checked_add(value.len()))
261            .ok_or(VectorIndexError::SizeOverflow)
262    })?;
263    let vector_bytes = descriptor
264        .dimension
265        .checked_mul(std::mem::size_of::<f32>())
266        .ok_or(VectorIndexError::SizeOverflow)?;
267    current
268        .checked_add(std::mem::size_of::<String>())
269        .and_then(|value| value.checked_add(std::mem::size_of::<BTreeMap<String, String>>()))
270        .and_then(|value| value.checked_add(id.len()))
271        .and_then(|value| value.checked_add(label_bytes))
272        .and_then(|value| value.checked_add(vector_bytes))
273        .ok_or(VectorIndexError::SizeOverflow)
274}
275
276fn prepare_vector(
277    mut vector: Vec<f32>,
278    descriptor: &VectorIndexDescriptor,
279    context: String,
280) -> VectorResult<Vec<f32>> {
281    validate_vector(&vector, descriptor, context.clone())?;
282    if descriptor.normalization == VectorNormalization::Unit {
283        normalize_unit(&mut vector);
284    }
285    Ok(vector)
286}
287
288fn validate_vector(
289    vector: &[f32],
290    descriptor: &VectorIndexDescriptor,
291    context: String,
292) -> VectorResult<()> {
293    if vector.len() != descriptor.dimension {
294        return Err(VectorIndexError::DimensionMismatch {
295            context,
296            expected: descriptor.dimension,
297            actual: vector.len(),
298        });
299    }
300    if let Some(element_index) = vector.iter().position(|value| !value.is_finite()) {
301        return Err(VectorIndexError::NonFiniteVector {
302            context,
303            element_index,
304        });
305    }
306    if descriptor.normalization == VectorNormalization::Unit {
307        let squared_norm = vector.iter().fold(0.0f64, |sum, value| {
308            let value = f64::from(*value);
309            sum + value * value
310        });
311        if squared_norm == 0.0 {
312            return Err(VectorIndexError::ZeroVector { context });
313        }
314    }
315    Ok(())
316}
317
318fn normalize_unit(vector: &mut [f32]) {
319    let norm = vector
320        .iter()
321        .fold(0.0f64, |sum, value| {
322            let value = f64::from(*value);
323            sum + value * value
324        })
325        .sqrt();
326    for value in vector {
327        *value = (f64::from(*value) / norm) as f32;
328    }
329}
330
331fn publish_partition(
332    inner: &IndexInner,
333    block: Arc<PartitionBlock>,
334) -> VectorResult<VectorIndexStatus> {
335    let mut published = write_unpoisoned(&inner.snapshot);
336    let current = Arc::clone(&published);
337    let existing = current.partitions.get(&block.name);
338
339    if block.record_count() == 0 && existing.is_none() {
340        return Ok(current.status());
341    }
342
343    let old_records = existing.map_or(0, |partition| partition.record_count());
344    let old_bytes = existing.map_or(0, |partition| partition.byte_count);
345    let record_count = current
346        .record_count
347        .checked_sub(old_records)
348        .and_then(|count| count.checked_add(block.record_count()))
349        .ok_or(VectorIndexError::SizeOverflow)?;
350    let retained_bytes = current
351        .byte_count
352        .checked_sub(old_bytes)
353        .ok_or(VectorIndexError::SizeOverflow)?;
354    let byte_count = if block.record_count() == 0 {
355        retained_bytes
356    } else {
357        retained_bytes
358            .checked_add(block.byte_count)
359            .ok_or(VectorIndexError::SizeOverflow)?
360    };
361    enforce_budgets(&inner.descriptor, record_count, byte_count)?;
362
363    let mut partitions = current.partitions.clone();
364    if block.record_count() == 0 {
365        partitions.remove(&block.name);
366    } else {
367        partitions.insert(block.name.clone(), block);
368    }
369    let next = Arc::new(IndexSnapshot {
370        revision: current.revision.next()?,
371        partitions,
372        record_count,
373        byte_count,
374    });
375    let status = next.status();
376    *published = next;
377    Ok(status)
378}
379
380fn remove_partition(inner: &IndexInner, partition: &str) -> VectorResult<VectorIndexStatus> {
381    let mut published = write_unpoisoned(&inner.snapshot);
382    let current = Arc::clone(&published);
383    let Some(existing) = current.partitions.get(partition) else {
384        return Ok(current.status());
385    };
386    let mut partitions = current.partitions.clone();
387    partitions.remove(partition);
388    let next = Arc::new(IndexSnapshot {
389        revision: current.revision.next()?,
390        partitions,
391        record_count: current
392            .record_count
393            .checked_sub(existing.record_count())
394            .ok_or(VectorIndexError::SizeOverflow)?,
395        byte_count: current
396            .byte_count
397            .checked_sub(existing.byte_count)
398            .ok_or(VectorIndexError::SizeOverflow)?,
399    });
400    let status = next.status();
401    *published = next;
402    Ok(status)
403}
404
405fn clear_index(inner: &IndexInner) -> VectorResult<VectorIndexStatus> {
406    let mut published = write_unpoisoned(&inner.snapshot);
407    let current = Arc::clone(&published);
408    if current.partitions.is_empty() {
409        return Ok(current.status());
410    }
411    let next = Arc::new(IndexSnapshot {
412        revision: current.revision.next()?,
413        ..IndexSnapshot::default()
414    });
415    let status = next.status();
416    *published = next;
417    Ok(status)
418}
419
420fn enforce_budgets(
421    descriptor: &VectorIndexDescriptor,
422    record_count: usize,
423    byte_count: usize,
424) -> VectorResult<()> {
425    if record_count > descriptor.max_records {
426        return Err(VectorIndexError::BudgetExceeded {
427            resource: VectorBudgetResource::Records,
428            limit: descriptor.max_records,
429            required: record_count,
430        });
431    }
432    if byte_count > descriptor.max_bytes {
433        return Err(VectorIndexError::BudgetExceeded {
434            resource: VectorBudgetResource::Bytes,
435            limit: descriptor.max_bytes,
436            required: byte_count,
437        });
438    }
439    Ok(())
440}
441
442fn read_unpoisoned<T>(lock: &RwLock<T>) -> RwLockReadGuard<'_, T> {
443    lock.read()
444        .unwrap_or_else(std::sync::PoisonError::into_inner)
445}
446
447fn write_unpoisoned<T>(lock: &RwLock<T>) -> RwLockWriteGuard<'_, T> {
448    lock.write()
449        .unwrap_or_else(std::sync::PoisonError::into_inner)
450}
451
452#[cfg(test)]
453mod lifetime_tests {
454    use super::*;
455
456    #[test]
457    fn last_index_handle_releases_the_complete_index_graph() {
458        let index = InMemoryVectorIndex::new(VectorIndexDescriptor::new(3)).unwrap();
459        let clone = index.clone();
460        let weak = Arc::downgrade(&index.inner);
461
462        drop(index);
463        assert!(weak.upgrade().is_some());
464        drop(clone);
465        assert!(weak.upgrade().is_none());
466    }
467}
a3s_memory/vector/in_memory.rs

a3s_memory/vector/
in_memory.rs
