datafusion_variant/
variant_to_json.rs

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

datafusion_variant/
variant_to_json.rs
