parquet_variant_compute/
to_json.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Module for transforming a batch of Variants into JSON strings.
19
20use crate::VariantArray;
21use arrow::array::{ArrayRef, BooleanBufferBuilder, StringArray};
22use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
23use arrow_schema::ArrowError;
24use parquet_variant_json::VariantToJson;
25
26/// Transform a [`VariantArray`] to a batch of JSON strings where nulls are preserved.
27pub fn variant_to_json(input: &ArrayRef) -> Result<StringArray, ArrowError> {
28    let array = VariantArray::try_new(input)?;
29
30    // Zero-copy builder
31    // The size per JSON string is assumed to be 128 bytes. If this holds true, resizing could be
32    // minimized for performance.
33    let mut json_buffer: Vec<u8> = Vec::with_capacity(array.len() * 128);
34    let mut offsets: Vec<i32> = Vec::with_capacity(array.len() + 1);
35    let mut validity = BooleanBufferBuilder::new(array.len());
36    let mut current_offset: i32 = 0;
37    offsets.push(current_offset);
38
39    for i in 0..array.len() {
40        if array.is_null(i) {
41            validity.append(false);
42            offsets.push(current_offset);
43        } else {
44            let variant = array.try_value(i)?;
45            let start_len = json_buffer.len();
46            variant.to_json(&mut json_buffer)?;
47            let written = (json_buffer.len() - start_len) as i32;
48            current_offset += written;
49            offsets.push(current_offset);
50            validity.append(true);
51        }
52    }
53
54    let offsets_buffer = OffsetBuffer::new(ScalarBuffer::from(offsets));
55    let value_buffer = Buffer::from_vec(json_buffer);
56    let null_buffer = NullBuffer::new(validity.finish());
57
58    StringArray::try_new(offsets_buffer, value_buffer, Some(null_buffer))
59}
60
61#[cfg(test)]
62mod test {
63    use crate::variant_to_json;
64    use arrow::array::{Array, ArrayRef, BinaryBuilder, BooleanBufferBuilder, StructArray};
65    use arrow::buffer::NullBuffer;
66    use arrow::datatypes::DataType;
67    use arrow::datatypes::Field;
68    use arrow_schema::Fields;
69    use std::sync::Arc;
70
71    #[test]
72    fn test_variant_to_json() {
73        let mut metadata_builder = BinaryBuilder::new();
74        let mut value_builder = BinaryBuilder::new();
75
76        // Row 0: [1, 0, 0], [12, 0]
77        metadata_builder.append_value([1, 0, 0]);
78        value_builder.append_value([12, 0]);
79
80        // Row 1: null
81        metadata_builder.append_null();
82        value_builder.append_null();
83
84        // Row 2: [1, 1, 0, 1, 97], [2, 1, 0, 0, 1, 32]
85        metadata_builder.append_value([1, 1, 0, 1, 97]);
86        value_builder.append_value([2, 1, 0, 0, 2, 12, 32]);
87
88        // Row 3: [1, 0, 0], [0]
89        metadata_builder.append_value([1, 0, 0]);
90        value_builder.append_value([0]);
91
92        // Row 4: null
93        metadata_builder.append_null();
94        value_builder.append_null();
95
96        let metadata_array = Arc::new(metadata_builder.finish()) as ArrayRef;
97        let value_array = Arc::new(value_builder.finish()) as ArrayRef;
98
99        let fields: Fields = vec![
100            Field::new("metadata", DataType::Binary, true),
101            Field::new("value", DataType::Binary, true),
102        ]
103        .into();
104
105        let mut validity = BooleanBufferBuilder::new(value_array.len());
106        for i in 0..value_array.len() {
107            let is_valid = value_array.is_valid(i) && metadata_array.is_valid(i);
108            validity.append(is_valid);
109        }
110        let null_buffer = NullBuffer::new(validity.finish());
111
112        let struct_array = StructArray::new(
113            fields,
114            vec![metadata_array.clone(), value_array.clone()],
115            Some(null_buffer), // Null bitmap (let Arrow infer from children)
116        );
117
118        let input = Arc::new(struct_array) as ArrayRef;
119
120        let result = variant_to_json(&input).unwrap();
121
122        // Expected output: ["0", null, "{\"a\":32}", "null", null]
123        let expected = vec![Some("0"), None, Some("{\"a\":32}"), Some("null"), None];
124
125        let result_vec: Vec<Option<&str>> = (0..result.len())
126            .map(|i| {
127                if result.is_null(i) {
128                    None
129                } else {
130                    Some(result.value(i))
131                }
132            })
133            .collect();
134
135        assert_eq!(result_vec, expected);
136    }
137}
parquet_variant_compute/to_json.rs

parquet_variant_compute/
to_json.rs
