diff --git a/parquet-variant-compute/src/shred_variant.rs b/parquet-variant-compute/src/shred_variant.rs index 7753fce0f674..a6198c9c1525 100644 --- a/parquet-variant-compute/src/shred_variant.rs +++ b/parquet-variant-compute/src/shred_variant.rs @@ -169,13 +169,16 @@ pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>( )?; VariantToShreddedVariantRowBuilder::Array(typed_value_builder) } - // Supported shredded primitive types, see Variant shredding spec: - // https://github.com/apache/parquet-format/blob/master/VariantShredding.md#shredded-value-types + // Supported Arrow canonical Variant primitive types: + // https://arrow.apache.org/docs/format/CanonicalExtensions.html#primitive-type-mappings DataType::Boolean | DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 | DataType::Float32 | DataType::Float64 | DataType::Decimal32(..) @@ -1459,7 +1462,7 @@ mod tests { let input = VariantArray::from_iter([Variant::from(42)]); let invalid_types = vec![ - DataType::UInt8, + DataType::UInt64, DataType::Float16, DataType::Decimal256(38, 10), DataType::Date64, @@ -1503,6 +1506,105 @@ mod tests { } } + #[test] + fn test_unsigned_shredding() { + let fallback = Variant::Int8(-1); + let cases = [ + ( + vec![ + Variant::Int16(0), + Variant::Int16(i16::from(u8::MAX)), + fallback.clone(), + ], + DataType::UInt8, + ), + ( + vec![ + Variant::Int32(0), + Variant::Int32(i32::from(u16::MAX)), + fallback.clone(), + ], + DataType::UInt16, + ), + ( + vec![ + Variant::Int64(0), + Variant::Int64(i64::from(u32::MAX)), + fallback.clone(), + ], + DataType::UInt32, + ), + ]; + + for (expected, data_type) in cases { + let input = VariantArray::from_iter(expected.iter().cloned()); + let shredded = shred_variant(&input, &data_type).unwrap(); + + assert_eq!( + shredded.typed_value_column().unwrap().data_type(), + &data_type + ); + let array = ArrayRef::from(shredded); + let shredded = VariantArray::try_new(array.as_ref()).unwrap(); + let unshredded = crate::unshred_variant(&shredded).unwrap(); + let fallback_index = expected.len() - 1; + for (index, expected) in expected[..fallback_index].iter().cloned().enumerate() { + assert_eq!(shredded.value(index), expected); + assert_eq!(unshredded.value(index), expected); + } + assert!( + shredded + .typed_value_column() + .unwrap() + .is_null(fallback_index) + ); + assert_eq!(shredded.value(fallback_index), fallback); + assert_eq!(unshredded.value(fallback_index), fallback); + } + } + + #[test] + fn test_nested_unsigned_shredding() { + let input = build_variant_array(vec![VariantRow::Object(vec![( + "items", + VariantValue::List(vec![VariantValue::from(Variant::Int16(i16::from(u8::MAX)))]), + )])]); + let list_type = DataType::List(Arc::new(Field::new("item", DataType::UInt8, true))); + let target = ShreddedSchemaBuilder::default() + .with_path("items", list_type) + .unwrap() + .build(); + + let shredded = shred_variant(&input, &target).unwrap(); + let object = shredded + .typed_value_column() + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let items = + ShreddedVariantFieldArray::try_new(object.column_by_name("items").unwrap()).unwrap(); + let items = items + .typed_value_column() + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let item = ShreddedVariantFieldArray::try_new(items.values().as_ref()).unwrap(); + assert_eq!( + item.typed_value_column().unwrap().data_type(), + &DataType::UInt8 + ); + let array = ArrayRef::from(shredded); + let shredded = VariantArray::try_new(&array).unwrap(); + let unshredded = crate::unshred_variant(&shredded).unwrap(); + let value = unshredded.value(0); + let object = value.as_object().unwrap(); + let items = object.get("items").unwrap(); + let items = items.as_list().unwrap(); + assert_eq!(items.get(0), Some(Variant::Int16(i16::from(u8::MAX)))); + } + #[test] fn test_array_shredding_as_list() { let input = build_variant_array(vec![ diff --git a/parquet-variant-compute/src/unshred_variant.rs b/parquet-variant-compute/src/unshred_variant.rs index 30f861ce1e94..bafe77c5896f 100644 --- a/parquet-variant-compute/src/unshred_variant.rs +++ b/parquet-variant-compute/src/unshred_variant.rs @@ -28,7 +28,8 @@ use arrow::array::{ use arrow::datatypes::{ ArrowPrimitiveType, DataType, Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, DecimalType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, - Time64MicrosecondType, TimeUnit, TimestampMicrosecondType, TimestampNanosecondType, + Time64MicrosecondType, TimeUnit, TimestampMicrosecondType, TimestampNanosecondType, UInt8Type, + UInt16Type, UInt32Type, }; use arrow::error::{ArrowError, Result}; use arrow::temporal_conversions::time64us_to_time; @@ -101,6 +102,9 @@ enum UnshredVariantRowBuilder<'a> { PrimitiveInt16(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), PrimitiveInt32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), PrimitiveInt64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveUInt8(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveUInt16(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), + PrimitiveUInt32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), PrimitiveFloat32(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), PrimitiveFloat64(UnshredPrimitiveRowBuilder<'a, PrimitiveArray>), Decimal32(DecimalUnshredRowBuilder<'a, Decimal32Type, VariantDecimal4>), @@ -146,6 +150,9 @@ impl<'a> UnshredVariantRowBuilder<'a> { Self::PrimitiveInt16(b) => b.append_row(builder, metadata, index), Self::PrimitiveInt32(b) => b.append_row(builder, metadata, index), Self::PrimitiveInt64(b) => b.append_row(builder, metadata, index), + Self::PrimitiveUInt8(b) => b.append_row(builder, metadata, index), + Self::PrimitiveUInt16(b) => b.append_row(builder, metadata, index), + Self::PrimitiveUInt32(b) => b.append_row(builder, metadata, index), Self::PrimitiveFloat32(b) => b.append_row(builder, metadata, index), Self::PrimitiveFloat64(b) => b.append_row(builder, metadata, index), Self::Decimal32(b) => b.append_row(builder, metadata, index), @@ -204,6 +211,9 @@ impl<'a> UnshredVariantRowBuilder<'a> { DataType::Int16 => primitive_builder!(PrimitiveInt16, as_primitive), DataType::Int32 => primitive_builder!(PrimitiveInt32, as_primitive), DataType::Int64 => primitive_builder!(PrimitiveInt64, as_primitive), + DataType::UInt8 => primitive_builder!(PrimitiveUInt8, as_primitive), + DataType::UInt16 => primitive_builder!(PrimitiveUInt16, as_primitive), + DataType::UInt32 => primitive_builder!(PrimitiveUInt32, as_primitive), DataType::Float32 => primitive_builder!(PrimitiveFloat32, as_primitive), DataType::Float64 => primitive_builder!(PrimitiveFloat64, as_primitive), DataType::Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => { @@ -443,6 +453,9 @@ impl_append_to_variant_builder!(PrimitiveArray); impl_append_to_variant_builder!(PrimitiveArray); impl_append_to_variant_builder!(PrimitiveArray); impl_append_to_variant_builder!(PrimitiveArray); +impl_append_to_variant_builder!(PrimitiveArray, |value| i16::from(value)); +impl_append_to_variant_builder!(PrimitiveArray, |value| i32::from(value)); +impl_append_to_variant_builder!(PrimitiveArray, |value| i64::from(value)); impl_append_to_variant_builder!(PrimitiveArray); impl_append_to_variant_builder!(PrimitiveArray); diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index 82cda7867f50..9fa13f6e0a78 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -28,7 +28,7 @@ use arrow::compute::cast; use arrow::datatypes::{ Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType, - TimestampMicrosecondType, TimestampNanosecondType, + TimestampMicrosecondType, TimestampNanosecondType, UInt8Type, UInt16Type, UInt32Type, }; use arrow::error::Result; use arrow_schema::extension::{ExtensionType, Uuid as UuidExtension}; @@ -996,6 +996,23 @@ fn typed_value_to_variant<'a>( DataType::Int64 => { primitive_conversion_single_value!(Int64Type, typed_value, index) } + DataType::UInt8 => { + generic_conversion_single_value!(UInt8Type, as_primitive, i16::from, typed_value, index) + } + DataType::UInt16 => generic_conversion_single_value!( + UInt16Type, + as_primitive, + i32::from, + typed_value, + index + ), + DataType::UInt32 => generic_conversion_single_value!( + UInt32Type, + as_primitive, + i64::from, + typed_value, + index + ), DataType::Float16 => { primitive_conversion_single_value!(Float16Type, typed_value, index) } @@ -1141,10 +1158,10 @@ fn canonicalize_and_verify_data_type(data_type: &DataType) -> Result borrow!(), - Int8 | Int16 | Int32 | Int64 | Float32 | Float64 => borrow!(), + Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | Float32 | Float64 => borrow!(), - // Unsigned integers and half-float are not allowed - UInt8 | UInt16 | UInt32 | UInt64 | Float16 => fail!(), + // UInt64 and half-float have no canonical Variant mapping + UInt64 | Float16 => fail!(), // Most decimal types are allowed, with restrictions on precision and scale // diff --git a/parquet-variant-compute/src/variant_get.rs b/parquet-variant-compute/src/variant_get.rs index 59407c617c39..64962a442233 100644 --- a/parquet-variant-compute/src/variant_get.rs +++ b/parquet-variant-compute/src/variant_get.rs @@ -497,6 +497,7 @@ mod test { LargeStringArray, ListArray, ListBuilder, ListViewArray, MapBuilder, NullArray, NullBuilder, StringArray, StringBuilder, StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, + UInt8Array, UInt16Array, UInt32Array, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::compute::{CastOptions, cast}; @@ -679,6 +680,21 @@ mod test { Int64Array, i64 ); + numeric_partially_shredded_variant_array_fn!( + partially_shredded_uint8_variant_array, + UInt8Array, + u8 + ); + numeric_partially_shredded_variant_array_fn!( + partially_shredded_uint16_variant_array, + UInt16Array, + u16 + ); + numeric_partially_shredded_variant_array_fn!( + partially_shredded_uint32_variant_array, + UInt32Array, + u32 + ); numeric_partially_shredded_variant_array_fn!( partially_shredded_float32_variant_array, Float32Array, @@ -739,6 +755,21 @@ mod test { numeric_partially_shredded_test!(f64, partially_shredded_float64_variant_array); } + #[test] + fn get_variant_partially_shredded_uint8_as_variant() { + numeric_partially_shredded_test!(i16, partially_shredded_uint8_variant_array); + } + + #[test] + fn get_variant_partially_shredded_uint16_as_variant() { + numeric_partially_shredded_test!(i32, partially_shredded_uint16_variant_array); + } + + #[test] + fn get_variant_partially_shredded_uint32_as_variant() { + numeric_partially_shredded_test!(i64, partially_shredded_uint32_variant_array); + } + #[test] fn get_variant_partially_shredded_bool_as_variant() { let array = partially_shredded_bool_variant_array(); diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 7dd7fb778909..b8d1b2e7b835 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1640,6 +1640,11 @@ fn write_leaf( let array = column.as_primitive::(); write_primitive(typed, array.values(), levels) } + ArrowDataType::UInt32 => { + let array: arrow_array::Int64Array = + column.as_primitive::().unary(i64::from); + write_primitive(typed, array.values(), levels) + } ArrowDataType::UInt64 => { let values = column.as_primitive::().values(); // follow C++ implementation and use overflow/reinterpret cast from u64 to i64 which will map diff --git a/parquet/src/arrow/schema/complex.rs b/parquet/src/arrow/schema/complex.rs index 161b26302825..167f846fa1e8 100644 --- a/parquet/src/arrow/schema/complex.rs +++ b/parquet/src/arrow/schema/complex.rs @@ -19,6 +19,8 @@ use std::collections::HashMap; use std::sync::Arc; use crate::arrow::schema::extension::try_add_extension_type; +#[cfg(feature = "variant_experimental")] +use crate::arrow::schema::extension::validate_variant_type; use crate::arrow::schema::primitive::convert_primitive; use crate::arrow::schema::virtual_type::{RowGroupIndex, RowNumber}; use crate::arrow::{PARQUET_FIELD_ID_META_KEY, ProjectionMask}; @@ -747,6 +749,13 @@ fn convert_field( match arrow_hint { Some(hint) => { + #[cfg(feature = "variant_experimental")] + if matches!( + parquet_type.get_basic_info().logical_type_ref(), + Some(crate::basic::LogicalType::Variant(_)) + ) { + validate_variant_type(parquet_type)?; + } // If the inferred type is a dictionary, preserve dictionary metadata #[allow(deprecated)] let field = match (&data_type, hint.dict_id(), hint.dict_is_ordered()) { diff --git a/parquet/src/arrow/schema/extension.rs b/parquet/src/arrow/schema/extension.rs index 4e1ec75084fc..ce80fceb815b 100644 --- a/parquet/src/arrow/schema/extension.rs +++ b/parquet/src/arrow/schema/extension.rs @@ -49,6 +49,7 @@ pub(crate) fn try_add_extension_type( Ok(match parquet_logical_type { #[cfg(feature = "variant_experimental")] LogicalType::Variant(_) => { + validate_variant_type(parquet_type)?; let mut arrow_field = arrow_field; arrow_field.try_with_extension_type(parquet_variant_compute::VariantType)?; arrow_field @@ -99,6 +100,39 @@ pub(crate) fn try_add_extension_type( }) } +#[cfg(feature = "variant_experimental")] +pub(crate) fn validate_variant_type(parquet_type: &Type) -> Result<(), ParquetError> { + use crate::basic::ConvertedType; + + match parquet_type { + Type::PrimitiveType { basic_info, .. } => { + let bit_width = match basic_info.logical_type_ref() { + Some(LogicalType::Integer(integer)) if !integer.is_signed => { + Some(integer.bit_width) + } + _ => match basic_info.converted_type() { + ConvertedType::UINT_8 => Some(8), + ConvertedType::UINT_16 => Some(16), + ConvertedType::UINT_32 => Some(32), + ConvertedType::UINT_64 => Some(64), + _ => None, + }, + }; + if let Some(bit_width) = bit_width { + return Err(ParquetError::General(format!( + "Illegal shredded value type: UInt{bit_width}" + ))); + } + } + Type::GroupType { fields, .. } => { + for field in fields { + validate_variant_type(field)?; + } + } + } + Ok(()) +} + /// Returns true if [`try_add_extension_type`] would add an extension type /// to the specified Parquet field. /// diff --git a/parquet/src/arrow/schema/mod.rs b/parquet/src/arrow/schema/mod.rs index f1d7ea7e887b..02cc5aced28a 100644 --- a/parquet/src/arrow/schema/mod.rs +++ b/parquet/src/arrow/schema/mod.rs @@ -39,6 +39,8 @@ pub mod virtual_type; use super::PARQUET_FIELD_ID_META_KEY; use crate::arrow::ProjectionMask; +#[cfg(feature = "variant_experimental")] +use crate::arrow::schema::extension::validate_variant_type; use crate::arrow::schema::extension::{ has_extension_type, logical_type_for_binary, logical_type_for_binary_view, logical_type_for_fixed_size_binary, logical_type_for_string, logical_type_for_struct, @@ -802,17 +804,33 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { if fields.is_empty() { return Err(arrow_err!("Parquet does not support writing empty structs",)); } + let logical_type = logical_type_for_struct(field); + #[cfg(feature = "variant_experimental")] + let is_variant = matches!(&logical_type, Some(LogicalType::Variant(_))); + #[cfg(feature = "variant_experimental")] + let variant_fields = if is_variant { + Some(variant_fields_for_parquet(fields)?) + } else { + None + }; + #[cfg(feature = "variant_experimental")] + let fields = variant_fields.as_ref().unwrap_or(fields); // recursively convert children to types/nodes let fields = fields .iter() .map(|f| arrow_to_parquet_type(f, coerce_types).map(Arc::new)) .collect::>()?; - Type::group_type_builder(name) + let parquet_type = Type::group_type_builder(name) .with_fields(fields) .with_repetition(repetition) .with_id(id) - .with_logical_type(logical_type_for_struct(field)) - .build() + .with_logical_type(logical_type) + .build()?; + #[cfg(feature = "variant_experimental")] + if is_variant { + validate_variant_type(&parquet_type)?; + } + Ok(parquet_type) } DataType::Map(field, _) => { if let DataType::Struct(struct_fields) = field.data_type() { @@ -867,6 +885,63 @@ fn arrow_to_parquet_type(field: &Field, coerce_types: bool) -> Result { } } +#[cfg(feature = "variant_experimental")] +fn variant_fields_for_parquet(fields: &Fields) -> Result { + fields + .iter() + .map(|field| { + if field.name() == "typed_value" { + variant_typed_value_for_parquet(field) + } else { + Ok(field.clone()) + } + }) + .collect() +} + +#[cfg(feature = "variant_experimental")] +fn variant_node_for_parquet(field: &FieldRef) -> Result { + let DataType::Struct(fields) = field.data_type() else { + return Err(arrow_err!( + "Invalid shredded Variant field: expected Struct, got {}", + field.data_type() + )); + }; + let fields = variant_fields_for_parquet(fields)?; + Ok(Arc::new( + field + .as_ref() + .clone() + .with_data_type(DataType::Struct(fields)), + )) +} + +#[cfg(feature = "variant_experimental")] +fn variant_typed_value_for_parquet(field: &FieldRef) -> Result { + let data_type = match field.data_type() { + DataType::UInt8 => DataType::Int16, + DataType::UInt16 => DataType::Int32, + DataType::UInt32 => DataType::Int64, + DataType::Struct(fields) => DataType::Struct( + fields + .iter() + .map(variant_node_for_parquet) + .collect::>()?, + ), + DataType::List(element) => DataType::List(variant_node_for_parquet(element)?), + DataType::LargeList(element) => DataType::LargeList(variant_node_for_parquet(element)?), + DataType::ListView(element) => DataType::ListView(variant_node_for_parquet(element)?), + DataType::LargeListView(element) => { + DataType::LargeListView(variant_node_for_parquet(element)?) + } + DataType::FixedSizeList(element, size) => { + DataType::FixedSizeList(variant_node_for_parquet(element)?, *size) + } + data_type => data_type.clone(), + }; + Ok(Arc::new(field.as_ref().clone().with_data_type(data_type))) +} + fn field_id(field: &Field) -> Option { let value = field.metadata().get(super::PARQUET_FIELD_ID_META_KEY)?; value.parse().ok() // Fail quietly if not a valid integer @@ -2369,6 +2444,87 @@ mod tests { Ok(()) } + #[test] + #[cfg(feature = "variant_experimental")] + fn variant_nested_unsigned() -> Result<()> { + use parquet_variant_compute::VariantType; + + let node = |data_type| { + DataType::Struct(Fields::from(vec![ + Field::new("value", DataType::Binary, true), + Field::new("typed_value", data_type, true), + ])) + }; + let list_element = Arc::new(Field::new("element", node(DataType::UInt8), false)); + let typed_value = DataType::Struct(Fields::from(vec![ + Field::new("number", node(DataType::UInt32), false), + Field::new("items", node(DataType::List(list_element)), false), + ])); + let variant = Field::new( + "variant", + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + Field::new("typed_value", typed_value, true), + ])), + true, + ) + .with_extension_type(VariantType); + + let parquet_schema = ArrowSchemaConverter::new().convert(&Schema::new(vec![variant]))?; + let number = parquet_schema.column(3); + let item = parquet_schema.column(6); + + assert_eq!(number.physical_type(), PhysicalType::INT64); + assert_eq!(number.logical_type_ref(), None); + assert_eq!(item.physical_type(), PhysicalType::INT32); + assert_eq!( + item.logical_type_ref(), + Some(&LogicalType::integer(16, true)) + ); + + Ok(()) + } + + #[test] + #[cfg(feature = "variant_experimental")] + fn reject_unsigned_variant_with_hint() { + use parquet_variant_compute::VariantType; + + let parquet_type = parse_message_type( + "message schema { + optional group variant (VARIANT) { + required binary metadata; + optional binary value; + optional int32 typed_value (INTEGER(32, false)); + } + }", + ) + .unwrap(); + let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_type)); + let variant = Field::new( + "variant", + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::Binary, false), + Field::new("value", DataType::Binary, true), + Field::new("typed_value", DataType::UInt32, true), + ])), + true, + ) + .with_extension_type(VariantType); + let metadata = vec![KeyValue::new( + crate::arrow::ARROW_SCHEMA_META_KEY.to_string(), + Some(encode_arrow_schema(&Schema::new(vec![variant]))), + )]; + + let error = parquet_to_arrow_schema(&parquet_schema, Some(&metadata)).unwrap_err(); + assert!( + error + .to_string() + .contains("Illegal shredded value type: UInt32") + ); + } + #[test] fn test_parquet_to_arrow_field_levels_with_virtual_rejects_non_virtual() { let message_type = " diff --git a/parquet/src/variant.rs b/parquet/src/variant.rs index 55df08673611..94b5a0eaacde 100644 --- a/parquet/src/variant.rs +++ b/parquet/src/variant.rs @@ -144,14 +144,15 @@ pub use parquet_variant_compute::*; mod tests { use crate::arrow::ArrowWriter; use crate::arrow::arrow_reader::ArrowReaderBuilder; + use crate::basic::{LogicalType, Type as PhysicalType}; use crate::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; use crate::file::reader::ChunkReader; use arrow::util::test_util::parquet_test_data; use arrow_array::{ArrayRef, RecordBatch}; - use arrow_schema::Schema; + use arrow_schema::{DataType, Schema}; use bytes::Bytes; use parquet_variant::{Variant, VariantBuilderExt}; - use parquet_variant_compute::{VariantArray, VariantArrayBuilder, VariantType}; + use parquet_variant_compute::{VariantArray, VariantArrayBuilder, VariantType, shred_variant}; use std::path::PathBuf; use std::sync::Arc; @@ -160,6 +161,54 @@ mod tests { roundtrip(variant_array()); } + #[test] + fn roundtrip_unsigned() { + let cases = [ + ( + vec![Variant::Int16(0), Variant::Int16(i16::from(u8::MAX))], + DataType::UInt8, + DataType::Int16, + PhysicalType::INT32, + Some(LogicalType::integer(16, true)), + ), + ( + vec![Variant::Int32(0), Variant::Int32(i32::from(u16::MAX))], + DataType::UInt16, + DataType::Int32, + PhysicalType::INT32, + None, + ), + ( + vec![Variant::Int64(0), Variant::Int64(i64::from(u32::MAX))], + DataType::UInt32, + DataType::Int64, + PhysicalType::INT64, + None, + ), + ]; + + for (expected, data_type, read_type, physical_type, logical_type) in cases { + let input = VariantArray::from_iter(expected.iter().cloned()); + let array = shred_variant(&input, &data_type).unwrap(); + assert_eq!(array.typed_value_column().unwrap().data_type(), &data_type); + + let batch = variant_array_to_batch(array); + let bytes = Bytes::from(write_to_buffer(&batch)); + let metadata = read_metadata(&bytes); + let typed_value = metadata.file_metadata().schema_descr().column(2); + assert_eq!(typed_value.physical_type(), physical_type); + assert_eq!(typed_value.logical_type_ref(), logical_type.as_ref()); + + let batch = read_to_batch(bytes); + let column = batch.column_by_name("data").unwrap(); + let result = VariantArray::try_new(column).unwrap(); + assert_eq!(result.typed_value_column().unwrap().data_type(), &read_type); + for (index, expected) in expected.into_iter().enumerate() { + assert_eq!(result.value(index), expected); + } + } + } + /// Ensure a file with Variant LogicalType, written by another writer in /// parquet-testing, can be read as a VariantArray #[test]