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

common_function/scalars/json/
json_to_string.rs
