common_function/scalars/json/
json_object.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::DataFusionError;
19use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, BinaryViewBuilder};
20use datafusion_common::arrow::compute;
21use datafusion_common::arrow::datatypes::{DataType, Float64Type, Int64Type, UInt64Type};
22use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, TypeSignature, Volatility};
23
24use crate::function::Function;
25
26const NAME: &str = "json_object";
27
28/// Builds a `JSONB` object from interleaved `(key, value, key, value, ...)`
29/// arguments, like MySQL's `JSON_OBJECT`; called with no arguments it returns
30/// `{}`. Values are written into the binary directly, so they need no JSON
31/// text escaping. Keys must be non-NULL and are converted to strings (so a
32/// numeric key like `1` becomes `"1"`, as in MySQL); values may be strings,
33/// integers, floats, booleans, or NULL (rendered as JSON null). Other types —
34/// including decimals, which JSONB numbers cannot represent exactly — are
35/// rejected; cast them explicitly. A duplicate key keeps the last value.
36#[derive(Clone, Debug)]
37pub(crate) struct JsonObjectFunction {
38    signature: Signature,
39}
40
41impl Default for JsonObjectFunction {
42    fn default() -> Self {
43        Self {
44            signature: Signature::one_of(
45                vec![TypeSignature::Nullary, TypeSignature::VariadicAny],
46                Volatility::Immutable,
47            ),
48        }
49    }
50}
51
52/// A value column normalized to the canonical arrow type its JSON rendering
53/// reads from.
54enum ValueColumn {
55    Null,
56    Bool(ArrayRef),
57    Int(ArrayRef),
58    UInt(ArrayRef),
59    Float(ArrayRef),
60    String(ArrayRef),
61}
62
63impl ValueColumn {
64    fn try_new(array: &ArrayRef) -> datafusion_common::Result<Self> {
65        let normalized = match array.data_type() {
66            DataType::Null => return Ok(ValueColumn::Null),
67            DataType::Boolean => ValueColumn::Bool(array.clone()),
68            DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => {
69                ValueColumn::Int(compute::cast(array, &DataType::Int64)?)
70            }
71            DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
72                ValueColumn::UInt(compute::cast(array, &DataType::UInt64)?)
73            }
74            DataType::Float16 | DataType::Float32 | DataType::Float64 => {
75                ValueColumn::Float(compute::cast(array, &DataType::Float64)?)
76            }
77            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
78                ValueColumn::String(compute::cast(array, &DataType::Utf8View)?)
79            }
80            other => {
81                return Err(DataFusionError::Execution(format!(
82                    "{NAME} does not support values of type {other}; cast the value to a string"
83                )));
84            }
85        };
86        Ok(normalized)
87    }
88
89    fn value(&self, row: usize) -> jsonb::Value<'_> {
90        let array = match self {
91            ValueColumn::Null => return jsonb::Value::Null,
92            ValueColumn::Bool(array)
93            | ValueColumn::Int(array)
94            | ValueColumn::UInt(array)
95            | ValueColumn::Float(array)
96            | ValueColumn::String(array) => array,
97        };
98        if !array.is_valid(row) {
99            return jsonb::Value::Null;
100        }
101        match self {
102            ValueColumn::Null => unreachable!(),
103            ValueColumn::Bool(array) => array.as_boolean().value(row).into(),
104            ValueColumn::Int(array) => array.as_primitive::<Int64Type>().value(row).into(),
105            ValueColumn::UInt(array) => array.as_primitive::<UInt64Type>().value(row).into(),
106            ValueColumn::Float(array) => array.as_primitive::<Float64Type>().value(row).into(),
107            ValueColumn::String(array) => array.as_string_view().value(row).into(),
108        }
109    }
110}
111
112impl Function for JsonObjectFunction {
113    fn name(&self) -> &str {
114        NAME
115    }
116
117    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
118        Ok(DataType::BinaryView)
119    }
120
121    fn signature(&self) -> &Signature {
122        &self.signature
123    }
124
125    fn invoke_with_args(
126        &self,
127        args: ScalarFunctionArgs,
128    ) -> datafusion_common::Result<ColumnarValue> {
129        let arrays = ColumnarValue::values_to_arrays(&args.args)?;
130        if arrays.len() % 2 != 0 {
131            return Err(DataFusionError::Execution(format!(
132                "{NAME} expects (key, value) argument pairs, got {} arguments",
133                arrays.len()
134            )));
135        }
136        let pairs = arrays
137            .chunks(2)
138            .map(|pair| {
139                let keys = compute::cast(&pair[0], &DataType::Utf8View).map_err(|_| {
140                    DataFusionError::Execution(format!(
141                        "{NAME} cannot convert keys of type {} to string",
142                        pair[0].data_type()
143                    ))
144                })?;
145                if keys.null_count() > 0 {
146                    return Err(DataFusionError::Execution(format!(
147                        "{NAME} does not allow NULL keys"
148                    )));
149                }
150                Ok((keys, ValueColumn::try_new(&pair[1])?))
151            })
152            .collect::<datafusion_common::Result<Vec<_>>>()?;
153
154        let rows = arrays.first().map_or(args.number_rows, |a| a.len());
155        let mut builder = BinaryViewBuilder::with_capacity(rows);
156        let mut buf = Vec::new();
157        for row in 0..rows {
158            let mut object = jsonb::Object::new();
159            for (keys, values) in &pairs {
160                object.insert(
161                    keys.as_string_view().value(row).to_string(),
162                    values.value(row),
163                );
164            }
165            buf.clear();
166            jsonb::Value::Object(object).write_to_vec(&mut buf);
167            builder.append_value(&buf);
168        }
169        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
170    }
171}
172
173impl Display for JsonObjectFunction {
174    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
175        write!(f, "JSON_OBJECT")
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use std::sync::Arc;
182
183    use arrow_schema::Field;
184    use datafusion_common::arrow::array::{
185        Int64Array, NullArray, StringArray, TimestampMillisecondArray,
186    };
187
188    use super::*;
189
190    fn invoke(args: Vec<ColumnarValue>, rows: usize) -> datafusion_common::Result<Vec<String>> {
191        let function = JsonObjectFunction::default();
192        let result = function.invoke_with_args(ScalarFunctionArgs {
193            args,
194            arg_fields: vec![],
195            number_rows: rows,
196            return_field: Arc::new(Field::new("x", DataType::BinaryView, true)),
197            config_options: Arc::new(Default::default()),
198        })?;
199        let array = result.to_array(rows)?;
200        let array = array.as_binary_view();
201        Ok((0..array.len())
202            .map(|i| jsonb::from_slice(array.value(i)).unwrap().to_string())
203            .collect())
204    }
205
206    fn key(name: &str) -> ColumnarValue {
207        ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(name.to_string())))
208    }
209
210    #[test]
211    fn test_builds_objects_from_mixed_types_without_escaping() {
212        let texts = invoke(
213            vec![
214                key("host"),
215                ColumnarValue::Array(Arc::new(StringArray::from(vec![
216                    Some("we\"ird\\\nhost"),
217                    None,
218                ]))),
219                key("pid"),
220                ColumnarValue::Array(Arc::new(Int64Array::from(vec![42, 7]))),
221                key("up"),
222                ColumnarValue::Scalar(datafusion_common::ScalarValue::Boolean(Some(true))),
223                key("v"),
224                ColumnarValue::Array(Arc::new(NullArray::new(2))),
225            ],
226            2,
227        )
228        .unwrap();
229        assert_eq!(
230            texts,
231            vec![
232                r#"{"host":"we\"ird\\\nhost","pid":42,"up":true,"v":null}"#,
233                r#"{"host":null,"pid":7,"up":true,"v":null}"#,
234            ]
235        );
236    }
237
238    #[test]
239    fn test_empty_call_builds_one_empty_object_per_row() {
240        assert_eq!(invoke(vec![], 3).unwrap(), vec!["{}", "{}", "{}"]);
241    }
242
243    #[test]
244    fn test_numeric_key_converts_to_string() {
245        let texts = invoke(
246            vec![
247                ColumnarValue::Scalar(datafusion_common::ScalarValue::Int64(Some(7))),
248                key("v"),
249            ],
250            1,
251        )
252        .unwrap();
253        assert_eq!(texts, vec![r#"{"7":"v"}"#]);
254    }
255
256    #[test]
257    fn test_rejects_odd_arguments_null_keys_and_unsupported_values() {
258        let err = invoke(vec![key("a")], 1).unwrap_err();
259        assert!(err.to_string().contains("(key, value) argument pairs"));
260
261        let err = invoke(
262            vec![
263                ColumnarValue::Array(Arc::new(StringArray::from(vec![None::<&str>]))),
264                key("v"),
265            ],
266            1,
267        )
268        .unwrap_err();
269        assert!(err.to_string().contains("NULL keys"));
270
271        let err = invoke(
272            vec![
273                key("ts"),
274                ColumnarValue::Array(Arc::new(TimestampMillisecondArray::from(vec![1_000]))),
275            ],
276            1,
277        )
278        .unwrap_err();
279        assert!(err.to_string().contains("does not support values"));
280    }
281}
common_function/scalars/json/json_object.rs

common_function/scalars/json/
json_object.rs
