a3s_vec/
types.rs

1//! Stable scalar, vector, metric, and operation types.
2
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6/// Data type of a collection field.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
8#[repr(u32)]
9pub enum DataType {
10    Undefined = 0,
11    Binary = 1,
12    String = 2,
13    Bool = 3,
14    Int32 = 4,
15    Int64 = 5,
16    Uint32 = 6,
17    Uint64 = 7,
18    Float = 8,
19    Double = 9,
20    VectorBinary32 = 20,
21    VectorBinary64 = 21,
22    VectorFp16 = 22,
23    VectorFp32 = 23,
24    VectorFp64 = 24,
25    VectorInt4 = 25,
26    VectorInt8 = 26,
27    VectorInt16 = 27,
28    SparseVectorFp16 = 30,
29    SparseVectorFp32 = 31,
30    ArrayBinary = 40,
31    ArrayString = 41,
32    ArrayBool = 42,
33    ArrayInt32 = 43,
34    ArrayInt64 = 44,
35    ArrayUint32 = 45,
36    ArrayUint64 = 46,
37    ArrayFloat = 47,
38    ArrayDouble = 48,
39}
40
41impl From<u32> for DataType {
42    fn from(value: u32) -> Self {
43        match value {
44            1 => Self::Binary,
45            2 => Self::String,
46            3 => Self::Bool,
47            4 => Self::Int32,
48            5 => Self::Int64,
49            6 => Self::Uint32,
50            7 => Self::Uint64,
51            8 => Self::Float,
52            9 => Self::Double,
53            20 => Self::VectorBinary32,
54            21 => Self::VectorBinary64,
55            22 => Self::VectorFp16,
56            23 => Self::VectorFp32,
57            24 => Self::VectorFp64,
58            25 => Self::VectorInt4,
59            26 => Self::VectorInt8,
60            27 => Self::VectorInt16,
61            30 => Self::SparseVectorFp16,
62            31 => Self::SparseVectorFp32,
63            40 => Self::ArrayBinary,
64            41 => Self::ArrayString,
65            42 => Self::ArrayBool,
66            43 => Self::ArrayInt32,
67            44 => Self::ArrayInt64,
68            45 => Self::ArrayUint32,
69            46 => Self::ArrayUint64,
70            47 => Self::ArrayFloat,
71            48 => Self::ArrayDouble,
72            _ => Self::Undefined,
73        }
74    }
75}
76
77impl From<DataType> for u32 {
78    fn from(value: DataType) -> Self {
79        value as u32
80    }
81}
82
83impl DataType {
84    pub fn is_vector(self) -> bool {
85        matches!(
86            self,
87            Self::VectorBinary32
88                | Self::VectorBinary64
89                | Self::VectorFp16
90                | Self::VectorFp32
91                | Self::VectorFp64
92                | Self::VectorInt4
93                | Self::VectorInt8
94                | Self::VectorInt16
95                | Self::SparseVectorFp16
96                | Self::SparseVectorFp32
97        )
98    }
99
100    pub fn is_dense_vector(self) -> bool {
101        self.is_vector() && !self.is_sparse_vector()
102    }
103
104    pub fn is_sparse_vector(self) -> bool {
105        matches!(self, Self::SparseVectorFp16 | Self::SparseVectorFp32)
106    }
107
108    pub fn is_array(self) -> bool {
109        matches!(
110            self,
111            Self::ArrayBinary
112                | Self::ArrayString
113                | Self::ArrayBool
114                | Self::ArrayInt32
115                | Self::ArrayInt64
116                | Self::ArrayUint32
117                | Self::ArrayUint64
118                | Self::ArrayFloat
119                | Self::ArrayDouble
120        )
121    }
122
123    pub fn is_scalar(self) -> bool {
124        !self.is_vector() && !self.is_array() && self != Self::Undefined
125    }
126}
127
128impl fmt::Display for DataType {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        write!(f, "{self:?}")
131    }
132}
133
134/// Index implementation selected for a field.
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
136#[repr(u32)]
137pub enum IndexType {
138    Undefined = 0,
139    Hnsw = 1,
140    Ivf = 2,
141    Flat = 3,
142    Diskann = 5,
143    /// Vamana is the graph-construction name used by `DiskANN`.
144    Vamana = 6,
145    IvfRabitq = 7,
146    HnswRabitq = 8,
147    Invert = 10,
148    Fts = 11,
149}
150
151impl From<u32> for IndexType {
152    fn from(value: u32) -> Self {
153        match value {
154            1 => Self::Hnsw,
155            2 => Self::Ivf,
156            3 => Self::Flat,
157            5 => Self::Diskann,
158            6 => Self::Vamana,
159            7 => Self::IvfRabitq,
160            8 => Self::HnswRabitq,
161            10 => Self::Invert,
162            11 => Self::Fts,
163            _ => Self::Undefined,
164        }
165    }
166}
167
168impl From<IndexType> for u32 {
169    fn from(value: IndexType) -> Self {
170        value as u32
171    }
172}
173
174impl fmt::Display for IndexType {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        write!(f, "{self:?}")
177    }
178}
179
180/// Similarity metric.  Search scores are always ordered from high to low;
181/// L2 is represented as negative squared distance at the API boundary.
182#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
183#[repr(u32)]
184pub enum MetricType {
185    Undefined = 0,
186    L2 = 1,
187    Ip = 2,
188    Cosine = 3,
189    MipsL2 = 4,
190}
191
192impl From<u32> for MetricType {
193    fn from(value: u32) -> Self {
194        match value {
195            1 => Self::L2,
196            2 => Self::Ip,
197            3 => Self::Cosine,
198            4 => Self::MipsL2,
199            _ => Self::Undefined,
200        }
201    }
202}
203
204impl From<MetricType> for u32 {
205    fn from(value: MetricType) -> Self {
206        value as u32
207    }
208}
209
210impl MetricType {
211    pub fn higher_is_better(self) -> bool {
212        true
213    }
214}
215
216impl fmt::Display for MetricType {
217    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
218        write!(f, "{self:?}")
219    }
220}
221
222/// Scalar/vector quantization mode.
223#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
224#[repr(u32)]
225pub enum QuantizeType {
226    Undefined = 0,
227    Fp16 = 1,
228    Int8 = 2,
229    Int4 = 3,
230    Rabitq = 4,
231    Pq = 5,
232}
233
234impl From<u32> for QuantizeType {
235    fn from(value: u32) -> Self {
236        match value {
237            1 => Self::Fp16,
238            2 => Self::Int8,
239            3 => Self::Int4,
240            4 => Self::Rabitq,
241            5 => Self::Pq,
242            _ => Self::Undefined,
243        }
244    }
245}
246
247impl From<QuantizeType> for u32 {
248    fn from(value: QuantizeType) -> Self {
249        value as u32
250    }
251}
252
253impl fmt::Display for QuantizeType {
254    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
255        write!(f, "{self:?}")
256    }
257}
258
259/// DML operation kind used by write-ahead records and telemetry.
260#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
261pub enum DocOperator {
262    Insert,
263    Update,
264    Upsert,
265    Delete,
266}
267
268impl fmt::Display for DocOperator {
269    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
270        write!(f, "{self:?}")
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::{DataType, DocOperator, IndexType, MetricType, QuantizeType};
277
278    #[test]
279    fn data_type_round_trips_and_classifies_every_variant() {
280        let variants = [
281            DataType::Undefined,
282            DataType::Binary,
283            DataType::String,
284            DataType::Bool,
285            DataType::Int32,
286            DataType::Int64,
287            DataType::Uint32,
288            DataType::Uint64,
289            DataType::Float,
290            DataType::Double,
291            DataType::VectorBinary32,
292            DataType::VectorBinary64,
293            DataType::VectorFp16,
294            DataType::VectorFp32,
295            DataType::VectorFp64,
296            DataType::VectorInt4,
297            DataType::VectorInt8,
298            DataType::VectorInt16,
299            DataType::SparseVectorFp16,
300            DataType::SparseVectorFp32,
301            DataType::ArrayBinary,
302            DataType::ArrayString,
303            DataType::ArrayBool,
304            DataType::ArrayInt32,
305            DataType::ArrayInt64,
306            DataType::ArrayUint32,
307            DataType::ArrayUint64,
308            DataType::ArrayFloat,
309            DataType::ArrayDouble,
310        ];
311        for variant in variants {
312            let encoded = u32::from(variant);
313            assert_eq!(DataType::from(encoded), variant);
314            assert!(!format!("{variant}").is_empty());
315            assert_eq!(variant.is_vector(), (20..40).contains(&encoded));
316            assert_eq!(
317                variant.is_sparse_vector(),
318                matches!(
319                    variant,
320                    DataType::SparseVectorFp16 | DataType::SparseVectorFp32
321                )
322            );
323            assert_eq!(
324                variant.is_dense_vector(),
325                variant.is_vector() && !variant.is_sparse_vector()
326            );
327            assert_eq!(variant.is_array(), encoded >= 40);
328            assert_eq!(
329                variant.is_scalar(),
330                !variant.is_vector() && !variant.is_array() && variant != DataType::Undefined
331            );
332        }
333        assert_eq!(DataType::from(999), DataType::Undefined);
334    }
335
336    #[test]
337    fn index_metric_quantize_and_operator_vocabularies_round_trip() {
338        for variant in [
339            IndexType::Undefined,
340            IndexType::Hnsw,
341            IndexType::Ivf,
342            IndexType::Flat,
343            IndexType::Diskann,
344            IndexType::Vamana,
345            IndexType::IvfRabitq,
346            IndexType::HnswRabitq,
347            IndexType::Invert,
348            IndexType::Fts,
349        ] {
350            assert_eq!(IndexType::from(u32::from(variant)), variant);
351            assert!(!format!("{variant}").is_empty());
352        }
353        assert_eq!(IndexType::from(999), IndexType::Undefined);
354
355        for variant in [
356            MetricType::Undefined,
357            MetricType::L2,
358            MetricType::Ip,
359            MetricType::Cosine,
360            MetricType::MipsL2,
361        ] {
362            assert_eq!(MetricType::from(u32::from(variant)), variant);
363            assert!(variant.higher_is_better());
364            assert!(!format!("{variant}").is_empty());
365        }
366        assert_eq!(MetricType::from(999), MetricType::Undefined);
367
368        for variant in [
369            QuantizeType::Undefined,
370            QuantizeType::Fp16,
371            QuantizeType::Int8,
372            QuantizeType::Int4,
373            QuantizeType::Rabitq,
374            QuantizeType::Pq,
375        ] {
376            assert_eq!(QuantizeType::from(u32::from(variant)), variant);
377            assert!(!format!("{variant}").is_empty());
378        }
379        assert_eq!(QuantizeType::from(999), QuantizeType::Undefined);
380
381        for variant in [
382            DocOperator::Insert,
383            DocOperator::Update,
384            DocOperator::Upsert,
385            DocOperator::Delete,
386        ] {
387            assert!(!format!("{variant}").is_empty());
388        }
389    }
390}
a3s_vec/types.rs

a3s_vec/
types.rs
