datafusion_functions_json/
common_union.rs

1use std::collections::HashMap;
2use std::sync::{Arc, OnceLock};
3
4use datafusion::arrow::array::{
5    Array, ArrayRef, AsArray, BooleanArray, Float64Array, Int64Array, NullArray, StringArray, UnionArray,
6};
7use datafusion::arrow::buffer::{Buffer, ScalarBuffer};
8use datafusion::arrow::datatypes::{DataType, Field, UnionFields, UnionMode};
9use datafusion::arrow::error::ArrowError;
10use datafusion::common::ScalarValue;
11
12pub fn is_json_union(data_type: &DataType) -> bool {
13    match data_type {
14        DataType::Union(fields, UnionMode::Sparse) => fields == &union_fields(),
15        _ => false,
16    }
17}
18
19/// Extract nested JSON from a `JsonUnion` `UnionArray`
20///
21/// # Arguments
22/// * `array` - The `UnionArray` to extract the nested JSON from
23/// * `object_lookup` - If `true`, extract from the "object" member of the union,
24///   otherwise extract from the "array" member
25pub(crate) fn nested_json_array(array: &ArrayRef, object_lookup: bool) -> Option<&StringArray> {
26    nested_json_array_ref(array, object_lookup).map(AsArray::as_string)
27}
28
29pub(crate) fn nested_json_array_ref(array: &ArrayRef, object_lookup: bool) -> Option<&ArrayRef> {
30    let union_array: &UnionArray = array.as_any().downcast_ref::<UnionArray>()?;
31    let type_id = if object_lookup { TYPE_ID_OBJECT } else { TYPE_ID_ARRAY };
32    Some(union_array.child(type_id))
33}
34
35/// Extract a JSON string from a `JsonUnion` scalar
36pub(crate) fn json_from_union_scalar<'a>(
37    type_id_value: Option<&'a (i8, Box<ScalarValue>)>,
38    fields: &UnionFields,
39) -> Option<&'a str> {
40    if let Some((type_id, value)) = type_id_value {
41        // we only want to take the ScalarValue string if the type_id indicates the value represents nested JSON
42        if fields == &union_fields() && (*type_id == TYPE_ID_ARRAY || *type_id == TYPE_ID_OBJECT) {
43            if let ScalarValue::Utf8(s) = value.as_ref() {
44                return s.as_deref();
45            }
46        }
47    }
48    None
49}
50
51#[derive(Debug)]
52pub(crate) struct JsonUnion {
53    bools: Vec<Option<bool>>,
54    ints: Vec<Option<i64>>,
55    floats: Vec<Option<f64>>,
56    strings: Vec<Option<String>>,
57    arrays: Vec<Option<String>>,
58    objects: Vec<Option<String>>,
59    type_ids: Vec<i8>,
60    index: usize,
61    length: usize,
62}
63
64impl JsonUnion {
65    pub fn new(length: usize) -> Self {
66        Self {
67            bools: vec![None; length],
68            ints: vec![None; length],
69            floats: vec![None; length],
70            strings: vec![None; length],
71            arrays: vec![None; length],
72            objects: vec![None; length],
73            type_ids: vec![TYPE_ID_NULL; length],
74            index: 0,
75            length,
76        }
77    }
78
79    pub fn data_type() -> DataType {
80        DataType::Union(union_fields(), UnionMode::Sparse)
81    }
82
83    pub fn push(&mut self, field: JsonUnionField) {
84        self.type_ids[self.index] = field.type_id();
85        match field {
86            JsonUnionField::JsonNull => (),
87            JsonUnionField::Bool(value) => self.bools[self.index] = Some(value),
88            JsonUnionField::Int(value) => self.ints[self.index] = Some(value),
89            JsonUnionField::Float(value) => self.floats[self.index] = Some(value),
90            JsonUnionField::Str(value) => self.strings[self.index] = Some(value),
91            JsonUnionField::Array(value) => self.arrays[self.index] = Some(value),
92            JsonUnionField::Object(value) => self.objects[self.index] = Some(value),
93        }
94        self.index += 1;
95        debug_assert!(self.index <= self.length);
96    }
97
98    pub fn push_none(&mut self) {
99        self.index += 1;
100        debug_assert!(self.index <= self.length);
101    }
102}
103
104/// So we can do `collect::<JsonUnion>()`
105impl FromIterator<Option<JsonUnionField>> for JsonUnion {
106    fn from_iter<I: IntoIterator<Item = Option<JsonUnionField>>>(iter: I) -> Self {
107        let inner = iter.into_iter();
108        let (lower, upper) = inner.size_hint();
109        let mut union = Self::new(upper.unwrap_or(lower));
110
111        for opt_field in inner {
112            if let Some(union_field) = opt_field {
113                union.push(union_field);
114            } else {
115                union.push_none();
116            }
117        }
118        union
119    }
120}
121
122impl TryFrom<JsonUnion> for UnionArray {
123    type Error = ArrowError;
124
125    fn try_from(value: JsonUnion) -> Result<Self, Self::Error> {
126        let children: Vec<Arc<dyn Array>> = vec![
127            Arc::new(NullArray::new(value.length)),
128            Arc::new(BooleanArray::from(value.bools)),
129            Arc::new(Int64Array::from(value.ints)),
130            Arc::new(Float64Array::from(value.floats)),
131            Arc::new(StringArray::from(value.strings)),
132            Arc::new(StringArray::from(value.arrays)),
133            Arc::new(StringArray::from(value.objects)),
134        ];
135        UnionArray::try_new(union_fields(), Buffer::from_vec(value.type_ids).into(), None, children)
136    }
137}
138
139#[derive(Debug)]
140pub(crate) enum JsonUnionField {
141    JsonNull,
142    Bool(bool),
143    Int(i64),
144    Float(f64),
145    Str(String),
146    Array(String),
147    Object(String),
148}
149
150pub(crate) const TYPE_ID_NULL: i8 = 0;
151const TYPE_ID_BOOL: i8 = 1;
152const TYPE_ID_INT: i8 = 2;
153const TYPE_ID_FLOAT: i8 = 3;
154const TYPE_ID_STR: i8 = 4;
155const TYPE_ID_ARRAY: i8 = 5;
156const TYPE_ID_OBJECT: i8 = 6;
157
158fn union_fields() -> UnionFields {
159    static FIELDS: OnceLock<UnionFields> = OnceLock::new();
160    FIELDS
161        .get_or_init(|| {
162            let json_metadata: HashMap<String, String> =
163                HashMap::from_iter(vec![("is_json".to_string(), "true".to_string())]);
164            UnionFields::from_iter([
165                (TYPE_ID_NULL, Arc::new(Field::new("null", DataType::Null, true))),
166                (TYPE_ID_BOOL, Arc::new(Field::new("bool", DataType::Boolean, false))),
167                (TYPE_ID_INT, Arc::new(Field::new("int", DataType::Int64, false))),
168                (TYPE_ID_FLOAT, Arc::new(Field::new("float", DataType::Float64, false))),
169                (TYPE_ID_STR, Arc::new(Field::new("str", DataType::Utf8, false))),
170                (
171                    TYPE_ID_ARRAY,
172                    Arc::new(Field::new("array", DataType::Utf8, false).with_metadata(json_metadata.clone())),
173                ),
174                (
175                    TYPE_ID_OBJECT,
176                    Arc::new(Field::new("object", DataType::Utf8, false).with_metadata(json_metadata.clone())),
177                ),
178            ])
179        })
180        .clone()
181}
182
183impl JsonUnionField {
184    fn type_id(&self) -> i8 {
185        match self {
186            Self::JsonNull => TYPE_ID_NULL,
187            Self::Bool(_) => TYPE_ID_BOOL,
188            Self::Int(_) => TYPE_ID_INT,
189            Self::Float(_) => TYPE_ID_FLOAT,
190            Self::Str(_) => TYPE_ID_STR,
191            Self::Array(_) => TYPE_ID_ARRAY,
192            Self::Object(_) => TYPE_ID_OBJECT,
193        }
194    }
195
196    pub fn scalar_value(f: Option<Self>) -> ScalarValue {
197        ScalarValue::Union(
198            f.map(|f| (f.type_id(), Box::new(f.into()))),
199            union_fields(),
200            UnionMode::Sparse,
201        )
202    }
203}
204
205impl From<JsonUnionField> for ScalarValue {
206    fn from(value: JsonUnionField) -> Self {
207        match value {
208            JsonUnionField::JsonNull => Self::Null,
209            JsonUnionField::Bool(b) => Self::Boolean(Some(b)),
210            JsonUnionField::Int(i) => Self::Int64(Some(i)),
211            JsonUnionField::Float(f) => Self::Float64(Some(f)),
212            JsonUnionField::Str(s) | JsonUnionField::Array(s) | JsonUnionField::Object(s) => Self::Utf8(Some(s)),
213        }
214    }
215}
216
217pub struct JsonUnionEncoder {
218    boolean: BooleanArray,
219    int: Int64Array,
220    float: Float64Array,
221    string: StringArray,
222    array: StringArray,
223    object: StringArray,
224    type_ids: ScalarBuffer<i8>,
225}
226
227impl JsonUnionEncoder {
228    #[must_use]
229    pub fn from_union(union: UnionArray) -> Option<Self> {
230        if is_json_union(union.data_type()) {
231            let (_, type_ids, _, c) = union.into_parts();
232            Some(Self {
233                boolean: c[1].as_boolean().clone(),
234                int: c[2].as_primitive().clone(),
235                float: c[3].as_primitive().clone(),
236                string: c[4].as_string().clone(),
237                array: c[5].as_string().clone(),
238                object: c[6].as_string().clone(),
239                type_ids,
240            })
241        } else {
242            None
243        }
244    }
245
246    #[must_use]
247    #[allow(clippy::len_without_is_empty)]
248    pub fn len(&self) -> usize {
249        self.type_ids.len()
250    }
251
252    /// Get the encodable value for a given index
253    ///
254    /// # Panics
255    ///
256    /// Panics if the idx is outside the union values or an invalid type id exists in the union.
257    #[must_use]
258    pub fn get_value(&self, idx: usize) -> JsonUnionValue {
259        let type_id = self.type_ids[idx];
260        match type_id {
261            TYPE_ID_NULL => JsonUnionValue::JsonNull,
262            TYPE_ID_BOOL => JsonUnionValue::Bool(self.boolean.value(idx)),
263            TYPE_ID_INT => JsonUnionValue::Int(self.int.value(idx)),
264            TYPE_ID_FLOAT => JsonUnionValue::Float(self.float.value(idx)),
265            TYPE_ID_STR => JsonUnionValue::Str(self.string.value(idx)),
266            TYPE_ID_ARRAY => JsonUnionValue::Array(self.array.value(idx)),
267            TYPE_ID_OBJECT => JsonUnionValue::Object(self.object.value(idx)),
268            _ => panic!("Invalid type_id: {type_id}, not a valid JSON type"),
269        }
270    }
271}
272
273#[derive(Debug, PartialEq)]
274pub enum JsonUnionValue<'a> {
275    JsonNull,
276    Bool(bool),
277    Int(i64),
278    Float(f64),
279    Str(&'a str),
280    Array(&'a str),
281    Object(&'a str),
282}
283
284#[cfg(test)]
285mod test {
286    use super::*;
287
288    #[test]
289    fn test_json_union() {
290        let json_union = JsonUnion::from_iter(vec![
291            Some(JsonUnionField::JsonNull),
292            Some(JsonUnionField::Bool(true)),
293            Some(JsonUnionField::Bool(false)),
294            Some(JsonUnionField::Int(42)),
295            Some(JsonUnionField::Float(42.0)),
296            Some(JsonUnionField::Str("foo".to_string())),
297            Some(JsonUnionField::Array("[42]".to_string())),
298            Some(JsonUnionField::Object(r#"{"foo": 42}"#.to_string())),
299            None,
300        ]);
301
302        let union_array = UnionArray::try_from(json_union).unwrap();
303        let encoder = JsonUnionEncoder::from_union(union_array).unwrap();
304
305        let values_after: Vec<_> = (0..encoder.len()).map(|idx| encoder.get_value(idx)).collect();
306        assert_eq!(
307            values_after,
308            vec![
309                JsonUnionValue::JsonNull,
310                JsonUnionValue::Bool(true),
311                JsonUnionValue::Bool(false),
312                JsonUnionValue::Int(42),
313                JsonUnionValue::Float(42.0),
314                JsonUnionValue::Str("foo"),
315                JsonUnionValue::Array("[42]"),
316                JsonUnionValue::Object(r#"{"foo": 42}"#),
317                JsonUnionValue::JsonNull,
318            ]
319        );
320    }
321}
datafusion_functions_json/common_union.rs

datafusion_functions_json/
common_union.rs
