parquet_variant_compute/
from_json.rs1use crate::{VariantArray, VariantArrayBuilder};
22use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray};
23use arrow_schema::ArrowError;
24use parquet_variant_json::JsonToVariant;
25
26macro_rules! string_array_to_variant {
28    ($input:expr, $array:expr, $builder:expr) => {{
29        for i in 0..$input.len() {
30            if $input.is_null(i) {
31                $builder.append_null();
32            } else {
33                $builder.append_json($array.value(i))?;
34            }
35        }
36    }};
37}
38
39pub fn json_to_variant(input: &ArrayRef) -> Result<VariantArray, ArrowError> {
48    let mut variant_array_builder = VariantArrayBuilder::new(input.len());
49
50    if let Some(string_array) = input.as_any().downcast_ref::<StringArray>() {
52        string_array_to_variant!(input, string_array, variant_array_builder);
53    } else if let Some(large_string_array) = input.as_any().downcast_ref::<LargeStringArray>() {
54        string_array_to_variant!(input, large_string_array, variant_array_builder);
55    } else if let Some(string_view_array) = input.as_any().downcast_ref::<StringViewArray>() {
56        string_array_to_variant!(input, string_view_array, variant_array_builder);
57    } else {
58        return Err(ArrowError::CastError(
59            "Expected reference to StringArray, LargeStringArray, or StringViewArray as input"
60                .into(),
61        ));
62    }
63
64    Ok(variant_array_builder.build())
65}
66
67#[cfg(test)]
68mod test {
69    use crate::json_to_variant;
70    use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray};
71    use arrow_schema::ArrowError;
72    use parquet_variant::{Variant, VariantBuilder};
73    use std::sync::Arc;
74
75    #[test]
76    fn test_json_to_variant() -> Result<(), ArrowError> {
77        let input = StringArray::from(vec![
78            Some("1"),
79            None,
80            Some("{\"a\": 32}"),
81            Some("null"),
82            None,
83        ]);
84        let array_ref: ArrayRef = Arc::new(input);
85        let variant_array = json_to_variant(&array_ref).unwrap();
86
87        let metadata_array = variant_array.metadata_field();
88        let value_array = variant_array.value_field().expect("value field");
89
90        assert!(!variant_array.is_null(0));
92        assert_eq!(variant_array.value(0), Variant::Int8(1));
93
94        assert!(variant_array.is_null(1));
96
97        assert!(!variant_array.is_null(2));
99        {
100            let mut vb = VariantBuilder::new();
101            let mut ob = vb.new_object();
102            ob.insert("a", Variant::Int8(32));
103            ob.finish();
104            let (object_metadata, object_value) = vb.finish();
105            let expected = Variant::new(&object_metadata, &object_value);
106            assert_eq!(variant_array.value(2), expected);
107        }
108
109        assert!(!variant_array.is_null(3));
111        assert_eq!(variant_array.value(3), Variant::Null);
112
113        assert!(variant_array.is_null(4));
115
116        assert!(!metadata_array.is_null(1));
118        assert!(!value_array.is_null(1));
119        assert!(!metadata_array.is_null(4));
120        assert!(!value_array.is_null(4));
121        Ok(())
122    }
123
124    #[test]
125    fn test_json_to_variant_large_string() -> Result<(), ArrowError> {
126        let input = LargeStringArray::from(vec![
127            Some("1"),
128            None,
129            Some("{\"a\": 32}"),
130            Some("null"),
131            None,
132        ]);
133        let array_ref: ArrayRef = Arc::new(input);
134        let variant_array = json_to_variant(&array_ref).unwrap();
135
136        let metadata_array = variant_array.metadata_field();
137        let value_array = variant_array.value_field().expect("value field");
138
139        assert!(!variant_array.is_null(0));
141        assert_eq!(variant_array.value(0), Variant::Int8(1));
142
143        assert!(variant_array.is_null(1));
145
146        assert!(!variant_array.is_null(2));
148        {
149            let mut vb = VariantBuilder::new();
150            let mut ob = vb.new_object();
151            ob.insert("a", Variant::Int8(32));
152            ob.finish();
153            let (object_metadata, object_value) = vb.finish();
154            let expected = Variant::new(&object_metadata, &object_value);
155            assert_eq!(variant_array.value(2), expected);
156        }
157
158        assert!(!variant_array.is_null(3));
160        assert_eq!(variant_array.value(3), Variant::Null);
161
162        assert!(variant_array.is_null(4));
164
165        assert!(!metadata_array.is_null(1));
167        assert!(!value_array.is_null(1));
168        assert!(!metadata_array.is_null(4));
169        assert!(!value_array.is_null(4));
170        Ok(())
171    }
172
173    #[test]
174    fn test_json_to_variant_string_view() -> Result<(), ArrowError> {
175        let input = StringViewArray::from(vec![
176            Some("1"),
177            None,
178            Some("{\"a\": 32}"),
179            Some("null"),
180            None,
181        ]);
182        let array_ref: ArrayRef = Arc::new(input);
183        let variant_array = json_to_variant(&array_ref).unwrap();
184
185        let metadata_array = variant_array.metadata_field();
186        let value_array = variant_array.value_field().expect("value field");
187
188        assert!(!variant_array.is_null(0));
190        assert_eq!(variant_array.value(0), Variant::Int8(1));
191
192        assert!(variant_array.is_null(1));
194
195        assert!(!variant_array.is_null(2));
197        {
198            let mut vb = VariantBuilder::new();
199            let mut ob = vb.new_object();
200            ob.insert("a", Variant::Int8(32));
201            ob.finish();
202            let (object_metadata, object_value) = vb.finish();
203            let expected = Variant::new(&object_metadata, &object_value);
204            assert_eq!(variant_array.value(2), expected);
205        }
206
207        assert!(!variant_array.is_null(3));
209        assert_eq!(variant_array.value(3), Variant::Null);
210
211        assert!(variant_array.is_null(4));
213
214        assert!(!metadata_array.is_null(1));
216        assert!(!value_array.is_null(1));
217        assert!(!metadata_array.is_null(4));
218        assert!(!value_array.is_null(4));
219        Ok(())
220    }
221}