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

common_function/scalars/json/
json_is.rs
