parquet_variant_compute/
to_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 Variants into JSON strings.
6
7use crate::VariantArray;
8use arrow::array::{ArrayRef, BooleanBufferBuilder, StringArray};
9use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
10use arrow_schema::ArrowError;
11use parquet_variant_json::VariantToJson;
12
13/// Transform a [`VariantArray`] to a batch of JSON strings where nulls are preserved.
14pub fn variant_to_json(input: &ArrayRef) -> Result<StringArray, ArrowError> {
15    let array = VariantArray::try_new(input)?;
16
17    // Zero-copy builder
18    // The size per JSON string is assumed to be 128 bytes. If this holds true, resizing could be
19    // minimized for performance.
20    let mut json_buffer: Vec<u8> = Vec::with_capacity(array.len() * 128);
21    let mut offsets: Vec<i32> = Vec::with_capacity(array.len() + 1);
22    let mut validity = BooleanBufferBuilder::new(array.len());
23    let mut current_offset: i32 = 0;
24    offsets.push(current_offset);
25
26    for i in 0..array.len() {
27        if array.is_null(i) {
28            validity.append(false);
29            offsets.push(current_offset);
30        } else {
31            let variant = array.try_value(i)?;
32            let start_len = json_buffer.len();
33            variant.to_json(&mut json_buffer)?;
34            let written = (json_buffer.len() - start_len) as i32;
35            current_offset += written;
36            offsets.push(current_offset);
37            validity.append(true);
38        }
39    }
40
41    let offsets_buffer = OffsetBuffer::new(ScalarBuffer::from(offsets));
42    let value_buffer = Buffer::from_vec(json_buffer);
43    let null_buffer = NullBuffer::new(validity.finish());
44
45    StringArray::try_new(offsets_buffer, value_buffer, Some(null_buffer))
46}
47
48#[cfg(test)]
49mod test {
50    use crate::variant_to_json;
51    use arrow::array::{Array, ArrayRef, BinaryBuilder, BooleanBufferBuilder, StructArray};
52    use arrow::buffer::NullBuffer;
53    use arrow::datatypes::DataType;
54    use arrow::datatypes::Field;
55    use arrow_schema::Fields;
56    use std::sync::Arc;
57
58    #[test]
59    fn test_variant_to_json() {
60        let mut metadata_builder = BinaryBuilder::new();
61        let mut value_builder = BinaryBuilder::new();
62
63        // Row 0: [1, 0, 0], [12, 0]
64        metadata_builder.append_value([1, 0, 0]);
65        value_builder.append_value([12, 0]);
66
67        // Row 1: null
68        metadata_builder.append_null();
69        value_builder.append_null();
70
71        // Row 2: [1, 1, 0, 1, 97], [2, 1, 0, 0, 1, 32]
72        metadata_builder.append_value([1, 1, 0, 1, 97]);
73        value_builder.append_value([2, 1, 0, 0, 2, 12, 32]);
74
75        // Row 3: [1, 0, 0], [0]
76        metadata_builder.append_value([1, 0, 0]);
77        value_builder.append_value([0]);
78
79        // Row 4: null
80        metadata_builder.append_null();
81        value_builder.append_null();
82
83        let metadata_array = Arc::new(metadata_builder.finish()) as ArrayRef;
84        let value_array = Arc::new(value_builder.finish()) as ArrayRef;
85
86        let fields: Fields = vec![
87            Field::new("metadata", DataType::Binary, true),
88            Field::new("value", DataType::Binary, true),
89        ]
90        .into();
91
92        let mut validity = BooleanBufferBuilder::new(value_array.len());
93        for i in 0..value_array.len() {
94            let is_valid = value_array.is_valid(i) && metadata_array.is_valid(i);
95            validity.append(is_valid);
96        }
97        let null_buffer = NullBuffer::new(validity.finish());
98
99        let struct_array = StructArray::new(
100            fields,
101            vec![metadata_array.clone(), value_array.clone()],
102            Some(null_buffer), // Null bitmap (let Arrow infer from children)
103        );
104
105        let input = Arc::new(struct_array) as ArrayRef;
106
107        let result = variant_to_json(&input).unwrap();
108
109        // Expected output: ["0", null, "{\"a\":32}", "null", null]
110        let expected = vec![Some("0"), None, Some("{\"a\":32}"), Some("null"), None];
111
112        let result_vec: Vec<Option<&str>> = (0..result.len())
113            .map(|i| {
114                if result.is_null(i) {
115                    None
116                } else {
117                    Some(result.value(i))
118                }
119            })
120            .collect();
121
122        assert_eq!(result_vec, expected);
123    }
124}
parquet_variant_compute/to_json.rs

parquet_variant_compute/
to_json.rs
