Skip to main content

weavatrix_search_vector/mutable/
lifecycle.rs

1use super::support::current_len;
2use super::{MutableSnapshot, MutableState, MutableVectorIndex, VectorRecord};
3use crate::config::IndexConfig;
4use crate::error::SearchError;
5use crate::hnsw::VectorIndex;
6use crate::metadata::MetadataIndex;
7use crate::simd::DistanceKernel;
8use std::collections::{BTreeMap, BTreeSet};
9use std::sync::{Arc, RwLock};
10
11impl MutableVectorIndex {
12    /// Builds a mutable index and preserves supplied metadata.
13    ///
14    /// # Errors
15    ///
16    /// Returns typed config, vector, allocation, capacity, or duplicate-key
17    /// errors.
18    pub fn build(config: IndexConfig, records: &[VectorRecord]) -> Result<Self, SearchError> {
19        let vectors = records
20            .iter()
21            .map(|record| (record.key, record.vector.as_slice()))
22            .collect::<Vec<_>>();
23        let base = Arc::new(VectorIndex::build(config.clone(), &vectors)?);
24        let mut metadata = MetadataIndex::new();
25        for record in records {
26            if !record.metadata.is_empty() {
27                metadata.insert(record.key, record.metadata.clone());
28            }
29        }
30        Ok(Self {
31            config,
32            state: RwLock::new(MutableState {
33                base,
34                sealed: None,
35                pending: BTreeMap::new(),
36                deleted: BTreeSet::new(),
37                metadata,
38                generation: 0,
39            }),
40            distance_kernel: DistanceKernel::detect(),
41        })
42    }
43
44    /// Wraps an existing immutable index with an empty mutable delta.
45    #[must_use]
46    pub fn from_index(index: VectorIndex) -> Self {
47        let config = index.config().clone();
48        Self {
49            config,
50            state: RwLock::new(MutableState {
51                base: Arc::new(index),
52                sealed: None,
53                pending: BTreeMap::new(),
54                deleted: BTreeSet::new(),
55                metadata: MetadataIndex::new(),
56                generation: 0,
57            }),
58            distance_kernel: DistanceKernel::detect(),
59        }
60    }
61
62    #[must_use]
63    pub const fn config(&self) -> &IndexConfig {
64        &self.config
65    }
66
67    #[must_use]
68    pub fn len(&self) -> usize {
69        let state = self
70            .state
71            .read()
72            .unwrap_or_else(std::sync::PoisonError::into_inner);
73        current_len(&state)
74    }
75
76    #[must_use]
77    pub fn is_empty(&self) -> bool {
78        self.len() == 0
79    }
80
81    #[must_use]
82    pub fn delta_len(&self) -> usize {
83        let state = self
84            .state
85            .read()
86            .unwrap_or_else(std::sync::PoisonError::into_inner);
87        state
88            .sealed
89            .as_ref()
90            .map_or(0, |index| index.len())
91            .saturating_add(state.pending.len())
92            .saturating_add(state.deleted.len())
93    }
94
95    /// Number of writes still using exact delta search.
96    #[must_use]
97    pub fn staged_len(&self) -> usize {
98        let state = self
99            .state
100            .read()
101            .unwrap_or_else(std::sync::PoisonError::into_inner);
102        state.pending.len()
103    }
104
105    /// Number of live vectors in the sealed HNSW delta.
106    #[must_use]
107    pub fn sealed_len(&self) -> usize {
108        let state = self
109            .state
110            .read()
111            .unwrap_or_else(std::sync::PoisonError::into_inner);
112        state.sealed.as_ref().map_or(0, |index| index.len())
113    }
114
115    #[must_use]
116    pub fn should_compact(&self, maximum_delta: usize) -> bool {
117        self.delta_len() >= maximum_delta
118    }
119
120    pub(crate) fn snapshot(&self) -> MutableSnapshot {
121        let state = self
122            .state
123            .read()
124            .unwrap_or_else(std::sync::PoisonError::into_inner);
125        MutableSnapshot {
126            config: self.config.clone(),
127            base: Arc::clone(&state.base),
128            sealed: state.sealed.clone(),
129            pending: state.pending.clone(),
130            deleted: state.deleted.clone(),
131            metadata: state.metadata.clone(),
132        }
133    }
134
135    pub(crate) fn from_snapshot(snapshot: MutableSnapshot) -> Self {
136        Self {
137            config: snapshot.config,
138            state: RwLock::new(MutableState {
139                base: snapshot.base,
140                sealed: snapshot.sealed,
141                pending: snapshot.pending,
142                deleted: snapshot.deleted,
143                metadata: snapshot.metadata,
144                generation: 0,
145            }),
146            distance_kernel: DistanceKernel::detect(),
147        }
148    }
149}