1use std::sync::Arc;
4
5use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray, StructArray};
6use arrow_schema::{DataType, Field, Fields};
7use datafusion::{
8 common::{exec_datafusion_err, exec_err},
9 error::Result,
10 logical_expr::{
11 ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
12 },
13 scalar::ScalarValue,
14};
15use parquet_variant_compute::{VariantArrayBuilder, VariantType};
16use parquet_variant_json::JsonToVariant as JsonToVariantExt;
17
18use crate::shared::{try_field_as_string, try_parse_string_scalar};
19
20#[derive(Debug, Hash, PartialEq, Eq)]
22pub struct JsonToVariantUdf {
23 signature: Signature,
24}
25
26impl Default for JsonToVariantUdf {
27 fn default() -> Self {
28 Self {
29 signature: Signature::new(
30 TypeSignature::Uniform(
31 1,
32 vec![DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View],
33 ),
34 datafusion::logical_expr::Volatility::Immutable,
35 ),
36 }
37 }
38}
39
40impl ScalarUDFImpl for JsonToVariantUdf {
41 fn as_any(&self) -> &dyn std::any::Any {
42 self
43 }
44
45 fn name(&self) -> &str {
46 "json_to_variant"
47 }
48
49 fn signature(&self) -> &Signature {
50 &self.signature
51 }
52
53 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
54 Ok(DataType::Struct(Fields::from(vec![
55 Field::new("metadata", DataType::BinaryView, false),
56 Field::new("value", DataType::BinaryView, false),
57 ])))
58 }
59
60 fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Arc<Field>> {
61 let data_type = self.return_type(
62 args.arg_fields
63 .iter()
64 .map(|f| f.data_type().clone())
65 .collect::<Vec<_>>()
66 .as_slice(),
67 )?;
68 Ok(Arc::new(
69 Field::new(self.name(), data_type, true).with_extension_type(VariantType),
70 ))
71 }
72
73 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
74 let arg_field = args
75 .arg_fields
76 .first()
77 .ok_or_else(|| exec_datafusion_err!("empty argument, expected 1 argument"))?;
78
79 try_field_as_string(arg_field.as_ref())?;
80
81 let arg = args
82 .args
83 .first()
84 .ok_or_else(|| exec_datafusion_err!("empty argument, expected 1 argument"))?;
85
86 let out = match arg {
87 ColumnarValue::Scalar(scalar_value) => {
88 let json_str = try_parse_string_scalar(scalar_value)?;
89
90 let mut builder = VariantArrayBuilder::new(1);
91
92 match json_str {
93 Some(json_str) => builder.append_json(json_str.as_str())?,
94 None => builder.append_null(),
95 }
96
97 let struct_array: StructArray = builder.build().into();
98 ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(struct_array)))
99 }
100 ColumnarValue::Array(arr) => match arr.data_type() {
101 DataType::Utf8 => ColumnarValue::Array(from_utf8_arr(arr)?),
102 DataType::LargeUtf8 => ColumnarValue::Array(from_large_utf8_arr(arr)?),
103 DataType::Utf8View => ColumnarValue::Array(from_utf8view_arr(arr)?),
104 _ => return exec_err!("Invalid data type {}", arr.data_type()),
105 },
106 };
107
108 Ok(out)
109 }
110}
111
112macro_rules! define_from_string_array {
113 ($fn_name:ident, $array_type:ty) => {
114 pub(crate) fn $fn_name(arr: &ArrayRef) -> Result<ArrayRef> {
115 let arr = arr
116 .as_any()
117 .downcast_ref::<$array_type>()
118 .ok_or(exec_datafusion_err!(
119 "Unable to downcast array as expected by type."
120 ))?;
121
122 let mut builder = VariantArrayBuilder::new(arr.len());
123
124 for v in arr {
125 match v {
126 Some(json_str) => builder.append_json(json_str)?,
127 None => builder.append_null(),
128 }
129 }
130
131 let variant_array: StructArray = builder.build().into();
132
133 Ok(Arc::new(variant_array) as ArrayRef)
134 }
135 };
136}
137
138define_from_string_array!(from_utf8_arr, StringArray);
139define_from_string_array!(from_utf8view_arr, StringViewArray);
140define_from_string_array!(from_large_utf8_arr, LargeStringArray);
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145 use datafusion::logical_expr::{ReturnFieldArgs, ScalarFunctionArgs};
146 use parquet_variant::{Variant, VariantBuilder};
147 use parquet_variant_compute::VariantArray;
148
149 #[test]
150 fn test_json_to_variant_udf_scalar_none() {
151 let json_input = ScalarValue::Utf8(None);
152
153 let udf = JsonToVariantUdf::default();
154 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
155
156 let return_field = udf
157 .return_field_from_args(ReturnFieldArgs {
158 arg_fields: &[arg_field.clone()],
159 scalar_arguments: &[],
160 })
161 .unwrap();
162
163 let args = ScalarFunctionArgs {
164 args: vec![ColumnarValue::Scalar(json_input)],
165 return_field,
166 arg_fields: vec![arg_field],
167 number_rows: Default::default(),
168 config_options: Default::default(),
169 };
170
171 let result = udf.invoke_with_args(args).unwrap();
172
173 match result {
174 ColumnarValue::Scalar(ScalarValue::Struct(sv)) => {
175 assert!(sv.is_null(0), "expected null struct");
176 }
177 _ => panic!("Expected null struct array result"),
178 }
179 }
180
181 #[test]
182 fn test_json_to_variant_udf_scalar_null() {
183 let json_input = ScalarValue::Utf8(Some("null".into()));
184
185 let udf = JsonToVariantUdf::default();
186 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
187 let return_field = udf
188 .return_field_from_args(ReturnFieldArgs {
189 arg_fields: &[arg_field.clone()],
190 scalar_arguments: &[],
191 })
192 .unwrap();
193
194 let args = ScalarFunctionArgs {
195 args: vec![ColumnarValue::Scalar(json_input)],
196 return_field,
197 arg_fields: vec![arg_field],
198 number_rows: Default::default(),
199 config_options: Default::default(),
200 };
201
202 let result = udf.invoke_with_args(args).unwrap();
203 match result {
204 ColumnarValue::Scalar(ScalarValue::Struct(v)) => {
205 let variant_array = VariantArray::try_new(v.as_ref()).unwrap();
206 let variant = variant_array.value(0);
207 assert_eq!(variant, Variant::from(()));
208 }
209 _ => panic!("Expected scalar BinaryView result"),
210 }
211 }
212
213 #[test]
214 fn test_json_to_variant_udf_scalar_complex() {
215 let json_input =
216 ScalarValue::Utf8(Some(r#"{"key": 123, "data": [4, 5, "str"]}"#.to_string()));
217
218 let udf = JsonToVariantUdf::default();
219
220 let (expected_m, expected_v) = {
221 let mut variant_builder = VariantBuilder::new();
222 let mut object_builder = variant_builder.new_object();
223
224 object_builder.insert("key", 123_u8);
225
226 let mut inner_array_builder = object_builder.new_list("data");
227
228 inner_array_builder.append_value(4u8);
229 inner_array_builder.append_value(5u8);
230 inner_array_builder.append_value("str");
231
232 inner_array_builder.finish();
233
234 object_builder.finish();
235
236 variant_builder.finish()
237 };
238
239 let expected_variant = Variant::try_new(&expected_m, &expected_v).unwrap();
240
241 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
242 let return_field = udf
243 .return_field_from_args(ReturnFieldArgs {
244 arg_fields: &[arg_field.clone()],
245 scalar_arguments: &[],
246 })
247 .unwrap();
248
249 let args = ScalarFunctionArgs {
250 args: vec![ColumnarValue::Scalar(json_input)],
251 return_field,
252 arg_fields: vec![arg_field],
253 number_rows: Default::default(),
254 config_options: Default::default(),
255 };
256
257 let result = udf.invoke_with_args(args).unwrap();
258
259 match result {
260 ColumnarValue::Scalar(ScalarValue::Struct(v)) => {
261 let variant_array = VariantArray::try_new(v.as_ref()).unwrap();
262 let variant = variant_array.value(0);
263 assert_eq!(variant, expected_variant);
264 }
265 _ => panic!("Expected scalar BinaryView result"),
266 }
267 }
268
269 #[test]
270 fn test_json_to_variant_udf_scalar_primitive() {
271 let json_input = ScalarValue::Utf8(Some("123".to_string()));
272
273 let udf = JsonToVariantUdf::default();
274 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
275 let return_field = udf
276 .return_field_from_args(ReturnFieldArgs {
277 arg_fields: &[arg_field.clone()],
278 scalar_arguments: &[],
279 })
280 .unwrap();
281
282 let args = ScalarFunctionArgs {
283 args: vec![ColumnarValue::Scalar(json_input)],
284 return_field,
285 arg_fields: vec![arg_field],
286 number_rows: Default::default(),
287 config_options: Default::default(),
288 };
289
290 let result = udf.invoke_with_args(args).unwrap();
291
292 match result {
293 ColumnarValue::Scalar(ScalarValue::Struct(v)) => {
294 let variant_array = VariantArray::try_new(v.as_ref()).unwrap();
295 let variant = variant_array.value(0);
296 assert_eq!(variant, Variant::from(123_u8));
297 }
298 _ => panic!("Expected scalar BinaryView result"),
299 }
300 }
301}