datatypes/types/
vector_type.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use arrow::datatypes::DataType as ArrowDataType;
16use common_base::bytes::Bytes;
17use serde::{Deserialize, Serialize};
18
19use crate::data_type::DataType;
20use crate::error::{InvalidVectorSnafu, Result};
21use crate::scalars::ScalarVectorBuilder;
22use crate::type_id::LogicalTypeId;
23use crate::value::Value;
24use crate::vectors::{BinaryVectorBuilder, MutableVector};
25
26/// `VectorType` is a data type for vector data with a fixed dimension.
27/// The type of items in the vector is float32.
28/// It is stored as binary data that contains the concatenated float32 values.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
30pub struct VectorType {
31    pub dim: u32,
32}
33
34impl VectorType {
35    pub fn new(dim: u32) -> Self {
36        Self { dim }
37    }
38}
39
40impl DataType for VectorType {
41    fn name(&self) -> String {
42        format!("Vector({})", self.dim)
43    }
44
45    fn logical_type_id(&self) -> LogicalTypeId {
46        LogicalTypeId::Vector
47    }
48
49    fn default_value(&self) -> Value {
50        Bytes::default().into()
51    }
52
53    fn as_arrow_type(&self) -> ArrowDataType {
54        ArrowDataType::Binary
55    }
56
57    fn create_mutable_vector(&self, capacity: usize) -> Box<dyn MutableVector> {
58        Box::new(BinaryVectorBuilder::with_capacity(capacity))
59    }
60
61    fn try_cast(&self, from: Value) -> Option<Value> {
62        match from {
63            Value::Binary(v) => Some(Value::Binary(v)),
64            _ => None,
65        }
66    }
67}
68
69/// Converts a vector type value to string
70/// for example: [1.0, 2.0, 3.0] -> "[1.0,2.0,3.0]"
71pub fn vector_type_value_to_string(val: &[u8], dim: u32) -> Result<String> {
72    let expected_len = dim as usize * std::mem::size_of::<f32>();
73    if val.len() != expected_len {
74        return InvalidVectorSnafu {
75            msg: format!(
76                "Failed to convert Vector value to string: wrong byte size, expected {}, got {}",
77                expected_len,
78                val.len()
79            ),
80        }
81        .fail();
82    }
83
84    if dim == 0 {
85        return Ok("[]".to_string());
86    }
87
88    let elements = val
89        .as_chunks::<{ std::mem::size_of::<f32>() }>()
90        .0
91        .iter()
92        .map(|e| f32::from_le_bytes(*e));
93
94    let mut s = String::from("[");
95    for (i, e) in elements.enumerate() {
96        if i > 0 {
97            s.push(',');
98        }
99        s.push_str(&e.to_string());
100    }
101    s.push(']');
102    Ok(s)
103}
104
105/// Parses a string to a vector type value
106/// Valid input format: "[1.0,2.0,3.0]", "[1.0, 2.0, 3.0]"
107pub fn parse_string_to_vector_type_value(s: &str, dim: Option<u32>) -> Result<Vec<u8>> {
108    // Trim the brackets
109    let trimmed = s.trim();
110    if !trimmed.starts_with('[') || !trimmed.ends_with(']') {
111        return InvalidVectorSnafu {
112            msg: format!("Failed to parse {s} to Vector value: not properly enclosed in brackets"),
113        }
114        .fail();
115    }
116    // Remove the brackets
117    let content = trimmed[1..trimmed.len() - 1].trim();
118
119    if content.is_empty() {
120        if dim.is_some_and(|d| d != 0) {
121            return InvalidVectorSnafu {
122                msg: format!("Failed to parse {s} to Vector value: wrong dimension"),
123            }
124            .fail();
125        }
126        return Ok(vec![]);
127    }
128
129    let elements = content
130        .split(',')
131        .map(|e| {
132            e.trim().parse::<f32>().map_err(|_| {
133                InvalidVectorSnafu {
134                    msg: format!(
135                        "Failed to parse {s} to Vector value: elements are not all float32"
136                    ),
137                }
138                .build()
139            })
140        })
141        .collect::<Result<Vec<f32>>>()?;
142
143    // Check dimension
144    if dim.is_some_and(|d| d as usize != elements.len()) {
145        return InvalidVectorSnafu {
146            msg: format!("Failed to parse {s} to Vector value: wrong dimension"),
147        }
148        .fail();
149    }
150
151    // Convert Vec<f32> to Vec<u8>
152    let bytes = if cfg!(target_endian = "little") {
153        unsafe {
154            std::slice::from_raw_parts(
155                elements.as_ptr() as *const u8,
156                elements.len() * std::mem::size_of::<f32>(),
157            )
158            .to_vec()
159        }
160    } else {
161        elements
162            .iter()
163            .flat_map(|e| e.to_le_bytes())
164            .collect::<Vec<u8>>()
165    };
166
167    Ok(bytes)
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn test_conversion_between_string_and_vector_type_value() {
176        let dim = 3;
177
178        let cases = [
179            ("[1.0,2.0,3]", "[1,2,3]"),
180            ("[0.0 , 0.0 , 0.0]", "[0,0,0]"),
181            (
182                "[3.4028235e38, -3.4028235e38, 1.1754944e-38]",
183                "[340282350000000000000000000000000000000,-340282350000000000000000000000000000000,0.000000000000000000000000000000000000011754944]",
184            ),
185        ];
186
187        for (s, expected) in cases.iter() {
188            let val = parse_string_to_vector_type_value(s, Some(dim)).unwrap();
189            let s = vector_type_value_to_string(&val, dim).unwrap();
190            assert_eq!(s, *expected);
191        }
192
193        let dim = 0;
194        let cases = [("[]", "[]"), ("[ ]", "[]"), ("[  ]", "[]")];
195        for (s, expected) in cases.iter() {
196            let val = parse_string_to_vector_type_value(s, Some(dim)).unwrap();
197            let s = vector_type_value_to_string(&val, dim).unwrap();
198            assert_eq!(s, *expected);
199        }
200    }
201
202    #[test]
203    fn test_vector_type_value_to_string_wrong_byte_size() {
204        let dim = 3;
205        let val = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
206        let res = vector_type_value_to_string(&val, dim);
207        assert!(res.is_err());
208
209        let dim = 0;
210        let val = vec![1];
211        let res = vector_type_value_to_string(&val, dim);
212        assert!(res.is_err());
213    }
214
215    #[test]
216    fn test_parse_string_to_vector_type_value_not_properly_enclosed_in_brackets() {
217        let dim = 3;
218        let s = "1.0,2.0,3.0";
219        let res = parse_string_to_vector_type_value(s, Some(dim));
220        assert!(res.is_err());
221
222        let s = "[1.0,2.0,3.0";
223        let res = parse_string_to_vector_type_value(s, Some(dim));
224        assert!(res.is_err());
225
226        let s = "1.0,2.0,3.0]";
227        let res = parse_string_to_vector_type_value(s, Some(dim));
228        assert!(res.is_err());
229    }
230
231    #[test]
232    fn test_parse_string_to_vector_type_value_wrong_dimension() {
233        let dim = 3;
234        let s = "[1.0,2.0]";
235        let res = parse_string_to_vector_type_value(s, Some(dim));
236        assert!(res.is_err());
237    }
238
239    #[test]
240    fn test_parse_string_to_vector_type_value_elements_are_not_all_float32() {
241        let dim = 3;
242        let s = "[1.0,2.0,ah]";
243        let res = parse_string_to_vector_type_value(s, Some(dim));
244        assert!(res.is_err());
245    }
246}
datatypes/types/vector_type.rs

datatypes/types/
vector_type.rs
