common_function/scalars/json/
json_is.rs1use 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
25macro_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 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 [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}