common_function/scalars/vector/
impl_conv.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 std::borrow::Cow;
16
17use common_query::error::{InvalidFuncArgsSnafu, Result};
18use datafusion_common::ScalarValue;
19
20/// Convert a string or binary literal to a vector literal.
21pub fn as_veclit(arg: &ScalarValue) -> Result<Option<Cow<'_, [f32]>>> {
22    match arg {
23        ScalarValue::Binary(b) | ScalarValue::BinaryView(b) => {
24            b.as_ref().map(|x| binlit_as_veclit(x)).transpose()
25        }
26        ScalarValue::Utf8(s) | ScalarValue::Utf8View(s) => s
27            .as_ref()
28            .map(|x| parse_veclit_from_strlit(x).map(Cow::Owned))
29            .transpose(),
30        _ => InvalidFuncArgsSnafu {
31            err_msg: format!("Unsupported data type: {:?}", arg.data_type()),
32        }
33        .fail(),
34    }
35}
36
37/// Convert a u8 slice to a vector literal.
38pub fn binlit_as_veclit(bytes: &[u8]) -> Result<Cow<'_, [f32]>> {
39    if !bytes.len().is_multiple_of(size_of::<f32>()) {
40        return InvalidFuncArgsSnafu {
41            err_msg: format!("Invalid binary length of vector: {}", bytes.len()),
42        }
43        .fail();
44    }
45
46    if cfg!(target_endian = "little") {
47        Ok(unsafe {
48            let vec = std::slice::from_raw_parts(
49                bytes.as_ptr() as *const f32,
50                bytes.len() / std::mem::size_of::<f32>(),
51            );
52            Cow::Borrowed(vec)
53        })
54    } else {
55        let v = bytes
56            .as_chunks::<{ std::mem::size_of::<f32>() }>()
57            .0
58            .iter()
59            .map(|chunk| f32::from_le_bytes(*chunk))
60            .collect::<Vec<f32>>();
61        Ok(Cow::Owned(v))
62    }
63}
64
65/// Parse a string literal to a vector literal.
66/// Valid inputs are strings like "[1.0, 2.0, 3.0]".
67pub fn parse_veclit_from_strlit(s: &str) -> Result<Vec<f32>> {
68    let trimmed = s.trim();
69    if !trimmed.starts_with('[') || !trimmed.ends_with(']') {
70        return InvalidFuncArgsSnafu {
71            err_msg: format!(
72                "Failed to parse {s} to Vector value: not properly enclosed in brackets"
73            ),
74        }
75        .fail();
76    }
77    let content = trimmed[1..trimmed.len() - 1].trim();
78    if content.is_empty() {
79        return Ok(Vec::new());
80    }
81
82    content
83        .split(',')
84        .map(|s| s.trim().parse::<f32>())
85        .collect::<std::result::Result<_, _>>()
86        .map_err(|e| {
87            InvalidFuncArgsSnafu {
88                err_msg: format!("Failed to parse {s} to Vector value: {e}"),
89            }
90            .build()
91        })
92}
93
94/// Convert a vector literal to a binary literal.
95pub fn veclit_to_binlit(vec: &[f32]) -> Vec<u8> {
96    if cfg!(target_endian = "little") {
97        unsafe {
98            std::slice::from_raw_parts(vec.as_ptr() as *const u8, std::mem::size_of_val(vec))
99                .to_vec()
100        }
101    } else {
102        let mut bytes = Vec::with_capacity(std::mem::size_of_val(vec));
103        for e in vec {
104            bytes.extend_from_slice(&e.to_le_bytes());
105        }
106        bytes
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn test_parse_veclit_from_strlit() {
116        let result = parse_veclit_from_strlit("[1.0, 2.0, 3.0]").unwrap();
117        assert_eq!(result, vec![1.0, 2.0, 3.0]);
118
119        let result = parse_veclit_from_strlit("[]").unwrap();
120        assert_eq!(result, Vec::<f32>::new());
121
122        let result = parse_veclit_from_strlit("[1.0, a, 3.0]");
123        assert!(result.is_err());
124    }
125
126    #[test]
127    fn test_binlit_as_veclit() {
128        let vec = &[1.0, 2.0, 3.0];
129        let bytes = veclit_to_binlit(vec);
130        let result = binlit_as_veclit(&bytes).unwrap();
131        assert_eq!(result.as_ref(), vec);
132
133        let invalid_bytes = [0, 0, 128];
134        let result = binlit_as_veclit(&invalid_bytes);
135        assert!(result.is_err());
136    }
137}
common_function/scalars/vector/impl_conv.rs

common_function/scalars/vector/
impl_conv.rs
