parquet_variant_compute/
from_json.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Module for transforming a batch of JSON strings into a batch of Variants represented as
6//! STRUCT<metadata: BINARY, value: BINARY>
7
8use crate::{VariantArray, VariantArrayBuilder};
9use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray};
10use arrow_schema::ArrowError;
11use parquet_variant_json::JsonToVariant;
12
13/// Macro to convert string array to variant array
14macro_rules! string_array_to_variant {
15    ($input:expr, $array:expr, $builder:expr) => {{
16        for i in 0..$input.len() {
17            if $input.is_null(i) {
18                $builder.append_null();
19            } else {
20                $builder.append_json($array.value(i))?;
21            }
22        }
23    }};
24}
25
26/// Parse a batch of JSON strings into a batch of Variants represented as
27/// STRUCT<metadata: BINARY, value: BINARY> where nulls are preserved. The JSON strings in the input
28/// must be valid.
29///
30/// Supports the following string array types:
31/// - [`StringArray`]
32/// - [`LargeStringArray`]
33/// - [`StringViewArray`]
34pub fn json_to_variant(input: &ArrayRef) -> Result<VariantArray, ArrowError> {
35    let mut variant_array_builder = VariantArrayBuilder::new(input.len());
36
37    // Try each string array type in sequence
38    if let Some(string_array) = input.as_any().downcast_ref::<StringArray>() {
39        string_array_to_variant!(input, string_array, variant_array_builder);
40    } else if let Some(large_string_array) = input.as_any().downcast_ref::<LargeStringArray>() {
41        string_array_to_variant!(input, large_string_array, variant_array_builder);
42    } else if let Some(string_view_array) = input.as_any().downcast_ref::<StringViewArray>() {
43        string_array_to_variant!(input, string_view_array, variant_array_builder);
44    } else {
45        return Err(ArrowError::CastError(
46            "Expected reference to StringArray, LargeStringArray, or StringViewArray as input"
47                .into(),
48        ));
49    }
50
51    Ok(variant_array_builder.build())
52}
53
54#[cfg(test)]
55mod test {
56    use crate::json_to_variant;
57    use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray};
58    use arrow_schema::ArrowError;
59    use parquet_variant::{Variant, VariantBuilder};
60    use std::sync::Arc;
61
62    #[test]
63    fn test_json_to_variant() -> Result<(), ArrowError> {
64        let input = StringArray::from(vec![
65            Some("1"),
66            None,
67            Some("{\"a\": 32}"),
68            Some("null"),
69            None,
70        ]);
71        let array_ref: ArrayRef = Arc::new(input);
72        let variant_array = json_to_variant(&array_ref).unwrap();
73
74        let metadata_array = variant_array.metadata_column();
75        let value_array = variant_array.value_column();
76
77        // Compare row 0
78        assert!(!variant_array.is_null(0));
79        assert_eq!(variant_array.value(0), Variant::Int8(1));
80
81        // Compare row 1
82        assert!(variant_array.is_null(1));
83
84        // Compare row 2
85        assert!(!variant_array.is_null(2));
86        {
87            let mut vb = VariantBuilder::new();
88            let mut ob = vb.new_object();
89            ob.insert("a", Variant::Int8(32));
90            ob.finish();
91            let (object_metadata, object_value) = vb.finish();
92            let expected = Variant::new(&object_metadata, &object_value);
93            assert_eq!(variant_array.value(2), expected);
94        }
95
96        // Compare row 3 (Note this is a variant NULL, not a null row)
97        assert!(!variant_array.is_null(3));
98        assert_eq!(variant_array.value(3), Variant::Null);
99
100        // Compare row 4
101        assert!(variant_array.is_null(4));
102
103        // Ensure that the subfields are not nullable
104        assert!(!metadata_array.is_null(1));
105        assert!(!value_array.is_null(1));
106        assert!(!metadata_array.is_null(4));
107        assert!(!value_array.is_null(4));
108        Ok(())
109    }
110
111    #[test]
112    fn test_json_to_variant_large_string() -> Result<(), ArrowError> {
113        let input = LargeStringArray::from(vec![
114            Some("1"),
115            None,
116            Some("{\"a\": 32}"),
117            Some("null"),
118            None,
119        ]);
120        let array_ref: ArrayRef = Arc::new(input);
121        let variant_array = json_to_variant(&array_ref).unwrap();
122
123        let metadata_array = variant_array.metadata_column();
124        let value_array = variant_array.value_column();
125
126        // Compare row 0
127        assert!(!variant_array.is_null(0));
128        assert_eq!(variant_array.value(0), Variant::Int8(1));
129
130        // Compare row 1
131        assert!(variant_array.is_null(1));
132
133        // Compare row 2
134        assert!(!variant_array.is_null(2));
135        {
136            let mut vb = VariantBuilder::new();
137            let mut ob = vb.new_object();
138            ob.insert("a", Variant::Int8(32));
139            ob.finish();
140            let (object_metadata, object_value) = vb.finish();
141            let expected = Variant::new(&object_metadata, &object_value);
142            assert_eq!(variant_array.value(2), expected);
143        }
144
145        // Compare row 3 (Note this is a variant NULL, not a null row)
146        assert!(!variant_array.is_null(3));
147        assert_eq!(variant_array.value(3), Variant::Null);
148
149        // Compare row 4
150        assert!(variant_array.is_null(4));
151
152        // Ensure that the subfields are not nullable
153        assert!(!metadata_array.is_null(1));
154        assert!(!value_array.is_null(1));
155        assert!(!metadata_array.is_null(4));
156        assert!(!value_array.is_null(4));
157        Ok(())
158    }
159
160    #[test]
161    fn test_json_to_variant_string_view() -> Result<(), ArrowError> {
162        let input = StringViewArray::from(vec![
163            Some("1"),
164            None,
165            Some("{\"a\": 32}"),
166            Some("null"),
167            None,
168        ]);
169        let array_ref: ArrayRef = Arc::new(input);
170        let variant_array = json_to_variant(&array_ref).unwrap();
171
172        let metadata_array = variant_array.metadata_column();
173        let value_array = variant_array.value_column();
174
175        // Compare row 0
176        assert!(!variant_array.is_null(0));
177        assert_eq!(variant_array.value(0), Variant::Int8(1));
178
179        // Compare row 1
180        assert!(variant_array.is_null(1));
181
182        // Compare row 2
183        assert!(!variant_array.is_null(2));
184        {
185            let mut vb = VariantBuilder::new();
186            let mut ob = vb.new_object();
187            ob.insert("a", Variant::Int8(32));
188            ob.finish();
189            let (object_metadata, object_value) = vb.finish();
190            let expected = Variant::new(&object_metadata, &object_value);
191            assert_eq!(variant_array.value(2), expected);
192        }
193
194        // Compare row 3 (Note this is a variant NULL, not a null row)
195        assert!(!variant_array.is_null(3));
196        assert_eq!(variant_array.value(3), Variant::Null);
197
198        // Compare row 4
199        assert!(variant_array.is_null(4));
200
201        // Ensure that the subfields are not nullable
202        assert!(!metadata_array.is_null(1));
203        assert!(!value_array.is_null(1));
204        assert!(!metadata_array.is_null(4));
205        assert!(!value_array.is_null(4));
206        Ok(())
207    }
208}
parquet_variant_compute/from_json.rs

parquet_variant_compute/
from_json.rs
