tcvectordb_rust/
index.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use crate::enums::{FieldType, IndexType, MetricType};
4use crate::error::{VectorDBError, Result};
5
6pub type SparseVector = Vec<Vec<f64>>;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct HNSWParams {
10    #[serde(rename = "M")]
11    pub m: u32,
12    #[serde(rename = "efConstruction")]
13    pub ef_construction: u32,
14}
15
16impl HNSWParams {
17    pub fn new(m: u32, ef_construction: u32) -> Self {
18        Self { m, ef_construction }
19    }
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct IVFFLATParams {
24    pub nlist: u32,
25}
26
27impl IVFFLATParams {
28    pub fn new(nlist: u32) -> Self {
29        Self { nlist }
30    }
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct IVFPQParams {
35    #[serde(rename = "M")]
36    pub m: u32,
37    pub nlist: u32,
38}
39
40impl IVFPQParams {
41    pub fn new(nlist: u32, m: u32) -> Self {
42        Self { m, nlist }
43    }
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct IVFSQ8Params {
48    pub nlist: u32,
49}
50
51impl IVFSQ8Params {
52    pub fn new(nlist: u32) -> Self {
53        Self { nlist }
54    }
55}
56
57pub type IVFSQ4Params = IVFSQ8Params;
58pub type IVFSQ16Params = IVFSQ8Params;
59
60#[derive(Debug, Clone, Serialize, Deserialize)]
61#[serde(untagged)]
62pub enum IndexParams {
63    HNSW(HNSWParams),
64    IVFFLAT(IVFFLATParams),
65    IVFPQ(IVFPQParams),
66    IVFSQ8(IVFSQ8Params),
67    IVFSQ4(IVFSQ4Params),
68    IVFSQ16(IVFSQ16Params),
69}
70
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct IndexField {
73    #[serde(rename = "fieldName")]
74    pub name: String,
75    #[serde(rename = "fieldType")]
76    pub field_type: FieldType,
77    #[serde(rename = "indexType", skip_serializing_if = "Option::is_none")]
78    pub index_type: Option<IndexType>,
79}
80
81impl IndexField {
82    pub fn new(name: impl Into<String>, field_type: FieldType, index_type: Option<IndexType>) -> Self {
83        Self {
84            name: name.into(),
85            field_type,
86            index_type,
87        }
88    }
89
90    pub fn is_vector_field(&self) -> bool {
91        matches!(
92            self.field_type,
93            FieldType::Vector | FieldType::BinaryVector | FieldType::Float16Vector | FieldType::BFloat16Vector
94        )
95    }
96
97    pub fn is_sparse_vector_field(&self) -> bool {
98        matches!(self.field_type, FieldType::SparseVector)
99    }
100
101    pub fn is_primary_key(&self) -> bool {
102        matches!(self.index_type, Some(IndexType::PRIMARY_KEY))
103    }
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct VectorIndex {
108    #[serde(flatten)]
109    pub base: IndexField,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub dimension: Option<u32>,
112    #[serde(rename = "metricType", skip_serializing_if = "Option::is_none")]
113    pub metric_type: Option<MetricType>,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub params: Option<IndexParams>,
116    #[serde(rename = "indexedCount", skip_serializing_if = "Option::is_none")]
117    pub indexed_count: Option<u64>,
118}
119
120impl VectorIndex {
121    pub fn new(
122        name: impl Into<String>,
123        dimension: u32,
124        index_type: IndexType,
125        metric_type: MetricType,
126        params: Option<IndexParams>,
127    ) -> Self {
128        Self {
129            base: IndexField::new(name, FieldType::Vector, Some(index_type)),
130            dimension: Some(dimension),
131            metric_type: Some(metric_type),
132            params,
133            indexed_count: None,
134        }
135    }
136
137    pub fn with_field_type(
138        name: impl Into<String>,
139        dimension: u32,
140        index_type: IndexType,
141        metric_type: MetricType,
142        field_type: FieldType,
143        params: Option<IndexParams>,
144    ) -> Self {
145        Self {
146            base: IndexField::new(name, field_type, Some(index_type)),
147            dimension: Some(dimension),
148            metric_type: Some(metric_type),
149            params,
150            indexed_count: None,
151        }
152    }
153}
154
155#[derive(Debug, Clone, Serialize, Deserialize)]
156pub struct FilterIndex {
157    #[serde(flatten)]
158    pub base: IndexField,
159    #[serde(rename = "autoId", skip_serializing_if = "Option::is_none")]
160    pub auto_id: Option<String>,
161}
162
163impl FilterIndex {
164    pub fn new(name: impl Into<String>, field_type: FieldType, index_type: IndexType) -> Self {
165        Self {
166            base: IndexField::new(name, field_type, Some(index_type)),
167            auto_id: None,
168        }
169    }
170
171    pub fn with_auto_id(mut self, auto_id: impl Into<String>) -> Self {
172        self.auto_id = Some(auto_id.into());
173        self
174    }
175}
176
177#[derive(Debug, Clone, Serialize, Deserialize)]
178pub struct SparseIndex {
179    #[serde(flatten)]
180    pub base: IndexField,
181    #[serde(rename = "metricType")]
182    pub metric_type: MetricType,
183}
184
185impl SparseIndex {
186    pub fn new(
187        name: impl Into<String>,
188        index_type: IndexType,
189        metric_type: MetricType,
190    ) -> Self {
191        Self {
192            base: IndexField::new(name, FieldType::SparseVector, Some(index_type)),
193            metric_type,
194        }
195    }
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize)]
199#[serde(untagged)]
200pub enum IndexDefinition {
201    Vector(VectorIndex),
202    Filter(FilterIndex),
203    Sparse(SparseIndex),
204}
205
206impl IndexDefinition {
207    pub fn name(&self) -> &str {
208        match self {
209            IndexDefinition::Vector(idx) => &idx.base.name,
210            IndexDefinition::Filter(idx) => &idx.base.name,
211            IndexDefinition::Sparse(idx) => &idx.base.name,
212        }
213    }
214
215    pub fn is_primary_key(&self) -> bool {
216        match self {
217            IndexDefinition::Vector(idx) => idx.base.is_primary_key(),
218            IndexDefinition::Filter(idx) => idx.base.is_primary_key(),
219            IndexDefinition::Sparse(idx) => idx.base.is_primary_key(),
220        }
221    }
222}
223
224#[derive(Debug, Clone, Default)]
225pub struct Index {
226    indexes: HashMap<String, IndexDefinition>,
227    primary_field: Option<String>,
228}
229
230impl Index {
231    pub fn new() -> Self {
232        Self::default()
233    }
234
235    pub fn add_vector_index(&mut self, index: VectorIndex) -> Result<&mut Self> {
236        let name = index.base.name.clone();
237        if self.indexes.contains_key(&name) {
238            return Err(VectorDBError::server_error(
239                15000,
240                "fieldName must exist and be unique",
241            ));
242        }
243
244        if index.base.is_primary_key() {
245            self.primary_field = Some(name.clone());
246        }
247
248        self.indexes.insert(name, IndexDefinition::Vector(index));
249        Ok(self)
250    }
251
252    pub fn add_filter_index(&mut self, index: FilterIndex) -> Result<&mut Self> {
253        let name = index.base.name.clone();
254        if self.indexes.contains_key(&name) {
255            return Err(VectorDBError::server_error(
256                15000,
257                "fieldName must exist and be unique",
258            ));
259        }
260
261        if index.base.is_primary_key() {
262            self.primary_field = Some(name.clone());
263        }
264
265        self.indexes.insert(name, IndexDefinition::Filter(index));
266        Ok(self)
267    }
268
269    pub fn add_sparse_index(&mut self, index: SparseIndex) -> Result<&mut Self> {
270        let name = index.base.name.clone();
271        if self.indexes.contains_key(&name) {
272            return Err(VectorDBError::server_error(
273                15000,
274                "fieldName must exist and be unique",
275            ));
276        }
277
278        if index.base.is_primary_key() {
279            self.primary_field = Some(name.clone());
280        }
281
282        self.indexes.insert(name, IndexDefinition::Sparse(index));
283        Ok(self)
284    }
285
286    pub fn remove(&mut self, index_name: &str) -> &mut Self {
287        self.indexes.remove(index_name);
288        self
289    }
290
291    pub fn get(&self, name: &str) -> Option<&IndexDefinition> {
292        self.indexes.get(name)
293    }
294
295    pub fn list(&self) -> Vec<&IndexDefinition> {
296        self.indexes.values().collect()
297    }
298
299    pub fn primary_field(&self) -> Option<&str> {
300        self.primary_field.as_deref()
301    }
302
303    pub fn is_empty(&self) -> bool {
304        self.indexes.is_empty()
305    }
306}
307
308impl Serialize for Index {
309    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
310    where
311        S: serde::Serializer,
312    {
313        self.indexes.serialize(serializer)
314    }
315}
316
317impl<'de> Deserialize<'de> for Index {
318    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
319    where
320        D: serde::Deserializer<'de>,
321    {
322        let indexes: HashMap<String, IndexDefinition> = HashMap::deserialize(deserializer)?;
323        let primary_field = indexes
324            .values()
325            .find(|idx| idx.is_primary_key())
326            .map(|idx| idx.name().to_string());
327
328        Ok(Self {
329            indexes,
330            primary_field,
331        })
332    }
333}
tcvectordb_rust/index.rs

tcvectordb_rust/
index.rs
