datafusion_variant/
variant_pretty.rs

1use std::sync::Arc;
2
3use arrow::array::StringViewArray;
4use arrow_schema::DataType;
5use datafusion::{
6    common::{exec_datafusion_err, exec_err},
7    error::Result,
8    logical_expr::{
9        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
10    },
11    scalar::ScalarValue,
12};
13use parquet_variant_compute::VariantArray;
14
15use crate::shared::try_field_as_variant_array;
16
17#[derive(Debug, Hash, PartialEq, Eq)]
18pub struct VariantPretty {
19    signature: Signature,
20}
21
22impl Default for VariantPretty {
23    fn default() -> Self {
24        Self {
25            signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable),
26        }
27    }
28}
29
30impl ScalarUDFImpl for VariantPretty {
31    fn as_any(&self) -> &dyn std::any::Any {
32        self
33    }
34
35    fn name(&self) -> &str {
36        "variant_pretty"
37    }
38
39    fn signature(&self) -> &Signature {
40        &self.signature
41    }
42
43    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
44        Ok(DataType::Utf8View)
45    }
46
47    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
48        let field = args
49            .arg_fields
50            .first()
51            .ok_or_else(|| exec_datafusion_err!("empty argument, expected 1 argument"))?;
52
53        try_field_as_variant_array(field.as_ref())?;
54
55        let arg = args
56            .args
57            .first()
58            .ok_or_else(|| exec_datafusion_err!("empty argument, expected 1 argument"))?;
59
60        let out = match arg {
61            ColumnarValue::Scalar(scalar) => {
62                let ScalarValue::Struct(variant_array) = scalar else {
63                    return exec_err!("Unsupported data type: {}", scalar.data_type());
64                };
65
66                let variant_array = VariantArray::try_new(variant_array.as_ref())?;
67                let v = variant_array.value(0);
68
69                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(format!("{:?}", v))))
70            }
71            ColumnarValue::Array(arr) => match arr.data_type() {
72                DataType::Struct(_) => {
73                    let variant_array = VariantArray::try_new(arr.as_ref())?;
74
75                    let out = variant_array
76                        .iter()
77                        .map(|v| v.map(|v| format!("{:?}", v)))
78                        .collect::<Vec<_>>();
79
80                    let out: StringViewArray = out.into();
81
82                    ColumnarValue::Array(Arc::new(out))
83                }
84                unsupported => return exec_err!("Invalid data type: {unsupported}"),
85            },
86        };
87
88        Ok(out)
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use arrow_schema::{Field, Fields};
95    use insta::assert_snapshot;
96    use parquet_variant_compute::VariantType;
97
98    use crate::shared::build_variant_array_from_json;
99
100    use super::*;
101
102    #[test]
103    fn test_pretty_primitive_scalar() {
104        let expected_json = serde_json::json!("norm");
105        let input = build_variant_array_from_json(&expected_json);
106
107        let variant_input = ScalarValue::Struct(Arc::new(input.into()));
108
109        let udf = VariantPretty::default();
110        let return_field = Arc::new(Field::new("result", DataType::Utf8View, true));
111        let arg_field = Arc::new(
112            Field::new("input", DataType::Struct(Fields::empty()), true)
113                .with_extension_type(VariantType),
114        );
115
116        let args = ScalarFunctionArgs {
117            args: vec![ColumnarValue::Scalar(variant_input)],
118            return_field,
119            arg_fields: vec![arg_field],
120            number_rows: Default::default(),
121            config_options: Default::default(),
122        };
123
124        let result = udf.invoke_with_args(args).unwrap();
125
126        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(j))) = result else {
127            panic!("expected valid json string")
128        };
129
130        assert_snapshot!(j, @r#"ShortString(ShortString("norm"))"#);
131    }
132
133    #[test]
134    fn test_pretty_variant_object() {
135        let expected_json = serde_json::json!({
136            "name": "norm",
137            "age": 50,
138            "list": [false, true, ()]
139        });
140
141        let input = build_variant_array_from_json(&expected_json);
142
143        let variant_input = ScalarValue::Struct(Arc::new(input.into()));
144
145        let udf = VariantPretty::default();
146
147        let return_field = Arc::new(Field::new("result", DataType::Utf8View, true));
148        let arg_field = Arc::new(
149            Field::new("input", DataType::Struct(Fields::empty()), true)
150                .with_extension_type(VariantType),
151        );
152
153        let args = ScalarFunctionArgs {
154            args: vec![ColumnarValue::Scalar(variant_input)],
155            return_field,
156            arg_fields: vec![arg_field],
157            number_rows: Default::default(),
158            config_options: Default::default(),
159        };
160
161        let result = udf.invoke_with_args(args).unwrap();
162
163        let ColumnarValue::Scalar(ScalarValue::Utf8View(Some(j))) = result else {
164            panic!("expected valid json string")
165        };
166
167        assert_snapshot!(j, @r#"{"age": Int8(50), "list": [BooleanFalse, BooleanTrue, Null], "name": ShortString(ShortString("norm"))}"#);
168    }
169}
datafusion_variant/variant_pretty.rs

datafusion_variant/
variant_pretty.rs
