1use std::sync::Arc;
2
3use arrow::array::{Array, ArrayRef, AsArray, StructArray};
4use arrow_schema::{DataType, Field};
5use datafusion::{
6 common::exec_err,
7 error::Result,
8 logical_expr::{
9 ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
10 TypeSignature, Volatility,
11 },
12 scalar::ScalarValue,
13};
14use parquet_variant::Variant;
15use parquet_variant_compute::{VariantArray, VariantArrayBuilder, cast_to_variant};
16
17use crate::shared::{try_parse_binary_columnar, try_parse_binary_scalar};
18
19#[derive(Debug, Hash, PartialEq, Eq)]
20pub struct CastToVariantUdf {
21 signature: Signature,
22}
23
24impl Default for CastToVariantUdf {
25 fn default() -> Self {
26 Self {
27 signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable),
28 }
29 }
30}
31
32impl CastToVariantUdf {
33 fn canonical_variant_data_type() -> DataType {
34 VariantArrayBuilder::new(0).build().data_type().clone()
35 }
36
37 fn canonical_return_field(name: &str) -> Field {
38 VariantArrayBuilder::new(0)
39 .build()
40 .field(name.to_string())
41 .with_nullable(true)
42 }
43
44 fn append_variant_or_null(
45 builder: &mut VariantArrayBuilder,
46 metadata: Option<&[u8]>,
47 value: Option<&[u8]>,
48 ) -> Result<()> {
49 match (metadata, value) {
50 (Some(m), Some(v)) if !m.is_empty() && !v.is_empty() => {
51 builder.append_variant(Variant::try_new(m, v)?);
52 }
53 _ => builder.append_null(),
54 }
55
56 Ok(())
57 }
58
59 fn from_metadata_value(
60 metadata_argument: &ColumnarValue,
61 variant_argument: &ColumnarValue,
62 ) -> Result<ColumnarValue> {
63 let out = match (metadata_argument, variant_argument) {
64 (ColumnarValue::Array(metadata_array), ColumnarValue::Array(value_array)) => {
65 if metadata_array.len() != value_array.len() {
66 return exec_err!(
67 "expected metadata array to be of same length as variant array"
68 );
69 }
70
71 let metadata_array = try_parse_binary_columnar(metadata_array)?;
72 let value_array = try_parse_binary_columnar(value_array)?;
73
74 let mut builder = VariantArrayBuilder::new(metadata_array.len());
75 for (m, v) in metadata_array.into_iter().zip(value_array) {
76 Self::append_variant_or_null(&mut builder, m, v)?;
77 }
78 let out: StructArray = builder.build().into();
79
80 ColumnarValue::Array(Arc::new(out) as ArrayRef)
81 }
82 (ColumnarValue::Scalar(metadata_value), ColumnarValue::Array(value_array)) => {
83 let metadata = try_parse_binary_scalar(metadata_value)?;
84 let value_array = try_parse_binary_columnar(value_array)?;
85
86 let mut builder = VariantArrayBuilder::new(value_array.len());
87 for v in value_array {
88 Self::append_variant_or_null(&mut builder, metadata.map(|m| m.as_slice()), v)?;
89 }
90 let arr: StructArray = builder.build().into();
91
92 ColumnarValue::Array(Arc::new(arr) as ArrayRef)
93 }
94 (ColumnarValue::Scalar(metadata_value), ColumnarValue::Scalar(value_scalar)) => {
95 let metadata = try_parse_binary_scalar(metadata_value)?;
96 let value = try_parse_binary_scalar(value_scalar)?;
97
98 match (metadata, value) {
99 (Some(m), Some(v)) if !m.is_empty() && !v.is_empty() => {
100 let mut b = VariantArrayBuilder::new(1);
101 b.append_variant(Variant::try_new(m.as_slice(), v.as_slice())?);
102 let arr: StructArray = b.build().into();
103
104 ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(arr)))
105 }
106 _ => ColumnarValue::Scalar(ScalarValue::Null),
107 }
108 }
109 (ColumnarValue::Array(metadata_array), ColumnarValue::Scalar(value_scalar)) => {
110 let metadata_array = try_parse_binary_columnar(metadata_array)?;
111 let value = try_parse_binary_scalar(value_scalar)?;
112
113 let mut builder = VariantArrayBuilder::new(metadata_array.len());
114 for m in metadata_array {
115 Self::append_variant_or_null(&mut builder, m, value.map(|v| v.as_slice()))?;
116 }
117 let arr: StructArray = builder.build().into();
118
119 ColumnarValue::Array(Arc::new(arr))
120 }
121 };
122
123 Ok(out)
124 }
125
126 fn from_array(array: &ArrayRef) -> Result<ColumnarValue> {
127 if let Some(struct_array) = array.as_struct_opt()
129 && VariantArray::try_new(struct_array).is_ok()
130 {
131 return Ok(ColumnarValue::Array(Arc::clone(array)));
132 }
133
134 let variant_array = cast_to_variant(array.as_ref())?;
135 let struct_array: StructArray = variant_array.into();
136
137 Ok(ColumnarValue::Array(Arc::new(struct_array)))
138 }
139
140 fn from_scalar_value(scalar_value: &ScalarValue) -> Result<ColumnarValue> {
141 if let ScalarValue::Struct(struct_array) = scalar_value
142 && VariantArray::try_new(struct_array.as_ref()).is_ok()
143 {
144 return Ok(ColumnarValue::Scalar(ScalarValue::Struct(
145 struct_array.clone(),
146 )));
147 }
148
149 let array = scalar_value.to_array_of_size(1)?;
150 let variant_array = cast_to_variant(array.as_ref())?;
151 let struct_array: StructArray = variant_array.into();
152
153 Ok(ColumnarValue::Scalar(ScalarValue::Struct(Arc::new(
154 struct_array,
155 ))))
156 }
157}
158
159impl ScalarUDFImpl for CastToVariantUdf {
160 fn as_any(&self) -> &dyn std::any::Any {
161 self
162 }
163
164 fn name(&self) -> &str {
165 "cast_to_variant"
166 }
167
168 fn signature(&self) -> &Signature {
169 &self.signature
170 }
171
172 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
173 Ok(Self::canonical_variant_data_type())
174 }
175
176 fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Arc<Field>> {
177 Ok(Arc::new(Self::canonical_return_field(self.name())))
178 }
179
180 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
181 match args.args.as_slice() {
182 [metadata_value, variant_value] => {
183 Self::from_metadata_value(metadata_value, variant_value)
184 }
185 [ColumnarValue::Scalar(scalar_value)] => Self::from_scalar_value(scalar_value),
186 [ColumnarValue::Array(array)] => Self::from_array(array),
187 _ => exec_err!("unrecognized argument"),
188 }
189 }
190}
191
192#[cfg(test)]
193mod tests {
194
195 use arrow::array::{FixedSizeBinaryBuilder, Int32Array, StringArray, StringViewArray};
196 use arrow_schema::Fields;
197 use parquet_variant::Variant;
198 use parquet_variant_compute::{VariantArray, VariantType};
199
200 use crate::shared::{build_variant_array_from_json, build_variant_array_from_json_array};
201
202 use super::*;
203
204 #[test]
205 fn test_scalar_float64() {
206 let udf = CastToVariantUdf::default();
207
208 let arg_field = Arc::new(Field::new("input", DataType::Float64, true));
209 let return_field = Arc::new(Field::new(
210 "res",
211 udf.return_type(&[DataType::Float64]).unwrap(),
212 true,
213 ));
214
215 let args = ScalarFunctionArgs {
216 args: vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(3.25)))],
217 return_field,
218 arg_fields: vec![arg_field],
219 number_rows: Default::default(),
220 config_options: Default::default(),
221 };
222
223 let res = udf.invoke_with_args(args).unwrap();
224
225 let ColumnarValue::Scalar(ScalarValue::Struct(variant_array)) = res else {
226 panic!("expected struct scalar")
227 };
228
229 let variant_array = VariantArray::try_new(variant_array.as_ref()).unwrap();
230
231 assert_eq!(variant_array.value(0), Variant::Double(3.25));
232 }
233
234 #[test]
235 fn test_array_int32() {
236 let udf = CastToVariantUdf::default();
237
238 let arg_field = Arc::new(Field::new("input", DataType::Int32, true));
239 let return_field = Arc::new(Field::new(
240 "res",
241 udf.return_type(&[DataType::Int32]).unwrap(),
242 true,
243 ));
244
245 let args = ScalarFunctionArgs {
246 args: vec![ColumnarValue::Array(Arc::new(Int32Array::from(vec![
247 Some(1),
248 None,
249 Some(-5),
250 ])) as ArrayRef)],
251 return_field,
252 arg_fields: vec![arg_field],
253 number_rows: Default::default(),
254 config_options: Default::default(),
255 };
256
257 let res = udf.invoke_with_args(args).unwrap();
258
259 let ColumnarValue::Array(arr) = res else {
260 panic!("expected array output")
261 };
262
263 let variant_array = VariantArray::try_new(arr.as_ref()).unwrap();
264
265 assert_eq!(variant_array.value(0), Variant::Int32(1));
266 assert!(variant_array.is_null(1));
267 assert_eq!(variant_array.value(2), Variant::Int32(-5));
268 }
269
270 #[test]
271 fn test_return_field_extension_type() {
272 let udf = CastToVariantUdf::default();
273
274 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
275
276 let return_field = udf
277 .return_field_from_args(ReturnFieldArgs {
278 arg_fields: std::slice::from_ref(&arg_field),
279 scalar_arguments: &[None],
280 })
281 .unwrap();
282
283 assert!(matches!(return_field.extension_type(), VariantType));
284 assert_eq!(
285 return_field.data_type(),
286 &DataType::Struct(Fields::from(vec![
287 Field::new("metadata", DataType::BinaryView, false),
288 Field::new("value", DataType::BinaryView, false),
289 ]))
290 );
291 }
292
293 #[test]
294 fn test_scalar_binary_views() {
295 let expected_variant_array = build_variant_array_from_json(&serde_json::json!({
296 "name": "norm",
297 }));
298
299 let (input_metadata, input_value) = {
300 let metadata = expected_variant_array.metadata_field().value(0);
301 let value = expected_variant_array.value_field().unwrap().value(0);
302
303 (metadata, value)
304 };
305
306 let udf = CastToVariantUdf::default();
307
308 let metadata_field = Arc::new(Field::new("metadata", DataType::BinaryView, true));
309 let variant_field = Arc::new(Field::new("value", DataType::BinaryView, true));
310
311 let return_field = Arc::new(Field::new(
312 "res",
313 udf.return_type(&[DataType::BinaryView, DataType::BinaryView])
314 .unwrap(),
315 true,
316 ));
317
318 let args = ScalarFunctionArgs {
319 args: vec![
320 ColumnarValue::Scalar(ScalarValue::BinaryView(Some(input_metadata.to_vec()))),
321 ColumnarValue::Scalar(ScalarValue::BinaryView(Some(input_value.to_vec()))),
322 ],
323 return_field,
324 arg_fields: vec![metadata_field, variant_field],
325 number_rows: Default::default(),
326 config_options: Default::default(),
327 };
328
329 let res = udf.invoke_with_args(args).unwrap();
330
331 let ColumnarValue::Scalar(ScalarValue::Struct(variant_array)) = res else {
332 panic!("expected scalar value struct array")
333 };
334
335 let variant_array = VariantArray::try_new(variant_array.as_ref()).unwrap();
336
337 assert_eq!(&variant_array, &expected_variant_array);
338 }
339
340 #[test]
341 fn test_array_string() {
342 let udf = CastToVariantUdf::default();
343
344 let arg_field = Arc::new(Field::new("input", DataType::Utf8, true));
345 let return_field = Arc::new(Field::new(
346 "res",
347 udf.return_type(&[DataType::Utf8]).unwrap(),
348 true,
349 ));
350
351 let args = ScalarFunctionArgs {
352 args: vec![ColumnarValue::Array(Arc::new(StringArray::from(vec![
353 Some("abcdefghijklmnop"),
354 None,
355 Some("hello world"),
356 ])) as ArrayRef)],
357 return_field,
358 arg_fields: vec![arg_field],
359 number_rows: Default::default(),
360 config_options: Default::default(),
361 };
362
363 let res = udf.invoke_with_args(args).unwrap();
364
365 let ColumnarValue::Array(arr) = res else {
366 panic!("expected array output")
367 };
368
369 let variant_array = VariantArray::try_new(arr.as_ref()).unwrap();
370
371 assert_eq!(variant_array.value(0), Variant::from("abcdefghijklmnop"));
372 assert!(variant_array.is_null(1));
373 assert_eq!(variant_array.value(2), Variant::from("hello world"));
374 }
375
376 #[test]
377 fn test_fixed_size_binary_uuid_like() {
378 let udf = CastToVariantUdf::default();
379
380 let arg_field = Arc::new(Field::new("input", DataType::FixedSizeBinary(16), true));
381 let return_field = Arc::new(Field::new(
382 "res",
383 udf.return_type(&[DataType::FixedSizeBinary(16)]).unwrap(),
384 true,
385 ));
386
387 let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 16);
388 builder.append_value([1u8; 16]).unwrap();
389 builder.append_null();
390 builder.append_value([2u8; 16]).unwrap();
391 let array = builder.finish();
392
393 let args = ScalarFunctionArgs {
394 args: vec![ColumnarValue::Array(Arc::new(array) as ArrayRef)],
395 return_field,
396 arg_fields: vec![arg_field],
397 number_rows: Default::default(),
398 config_options: Default::default(),
399 };
400
401 let res = udf.invoke_with_args(args).unwrap();
402
403 let ColumnarValue::Array(arr) = res else {
404 panic!("expected array output")
405 };
406
407 let variant_array = VariantArray::try_new(arr.as_ref()).unwrap();
408
409 assert_eq!(variant_array.value(0), Variant::Binary(&[1u8; 16]));
410 assert!(variant_array.is_null(1));
411 assert_eq!(variant_array.value(2), Variant::Binary(&[2u8; 16]));
412 }
413
414 #[test]
415 fn test_array_binary_views() {
416 let expected_variant_array = build_variant_array_from_json_array(&[
417 Some(serde_json::json!({
418 "name": "norm",
419 })),
420 None,
421 None,
422 Some(serde_json::json!({
423 "id": 1,
424 "parent_id": 0,
425 "child_ids": [2, 3, 4, 5]
426 })),
427 ]);
428
429 let (input_metadata_array, input_value_array) = {
430 let metadata = expected_variant_array.metadata_field().clone();
431 let value = expected_variant_array.value_field().unwrap().clone();
432
433 (metadata, value)
434 };
435
436 let udf = CastToVariantUdf::default();
437
438 let metadata_field = Arc::new(Field::new("metadata", DataType::BinaryView, true));
439 let variant_field = Arc::new(Field::new("value", DataType::BinaryView, true));
440
441 let return_field = Arc::new(Field::new(
442 "res",
443 udf.return_type(&[DataType::BinaryView, DataType::BinaryView])
444 .unwrap(),
445 true,
446 ));
447
448 let args = ScalarFunctionArgs {
449 args: vec![
450 ColumnarValue::Array(Arc::new(input_metadata_array) as ArrayRef),
451 ColumnarValue::Array(Arc::new(input_value_array) as ArrayRef),
452 ],
453 return_field,
454 arg_fields: vec![metadata_field, variant_field],
455 number_rows: Default::default(),
456 config_options: Default::default(),
457 };
458
459 let res = udf.invoke_with_args(args).unwrap();
460
461 let ColumnarValue::Array(variant_array) = res else {
462 panic!("expected scalar value struct array")
463 };
464
465 let variant_array = VariantArray::try_new(variant_array.as_ref()).unwrap();
466
467 assert_eq!(&variant_array, &expected_variant_array);
468 }
469
470 #[test]
471 fn test_array_string_view() {
472 let udf = CastToVariantUdf::default();
473
474 let arg_field = Arc::new(Field::new("input", DataType::Utf8View, true));
475 let return_field = Arc::new(Field::new(
476 "res",
477 udf.return_type(&[DataType::Utf8View]).unwrap(),
478 true,
479 ));
480
481 let args = ScalarFunctionArgs {
482 args: vec![ColumnarValue::Array(Arc::new(StringViewArray::from(vec![
483 Some("short"),
484 None,
485 Some("another"),
486 ])) as ArrayRef)],
487 return_field,
488 arg_fields: vec![arg_field],
489 number_rows: Default::default(),
490 config_options: Default::default(),
491 };
492
493 let res = udf.invoke_with_args(args).unwrap();
494
495 let ColumnarValue::Array(arr) = res else {
496 panic!("expected array output")
497 };
498
499 let variant_array = VariantArray::try_new(arr.as_ref()).unwrap();
500
501 assert_eq!(variant_array.value(0), Variant::from("short"));
502 assert!(variant_array.is_null(1));
503 assert_eq!(variant_array.value(2), Variant::from("another"));
504 }
505}