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
19pub(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
35pub(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 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
104impl 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 #[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}