common_function/scalars/json/
json_path_exists.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, TypeSignature};
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/// Check if the given JSON data contains the given JSON path.
29#[derive(Clone, Debug, Default)]
30pub struct JsonPathExistsFunction;
31
32const NAME: &str = "json_path_exists";
33
34impl Function for JsonPathExistsFunction {
35    fn name(&self) -> &str {
36        NAME
37    }
38
39    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
40        Ok(ConcreteDataType::boolean_datatype())
41    }
42
43    fn signature(&self) -> Signature {
44        Signature::one_of(
45            vec![
46                TypeSignature::Exact(vec![
47                    ConcreteDataType::json_datatype(),
48                    ConcreteDataType::string_datatype(),
49                ]),
50                TypeSignature::Exact(vec![
51                    ConcreteDataType::null_datatype(),
52                    ConcreteDataType::string_datatype(),
53                ]),
54                TypeSignature::Exact(vec![
55                    ConcreteDataType::json_datatype(),
56                    ConcreteDataType::null_datatype(),
57                ]),
58                TypeSignature::Exact(vec![
59                    ConcreteDataType::null_datatype(),
60                    ConcreteDataType::null_datatype(),
61                ]),
62            ],
63            Volatility::Immutable,
64        )
65    }
66
67    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
68        ensure!(
69            columns.len() == 2,
70            InvalidFuncArgsSnafu {
71                err_msg: format!(
72                    "The length of the args is not correct, expect exactly two, have: {}",
73                    columns.len()
74                ),
75            }
76        );
77        let jsons = &columns[0];
78        let paths = &columns[1];
79
80        let size = jsons.len();
81        let mut results = BooleanVectorBuilder::with_capacity(size);
82
83        match (jsons.data_type(), paths.data_type()) {
84            (ConcreteDataType::Binary(_), ConcreteDataType::String(_)) => {
85                for i in 0..size {
86                    let result = match (jsons.get_ref(i).as_binary(), paths.get_ref(i).as_string())
87                    {
88                        (Ok(Some(json)), Ok(Some(path))) => {
89                            // Get `JsonPath`.
90                            let json_path = match jsonb::jsonpath::parse_json_path(path.as_bytes())
91                            {
92                                Ok(json_path) => json_path,
93                                Err(_) => {
94                                    return InvalidFuncArgsSnafu {
95                                        err_msg: format!("Illegal json path: {:?}", path),
96                                    }
97                                    .fail();
98                                }
99                            };
100                            jsonb::path_exists(json, json_path).ok()
101                        }
102                        _ => None,
103                    };
104
105                    results.push(result);
106                }
107            }
108
109            // Any null args existence causes the result to be NULL.
110            (ConcreteDataType::Null(_), ConcreteDataType::String(_)) => results.push_nulls(size),
111            (ConcreteDataType::Binary(_), ConcreteDataType::Null(_)) => results.push_nulls(size),
112            (ConcreteDataType::Null(_), ConcreteDataType::Null(_)) => results.push_nulls(size),
113
114            _ => {
115                return UnsupportedInputDataTypeSnafu {
116                    function: NAME,
117                    datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
118                }
119                .fail();
120            }
121        }
122
123        Ok(results.to_vector())
124    }
125}
126
127impl Display for JsonPathExistsFunction {
128    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
129        write!(f, "JSON_PATH_EXISTS")
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use std::sync::Arc;
136
137    use common_query::prelude::TypeSignature;
138    use datatypes::prelude::ScalarVector;
139    use datatypes::vectors::{BinaryVector, NullVector, StringVector};
140
141    use super::*;
142
143    #[test]
144    fn test_json_path_exists_function() {
145        let json_path_exists = JsonPathExistsFunction;
146
147        assert_eq!("json_path_exists", json_path_exists.name());
148        assert_eq!(
149            ConcreteDataType::boolean_datatype(),
150            json_path_exists
151                .return_type(&[ConcreteDataType::json_datatype()])
152                .unwrap()
153        );
154
155        assert!(matches!(json_path_exists.signature(),
156                         Signature {
157                             type_signature: TypeSignature::OneOf(valid_types),
158                             volatility: Volatility::Immutable
159                         } if valid_types ==
160            vec![
161                TypeSignature::Exact(vec![
162                    ConcreteDataType::json_datatype(),
163                    ConcreteDataType::string_datatype(),
164                ]),
165                TypeSignature::Exact(vec![
166                    ConcreteDataType::null_datatype(),
167                    ConcreteDataType::string_datatype(),
168                ]),
169                TypeSignature::Exact(vec![
170                    ConcreteDataType::json_datatype(),
171                    ConcreteDataType::null_datatype(),
172                ]),
173                TypeSignature::Exact(vec![
174                    ConcreteDataType::null_datatype(),
175                    ConcreteDataType::null_datatype(),
176                ]),
177            ],
178        ));
179
180        let json_strings = [
181            r#"{"a": {"b": 2}, "b": 2, "c": 3}"#,
182            r#"{"a": 4, "b": {"c": 6}, "c": 6}"#,
183            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
184            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
185            r#"[1, 2, 3]"#,
186            r#"null"#,
187            r#"{"a": 7, "b": 8, "c": {"a": 7}}"#,
188            r#"null"#,
189        ];
190        let paths = vec![
191            "$.a.b.c", "$.b", "$.c.a", ".d", "$[0]", "$.a", "null", "null",
192        ];
193        let expected = [false, true, true, false, true, false, false, false];
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
203        let json_vector = BinaryVector::from_vec(jsonbs);
204        let path_vector = StringVector::from_vec(paths);
205        let args: Vec<VectorRef> = vec![Arc::new(json_vector), Arc::new(path_vector)];
206        let vector = json_path_exists
207            .eval(&FunctionContext::default(), &args)
208            .unwrap();
209
210        // Test for non-nulls.
211        assert_eq!(8, vector.len());
212        for (i, real) in expected.iter().enumerate() {
213            let result = vector.get_ref(i);
214            assert!(!result.is_null());
215            let val = result.as_boolean().unwrap().unwrap();
216            assert_eq!(val, *real);
217        }
218
219        // Test for path error.
220        let json_bytes = jsonb::parse_value("{}".as_bytes()).unwrap().to_vec();
221        let json = BinaryVector::from_vec(vec![json_bytes]);
222        let illegal_path = StringVector::from_vec(vec!["$..a"]);
223
224        let args: Vec<VectorRef> = vec![Arc::new(json), Arc::new(illegal_path)];
225        let err = json_path_exists.eval(&FunctionContext::default(), &args);
226        assert!(err.is_err());
227
228        // Test for nulls.
229        let json_bytes = jsonb::parse_value("{}".as_bytes()).unwrap().to_vec();
230        let json = BinaryVector::from_vec(vec![json_bytes]);
231        let null_json = NullVector::new(1);
232
233        let path = StringVector::from_vec(vec!["$.a"]);
234        let null_path = NullVector::new(1);
235
236        let args: Vec<VectorRef> = vec![Arc::new(null_json), Arc::new(path)];
237        let result1 = json_path_exists
238            .eval(&FunctionContext::default(), &args)
239            .unwrap();
240        let args: Vec<VectorRef> = vec![Arc::new(json), Arc::new(null_path)];
241        let result2 = json_path_exists
242            .eval(&FunctionContext::default(), &args)
243            .unwrap();
244
245        assert_eq!(result1.len(), 1);
246        assert!(result1.get_ref(0).is_null());
247        assert_eq!(result2.len(), 1);
248        assert!(result2.get_ref(0).is_null());
249    }
250}
common_function/scalars/json/json_path_exists.rs

common_function/scalars/json/
json_path_exists.rs
