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

common_function/scalars/json/
json_is.rs
