datafusion_functions_json/
json_union_to_text.rs

1use std::sync::Arc;
2
3use datafusion::arrow::array::{Array, ArrayRef, StringViewBuilder, UnionArray};
4use datafusion::arrow::datatypes::{DataType, Field, FieldRef};
5use datafusion::common::{exec_datafusion_err, exec_err, plan_err, Result as DataFusionResult};
6use datafusion::logical_expr::{
7    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
8};
9
10use crate::common_macros::make_udf_function;
11use crate::common_union::{is_json_union, json_field_metadata, JsonUnionEncoder, JsonUnionValue, JSON_UNION_DATA_TYPE};
12
13make_udf_function!(
14    JsonUnionToText,
15    json_union_to_text,
16    json_union,
17    "Flatten a JSON union value (produced by `json_get`) into its canonical JSON text"
18);
19
20/// Flattens the heterogeneous JSON union that `json_get` produces into a single
21/// `Utf8View` column of canonical JSON text: scalars render as `true` / `42` /
22/// `1.5`, strings are JSON-quoted and escaped, and array/object arms (already raw
23/// JSON text) pass through. A JSON `null` arm becomes a SQL `NULL`.
24///
25/// Useful when a JSON-union column must be materialized somewhere that can't
26/// represent an Arrow `Union` — e.g. the Parquet writer, which rejects unions
27/// (`arrow_to_parquet_schema` panics with "See ARROW-8817.").
28#[derive(Debug, PartialEq, Eq, Hash)]
29pub(super) struct JsonUnionToText {
30    signature: Signature,
31    aliases: [String; 1],
32}
33
34impl Default for JsonUnionToText {
35    fn default() -> Self {
36        Self {
37            // Exactly the JSON union — any other argument type is a planning error.
38            signature: Signature::exact(vec![JSON_UNION_DATA_TYPE.clone()], Volatility::Immutable),
39            aliases: ["json_union_to_text".to_string()],
40        }
41    }
42}
43
44impl ScalarUDFImpl for JsonUnionToText {
45    fn name(&self) -> &str {
46        self.aliases[0].as_str()
47    }
48
49    fn signature(&self) -> &Signature {
50        &self.signature
51    }
52
53    fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> {
54        match arg_types {
55            [t] if is_json_union(t) => Ok(DataType::Utf8View),
56            _ => plan_err!("json_union_to_text expects a single JSON-union argument, got {arg_types:?}"),
57        }
58    }
59
60    fn return_field_from_args(&self, args: ReturnFieldArgs) -> DataFusionResult<FieldRef> {
61        let arg_types: Vec<DataType> = args.arg_fields.iter().map(|f| f.data_type().clone()).collect();
62        let return_type = self.return_type(&arg_types)?;
63        Ok(Arc::new(
64            Field::new(self.name(), return_type, true).with_metadata(json_field_metadata()),
65        ))
66    }
67
68    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
69        let Some(arg) = args.args.into_iter().next() else {
70            return exec_err!("json_union_to_text expects one argument");
71        };
72        let array = arg.into_array(args.number_rows)?;
73        Ok(ColumnarValue::Array(json_union_to_text_array(&array)?))
74    }
75
76    fn aliases(&self) -> &[String] {
77        &self.aliases
78    }
79}
80
81/// Encode a JSON-union array into a `Utf8View` array of canonical JSON text.
82fn json_union_to_text_array(array: &ArrayRef) -> DataFusionResult<ArrayRef> {
83    let Some(union) = array.as_any().downcast_ref::<UnionArray>() else {
84        return exec_err!("json_union_to_text expects a UnionArray argument");
85    };
86    let Some(encoder) = JsonUnionEncoder::from_union(union.clone()) else {
87        return exec_err!("json_union_to_text argument is not the JSON union type");
88    };
89
90    let mut builder = StringViewBuilder::with_capacity(encoder.len());
91    // Scalar arms are JSON-encoded with serde_json (string escaping, float
92    // formatting, …); the array/object arms already hold raw JSON text and pass
93    // through verbatim.
94    let mut scratch: Vec<u8> = Vec::new();
95    for idx in 0..encoder.len() {
96        scratch.clear();
97        let write_result = match encoder.get_value(idx) {
98            JsonUnionValue::JsonNull => {
99                builder.append_null();
100                continue;
101            }
102            JsonUnionValue::Bool(b) => serde_json::to_writer(&mut scratch, &b),
103            JsonUnionValue::Int(i) => serde_json::to_writer(&mut scratch, &i),
104            JsonUnionValue::Float(f) => serde_json::to_writer(&mut scratch, &f),
105            JsonUnionValue::Str(s) => serde_json::to_writer(&mut scratch, s),
106            JsonUnionValue::Array(s) | JsonUnionValue::Object(s) => {
107                builder.append_value(s);
108                continue;
109            }
110        };
111        write_result.map_err(|e| exec_datafusion_err!("json_union_to_text: failed to encode JSON value: {e}"))?;
112        // `serde_json` always emits valid UTF-8.
113        let text = std::str::from_utf8(&scratch)
114            .map_err(|e| exec_datafusion_err!("json_union_to_text: encoded value was not UTF-8: {e}"))?;
115        builder.append_value(text);
116    }
117    Ok(Arc::new(builder.finish()))
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::common_union::{JsonUnion, JsonUnionField};
124    use datafusion::arrow::array::StringViewArray;
125
126    #[test]
127    fn flattens_each_arm_to_json_text() {
128        let union = JsonUnion::from_iter(vec![
129            Some(JsonUnionField::JsonNull),
130            Some(JsonUnionField::Bool(true)),
131            Some(JsonUnionField::Int(42)),
132            Some(JsonUnionField::Float(1.5)),
133            Some(JsonUnionField::Str("foo\"bar\n\u{1}".to_string())),
134            Some(JsonUnionField::Array("[1,2]".to_string())),
135            Some(JsonUnionField::Object(r#"{"a":1}"#.to_string())),
136            None,
137        ]);
138        let array: ArrayRef = Arc::new(UnionArray::try_from(union).unwrap());
139
140        let out = json_union_to_text_array(&array).unwrap();
141        let strings = out.as_any().downcast_ref::<StringViewArray>().unwrap();
142        let got: Vec<Option<&str>> = (0..strings.len())
143            .map(|i| (!strings.is_null(i)).then(|| strings.value(i)))
144            .collect();
145        assert_eq!(
146            got,
147            vec![
148                None,                             // JsonNull
149                Some("true"),                     // Bool
150                Some("42"),                       // Int
151                Some("1.5"),                      // Float
152                Some("\"foo\\\"bar\\n\\u0001\""), // Str: JSON-quoted + escaped (quote, newline, control char)
153                Some("[1,2]"),                    // Array (passthrough)
154                Some(r#"{"a":1}"#),               // Object (passthrough)
155                None,                             // None
156            ]
157        );
158    }
159
160    #[test]
161    fn output_field_is_marked_as_json() {
162        let udf = JsonUnionToText::default();
163        let arg = Arc::new(Field::new("j", JSON_UNION_DATA_TYPE.clone(), true));
164        let field = udf
165            .return_field_from_args(ReturnFieldArgs {
166                arg_fields: std::slice::from_ref(&arg),
167                scalar_arguments: &[],
168            })
169            .unwrap();
170        assert_eq!(field.data_type(), &DataType::Utf8View);
171        assert_eq!(
172            field.metadata().get("ARROW:extension:name").map(String::as_str),
173            Some("arrow.json")
174        );
175    }
176}
datafusion_functions_json/json_union_to_text.rs

datafusion_functions_json/
json_union_to_text.rs
