common_function/scalars/json/
json_is.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::arrow::array::{Array, AsArray, BooleanBuilder};
19use datafusion_common::arrow::compute;
20use datafusion_common::arrow::datatypes::DataType;
21use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
22
23use crate::function::{Function, extract_args};
24
25/// Checks if the input is a JSON object of the given type.
26macro_rules! json_is {
27    ($name:ident, $json_type:ident, $doc:expr) => {
28        paste::paste! {
29            #[derive(Clone, Debug)]
30            pub(crate) struct $name {
31                signature: Signature,
32            }
33
34            impl Default for $name {
35                fn default() -> Self {
36                    Self {
37                        // TODO(LFC): Use a more clear type here instead of "Binary" for Json input, once we have a "Json" type.
38                        signature: Signature::uniform(
39                            1,
40                            vec![DataType::Binary, DataType::BinaryView],
41                            Volatility::Immutable,
42                        ),
43                    }
44                }
45            }
46
47            impl Function for $name {
48                fn name(&self) -> &str {
49                    stringify!([<$name:snake>])
50                }
51
52                fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
53                    Ok(DataType::Boolean)
54                }
55
56                fn signature(&self) -> &Signature {
57                    &self.signature
58                }
59
60                fn invoke_with_args(
61                    &self,
62                    args: ScalarFunctionArgs,
63                ) -> datafusion_common::Result<ColumnarValue> {
64                    let [arg0] = extract_args(self.name(), &args)?;
65
66                    let arg0 = compute::cast(&arg0, &DataType::BinaryView)?;
67                    let jsons = arg0.as_binary_view();
68                    let size = jsons.len();
69                    let mut builder = BooleanBuilder::with_capacity(size);
70
71                    for i in 0..size {
72                        let json = jsons.is_valid(i).then(|| jsons.value(i));
73                        let result = match json {
74                            Some(json) => {
75                                Some(jsonb::[<is_ $json_type>](json))
76                            }
77                            _ => None,
78                        };
79                        builder.append_option(result);
80                    }
81
82                    Ok(ColumnarValue::Array(Arc::new(builder.finish())))
83                }
84            }
85
86            impl Display for $name {
87                fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
88                    write!(f, "{}", stringify!([<$name:snake>]).to_ascii_uppercase())
89                }
90            }
91        }
92    };
93}
94
95json_is!(JsonIsNull, null, "Checks if the input JSONB is null");
96json_is!(
97    JsonIsBool,
98    boolean,
99    "Checks if the input JSONB is a boolean type JSON value"
100);
101json_is!(
102    JsonIsInt,
103    i64,
104    "Checks if the input JSONB is a integer type JSON value"
105);
106json_is!(
107    JsonIsFloat,
108    number,
109    "Checks if the input JSONB is a JSON float"
110);
111json_is!(
112    JsonIsString,
113    string,
114    "Checks if the input JSONB is a JSON string"
115);
116json_is!(
117    JsonIsArray,
118    array,
119    "Checks if the input JSONB is a JSON array"
120);
121json_is!(
122    JsonIsObject,
123    object,
124    "Checks if the input JSONB is a JSON object"
125);
126
127#[cfg(test)]
128mod tests {
129    use std::sync::Arc;
130
131    use arrow_schema::Field;
132    use datafusion_common::arrow::array::{AsArray, BinaryArray};
133
134    use super::*;
135
136    #[test]
137    fn test_json_is_functions() {
138        let json_is_functions: [&dyn Function; 6] = [
139            &JsonIsBool::default(),
140            &JsonIsInt::default(),
141            &JsonIsFloat::default(),
142            &JsonIsString::default(),
143            &JsonIsArray::default(),
144            &JsonIsObject::default(),
145        ];
146        let expected_names = [
147            "json_is_bool",
148            "json_is_int",
149            "json_is_float",
150            "json_is_string",
151            "json_is_array",
152            "json_is_object",
153        ];
154        for (func, expected_name) in json_is_functions.iter().zip(expected_names.iter()) {
155            assert_eq!(func.name(), *expected_name);
156            assert_eq!(
157                func.return_type(&[DataType::Binary]).unwrap(),
158                DataType::Boolean
159            );
160        }
161
162        let json_strings = [
163            r#"true"#,
164            r#"1"#,
165            r#"1.0"#,
166            r#""The pig fly through a castle, and has been attracted by the princess.""#,
167            r#"[1, 2]"#,
168            r#"{"a": 1}"#,
169        ];
170        let expected_results = [
171            [true, false, false, false, false, false],
172            [false, true, false, false, false, false],
173            // Integers are also floats
174            [false, true, true, false, false, false],
175            [false, false, false, true, false, false],
176            [false, false, false, false, true, false],
177            [false, false, false, false, false, true],
178        ];
179
180        let jsonbs = json_strings
181            .iter()
182            .map(|s| {
183                let value = jsonb::parse_value(s.as_bytes()).unwrap();
184                value.to_vec()
185            })
186            .collect::<Vec<_>>();
187        let args = ScalarFunctionArgs {
188            args: vec![ColumnarValue::Array(Arc::new(
189                BinaryArray::from_iter_values(jsonbs),
190            ))],
191            arg_fields: vec![],
192            number_rows: 6,
193            return_field: Arc::new(Field::new("", DataType::Boolean, false)),
194            config_options: Arc::new(Default::default()),
195        };
196
197        for (func, expected_result) in json_is_functions.iter().zip(expected_results.iter()) {
198            let result = func
199                .invoke_with_args(args.clone())
200                .and_then(|x| x.to_array(6))
201                .unwrap();
202            let vector = result.as_boolean();
203            assert_eq!(vector.len(), json_strings.len());
204
205            for (i, expected) in expected_result.iter().enumerate() {
206                let result = vector.value(i);
207                assert_eq!(result, *expected);
208            }
209        }
210    }
211}
common_function/scalars/json/json_is.rs

common_function/scalars/json/
json_is.rs
