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

common_function/scalars/json/
json_to_string.rs
