1use std::sync::Arc;
2
3use arrow::{
4 array::{Array, ArrayRef, StructArray},
5 compute::concat,
6};
7use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields};
8use datafusion::{
9 common::{arrow_datafusion_err, exec_datafusion_err, exec_err},
10 error::{DataFusionError, Result},
11 logical_expr::{
12 ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
13 },
14 scalar::ScalarValue,
15};
16use parquet_variant::VariantPath;
17use parquet_variant_compute::{GetOptions, VariantArray, VariantType, variant_get};
18
19use crate::shared::{
20 try_field_as_variant_array, try_parse_string_columnar, try_parse_string_scalar,
21};
22
23fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<FieldRef> {
24 let type_name = match scalar {
25 ScalarValue::Utf8(Some(value))
26 | ScalarValue::Utf8View(Some(value))
27 | ScalarValue::LargeUtf8(Some(value)) => value.as_str(),
28 other => {
29 return exec_err!(
30 "type hint must be a non-null UTF8 literal, got {}",
31 other.data_type()
32 );
33 }
34 };
35
36 let casted_type = match type_name.parse::<DataType>() {
37 Ok(data_type) => Ok(data_type),
38 Err(ArrowError::ParseError(e)) => Err(exec_datafusion_err!("{e}")),
39 Err(e) => Err(arrow_datafusion_err!(e)),
40 }?;
41
42 Ok(Arc::new(Field::new(field_name, casted_type, true)))
43}
44
45fn type_hint_from_value(field_name: &str, arg: &ColumnarValue) -> Result<FieldRef> {
46 match arg {
47 ColumnarValue::Scalar(value) => type_hint_from_scalar(field_name, value),
48 ColumnarValue::Array(_) => {
49 exec_err!("type hint argument must be a scalar UTF8 literal")
50 }
51 }
52}
53
54fn build_get_options<'a>(path: VariantPath<'a>, as_type: &Option<FieldRef>) -> GetOptions<'a> {
55 match as_type {
56 Some(field) => GetOptions::new_with_path(path).with_as_type(Some(field.clone())),
57 None => GetOptions::new_with_path(path),
58 }
59}
60
61#[derive(Debug, Hash, PartialEq, Eq)]
62pub struct VariantGetUdf {
63 signature: Signature,
64}
65
66impl Default for VariantGetUdf {
67 fn default() -> Self {
68 Self {
69 signature: Signature::new(
70 TypeSignature::OneOf(vec![TypeSignature::Any(2), TypeSignature::Any(3)]),
71 Volatility::Immutable,
72 ),
73 }
74 }
75}
76
77impl ScalarUDFImpl for VariantGetUdf {
78 fn as_any(&self) -> &dyn std::any::Any {
79 self
80 }
81
82 fn name(&self) -> &str {
83 "variant_get"
84 }
85
86 fn signature(&self) -> &Signature {
87 &self.signature
88 }
89
90 fn return_type(&self, _arg_types: &[arrow_schema::DataType]) -> Result<arrow_schema::DataType> {
91 Err(DataFusionError::Internal(
92 "implemented return_field_from_args instead".into(),
93 ))
94 }
95
96 fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Arc<Field>> {
97 if let Some(maybe_scalar) = args.scalar_arguments.get(2) {
98 let scalar = maybe_scalar.ok_or_else(|| {
99 exec_datafusion_err!("type hint argument to variant_get must be a literal")
100 })?;
101 return type_hint_from_scalar(self.name(), scalar);
102 }
103
104 let data_type = DataType::Struct(Fields::from(vec![
105 Field::new("metadata", DataType::BinaryView, false),
106 Field::new("value", DataType::BinaryView, true),
107 ]));
108
109 Ok(Arc::new(
110 Field::new(self.name(), data_type, true).with_extension_type(VariantType),
111 ))
112 }
113
114 fn invoke_with_args(
115 &self,
116 args: datafusion::logical_expr::ScalarFunctionArgs,
117 ) -> Result<ColumnarValue> {
118 let (variant_arg, variant_path, type_arg) = match args.args.as_slice() {
119 [variant_arg, variant_path] => (variant_arg, variant_path, None),
120 [variant_arg, variant_path, type_arg] => (variant_arg, variant_path, Some(type_arg)),
121 _ => return exec_err!("expected 2 or 3 arguments"),
122 };
123
124 let variant_field = args
125 .arg_fields
126 .first()
127 .ok_or_else(|| exec_datafusion_err!("expected argument field"))?;
128
129 try_field_as_variant_array(variant_field.as_ref())?;
130
131 let type_field = type_arg
132 .map(|arg| type_hint_from_value(self.name(), arg))
133 .transpose()?;
134
135 let out = match (variant_arg, variant_path) {
136 (ColumnarValue::Array(variant_array), ColumnarValue::Scalar(variant_path)) => {
137 let variant_path = try_parse_string_scalar(variant_path)?
138 .map(|s| s.as_str())
139 .unwrap_or_default();
140
141 let res = variant_get(
142 variant_array,
143 build_get_options(VariantPath::from(variant_path), &type_field),
144 )?;
145
146 ColumnarValue::Array(res)
147 }
148 (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Scalar(variant_path)) => {
149 let ScalarValue::Struct(variant_array) = scalar_variant else {
150 return exec_err!("expected struct array");
151 };
152
153 let variant_array = Arc::clone(variant_array) as ArrayRef;
154
155 let variant_path = try_parse_string_scalar(variant_path)?
156 .map(|s| s.as_str())
157 .unwrap_or_default();
158
159 let res = variant_get(
160 &variant_array,
161 build_get_options(VariantPath::from(variant_path), &type_field),
162 )?;
163
164 let scalar = ScalarValue::try_from_array(res.as_ref(), 0)?;
165 ColumnarValue::Scalar(scalar)
166 }
167 (ColumnarValue::Array(variant_array), ColumnarValue::Array(variant_paths)) => {
168 if variant_array.len() != variant_paths.len() {
169 return exec_err!(
170 "expected variant_array and variant paths to be of same length"
171 );
172 }
173
174 let variant_paths = try_parse_string_columnar(variant_paths)?;
175 let variant_array = VariantArray::try_new(variant_array.as_ref())?;
176
177 let mut out = Vec::with_capacity(variant_array.len());
178
179 for (i, path) in variant_paths.iter().enumerate() {
180 let v = variant_array.value(i);
181 let singleton_variant_array: StructArray = VariantArray::from_iter([v]).into();
183
184 let arr = Arc::new(singleton_variant_array) as ArrayRef;
185
186 let res = variant_get(
187 &arr,
188 build_get_options(VariantPath::from(path.unwrap_or_default()), &type_field),
189 )?;
190
191 out.push(res);
192 }
193
194 let out_refs: Vec<&dyn Array> = out.iter().map(|a| a.as_ref()).collect();
195 ColumnarValue::Array(concat(&out_refs)?)
196 }
197 (ColumnarValue::Scalar(scalar_variant), ColumnarValue::Array(variant_paths)) => {
198 let ScalarValue::Struct(variant_array) = scalar_variant else {
199 return exec_err!("expected struct array");
200 };
201
202 let variant_array = Arc::clone(variant_array) as ArrayRef;
203 let variant_paths = try_parse_string_columnar(variant_paths)?;
204
205 let mut out = Vec::with_capacity(variant_paths.len());
206
207 for path in variant_paths {
208 let path = path.unwrap_or_default();
209 let res = variant_get(
210 &variant_array,
211 build_get_options(VariantPath::from(path), &type_field),
212 )?;
213
214 out.push(res);
215 }
216
217 let out_refs: Vec<&dyn Array> = out.iter().map(|a| a.as_ref()).collect();
218 ColumnarValue::Array(concat(&out_refs)?)
219 }
220 };
221
222 Ok(out)
223 }
224}
225
226#[cfg(test)]
227mod tests {
228 use arrow::array::{Array, BinaryViewArray, Int64Array};
229 use arrow_schema::{Field, Fields};
230 use datafusion::logical_expr::{ReturnFieldArgs, ScalarFunctionArgs};
231 use parquet_variant::Variant;
232 use parquet_variant_compute::{VariantArrayBuilder, VariantType};
233 use parquet_variant_json::JsonToVariant;
234
235 use super::*;
236
237 fn variant_scalar_from_json(json: serde_json::Value) -> ScalarValue {
238 let mut builder = VariantArrayBuilder::new(1);
239 builder.append_json(json.to_string().as_str()).unwrap();
240 ScalarValue::Struct(Arc::new(builder.build().into()))
241 }
242
243 fn variant_array_from_json_rows(json_rows: &[serde_json::Value]) -> ArrayRef {
244 let mut builder = VariantArrayBuilder::new(json_rows.len());
245 for value in json_rows {
246 builder.append_json(value.to_string().as_str()).unwrap();
247 }
248 let variant_array: StructArray = builder.build().into();
249 Arc::new(variant_array) as ArrayRef
250 }
251
252 fn standard_arg_fields(with_type_hint: bool) -> Vec<FieldRef> {
253 let mut fields = vec![
254 Arc::new(
255 Field::new("input", DataType::Struct(Fields::empty()), true)
256 .with_extension_type(VariantType),
257 ),
258 Arc::new(Field::new("path", DataType::Utf8, true)),
259 ];
260 if with_type_hint {
261 fields.push(Arc::new(Field::new("type", DataType::Utf8, true)));
262 }
263 fields
264 }
265
266 fn get_return_field(
267 udf: &VariantGetUdf,
268 arg_fields: &[FieldRef],
269 type_hint_value: Option<&ScalarValue>,
270 ) -> FieldRef {
271 let scalar_arguments: Vec<Option<&ScalarValue>> = if let Some(hint) = type_hint_value {
272 vec![None, None, Some(hint)]
273 } else {
274 vec![]
275 };
276
277 udf.return_field_from_args(ReturnFieldArgs {
278 arg_fields,
279 scalar_arguments: &scalar_arguments,
280 })
281 .unwrap()
282 }
283
284 fn build_scalar_function_args(
285 variant_input: ColumnarValue,
286 path: &str,
287 arg_fields: Vec<FieldRef>,
288 return_field: FieldRef,
289 type_hint: Option<ScalarValue>,
290 ) -> ScalarFunctionArgs {
291 let mut args = vec![
292 variant_input,
293 ColumnarValue::Scalar(ScalarValue::Utf8(Some(path.to_string()))),
294 ];
295 if let Some(hint) = type_hint {
296 args.push(ColumnarValue::Scalar(hint));
297 }
298
299 ScalarFunctionArgs {
300 args,
301 return_field,
302 arg_fields,
303 number_rows: Default::default(),
304 config_options: Default::default(),
305 }
306 }
307
308 #[test]
309 fn test_get_variant_scalar() {
310 let variant_input = variant_scalar_from_json(serde_json::json!({
311 "name": "norm",
312 "age": 50,
313 "list": [false, true, ()]
314 }));
315
316 let udf = VariantGetUdf::default();
317 let arg_fields = standard_arg_fields(false);
318 let return_field = get_return_field(&udf, &arg_fields, None);
319
320 let args = build_scalar_function_args(
321 ColumnarValue::Scalar(variant_input),
322 "name",
323 arg_fields,
324 return_field,
325 None,
326 );
327
328 let result = udf.invoke_with_args(args).unwrap();
329
330 let ColumnarValue::Scalar(ScalarValue::Struct(struct_arr)) = result else {
331 panic!("expected ScalarValue struct");
332 };
333
334 assert_eq!(struct_arr.len(), 1);
335
336 let metadata_arr = struct_arr
337 .column(0)
338 .as_any()
339 .downcast_ref::<BinaryViewArray>()
340 .unwrap();
341 let value_arr = struct_arr
342 .column(1)
343 .as_any()
344 .downcast_ref::<BinaryViewArray>()
345 .unwrap();
346
347 let metadata = metadata_arr.value(0);
348 let value = value_arr.value(0);
349 let v = Variant::try_new(metadata, value).unwrap();
350
351 assert_eq!(v, Variant::from("norm"))
352 }
353
354 #[test]
355 fn test_return_field_with_type_hint() {
356 let udf = VariantGetUdf::default();
357 let arg_fields = standard_arg_fields(true);
358 let type_hint = ScalarValue::Utf8(Some("Int64".to_string()));
359 let return_field = get_return_field(&udf, &arg_fields, Some(&type_hint));
360
361 assert_eq!(return_field.data_type(), &DataType::Int64);
362 }
363
364 #[test]
365 fn test_get_variant_scalar_with_type_hint() {
366 let variant_input = variant_scalar_from_json(serde_json::json!({
367 "name": "norm",
368 "age": 50,
369 }));
370
371 let udf = VariantGetUdf::default();
372 let arg_fields = standard_arg_fields(true);
373 let type_hint = ScalarValue::Utf8(Some("Int64".to_string()));
374 let return_field = get_return_field(&udf, &arg_fields, Some(&type_hint));
375
376 let args = build_scalar_function_args(
377 ColumnarValue::Scalar(variant_input),
378 "age",
379 arg_fields,
380 return_field,
381 Some(type_hint),
382 );
383
384 let result = udf.invoke_with_args(args).unwrap();
385
386 let ColumnarValue::Scalar(ScalarValue::Int64(Some(value))) = result else {
387 panic!("expected ScalarValue Int64");
388 };
389
390 assert_eq!(value, 50);
391 }
392
393 #[test]
394 fn test_get_variant_array_with_type_hint() {
395 let json_rows = vec![
396 serde_json::json!({ "age": 50 }),
397 serde_json::json!({ "age": 60 }),
398 ];
399
400 let variant_array = variant_array_from_json_rows(&json_rows);
401
402 let udf = VariantGetUdf::default();
403 let arg_fields = standard_arg_fields(true);
404 let type_hint = ScalarValue::Utf8(Some("Int64".to_string()));
405 let return_field = get_return_field(&udf, &arg_fields, Some(&type_hint));
406
407 let args = build_scalar_function_args(
408 ColumnarValue::Array(variant_array),
409 "age",
410 arg_fields,
411 return_field,
412 Some(type_hint),
413 );
414
415 let result = udf.invoke_with_args(args).unwrap();
416
417 let ColumnarValue::Array(array) = result else {
418 panic!("expected array output");
419 };
420
421 let values = array.as_any().downcast_ref::<Int64Array>().unwrap();
422 assert_eq!(values.len(), 2);
423 assert_eq!(values.value(0), 50);
424 assert_eq!(values.value(1), 60);
425 }
426}