common_function/scalars/json/
json_object_keys.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 arrow::array::{Array, AsArray, ListBuilder, StringViewBuilder};
19use arrow::compute;
20use arrow::datatypes::{DataType, Field};
21use datafusion_common::DataFusionError;
22use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, Volatility};
23
24use crate::function::{Function, extract_args};
25
26/// Returns the keys of the outermost JSON object as a list of strings.
27#[derive(Clone, Debug)]
28pub(crate) struct JsonObjectKeysFunction {
29    signature: Signature,
30}
31
32impl Default for JsonObjectKeysFunction {
33    fn default() -> Self {
34        Self {
35            signature: Signature::uniform(
36                1,
37                vec![
38                    DataType::Binary,
39                    DataType::LargeBinary,
40                    DataType::BinaryView,
41                    DataType::Null,
42                ],
43                Volatility::Immutable,
44            ),
45        }
46    }
47}
48
49const NAME: &str = "json_object_keys";
50
51impl Function for JsonObjectKeysFunction {
52    fn name(&self) -> &str {
53        NAME
54    }
55
56    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
57        Ok(DataType::List(Arc::new(Field::new(
58            "item",
59            DataType::Utf8View,
60            true,
61        ))))
62    }
63
64    fn signature(&self) -> &Signature {
65        &self.signature
66    }
67
68    fn invoke_with_args(
69        &self,
70        args: ScalarFunctionArgs,
71    ) -> datafusion_common::Result<ColumnarValue> {
72        let [jsons] = extract_args(self.name(), &args)?;
73        let jsons = compute::cast(&jsons, &DataType::BinaryView)?;
74        let jsons = jsons.as_binary_view();
75
76        let size = jsons.len();
77        let mut builder = ListBuilder::with_capacity(StringViewBuilder::new(), size);
78
79        for i in 0..size {
80            let Some(json) = jsons.is_valid(i).then(|| jsons.value(i)) else {
81                builder.append_null();
82                continue;
83            };
84
85            match jsonb::from_slice(json) {
86                Ok(jsonb::Value::Object(object)) => {
87                    for key in object.keys() {
88                        builder.values().append_value(key);
89                    }
90                    builder.append(true);
91                }
92                Ok(_) => builder.append_null(),
93                Err(e) => {
94                    return Err(DataFusionError::Execution(format!(
95                        "invalid json binary: {e}"
96                    )));
97                }
98            }
99        }
100
101        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
102    }
103}
104
105impl Display for JsonObjectKeysFunction {
106    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
107        write!(f, "JSON_OBJECT_KEYS")
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use std::sync::Arc;
114
115    use arrow::array::{BinaryArray, NullArray};
116    use arrow_schema::Field;
117
118    use super::*;
119
120    #[test]
121    fn test_json_object_keys_function() {
122        let json_object_keys = JsonObjectKeysFunction::default();
123        let return_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8View, true)));
124
125        assert_eq!("json_object_keys", json_object_keys.name());
126        assert_eq!(
127            return_type,
128            json_object_keys.return_type(&[DataType::Binary]).unwrap()
129        );
130
131        let json_strings = [
132            Some(r#"{"b": 2, "a": 1}"#),
133            Some("{}"),
134            Some(r#"{"outer": {"inner": 1}, "value": 2}"#),
135            Some("[1, 2]"),
136            Some("42"),
137            Some("null"),
138            None,
139        ];
140
141        let results = [
142            Some(vec!["a", "b"]),
143            Some(vec![]),
144            Some(vec!["outer", "value"]),
145            None,
146            None,
147            None,
148            None,
149        ];
150
151        let jsonbs = json_strings
152            .into_iter()
153            .map(|s| s.map(|json| jsonb::parse_value(json.as_bytes()).unwrap().to_vec()))
154            .collect::<Vec<_>>();
155
156        let args = ScalarFunctionArgs {
157            args: vec![ColumnarValue::Array(Arc::new(BinaryArray::from_iter(
158                jsonbs,
159            )))],
160            arg_fields: vec![],
161            number_rows: 7,
162            return_field: Arc::new(Field::new("x", return_type.clone(), true)),
163            config_options: Arc::new(Default::default()),
164        };
165        let result = json_object_keys
166            .invoke_with_args(args)
167            .and_then(|x| x.to_array(7))
168            .unwrap();
169        let vector = result.as_list::<i32>();
170
171        assert_eq!(7, vector.len());
172        for (i, expected) in results.iter().enumerate() {
173            match expected {
174                Some(expected) => {
175                    let values = vector.value(i);
176                    let values = values.as_string_view();
177                    let keys = values.iter().flatten().collect::<Vec<_>>();
178                    assert_eq!(expected, &keys);
179                }
180                None => assert!(vector.is_null(i)),
181            }
182        }
183
184        let invalid_jsonb = vec![b"invalid json"];
185        let args = ScalarFunctionArgs {
186            args: vec![ColumnarValue::Array(Arc::new(
187                BinaryArray::from_iter_values(invalid_jsonb),
188            ))],
189            arg_fields: vec![],
190            number_rows: 1,
191            return_field: Arc::new(Field::new("x", return_type.clone(), true)),
192            config_options: Arc::new(Default::default()),
193        };
194        let result = json_object_keys.invoke_with_args(args);
195        assert!(result.is_err());
196
197        let args = ScalarFunctionArgs {
198            args: vec![ColumnarValue::Array(Arc::new(NullArray::new(1)))],
199            arg_fields: vec![],
200            number_rows: 1,
201            return_field: Arc::new(Field::new("x", return_type, true)),
202            config_options: Arc::new(Default::default()),
203        };
204        let result = json_object_keys
205            .invoke_with_args(args)
206            .and_then(|x| x.to_array(1))
207            .unwrap();
208        let vector = result.as_list::<i32>();
209        assert!(vector.is_null(0));
210    }
211}
common_function/scalars/json/json_object_keys.rs

common_function/scalars/json/
json_object_keys.rs
