datatypes/vectors/json/
variant.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use arrow_array::ArrayRef;
18use arrow_schema::{ArrowError, DataType, Field};
19use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16};
20#[cfg(test)]
21use parquet_variant_compute::VariantArrayBuilder;
22use parquet_variant_compute::{VariantArray, VariantType};
23use parquet_variant_json::VariantToJson;
24use snafu::ResultExt;
25
26use crate::error::{ArrowComputeSnafu, Result};
27use crate::json::value::{JsonNumber, JsonVariant, JsonVariantRef, decode_json_variant};
28
29/// Returns the canonical Arrow field for an unshredded Parquet Variant array.
30pub fn variant_field(name: impl Into<String>, nullable: bool) -> Field {
31    Field::new_struct(
32        name,
33        [
34            Arc::new(Field::new("metadata", DataType::BinaryView, false)),
35            Arc::new(Field::new("value", DataType::BinaryView, false)),
36        ],
37        nullable,
38    )
39    .with_extension_type(VariantType)
40}
41
42/// Encodes JSON values as an unshredded Parquet Variant array.
43///
44/// `None` represents an Arrow null while `Some(Value::Null)` represents a JSON
45/// null, preserving the distinction required by JSON2.
46#[cfg(test)]
47pub(crate) fn json_values_to_variant(values: &[Option<serde_json::Value>]) -> Result<ArrayRef> {
48    let mut builder = VariantArrayBuilder::new(values.len());
49    for value in values {
50        match value {
51            Some(value) => append_json_value(&mut builder, value).context(ArrowComputeSnafu)?,
52            None => builder.append_null(),
53        }
54    }
55    Ok(ArrayRef::from(builder.build()))
56}
57
58/// Encodes JSON variants as an unshredded Parquet Variant array.
59#[cfg(test)]
60pub(crate) fn json_variants_to_variant(values: &[Option<JsonVariant>]) -> Result<ArrayRef> {
61    let mut builder = VariantArrayBuilder::new(values.len());
62    for value in values {
63        match value {
64            Some(value) => append_json_variant(&mut builder, value).context(ArrowComputeSnafu)?,
65            None => builder.append_null(),
66        }
67    }
68    Ok(ArrayRef::from(builder.build()))
69}
70
71pub(super) fn append_json_variant(
72    builder: &mut impl VariantBuilderExt,
73    value: &JsonVariant,
74) -> std::result::Result<(), ArrowError> {
75    match value {
76        JsonVariant::Null => builder.append_value(Variant::Null),
77        JsonVariant::Bool(value) => builder.append_value(*value),
78        JsonVariant::Number(JsonNumber::PosInt(value)) => {
79            if let Ok(value) = i64::try_from(*value) {
80                builder.append_value(value);
81            } else {
82                append_large_u64(builder, *value)?;
83            }
84        }
85        JsonVariant::Number(JsonNumber::NegInt(value)) => builder.append_value(*value),
86        JsonVariant::Number(JsonNumber::Float(value)) => {
87            if value.0.is_finite() {
88                builder.append_value(value.0)
89            } else {
90                builder.append_value("NaN")
91            }
92        }
93        JsonVariant::String(value) => builder.append_value(value.as_str()),
94        JsonVariant::Array(values) => {
95            let mut list = builder.try_new_list()?;
96            for value in values {
97                append_json_variant(&mut list, value)?;
98            }
99            list.finish();
100        }
101        JsonVariant::Object(values) => {
102            let mut object = builder.try_new_object()?;
103            for (name, value) in values {
104                append_json_variant(&mut ObjectFieldBuilder::new(name, &mut object), value)?;
105            }
106            object.finish();
107        }
108        JsonVariant::Variant(value) => {
109            let value = decode_json_variant(value)
110                .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?;
111            append_json_value(builder, &value)?;
112        }
113    }
114    Ok(())
115}
116
117pub(super) fn append_json_variant_ref(
118    builder: &mut impl VariantBuilderExt,
119    value: &JsonVariantRef<'_>,
120) -> std::result::Result<(), ArrowError> {
121    match value {
122        JsonVariantRef::Null => builder.append_value(Variant::Null),
123        JsonVariantRef::Bool(value) => builder.append_value(*value),
124        JsonVariantRef::Number(JsonNumber::PosInt(value)) => {
125            if let Ok(value) = i64::try_from(*value) {
126                builder.append_value(value);
127            } else {
128                append_large_u64(builder, *value)?;
129            }
130        }
131        JsonVariantRef::Number(JsonNumber::NegInt(value)) => builder.append_value(*value),
132        JsonVariantRef::Number(JsonNumber::Float(value)) => {
133            if value.0.is_finite() {
134                builder.append_value(value.0)
135            } else {
136                builder.append_value("NaN")
137            }
138        }
139        JsonVariantRef::String(value) => builder.append_value(*value),
140        JsonVariantRef::Array(values) => {
141            let mut list = builder.try_new_list()?;
142            for value in values {
143                append_json_variant_ref(&mut list, value)?;
144            }
145            list.finish();
146        }
147        JsonVariantRef::Object(values) => {
148            let mut object = builder.try_new_object()?;
149            for (name, value) in values {
150                append_json_variant_ref(&mut ObjectFieldBuilder::new(name, &mut object), value)?;
151            }
152            object.finish();
153        }
154        JsonVariantRef::Variant(value) => {
155            let value = decode_json_variant(value)
156                .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?;
157            append_json_value(builder, &value)?;
158        }
159    }
160    Ok(())
161}
162
163fn append_json_value(
164    builder: &mut impl VariantBuilderExt,
165    value: &serde_json::Value,
166) -> std::result::Result<(), ArrowError> {
167    match value {
168        serde_json::Value::Null => builder.append_value(Variant::Null),
169        serde_json::Value::Bool(value) => builder.append_value(*value),
170        serde_json::Value::Number(value) => {
171            if let Some(value) = value.as_i64() {
172                builder.append_value(value);
173            } else if let Some(value) = value.as_u64() {
174                append_large_u64(builder, value)?;
175            } else if let Some(value) = value.as_f64() {
176                builder.append_value(value);
177            } else {
178                return Err(ArrowError::InvalidArgumentError(format!(
179                    "Failed to encode JSON number as Variant: {value}"
180                )));
181            }
182        }
183        serde_json::Value::String(value) => builder.append_value(value.as_str()),
184        serde_json::Value::Array(values) => {
185            let mut list = builder.try_new_list()?;
186            for value in values {
187                append_json_value(&mut list, value)?;
188            }
189            list.finish();
190        }
191        serde_json::Value::Object(values) => {
192            let mut object = builder.try_new_object()?;
193            for (name, value) in values {
194                append_json_value(&mut ObjectFieldBuilder::new(name, &mut object), value)?;
195            }
196            object.finish();
197        }
198    }
199    Ok(())
200}
201
202/// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16
203/// to represent large (larger than i64::MAX) u64.
204fn append_large_u64(
205    builder: &mut impl VariantBuilderExt,
206    value: u64,
207) -> std::result::Result<(), ArrowError> {
208    // Parquet Variant has no unsigned integer primitive. Decimal16 preserves the full u64 range.
209    let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| {
210        ArrowError::InvalidArgumentError(format!(
211            "Failed to encode JSON large integer as Variant Decimal16: {e}"
212        ))
213    })?;
214    builder.append_value(value);
215    Ok(())
216}
217
218/// Decodes an unshredded Parquet Variant array into JSON values.
219pub fn variant_to_json_values(array: &ArrayRef) -> Result<Vec<Option<serde_json::Value>>> {
220    let variants = VariantArray::try_new(array.as_ref()).context(ArrowComputeSnafu)?;
221    (0..variants.len())
222        .map(|i| {
223            if variants.is_null(i) {
224                Ok(None)
225            } else {
226                variants
227                    .try_value(i)
228                    .and_then(|x| x.to_json_value())
229                    .context(ArrowComputeSnafu)
230                    .map(Some)
231            }
232        })
233        .collect()
234}
235
236#[cfg(test)]
237mod tests {
238    use serde_json::json;
239
240    use super::*;
241
242    #[test]
243    fn test_variant_field_matches_canonical_layout() {
244        let expected = VariantArrayBuilder::new(0)
245            .build()
246            .field("remainder")
247            .with_nullable(true);
248        assert_eq!(expected, variant_field("remainder", true));
249    }
250
251    #[test]
252    fn test_variant_json_round_trip() -> Result<()> {
253        let values = vec![
254            None,
255            Some(serde_json::Value::Null),
256            Some(json!({})),
257            Some(json!({"nested": {"items": [1, "two", null]}})),
258            Some(json!({"unicode": "\u{503c}"})),
259            Some(json!({"max_u64": u64::MAX})),
260        ];
261
262        let array = json_values_to_variant(&values)?;
263        assert_eq!(values, variant_to_json_values(&array)?);
264
265        let variants = values
266            .clone()
267            .into_iter()
268            .map(|x| x.map(JsonVariant::from))
269            .collect::<Vec<_>>();
270        let array = json_variants_to_variant(&variants)?;
271        assert_eq!(values, variant_to_json_values(&array)?);
272
273        let variants = [Some(JsonVariant::Variant(
274            jsonb::parse_value(br#"{"nested": true}"#).unwrap().to_vec(),
275        ))];
276        let array = json_variants_to_variant(&variants)?;
277        assert_eq!(
278            vec![Some(json!({"nested": true}))],
279            variant_to_json_values(&array)?
280        );
281        Ok(())
282    }
283}
datatypes/vectors/json/variant.rs

datatypes/vectors/json/
variant.rs
