common_function/scalars/vector/convert/
vector_to_string.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::fmt::Display;
16
17use common_query::error::{InvalidFuncArgsSnafu, Result};
18use datafusion::arrow::datatypes::DataType;
19use datafusion_expr::type_coercion::aggregates::BINARYS;
20use datafusion_expr::{Signature, Volatility};
21use datatypes::scalars::ScalarVectorBuilder;
22use datatypes::types::vector_type_value_to_string;
23use datatypes::value::Value;
24use datatypes::vectors::{MutableVector, StringVectorBuilder, VectorRef};
25use snafu::ensure;
26
27use crate::function::{Function, FunctionContext};
28
29const NAME: &str = "vec_to_string";
30
31#[derive(Debug, Clone, Default)]
32pub struct VectorToStringFunction;
33
34impl Function for VectorToStringFunction {
35    fn name(&self) -> &str {
36        NAME
37    }
38
39    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
40        Ok(DataType::Utf8)
41    }
42
43    fn signature(&self) -> Signature {
44        Signature::uniform(1, BINARYS.to_vec(), Volatility::Immutable)
45    }
46
47    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
48        ensure!(
49            columns.len() == 1,
50            InvalidFuncArgsSnafu {
51                err_msg: format!(
52                    "The length of the args is not correct, expect exactly one, have: {}",
53                    columns.len()
54                ),
55            }
56        );
57
58        let column = &columns[0];
59        let size = column.len();
60
61        let mut result = StringVectorBuilder::with_capacity(size);
62        for i in 0..size {
63            let value = column.get(i);
64            match value {
65                Value::Binary(bytes) => {
66                    let len = bytes.len();
67                    if len % std::mem::size_of::<f32>() != 0 {
68                        return InvalidFuncArgsSnafu {
69                            err_msg: format!("Invalid binary length of vector: {}", len),
70                        }
71                        .fail();
72                    }
73
74                    let dim = len / std::mem::size_of::<f32>();
75                    // Safety: `dim` is calculated from the length of `bytes` and is guaranteed to be valid
76                    let res = vector_type_value_to_string(&bytes, dim as _).unwrap();
77                    result.push(Some(&res));
78                }
79                Value::Null => {
80                    result.push_null();
81                }
82                _ => {
83                    return InvalidFuncArgsSnafu {
84                        err_msg: format!("Invalid value type: {:?}", value.data_type()),
85                    }
86                    .fail();
87                }
88            }
89        }
90
91        Ok(result.to_vector())
92    }
93}
94
95impl Display for VectorToStringFunction {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        write!(f, "{}", NAME.to_ascii_uppercase())
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use datatypes::value::Value;
104    use datatypes::vectors::BinaryVectorBuilder;
105
106    use super::*;
107
108    #[test]
109    fn test_vector_to_string() {
110        let func = VectorToStringFunction;
111
112        let mut builder = BinaryVectorBuilder::with_capacity(3);
113        builder.push(Some(
114            [1.0f32, 2.0, 3.0]
115                .iter()
116                .flat_map(|e| e.to_le_bytes())
117                .collect::<Vec<_>>()
118                .as_slice(),
119        ));
120        builder.push(Some(
121            [4.0f32, 5.0, 6.0]
122                .iter()
123                .flat_map(|e| e.to_le_bytes())
124                .collect::<Vec<_>>()
125                .as_slice(),
126        ));
127        builder.push_null();
128        let vector = builder.to_vector();
129
130        let result = func.eval(&FunctionContext::default(), &[vector]).unwrap();
131
132        assert_eq!(result.len(), 3);
133        assert_eq!(result.get(0), Value::String("[1,2,3]".to_string().into()));
134        assert_eq!(result.get(1), Value::String("[4,5,6]".to_string().into()));
135        assert_eq!(result.get(2), Value::Null);
136    }
137}