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};
16use std::sync::Arc;
17
18use datafusion_common::DataFusionError;
19use datafusion_common::arrow::array::{Array, AsArray, StringViewBuilder};
20use datafusion_common::arrow::datatypes::DataType;
21use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
22
23use crate::function::{Function, extract_args};
24
25/// Converts the `JSONB` into `String`. It's useful for displaying JSONB content.
26#[derive(Clone, Debug)]
27pub(crate) struct JsonToStringFunction {
28    signature: Signature,
29}
30
31impl Default for JsonToStringFunction {
32    fn default() -> Self {
33        Self {
34            // TODO(LFC): Use a more clear type here instead of "Binary" for Json input, once we have a "Json" type.
35            signature: Signature::uniform(
36                1,
37                vec![
38                    DataType::Binary,
39                    DataType::LargeBinary,
40                    DataType::BinaryView,
41                ],
42                Volatility::Immutable,
43            ),
44        }
45    }
46}
47
48const NAME: &str = "json_to_string";
49
50impl Function for JsonToStringFunction {
51    fn name(&self) -> &str {
52        NAME
53    }
54
55    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
56        Ok(DataType::Utf8View)
57    }
58
59    fn signature(&self) -> &Signature {
60        &self.signature
61    }
62
63    fn invoke_with_args(
64        &self,
65        args: ScalarFunctionArgs,
66    ) -> datafusion_common::Result<ColumnarValue> {
67        let [arg0] = extract_args(self.name(), &args)?;
68        let arg0 = arrow::compute::cast(&arg0, &DataType::BinaryView)?;
69        let jsons = arg0.as_binary_view();
70
71        let size = jsons.len();
72        let mut builder = StringViewBuilder::with_capacity(size);
73
74        for i in 0..size {
75            let json = jsons.is_valid(i).then(|| jsons.value(i));
76            let result = json
77                .map(|json| jsonb::from_slice(json).map(|x| x.to_string()))
78                .transpose()
79                .map_err(|e| DataFusionError::Execution(format!("invalid json binary: {e}")))?;
80
81            builder.append_option(result.as_deref());
82        }
83
84        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
85    }
86}
87
88impl Display for JsonToStringFunction {
89    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
90        write!(f, "JSON_TO_STRING")
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::sync::Arc;
97
98    use arrow_schema::Field;
99    use datafusion_common::arrow::array::BinaryArray;
100
101    use super::*;
102
103    #[test]
104    fn test_json_to_string_function() {
105        let json_to_string = JsonToStringFunction::default();
106
107        assert_eq!("json_to_string", json_to_string.name());
108        assert_eq!(
109            DataType::Utf8View,
110            json_to_string.return_type(&[DataType::Binary]).unwrap()
111        );
112
113        let json_strings = [
114            r#"{"a": {"b": 2}, "b": 2, "c": 3}"#,
115            r#"{"a": 4, "b": {"c": 6}, "c": 6}"#,
116            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
117        ];
118
119        let jsonbs = json_strings
120            .iter()
121            .map(|s| {
122                let value = jsonb::parse_value(s.as_bytes()).unwrap();
123                value.to_vec()
124            })
125            .collect::<Vec<_>>();
126
127        let args = ScalarFunctionArgs {
128            args: vec![ColumnarValue::Array(Arc::new(
129                BinaryArray::from_iter_values(jsonbs),
130            ))],
131            arg_fields: vec![],
132            number_rows: 3,
133            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
134            config_options: Arc::new(Default::default()),
135        };
136        let result = json_to_string
137            .invoke_with_args(args)
138            .and_then(|x| x.to_array(1))
139            .unwrap();
140        let vector = result.as_string_view();
141
142        assert_eq!(3, vector.len());
143        for (i, gt) in json_strings.iter().enumerate() {
144            let result = vector.value(i);
145            // remove whitespaces
146            assert_eq!(gt.replace(" ", ""), result);
147        }
148
149        let invalid_jsonb = vec![b"invalid json"];
150        let args = ScalarFunctionArgs {
151            args: vec![ColumnarValue::Array(Arc::new(
152                BinaryArray::from_iter_values(invalid_jsonb),
153            ))],
154            arg_fields: vec![],
155            number_rows: 1,
156            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
157            config_options: Arc::new(Default::default()),
158        };
159        let vector = json_to_string.invoke_with_args(args);
160        assert!(vector.is_err());
161    }
162}
common_function/scalars/json/json_to_string.rs

common_function/scalars/json/
json_to_string.rs
