common_function/scalars/json/
json_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::{self, Display};
16
17use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
18use datafusion_expr::{Signature, Volatility};
19use datatypes::arrow::datatypes::DataType;
20use datatypes::data_type::ConcreteDataType;
21use datatypes::prelude::VectorRef;
22use datatypes::scalars::ScalarVectorBuilder;
23use datatypes::vectors::{MutableVector, StringVectorBuilder};
24use snafu::ensure;
25
26use crate::function::{Function, FunctionContext};
27
28/// Converts the `JSONB` into `String`. It's useful for displaying JSONB content.
29#[derive(Clone, Debug, Default)]
30pub struct JsonToStringFunction;
31
32const NAME: &str = "json_to_string";
33
34impl Function for JsonToStringFunction {
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        // TODO(LFC): Use a more clear type here instead of "Binary" for Json input, once we have a "Json" type.
45        Signature::exact(vec![DataType::Binary], Volatility::Immutable)
46    }
47
48    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
49        ensure!(
50            columns.len() == 1,
51            InvalidFuncArgsSnafu {
52                err_msg: format!(
53                    "The length of the args is not correct, expect exactly one, have: {}",
54                    columns.len()
55                ),
56            }
57        );
58        let jsons = &columns[0];
59
60        let size = jsons.len();
61        let datatype = jsons.data_type();
62        let mut results = StringVectorBuilder::with_capacity(size);
63
64        match datatype {
65            // JSON data type uses binary vector
66            ConcreteDataType::Binary(_) => {
67                for i in 0..size {
68                    let json = jsons.get_ref(i);
69
70                    let json = json.as_binary();
71                    let result = match json {
72                        Ok(Some(json)) => match jsonb::from_slice(json) {
73                            Ok(json) => {
74                                let json = json.to_string();
75                                Some(json)
76                            }
77                            Err(_) => {
78                                return InvalidFuncArgsSnafu {
79                                    err_msg: format!("Illegal json binary: {:?}", json),
80                                }
81                                .fail();
82                            }
83                        },
84                        _ => None,
85                    };
86
87                    results.push(result.as_deref());
88                }
89            }
90            _ => {
91                return UnsupportedInputDataTypeSnafu {
92                    function: NAME,
93                    datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
94                }
95                .fail();
96            }
97        }
98
99        Ok(results.to_vector())
100    }
101}
102
103impl Display for JsonToStringFunction {
104    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
105        write!(f, "JSON_TO_STRING")
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use std::sync::Arc;
112
113    use datafusion_expr::TypeSignature;
114    use datatypes::scalars::ScalarVector;
115    use datatypes::vectors::BinaryVector;
116
117    use super::*;
118
119    #[test]
120    fn test_json_to_string_function() {
121        let json_to_string = JsonToStringFunction;
122
123        assert_eq!("json_to_string", json_to_string.name());
124        assert_eq!(
125            DataType::Utf8,
126            json_to_string.return_type(&[DataType::Binary]).unwrap()
127        );
128
129        assert!(matches!(json_to_string.signature(),
130                         Signature {
131                             type_signature: TypeSignature::Exact(valid_types),
132                             volatility: Volatility::Immutable
133                         } if  valid_types == vec![DataType::Binary]
134        ));
135
136        let json_strings = [
137            r#"{"a": {"b": 2}, "b": 2, "c": 3}"#,
138            r#"{"a": 4, "b": {"c": 6}, "c": 6}"#,
139            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
140        ];
141
142        let jsonbs = json_strings
143            .iter()
144            .map(|s| {
145                let value = jsonb::parse_value(s.as_bytes()).unwrap();
146                value.to_vec()
147            })
148            .collect::<Vec<_>>();
149
150        let json_vector = BinaryVector::from_vec(jsonbs);
151        let args: Vec<VectorRef> = vec![Arc::new(json_vector)];
152        let vector = json_to_string
153            .eval(&FunctionContext::default(), &args)
154            .unwrap();
155
156        assert_eq!(3, vector.len());
157        for (i, gt) in json_strings.iter().enumerate() {
158            let result = vector.get_ref(i);
159            let result = result.as_string().unwrap().unwrap();
160            // remove whitespaces
161            assert_eq!(gt.replace(" ", ""), result);
162        }
163
164        let invalid_jsonb = vec![b"invalid json"];
165        let invalid_json_vector = BinaryVector::from_vec(invalid_jsonb);
166        let args: Vec<VectorRef> = vec![Arc::new(invalid_json_vector)];
167        let vector = json_to_string.eval(&FunctionContext::default(), &args);
168        assert!(vector.is_err());
169    }
170}
common_function/scalars/json/json_to_string.rs

common_function/scalars/json/
json_to_string.rs
