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

common_function/scalars/json/
json_is.rs
