diff --git a/Cargo.lock b/Cargo.lock index 0f473d4e03..16b4c6b104 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4554,6 +4554,9 @@ dependencies = [ "num", "num-traits", "ordered-float 4.6.0", + "parquet-variant", + "parquet-variant-compute", + "parquet-variant-json", "paste", "regex", "serde", @@ -10259,6 +10262,53 @@ dependencies = [ "zstd", ] +[[package]] +name = "parquet-variant" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74c8db065291f088a2aad8ab831853eae1871c0d311c8d0b83bbc3b7e735d0fc" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "num-traits", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-compute" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a530e8d5b5e14efcb39c9a6ec55432ad11f6afb7dc4455a79be0dc615fe3cc31" +dependencies = [ + "arrow 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "half", + "indexmap 2.13.0", + "parquet-variant", + "parquet-variant-json", + "serde_json", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00ed89908289f67caa2ca078f9ff9aacd6229a313ec92b12bf4f48f613dc2b97" +dependencies = [ + "arrow-schema 58.3.0", + "base64 0.22.1", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "parse-size" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index daee1bd409..777d8b302b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -195,6 +195,9 @@ otel-arrow-rust = { git = "https://github.com/GreptimeTeam/otel-arrow", rev = "5 ] } parking_lot = "0.12" parquet = { version = "58.3", default-features = false, features = ["arrow", "async", "object_store"] } +parquet-variant = "58.3" +parquet-variant-compute = "58.3" +parquet-variant-json = "58.3" paste = "1.0" pin-project = "1.0" pretty_assertions = "1.4.0" diff --git a/src/datatypes/Cargo.toml b/src/datatypes/Cargo.toml index 3bf2a5d8ee..124b1e0ea6 100644 --- a/src/datatypes/Cargo.toml +++ b/src/datatypes/Cargo.toml @@ -27,6 +27,9 @@ jsonb.workspace = true num = "0.4" num-traits = "0.2" ordered-float.workspace = true +parquet-variant.workspace = true +parquet-variant-compute.workspace = true +parquet-variant-json.workspace = true paste.workspace = true regex.workspace = true serde.workspace = true diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 109aa853b2..18aee22a96 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -15,34 +15,154 @@ use std::collections::HashMap; use std::sync::Arc; +#[cfg(test)] +use arrow_schema::FieldRef; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; use arrow_schema::{ArrowError, DataType, Field}; +#[cfg(test)] +use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; -use snafu::ResultExt; +#[cfg(test)] +use snafu::OptionExt; +use snafu::{ResultExt, ensure}; +use crate::error::InvalidJson2LayoutSnafu; pub use crate::json::JSON2_REMAINDER_FIELD_NAME; use crate::json::JsonSettings; const LEGACY_JSON_STRUCTURE_SETTINGS_KEY: &str = "json_structure_settings"; +const JSON2_LAYOUT_V1: u8 = 1; +const JSON2_LAYOUT_V2: u8 = 2; -#[derive(Debug, Clone, Serialize, Deserialize, Default)] +/// Parsed physical layout of a JSON2 Arrow root field. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Json2PhysicalLayout { + version: u8, +} + +impl Json2PhysicalLayout { + /// Parses and validates the JSON2 layout version from a root field. + pub fn try_from_root(field: &Field) -> crate::error::Result { + ensure!( + is_json2_extension_type(field), + InvalidJson2LayoutSnafu { + reason: format!("field '{}' is not a JSON2 extension", field.name()), + } + ); + + let version = match field.extension_type_name() { + Some(Json2ExtensionType::NAME) => field + .metadata() + .get(EXTENSION_TYPE_METADATA_KEY) + .map(|x| parse_version(x)) + .transpose()? + .flatten() + .unwrap_or(JSON2_LAYOUT_V1), + _ => JSON2_LAYOUT_V1, + }; + ensure!( + matches!(version, JSON2_LAYOUT_V1 | JSON2_LAYOUT_V2), + InvalidJson2LayoutSnafu { + reason: format!("unsupported JSON2 layout version: {version}"), + } + ); + Ok(Self { version }) + } + + /// Returns whether this is the JSON2 physical layout version 2. + pub fn is_version_2(&self) -> bool { + self.version == JSON2_LAYOUT_V2 + } +} + +fn parse_version(metadata: &str) -> crate::error::Result> { + serde_json::from_str::(metadata) + .map(|x| x.layout_version) + .map_err(|e| { + InvalidJson2LayoutSnafu { + reason: format!(r#"invalid extension metadata: "{metadata}", error: {e}"#), + } + .build() + }) +} + +/// Returns the remainder field of a JSON2 v2 root. +#[cfg(test)] +fn json2_remainder_field(field: &Field) -> crate::error::Result<&FieldRef> { + if !Json2PhysicalLayout::try_from_root(field)?.is_version_2() { + return InvalidJson2LayoutSnafu { + reason: format!("JSON2 root '{}' is not layout v2", field.name()), + } + .fail(); + } + + let DataType::Struct(fields) = field.data_type() else { + return InvalidJson2LayoutSnafu { + reason: format!( + "expecting the Struct datatype, actual: '{}'", + field.data_type(), + ), + } + .fail(); + }; + let remainder = fields + .iter() + .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME) + .context(InvalidJson2LayoutSnafu { + reason: "remainder field not found", + })?; + let _ = remainder.try_extension_type::().map_err(|e| { + InvalidJson2LayoutSnafu { + reason: e.to_string(), + } + .build() + })?; + Ok(remainder) +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct JsonMetadata { /// JSON2 settings stored in Arrow extension metadata. json_settings: JsonSettings, + /// Physical JSON2 layout used by this Arrow field. + #[serde(default, skip_serializing_if = "Option::is_none")] + layout_version: Option, } impl JsonMetadata { - /// Creates JSON2 extension metadata. + /// Creates metadata for the legacy JSON2 layout. pub fn new(json_settings: JsonSettings) -> Self { - Self { json_settings } + Self { + json_settings, + layout_version: None, + } + } + + /// Creates metadata for the new JSON2 physical v2 layout. + pub fn new_v2(json_settings: JsonSettings) -> Self { + Self { + json_settings, + layout_version: Some(JSON2_LAYOUT_V2), + } } /// Returns the JSON2 settings. pub fn json_settings(&self) -> &JsonSettings { &self.json_settings } + + /// Returns whether this metadata describes JSON2 layout version 2. + pub fn is_version_2(&self) -> bool { + self.layout_version == Some(JSON2_LAYOUT_V2) + } +} + +impl Default for JsonMetadata { + fn default() -> Self { + Self::new(JsonSettings::default()) + } } /// Arrow extension type for legacy JSONB columns. @@ -227,6 +347,7 @@ mod tests { use arrow_schema::{Field, Fields}; use super::*; + use crate::vectors::json::variant::variant_field; #[test] fn test_json2_extension_type_detection() { @@ -282,4 +403,78 @@ mod tests { assert!(JsonExtensionType::try_new(&DataType::Null, ()).is_ok()); assert!(JsonExtensionType::try_new(&DataType::Struct(Fields::empty()), ()).is_err()); } + + #[test] + fn test_json_metadata_layout_version_compatibility() -> serde_json::Result<()> { + let legacy: JsonMetadata = serde_json::from_str(r#"{"json_settings":{}}"#)?; + assert!(!legacy.is_version_2()); + + let metadata = JsonMetadata::new_v2(JsonSettings::default()); + assert!(metadata.is_version_2()); + let serialized = serde_json::to_string(&metadata)?; + let deserialized: JsonMetadata = serde_json::from_str(&serialized)?; + assert!(deserialized.is_version_2()); + assert_eq!(deserialized, metadata); + Ok(()) + } + + #[test] + fn test_parse_json2_physical_layout() -> crate::error::Result<()> { + let legacy = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::default()); + assert!(!Json2PhysicalLayout::try_from_root(&legacy)?.is_version_2()); + + let v2 = Field::new( + "data", + DataType::Struct( + vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ), + true, + ) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2( + JsonSettings::default(), + )))); + assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2()); + let remainder = json2_remainder_field(&v2)?; + assert_eq!(JSON2_REMAINDER_FIELD_NAME, remainder.name()); + Ok(()) + } + + #[test] + fn test_reject_invalid_json2_physical_layout() { + let field = Field::new("data", DataType::Struct(Fields::empty()), true); + assert!(Json2PhysicalLayout::try_from_root(&field).is_err()); + + let metadata = + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))); + let missing = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(metadata.clone()); + assert!(json2_remainder_field(&missing).is_err()); + + let invalid = Field::new( + "data", + DataType::Struct( + vec![Arc::new(Field::new( + JSON2_REMAINDER_FIELD_NAME, + DataType::Binary, + true, + ))] + .into(), + ), + true, + ) + .with_extension_type(metadata); + assert!(json2_remainder_field(&invalid).is_err()); + + let future = Field::new("data", DataType::Struct(Fields::empty()), true) + .with_extension_type(Json2ExtensionType::new(Arc::new(JsonMetadata { + json_settings: JsonSettings::default(), + layout_version: Some(JSON2_LAYOUT_V2 + 1), + }))); + assert!(Json2PhysicalLayout::try_from_root(&future).is_err()); + } } diff --git a/src/datatypes/src/vectors/json.rs b/src/datatypes/src/vectors/json.rs index 3dcb8d7bc0..c68e6cf3d9 100644 --- a/src/datatypes/src/vectors/json.rs +++ b/src/datatypes/src/vectors/json.rs @@ -14,3 +14,4 @@ pub mod array; pub(crate) mod builder; +pub mod variant; diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index afe9f55305..9654f84c6d 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -17,9 +17,12 @@ use std::sync::Arc; use arrow::compute::{can_cast_types, cast}; use arrow_array::cast::AsArray; -use arrow_array::types::{Float64Type, Int64Type, UInt64Type}; +use arrow_array::types::{ + Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, + UInt32Type, UInt64Type, +}; use arrow_array::{Array, ArrayRef, GenericListArray, ListArray, StructArray, new_null_array}; -use arrow_schema::{DataType, FieldRef}; +use arrow_schema::{DataType, Field, FieldRef}; use common_telemetry::trace; use serde_json::Value; use snafu::{OptionExt, ResultExt}; @@ -29,9 +32,11 @@ use crate::data_type::ConcreteDataType; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; +use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2PhysicalLayout}; use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; use crate::prelude::{DataType as _, Value as GreptimeValue}; use crate::value::{ListValue, StructValue}; +use crate::vectors::json::variant::variant_to_json_values; pub struct JsonArray<'a> { inner: &'a ArrayRef, @@ -48,8 +53,15 @@ impl JsonArray<'_> { let value = match array.data_type() { DataType::Null => Value::Null, DataType::Boolean => Value::Bool(array.as_boolean().value(i)), + DataType::Int8 => Value::from(array.as_primitive::().value(i)), + DataType::Int16 => Value::from(array.as_primitive::().value(i)), + DataType::Int32 => Value::from(array.as_primitive::().value(i)), DataType::Int64 => Value::from(array.as_primitive::().value(i)), + DataType::UInt8 => Value::from(array.as_primitive::().value(i)), + DataType::UInt16 => Value::from(array.as_primitive::().value(i)), + DataType::UInt32 => Value::from(array.as_primitive::().value(i)), DataType::UInt64 => Value::from(array.as_primitive::().value(i)), + DataType::Float32 => Value::from(array.as_primitive::().value(i)), DataType::Float64 => Value::from(array.as_primitive::().value(i)), DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { Value::String(string_array_value(array, i).to_string()) @@ -92,6 +104,69 @@ impl JsonArray<'_> { Ok(value) } + /// Projects a physical JSON2 array to a logical query type. + pub fn project_json2(&self, field: &Field, target: &DataType) -> Result { + let layout = Json2PhysicalLayout::try_from_root(field)?; + if !layout.is_version_2() { + return self.project_to(target); + } + + let has_remainder = match field.data_type() { + DataType::Struct(fields) => fields + .iter() + .any(|x| x.name() == JSON2_REMAINDER_FIELD_NAME), + _ => false, + }; + if !has_remainder { + return self.project_to(target); + } + + project_json_values(self.json2_values()?, target) + } + + fn json2_values(&self) -> Result> { + let structs = self.inner.as_struct_opt().context(AlignJsonArraySnafu { + reason: "JSON2 layout v2 root array must be a struct", + })?; + let remainder = structs.column_by_name(JSON2_REMAINDER_FIELD_NAME); + let mut remainders = if let Some(remainder) = remainder { + variant_to_json_values(remainder)? + } else { + vec![None; structs.len()] + }; + let mut values = Vec::with_capacity(structs.len()); + let mut path = Vec::new(); + + for (i, remainder) in remainders.iter_mut().enumerate() { + if structs.is_null(i) { + values.push(Value::Null); + continue; + } + + let mut object = match remainder.take() { + None => serde_json::Map::new(), + Some(Value::Object(object)) => object, + Some(value) => { + return InvalidJsonSnafu { + value: format!("JSON2 layout v2 remainder must be an object, got {value}"), + } + .fail(); + } + }; + + for (child, column) in structs.fields().iter().zip(structs.columns()) { + if child.name() == JSON2_REMAINDER_FIELD_NAME { + continue; + } + let value = JsonArray::from(column).try_get_value(i)?; + merge_explicit_value(&mut object, child.name().clone(), value, &mut path)?; + } + values.push(Value::Object(object)); + } + + Ok(values) + } + /// Normalizes a JSON2 array to the wider `expect` data type without losing /// information. /// @@ -326,6 +401,34 @@ impl JsonArray<'_> { } } +fn merge_explicit_value( + remainder: &mut serde_json::Map, + key: String, + explicit: Value, + path: &mut Vec, +) -> Result<()> { + let Some(existing) = remainder.get_mut(&key) else { + remainder.insert(key, explicit); + return Ok(()); + }; + path.push(key); + + let (Value::Object(remainder), Value::Object(explicit)) = (existing, explicit) else { + return InvalidJsonSnafu { + value: format!( + "cannot merge '{}' in explicit fields and remainder: not both objects", + path.join("."), + ), + } + .fail(); + }; + for (key, value) in explicit { + merge_explicit_value(remainder, key, value, path)?; + } + path.pop(); + Ok(()) +} + /// Returns whether Arrow can cast between the types without JSON-aware projection. /// Binary and nested types require JSONB decoding or recursive projection. fn can_fast_cast_types(from_type: &DataType, to_type: &DataType) -> bool { @@ -453,13 +556,16 @@ impl<'a> From<&'a ArrayRef> for JsonArray<'a> { mod test { use arrow_array::types::Int64Type; use arrow_array::{ - BinaryArray, BooleanArray, Float64Array, Int32Array, Int64Array, ListArray, StringArray, - UInt64Array, + BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, + Int64Array, ListArray, StringArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array, }; use arrow_schema::{Field, Fields}; use serde_json::json; use super::*; + use crate::extension::json::{Json2ExtensionType, JsonMetadata}; + use crate::json::JsonSettings; + use crate::vectors::json::variant::{json_values_to_variant, variant_field}; #[test] fn test_try_get_value() -> Result<()> { @@ -474,6 +580,20 @@ mod test { assert_eq!(JsonArray::from(&ints).try_get_value(0)?, json!(-7)); assert_eq!(JsonArray::from(&ints).try_get_value(1)?, Value::Null); + macro_rules! assert_number { + ($array:expr, $expected:expr) => {{ + let array: ArrayRef = Arc::new($array); + assert_eq!(JsonArray::from(&array).try_get_value(0)?, json!($expected)); + }}; + } + assert_number!(Int8Array::from(vec![-8]), -8); + assert_number!(Int16Array::from(vec![-16]), -16); + assert_number!(Int32Array::from(vec![-32]), -32); + assert_number!(UInt8Array::from(vec![8]), 8); + assert_number!(UInt16Array::from(vec![16]), 16); + assert_number!(UInt32Array::from(vec![32]), 32); + assert_number!(Float32Array::from(vec![1.25]), 1.25); + let floats: ArrayRef = Arc::new(Float64Array::from(vec![Some(1.5)])); assert_eq!(JsonArray::from(&floats).try_get_value(0)?, json!(1.5)); @@ -529,15 +649,6 @@ mod test { json!({"flag": null, "items": [2]}) ); - let unsupported: ArrayRef = Arc::new(Int32Array::from(vec![1])); - assert_eq!( - JsonArray::from(&unsupported) - .try_get_value(0) - .unwrap_err() - .to_string(), - "Invalid JSON: unknown JSON type Int32" - ); - Ok(()) } @@ -915,4 +1026,124 @@ mod test { Ok(()) } + + #[test] + fn test_reconstruct_json2_v2_value() -> Result<()> { + let remainders = json_values_to_variant(&[ + Some(json!({"cold": 1, "nested": {"right": true}})), + Some(json!({"!__remainder__!": "user value"})), + ])?; + let remainder = Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)); + let nested = Arc::new(Field::new_struct( + "nested", + [Arc::new(Field::new("left", DataType::Utf8, true))], + true, + )); + let nested_values: ArrayRef = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("left", DataType::Utf8, true)), + Arc::new(StringArray::from(vec![Some("value"), None])) as ArrayRef, + )])); + let fields = Fields::from(vec![ + remainder, + Arc::new(Field::new("count", DataType::Int64, true)), + nested, + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![ + remainders, + Arc::new(Int64Array::from(vec![Some(42), None])), + nested_values, + ], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + assert_eq!( + json!({ + "cold": 1, + "count": 42, + "nested": {"left": "value", "right": true} + }), + JsonArray::from(&array).json2_values()?[0] + ); + assert_eq!( + json!({ + "!__remainder__!": "user value", + "count": null, + "nested": {"left": null} + }), + JsonArray::from(&array).json2_values()?[1] + ); + + let target = DataType::Struct( + vec![ + Arc::new(Field::new("cold", DataType::UInt64, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ] + .into(), + ); + let projected = JsonArray::from(&array).project_json2(&field, &target)?; + assert_eq!( + json!({"cold": 1, "count": 42}), + JsonArray::from(&projected).try_get_value(0)? + ); + assert_eq!( + json!({"cold": null, "count": null}), + JsonArray::from(&projected).try_get_value(1)? + ); + Ok(()) + } + + #[test] + fn test_project_partial_json2_v2_without_remainder() -> Result<()> { + let fields = Fields::from(vec![Arc::new(Field::new("hot", DataType::Int64, true))]); + let array: ArrayRef = Arc::new(StructArray::new( + fields.clone(), + vec![Arc::new(Int64Array::from(vec![1, 2]))], + None, + )); + let field = Field::new("data", DataType::Struct(fields), true).with_extension_type( + Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), + ); + + let projected = JsonArray::from(&array).project_json2(&field, field.data_type())?; + assert!(Arc::ptr_eq(&array, &projected)); + Ok(()) + } + + #[test] + fn test_reject_conflict_json2_v2_path() -> Result<()> { + let remainders = json_values_to_variant(&[Some(json!({"count": 1}))])?; + let fields = Fields::from(vec![ + Arc::new(variant_field(JSON2_REMAINDER_FIELD_NAME, true)), + Arc::new(Field::new("count", DataType::Int64, true)), + ]); + let array: ArrayRef = Arc::new(StructArray::new( + fields, + vec![remainders, Arc::new(Int64Array::from(vec![2]))], + None, + )); + let error = JsonArray::from(&array).json2_values().unwrap_err(); + assert!( + error.to_string().contains( + "cannot merge 'count' in explicit fields and remainder: not both objects" + ) + ); + + let Value::Object(mut remainder) = json!({"nested": {"count": 1}}) else { + unreachable!(); + }; + let error = merge_explicit_value( + &mut remainder, + "nested".to_string(), + json!({"count": 2}), + &mut Vec::new(), + ) + .unwrap_err(); + assert!(error.to_string().contains("cannot merge 'nested.count'")); + Ok(()) + } } diff --git a/src/datatypes/src/vectors/json/variant.rs b/src/datatypes/src/vectors/json/variant.rs new file mode 100644 index 0000000000..1a4b8a861d --- /dev/null +++ b/src/datatypes/src/vectors/json/variant.rs @@ -0,0 +1,239 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use arrow_array::ArrayRef; +#[cfg(test)] +use arrow_schema::ArrowError; +use arrow_schema::{DataType, Field}; +#[cfg(test)] +use parquet_variant::{ObjectFieldBuilder, Variant, VariantBuilderExt, VariantDecimal16}; +#[cfg(test)] +use parquet_variant_compute::VariantArrayBuilder; +use parquet_variant_compute::{VariantArray, VariantType}; +use parquet_variant_json::VariantToJson; +use snafu::ResultExt; + +use crate::error::{ArrowComputeSnafu, Result}; +#[cfg(test)] +use crate::json::value::{JsonNumber, JsonVariant, decode_json_variant}; + +/// Returns the canonical Arrow field for an unshredded Parquet Variant array. +pub fn variant_field(name: impl Into, nullable: bool) -> Field { + Field::new_struct( + name, + [ + Arc::new(Field::new("metadata", DataType::BinaryView, false)), + Arc::new(Field::new("value", DataType::BinaryView, false)), + ], + nullable, + ) + .with_extension_type(VariantType) +} + +#[cfg(test)] +pub(crate) fn json_values_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_value(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +/// Encodes JSON variants as an unshredded Parquet Variant array. +#[cfg(test)] +fn json_variants_to_variant(values: &[Option]) -> Result { + let mut builder = VariantArrayBuilder::new(values.len()); + for value in values { + match value { + Some(value) => append_json_variant(&mut builder, value).context(ArrowComputeSnafu)?, + None => builder.append_null(), + } + } + Ok(ArrayRef::from(builder.build())) +} + +#[cfg(test)] +fn append_json_variant( + builder: &mut impl VariantBuilderExt, + value: &JsonVariant, +) -> std::result::Result<(), ArrowError> { + match value { + JsonVariant::Null => builder.append_value(Variant::Null), + JsonVariant::Bool(value) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::PosInt(value)) => { + if let Ok(value) = i64::try_from(*value) { + builder.append_value(value); + } else { + append_large_u64(builder, *value)?; + } + } + JsonVariant::Number(JsonNumber::NegInt(value)) => builder.append_value(*value), + JsonVariant::Number(JsonNumber::Float(value)) => { + if value.0.is_finite() { + builder.append_value(value.0) + } else { + builder.append_value("NaN") + } + } + JsonVariant::String(value) => builder.append_value(value.as_str()), + JsonVariant::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_variant(&mut list, value)?; + } + list.finish(); + } + JsonVariant::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_variant(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + JsonVariant::Variant(value) => { + let value = decode_json_variant(value) + .map_err(|e| ArrowError::JsonError(format!("Failed to decode JSONB: {e}")))?; + append_json_value(builder, &value)?; + } + } + Ok(()) +} + +#[cfg(test)] +fn append_json_value( + builder: &mut impl VariantBuilderExt, + value: &serde_json::Value, +) -> std::result::Result<(), ArrowError> { + match value { + serde_json::Value::Null => builder.append_value(Variant::Null), + serde_json::Value::Bool(value) => builder.append_value(*value), + serde_json::Value::Number(value) => { + if let Some(value) = value.as_i64() { + builder.append_value(value); + } else if let Some(value) = value.as_u64() { + append_large_u64(builder, value)?; + } else if let Some(value) = value.as_f64() { + builder.append_value(value); + } else { + return Err(ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON number as Variant: {value}" + ))); + } + } + serde_json::Value::String(value) => builder.append_value(value.as_str()), + serde_json::Value::Array(values) => { + let mut list = builder.try_new_list()?; + for value in values { + append_json_value(&mut list, value)?; + } + list.finish(); + } + serde_json::Value::Object(values) => { + let mut object = builder.try_new_object()?; + for (name, value) in values { + append_json_value(&mut ObjectFieldBuilder::new(name, &mut object), value)?; + } + object.finish(); + } + } + Ok(()) +} + +/// Parquet Variant has no unsigned integer primitive. Treat u64 as i64 first, then use Decimal16 +/// to represent large (larger than i64::MAX) u64. +#[cfg(test)] +fn append_large_u64( + builder: &mut impl VariantBuilderExt, + value: u64, +) -> std::result::Result<(), ArrowError> { + let value = VariantDecimal16::try_new(value as i128, 0).map_err(|e| { + ArrowError::InvalidArgumentError(format!( + "Failed to encode JSON large integer as Variant Decimal16: {e}" + )) + })?; + builder.append_value(value); + Ok(()) +} + +/// Decodes an unshredded Parquet Variant array into JSON values. +pub fn variant_to_json_values(array: &ArrayRef) -> Result>> { + let variants = VariantArray::try_new(array.as_ref()).context(ArrowComputeSnafu)?; + (0..variants.len()) + .map(|i| { + if variants.is_null(i) { + Ok(None) + } else { + variants + .try_value(i) + .and_then(|x| x.to_json_value()) + .context(ArrowComputeSnafu) + .map(Some) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn test_variant_field_matches_canonical_layout() { + let expected = VariantArrayBuilder::new(0) + .build() + .field("remainder") + .with_nullable(true); + assert_eq!(expected, variant_field("remainder", true)); + } + + #[test] + fn test_variant_json_round_trip() -> Result<()> { + let values = vec![ + None, + Some(serde_json::Value::Null), + Some(json!({})), + Some(json!({"nested": {"items": [1, "two", null]}})), + Some(json!({"unicode": "\u{503c}"})), + Some(json!({"max_u64": u64::MAX})), + ]; + + let array = json_values_to_variant(&values)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = values + .clone() + .into_iter() + .map(|x| x.map(JsonVariant::from)) + .collect::>(); + let array = json_variants_to_variant(&variants)?; + assert_eq!(values, variant_to_json_values(&array)?); + + let variants = [Some(JsonVariant::Variant( + jsonb::parse_value(br#"{"nested": true}"#).unwrap().to_vec(), + ))]; + let array = json_variants_to_variant(&variants)?; + assert_eq!( + vec![Some(json!({"nested": true}))], + variant_to_json_values(&array)? + ); + Ok(()) + } +}