milvus/
value.rs

1use std::borrow::Cow;
2
3use crate::{
4    proto::schema::{
5        field_data::Field, scalar_field::Data as ScalarData, vector_field::Data as VectorData,
6        DataType,
7    },
8};
9
10pub enum Value<'a> {
11    None,
12    Bool(bool),
13    Int8(i8),
14    Int16(i16),
15    Int32(i32),
16    Long(i64),
17    Float(f32),
18    Double(f64),
19    FloatArray(Cow<'a, [f32]>),
20    Binary(Cow<'a, [u8]>),
21    String(Cow<'a, str>),
22}
23
24macro_rules! impl_from_for_field_data_column {
25    ( $($t: ty, $o: ident ),+ ) => {$(
26        impl From<$t> for Value<'static> {
27            fn from(v: $t) -> Self {
28                Self::$o(v as _)
29            }
30        }
31    )*};
32}
33
34impl_from_for_field_data_column! {
35    bool, Bool,
36    i8,  Int8,
37    i16, Int16,
38    i32, Int32,
39    i64, Long,
40    f32, Float,
41    f64, Double
42}
43
44impl Value<'_> {
45    pub fn data_type(&self) -> DataType {
46        match self {
47            Value::None => DataType::None,
48            Value::Bool(_) => DataType::Bool,
49            Value::Int8(_) => DataType::Int8,
50            Value::Int16(_) => DataType::Int16,
51            Value::Int32(_) => DataType::Int32,
52            Value::Long(_) => DataType::Int64,
53            Value::Float(_) => DataType::Float,
54            Value::Double(_) => DataType::Double,
55            Value::String(_) => DataType::String,
56            Value::FloatArray(_) => DataType::FloatVector,
57            Value::Binary(_) => DataType::BinaryVector,
58        }
59    }
60}
61
62impl From<String> for Value<'static> {
63    fn from(v: String) -> Self {
64        Self::String(Cow::Owned(v))
65    }
66}
67
68impl<'a> From<&'a str> for Value<'a> {
69    fn from(v: &'a str) -> Self {
70        Self::String(Cow::Borrowed(v))
71    }
72}
73
74impl<'a> From<&'a [u8]> for Value<'a> {
75    fn from(v: &'a [u8]) -> Self {
76        Self::Binary(Cow::Borrowed(v))
77    }
78}
79
80impl From<Vec<u8>> for Value<'static> {
81    fn from(v: Vec<u8>) -> Self {
82        Self::Binary(Cow::Owned(v))
83    }
84}
85
86impl<'a> From<&'a [f32]> for Value<'a> {
87    fn from(v: &'a [f32]) -> Self {
88        Self::FloatArray(Cow::Borrowed(v))
89    }
90}
91
92impl From<Vec<f32>> for Value<'static> {
93    fn from(v: Vec<f32>) -> Self {
94        Self::FloatArray(Cow::Owned(v))
95    }
96}
97
98#[derive(Debug, Clone)]
99pub enum ValueVec {
100    None,
101    Bool(Vec<bool>),
102    Int(Vec<i32>),
103    Long(Vec<i64>),
104    Float(Vec<f32>),
105    Double(Vec<f64>),
106    Binary(Vec<u8>),
107    String(Vec<String>),
108}
109
110macro_rules! impl_from_for_value_vec {
111    ( $($t: ty, $o: ident ),+ ) => {$(
112        impl From<$t> for ValueVec {
113            fn from(v: $t) -> Self {
114                Self::$o(v)
115            }
116        }
117    )*};
118}
119
120impl_from_for_value_vec! {
121    Vec<bool>, Bool,
122    Vec<i32>, Int,
123    Vec<i64>, Long,
124    Vec<String>, String,
125    Vec<u8>, Binary,
126    Vec<f32>, Float,
127    Vec<f64>, Double
128}
129
130macro_rules! impl_try_from_for_value_vec {
131    ( $($o: ident, $t: ty ),+ ) => {$(
132        impl TryFrom<ValueVec> for $t {
133            type Error = crate::error::Error;
134            fn try_from(value: ValueVec) -> Result<Self, Self::Error> {
135                match value {
136                    ValueVec::$o(v) => Ok(v),
137                    _ => Err(crate::error::Error::Conversion),
138                }
139            }
140        }
141    )*};
142}
143
144impl_try_from_for_value_vec! {
145    Bool, Vec<bool>,
146    Int, Vec<i32>,
147    Long, Vec<i64>,
148    String, Vec<String>,
149    Binary, Vec<u8>,
150    Float, Vec<f32>,
151    Double, Vec<f64>
152}
153
154impl From<Vec<i8>> for ValueVec {
155    fn from(v: Vec<i8>) -> Self {
156        Self::Int(v.into_iter().map(Into::into).collect())
157    }
158}
159
160impl From<Vec<i16>> for ValueVec {
161    fn from(v: Vec<i16>) -> Self {
162        Self::Int(v.into_iter().map(Into::into).collect())
163    }
164}
165
166impl ValueVec {
167    pub fn new(dtype: DataType) -> Self {
168        match dtype {
169            DataType::None => Self::None,
170            DataType::Bool => Self::Bool(Vec::new()),
171            DataType::Int8 => Self::Int(Vec::new()),
172            DataType::Int16 => Self::Int(Vec::new()),
173            DataType::Int32 => Self::Int(Vec::new()),
174            DataType::Int64 => Self::Long(Vec::new()),
175            DataType::Float => Self::Float(Vec::new()),
176            DataType::Double => Self::Double(Vec::new()),
177            DataType::String => Self::String(Vec::new()),
178            DataType::VarChar => Self::String(Vec::new()),
179            DataType::BinaryVector => Self::Binary(Vec::new()),
180            DataType::FloatVector => Self::Float(Vec::new()),
181        }
182    }
183
184    pub fn check_dtype(&self, dtype: DataType) -> bool {
185        match (self, dtype) {
186            (ValueVec::Binary(..), DataType::BinaryVector)
187            | (ValueVec::Float(..), DataType::FloatVector)
188            | (ValueVec::Float(..), DataType::Float)
189            | (ValueVec::Int(..), DataType::Int8)
190            | (ValueVec::Int(..), DataType::Int16)
191            | (ValueVec::Int(..), DataType::Int32)
192            | (ValueVec::Long(..), DataType::Int64)
193            | (ValueVec::Bool(..), DataType::Bool)
194            | (ValueVec::String(..), DataType::String)
195            | (ValueVec::String(..), DataType::VarChar)
196            | (ValueVec::None, _)
197            | (ValueVec::Double(..), DataType::Double) => true,
198            _ => false,
199        }
200    }
201
202    #[inline]
203    pub fn is_empty(&self) -> bool {
204        self.len() == 0
205    }
206
207    pub fn len(&self) -> usize {
208        match self {
209            ValueVec::None => 0,
210            ValueVec::Bool(v) => v.len(),
211            ValueVec::Int(v) => v.len(),
212            ValueVec::Long(v) => v.len(),
213            ValueVec::Float(v) => v.len(),
214            ValueVec::Double(v) => v.len(),
215            ValueVec::Binary(v) => v.len(),
216            ValueVec::String(v) => v.len(),
217        }
218    }
219
220    pub fn clear(&mut self) {
221        match self {
222            ValueVec::None => (),
223            ValueVec::Bool(v) => v.clear(),
224            ValueVec::Int(v) => v.clear(),
225            ValueVec::Long(v) => v.clear(),
226            ValueVec::Float(v) => v.clear(),
227            ValueVec::Double(v) => v.clear(),
228            ValueVec::Binary(v) => v.clear(),
229            ValueVec::String(v) => v.clear(),
230        }
231    }
232}
233
234impl From<Field> for ValueVec {
235    fn from(f: Field) -> Self {
236        match f {
237            Field::Scalars(s) => match s.data {
238                Some(x) => match x {
239                    ScalarData::BoolData(v) => Self::Bool(v.data),
240                    ScalarData::IntData(v) => Self::Int(v.data),
241                    ScalarData::LongData(v) => Self::Long(v.data),
242                    ScalarData::FloatData(v) => Self::Float(v.data),
243                    ScalarData::DoubleData(v) => Self::Double(v.data),
244                    ScalarData::StringData(v) => Self::String(v.data),
245                    ScalarData::BytesData(_) => unimplemented!(), // Self::Bytes(v.data),
246                },
247                None => Self::None,
248            },
249
250            Field::Vectors(arr) => match arr.data {
251                Some(x) => match x {
252                    VectorData::FloatVector(v) => Self::Float(v.data),
253                    VectorData::BinaryVector(v) => Self::Binary(v),
254                },
255                None => Self::None,
256            },
257        }
258    }
259}
260
261macro_rules! impl_try_from_for_value_column {
262    ( $($o: ident,$t: ty ),+ ) => {$(
263        impl TryFrom<Value<'_>> for $t {
264            type Error = crate::error::Error;
265            fn try_from(value: Value<'_>) -> Result<Self, Self::Error> {
266                match value {
267                    Value::$o(v) => Ok(v),
268                    _ => Err(crate::error::Error::Conversion),
269                }
270            }
271        }
272    )*};
273}
274
275impl_try_from_for_value_column! {
276    Bool,bool,
277    Int8,i8,
278    Int16,i16,
279    Int32,i32,
280    Long,i64,
281    Float,f32,
282    Double,f64
283}
284
285#[cfg(test)]
286mod test {
287    use crate::{
288        error::Error,
289        value::{Value, ValueVec},
290    };
291
292    #[test]
293    fn test_try_from_for_value_column() {
294        // bool
295        let b = Value::Bool(false);
296        let b: Result<bool, Error> = b.try_into();
297        assert!(b.is_ok());
298        assert!(!b.unwrap());
299        //i8
300        let int8 = Value::Int8(12);
301        let r: Result<i8, Error> = int8.try_into();
302        assert!(r.is_ok());
303        assert_eq!(12, r.unwrap());
304        //i16
305        let int16 = Value::Int16(1225);
306        let r: Result<i16, Error> = int16.try_into();
307        assert!(r.is_ok());
308        assert_eq!(1225, r.unwrap());
309        //i32
310        let int32 = Value::Int32(37360798);
311        let r: Result<i32, Error> = int32.try_into();
312        assert!(r.is_ok());
313        assert_eq!(37360798, r.unwrap());
314        //i64
315        let long = Value::Long(123);
316        let r: Result<i64, Error> = long.try_into();
317        assert!(r.is_ok());
318        assert_eq!(123, r.unwrap());
319
320        let float = Value::Float(22104f32);
321        let r: Result<f32, Error> = float.try_into();
322        assert!(r.is_ok());
323        assert_eq!(22104f32, r.unwrap());
324
325        let double = Value::Double(22104f64);
326        let r: Result<f64, Error> = double.try_into();
327        assert!(r.is_ok());
328        assert_eq!(22104f64, r.unwrap());
329    }
330
331    #[test]
332    fn test_try_from_for_value_vec() {
333        let b = ValueVec::Bool(vec![false, false]);
334        let b: Result<Vec<bool>, Error> = b.try_into();
335        assert!(b.is_ok());
336        assert_eq!(vec![false, false], b.unwrap());
337
338        let b = ValueVec::Int(vec![1, 2]);
339        let b: Result<Vec<i32>, Error> = b.try_into();
340        assert!(b.is_ok());
341        assert_eq!(vec![1, 2], b.unwrap());
342
343        let v: Vec<i64> = vec![4095291003, 2581116377, 3892395808];
344        let b = ValueVec::Long(v.clone());
345        let b: Result<Vec<i64>, Error> = b.try_into();
346        assert!(b.is_ok());
347        assert_eq!(v, b.unwrap());
348
349        let v: Vec<f32> = vec![11., 7., 754., 68., 34.];
350        let b = ValueVec::Float(v.clone());
351        let b: Result<Vec<f32>, Error> = b.try_into();
352        assert!(b.is_ok());
353        assert_eq!(v, b.unwrap());
354
355        let v: Vec<f64> = vec![28., 9., 92., 6099786761., 64.];
356        let b = ValueVec::Double(v.clone());
357        let b_err: Result<Vec<f32>, Error> = b.clone().try_into();
358        assert!(b_err.is_err());
359        let b: Result<Vec<f64>, Error> = b.try_into();
360        assert_eq!(v, b.unwrap());
361
362        let v = vec![28, 5, 70, 62, 57, 103, 68];
363        let b = ValueVec::Binary(v.clone());
364        let b: Result<Vec<u8>, Error> = b.try_into();
365        assert!(b.is_ok());
366        assert_eq!(v, b.unwrap());
367
368        let v: Vec<String> = vec!["Janoato", "Samoa", "opde@tuwuv.yt"]
369            .into_iter()
370            .map(|a| a.into())
371            .collect();
372        let b = ValueVec::String(v.clone());
373        let b: Result<Vec<String>, Error> = b.try_into();
374        assert!(b.is_ok());
375        assert_eq!(v, b.unwrap());
376    }
377}