1use crate::VariantArrayBuilder;
21use crate::type_conversion::{
22 generic_conversion_single_value, generic_conversion_single_value_with_result,
23 primitive_conversion_single_value,
24};
25use arrow::array::{Array, ArrayRef, AsArray, StructArray, new_null_array};
26use arrow::buffer::NullBuffer;
27use arrow::compute::cast;
28use arrow::datatypes::{
29 Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, Float16Type, Float32Type,
30 Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType,
31 TimestampMicrosecondType, TimestampNanosecondType,
32};
33use arrow::error::Result;
34use arrow_schema::extension::{ExtensionType, Uuid as UuidExtension};
35use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, TimeUnit};
36use chrono::{DateTime, NaiveTime};
37use parquet_variant::{
38 Uuid, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16, VariantDecimalType as _,
39};
40
41use std::borrow::Cow;
42use std::sync::Arc;
43
44pub(crate) fn binary_array_value(array: &dyn Array, index: usize) -> Option<&[u8]> {
46 match array.data_type() {
47 DataType::Binary => Some(array.as_binary::<i32>().value(index)),
48 DataType::LargeBinary => Some(array.as_binary::<i64>().value(index)),
49 DataType::BinaryView => Some(array.as_binary_view().value(index)),
50 _ => None,
51 }
52}
53
54pub(crate) fn variant_from_arrays_at<'m, 'v>(
57 metadata: &'m dyn Array,
58 value: &'v dyn Array,
59 index: usize,
60) -> Option<Variant<'m, 'v>> {
61 let metadata = binary_array_value(metadata, index)?;
62 let value = binary_array_value(value, index)?;
63 Some(Variant::new(metadata, value))
64}
65
66pub(crate) fn all_null_value_column(len: usize) -> ArrayRef {
72 new_null_array(&DataType::BinaryView, len)
73}
74
75pub(crate) fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> {
77 match array.data_type() {
78 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()),
79 _ => Err(ArrowError::InvalidArgumentError(format!(
80 "VariantArray '{field_name}' field must be Binary, LargeBinary, or BinaryView, got {}",
81 array.data_type()
82 ))),
83 }
84}
85
86pub struct VariantType;
91
92impl ExtensionType for VariantType {
93 const NAME: &'static str = "arrow.parquet.variant";
94
95 type Metadata = &'static str;
98
99 fn metadata(&self) -> &Self::Metadata {
100 &""
101 }
102
103 fn serialize_metadata(&self) -> Option<String> {
104 Some(String::new())
105 }
106
107 fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata> {
108 Ok("")
109 }
110
111 fn supports_data_type(&self, data_type: &DataType) -> Result<()> {
112 if matches!(data_type, DataType::Struct(_)) {
113 Ok(())
114 } else {
115 Err(ArrowError::InvalidArgumentError(format!(
116 "VariantType only supports StructArray, got {data_type}"
117 )))
118 }
119 }
120
121 fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self> {
122 Self.supports_data_type(data_type)?;
123 Ok(Self)
124 }
125
126 fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<()> {
127 Self.supports_data_type(data_type)
128 }
129}
130
131#[derive(Debug, Clone)]
262pub struct VariantArray {
263 inner: StructArray,
265
266 metadata: ArrayRef,
268
269 shredding_state: ShreddingState,
271}
272
273impl VariantArray {
274 pub fn try_new(inner: &dyn Array) -> Result<Self> {
301 let inner = canonicalize_shredded_types(inner)?;
303
304 let Some(inner) = inner.as_struct_opt() else {
305 return Err(ArrowError::InvalidArgumentError(
306 "Invalid VariantArray: requires StructArray as input".to_string(),
307 ));
308 };
309
310 let Some(metadata_col) = inner.column_by_name("metadata") else {
314 return Err(ArrowError::InvalidArgumentError(
315 "Invalid VariantArray: StructArray must contain a 'metadata' field".to_string(),
316 ));
317 };
318 validate_binary_array(metadata_col.as_ref(), "metadata")?;
319
320 let shredding_state = ShreddingState::try_from(inner)?;
321
322 if inner.column_by_name("value").is_none() {
325 return Ok(Self::from_parts(
326 metadata_col.clone(),
327 shredding_state.value_column().clone(),
328 shredding_state.typed_value_column().cloned(),
329 inner.nulls().cloned(),
330 ));
331 }
332
333 Ok(Self {
335 inner: inner.clone(),
336 metadata: metadata_col.clone(),
337 shredding_state,
338 })
339 }
340
341 pub(crate) fn from_parts(
342 metadata: ArrayRef,
343 value: ArrayRef,
344 typed_value: Option<ArrayRef>,
345 nulls: Option<NullBuffer>,
346 ) -> Self {
347 let mut builder = StructArrayBuilder::new()
348 .with_field("metadata", metadata.clone(), false)
349 .with_field("value", value.clone(), true);
350 if let Some(typed_value) = typed_value.clone() {
351 builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
352 }
353 if let Some(nulls) = nulls {
354 builder = builder.with_nulls(nulls);
355 }
356
357 Self {
358 inner: builder.build(),
359 metadata,
360 shredding_state: ShreddingState::new(value, typed_value),
361 }
362 }
363
364 pub fn inner(&self) -> &StructArray {
366 &self.inner
367 }
368
369 pub fn into_inner(self) -> StructArray {
371 self.inner
372 }
373
374 pub fn shredding_state(&self) -> &ShreddingState {
376 &self.shredding_state
377 }
378
379 pub fn value(&self, index: usize) -> Variant<'_, '_> {
389 self.try_value(index).unwrap()
390 }
391
392 pub fn try_value(&self, index: usize) -> Result<Variant<'_, '_>> {
420 let value = self.value_column();
421 match self.typed_value_column() {
422 Some(typed_value) if typed_value.is_valid(index) => {
424 if !matches!(typed_value.data_type(), DataType::Struct(_)) && value.is_valid(index) {
425 panic!("Invalid variant, conflicting value and typed_value");
427 }
428 typed_value_to_variant(typed_value, index)
429 }
430 _ if value.is_valid(index) => variant_from_arrays_at(&self.metadata, value, index)
432 .ok_or_else(|| {
433 ArrowError::InvalidArgumentError(format!(
434 "metadata and value fields must be binary-like arrays, instead got {} and {}",
435 self.metadata.data_type(),
436 value.data_type()
437 ))
438 }),
439 _ => Ok(Variant::Null),
442 }
443 }
444
445 pub fn metadata_column(&self) -> &ArrayRef {
447 &self.metadata
448 }
449
450 pub fn value_column(&self) -> &ArrayRef {
452 self.shredding_state.value_column()
453 }
454
455 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
457 self.shredding_state.typed_value_column()
458 }
459
460 pub fn field(&self, name: impl Into<String>) -> Field {
463 Field::new(
464 name.into(),
465 self.data_type().clone(),
466 self.inner.is_nullable(),
467 )
468 .with_extension_type(VariantType)
469 }
470
471 pub fn data_type(&self) -> &DataType {
473 self.inner.data_type()
474 }
475
476 pub fn slice(&self, offset: usize, length: usize) -> Self {
477 let inner = self.inner.slice(offset, length);
478 let metadata = self.metadata.slice(offset, length);
479 let shredding_state = self.shredding_state.slice(offset, length);
480 Self {
481 inner,
482 metadata,
483 shredding_state,
484 }
485 }
486
487 pub fn len(&self) -> usize {
488 self.inner.len()
489 }
490
491 pub fn is_empty(&self) -> bool {
492 self.inner.is_empty()
493 }
494
495 pub fn nulls(&self) -> Option<&NullBuffer> {
496 self.inner.nulls()
497 }
498
499 pub fn is_null(&self, index: usize) -> bool {
501 self.nulls().is_some_and(|n| n.is_null(index))
502 }
503
504 pub fn is_valid(&self, index: usize) -> bool {
506 !self.is_null(index)
507 }
508
509 pub fn iter(&self) -> VariantArrayIter<'_> {
511 VariantArrayIter::new(self)
512 }
513}
514
515impl PartialEq for VariantArray {
516 fn eq(&self, other: &Self) -> bool {
517 self.inner == other.inner
518 }
519}
520
521impl From<VariantArray> for StructArray {
522 fn from(variant_array: VariantArray) -> Self {
523 variant_array.into_inner()
524 }
525}
526
527impl From<VariantArray> for ArrayRef {
528 fn from(variant_array: VariantArray) -> Self {
529 Arc::new(variant_array.into_inner())
530 }
531}
532
533impl<'m, 'v> FromIterator<Option<Variant<'m, 'v>>> for VariantArray {
534 fn from_iter<T: IntoIterator<Item = Option<Variant<'m, 'v>>>>(iter: T) -> Self {
535 let iter = iter.into_iter();
536
537 let mut b = VariantArrayBuilder::new(iter.size_hint().0);
538 b.extend(iter);
539 b.build()
540 }
541}
542
543impl<'m, 'v> FromIterator<Variant<'m, 'v>> for VariantArray {
544 fn from_iter<T: IntoIterator<Item = Variant<'m, 'v>>>(iter: T) -> Self {
545 Self::from_iter(iter.into_iter().map(Some))
546 }
547}
548
549#[derive(Debug)]
574pub struct VariantArrayIter<'a> {
575 array: &'a VariantArray,
576 head_i: usize,
577 tail_i: usize,
578}
579
580impl<'a> VariantArrayIter<'a> {
581 pub fn new(array: &'a VariantArray) -> Self {
583 Self {
584 array,
585 head_i: 0,
586 tail_i: array.len(),
587 }
588 }
589
590 fn value_opt(&self, i: usize) -> Option<Variant<'a, 'a>> {
591 self.array.is_valid(i).then(|| self.array.value(i))
592 }
593}
594
595impl<'a> Iterator for VariantArrayIter<'a> {
596 type Item = Option<Variant<'a, 'a>>;
597
598 #[inline]
599 fn next(&mut self) -> Option<Self::Item> {
600 if self.head_i == self.tail_i {
601 return None;
602 }
603
604 let out = self.value_opt(self.head_i);
605
606 self.head_i += 1;
607
608 Some(out)
609 }
610
611 fn size_hint(&self) -> (usize, Option<usize>) {
612 let remainder = self.tail_i - self.head_i;
613
614 (remainder, Some(remainder))
615 }
616}
617
618impl<'a> DoubleEndedIterator for VariantArrayIter<'a> {
619 fn next_back(&mut self) -> Option<Self::Item> {
620 if self.head_i == self.tail_i {
621 return None;
622 }
623
624 self.tail_i -= 1;
625
626 Some(self.value_opt(self.tail_i))
627 }
628}
629
630impl<'a> ExactSizeIterator for VariantArrayIter<'a> {}
631
632#[derive(Debug)]
667pub struct ShreddedVariantFieldArray {
668 inner: StructArray,
670 shredding_state: ShreddingState,
671}
672
673#[allow(unused)]
674impl ShreddedVariantFieldArray {
675 pub fn try_new(inner: &dyn Array) -> Result<Self> {
695 let Some(inner_struct) = inner.as_struct_opt() else {
696 return Err(ArrowError::InvalidArgumentError(
697 "Invalid ShreddedVariantFieldArray: requires StructArray as input".to_string(),
698 ));
699 };
700
701 let shredding_state = ShreddingState::try_from(inner_struct)?;
702
703 if inner_struct.column_by_name("value").is_none() {
706 return Ok(Self::from_parts(
707 shredding_state.value_column().clone(),
708 shredding_state.typed_value_column().cloned(),
709 inner_struct.nulls().cloned(),
710 ));
711 }
712
713 Ok(Self {
715 inner: inner_struct.clone(),
716 shredding_state,
717 })
718 }
719
720 pub fn shredding_state(&self) -> &ShreddingState {
722 &self.shredding_state
723 }
724
725 pub fn value_column(&self) -> &ArrayRef {
727 self.shredding_state.value_column()
728 }
729
730 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
732 self.shredding_state.typed_value_column()
733 }
734
735 pub fn inner(&self) -> &StructArray {
737 &self.inner
738 }
739
740 pub(crate) fn from_parts(
741 value: ArrayRef,
742 typed_value: Option<ArrayRef>,
743 nulls: Option<NullBuffer>,
744 ) -> Self {
745 let mut builder = StructArrayBuilder::new().with_field("value", value.clone(), true);
746 if let Some(typed_value) = typed_value.clone() {
747 builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
748 }
749 if let Some(nulls) = nulls {
750 builder = builder.with_nulls(nulls);
751 }
752
753 Self {
754 inner: builder.build(),
755 shredding_state: ShreddingState::new(value, typed_value),
756 }
757 }
758
759 pub fn into_inner(self) -> StructArray {
761 self.inner
762 }
763
764 pub fn data_type(&self) -> &DataType {
765 self.inner.data_type()
766 }
767
768 pub fn len(&self) -> usize {
769 self.inner.len()
770 }
771
772 pub fn is_empty(&self) -> bool {
773 self.inner.is_empty()
774 }
775
776 pub fn offset(&self) -> usize {
777 self.inner.offset()
778 }
779
780 pub fn nulls(&self) -> Option<&NullBuffer> {
781 None
785 }
786 pub fn is_null(&self, index: usize) -> bool {
788 self.nulls().is_some_and(|n| n.is_null(index))
789 }
790
791 pub fn is_valid(&self, index: usize) -> bool {
793 !self.is_null(index)
794 }
795}
796
797impl From<ShreddedVariantFieldArray> for ArrayRef {
798 fn from(array: ShreddedVariantFieldArray) -> Self {
799 Arc::new(array.into_inner())
800 }
801}
802
803impl From<ShreddedVariantFieldArray> for StructArray {
804 fn from(array: ShreddedVariantFieldArray) -> Self {
805 array.into_inner()
806 }
807}
808
809#[derive(Debug, Clone)]
846pub struct ShreddingState {
847 value: ArrayRef,
848 typed_value: Option<ArrayRef>,
849}
850
851impl ShreddingState {
852 pub fn new(value: ArrayRef, typed_value: Option<ArrayRef>) -> Self {
867 Self { value, typed_value }
868 }
869
870 pub fn value_column(&self) -> &ArrayRef {
872 &self.value
873 }
874
875 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
877 self.typed_value.as_ref()
878 }
879
880 pub fn slice(&self, offset: usize, length: usize) -> Self {
882 Self {
883 value: self.value.slice(offset, length),
884 typed_value: self.typed_value.as_ref().map(|tv| tv.slice(offset, length)),
885 }
886 }
887}
888
889impl TryFrom<&StructArray> for ShreddingState {
890 type Error = ArrowError;
891
892 fn try_from(inner_struct: &StructArray) -> Result<Self> {
893 let typed_value = inner_struct.column_by_name("typed_value").cloned();
894 let value = match inner_struct.column_by_name("value") {
895 Some(value) => {
896 validate_binary_array(value.as_ref(), "value")?;
897 value.clone()
898 }
899 None if typed_value.is_some() => all_null_value_column(inner_struct.len()),
902 None => {
903 return Err(ArrowError::InvalidArgumentError(
904 "Invalid VariantArray: StructArray must contain a 'value' field".to_string(),
905 ));
906 }
907 };
908 Ok(ShreddingState::new(value, typed_value))
909 }
910}
911
912fn typed_value_field(array: &ArrayRef) -> FieldRef {
918 let mut field = Field::new("typed_value", array.data_type().clone(), true);
919 if matches!(array.data_type(), DataType::FixedSizeBinary(16)) {
920 field = field.with_extension_type(UuidExtension);
921 }
922 Arc::new(field)
923}
924
925#[derive(Debug, Default, Clone)]
929pub(crate) struct StructArrayBuilder {
930 fields: Vec<FieldRef>,
931 arrays: Vec<ArrayRef>,
932 nulls: Option<NullBuffer>,
933}
934
935impl StructArrayBuilder {
936 pub fn new() -> Self {
937 Default::default()
938 }
939
940 pub fn with_field(mut self, field_name: &str, array: ArrayRef, nullable: bool) -> Self {
942 let field = Field::new(field_name, array.data_type().clone(), nullable);
943 self.fields.push(Arc::new(field));
944 self.arrays.push(array);
945 self
946 }
947
948 pub fn with_field_ref(mut self, field: FieldRef, array: ArrayRef) -> Self {
953 self.fields.push(field);
954 self.arrays.push(array);
955 self
956 }
957
958 pub fn with_nulls(mut self, nulls: NullBuffer) -> Self {
960 self.nulls = Some(nulls);
961 self
962 }
963
964 pub fn build(self) -> StructArray {
965 let Self {
966 fields,
967 arrays,
968 nulls,
969 } = self;
970 StructArray::new(Fields::from(fields), arrays, nulls)
971 }
972}
973
974fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Result<Variant<'_, '_>> {
976 let data_type = typed_value.data_type();
977 match data_type {
978 DataType::Null => Ok(Variant::Null),
979 DataType::Boolean => {
980 let boolean_array = typed_value.as_boolean();
981 let value = boolean_array.value(index);
982 Ok(Variant::from(value))
983 }
984 DataType::FixedSizeBinary(16) => {
986 let array = typed_value.as_fixed_size_binary();
987 let value = array.value(index);
988 Ok(Uuid::from_slice(value).unwrap().into()) }
990 DataType::Binary => {
991 let array = typed_value.as_binary::<i32>();
992 let value = array.value(index);
993 Ok(Variant::from(value))
994 }
995 DataType::LargeBinary => {
996 let array = typed_value.as_binary::<i64>();
997 let value = array.value(index);
998 Ok(Variant::from(value))
999 }
1000 DataType::BinaryView => {
1001 let array = typed_value.as_binary_view();
1002 let value = array.value(index);
1003 Ok(Variant::from(value))
1004 }
1005 DataType::Utf8 => {
1006 let array = typed_value.as_string::<i32>();
1007 let value = array.value(index);
1008 Ok(Variant::from(value))
1009 }
1010 DataType::LargeUtf8 => {
1011 let array = typed_value.as_string::<i64>();
1012 let value = array.value(index);
1013 Ok(Variant::from(value))
1014 }
1015 DataType::Utf8View => {
1016 let array = typed_value.as_string_view();
1017 let value = array.value(index);
1018 Ok(Variant::from(value))
1019 }
1020 DataType::Int8 => {
1021 primitive_conversion_single_value!(Int8Type, typed_value, index)
1022 }
1023 DataType::Int16 => {
1024 primitive_conversion_single_value!(Int16Type, typed_value, index)
1025 }
1026 DataType::Int32 => {
1027 primitive_conversion_single_value!(Int32Type, typed_value, index)
1028 }
1029 DataType::Int64 => {
1030 primitive_conversion_single_value!(Int64Type, typed_value, index)
1031 }
1032 DataType::Float16 => {
1033 primitive_conversion_single_value!(Float16Type, typed_value, index)
1034 }
1035 DataType::Float32 => {
1036 primitive_conversion_single_value!(Float32Type, typed_value, index)
1037 }
1038 DataType::Float64 => {
1039 primitive_conversion_single_value!(Float64Type, typed_value, index)
1040 }
1041 DataType::Decimal32(_, s) => {
1042 generic_conversion_single_value_with_result!(
1043 Decimal32Type,
1044 as_primitive,
1045 |v| VariantDecimal4::try_new(v, *s as u8),
1046 typed_value,
1047 index
1048 )
1049 }
1050 DataType::Decimal64(_, s) => {
1051 generic_conversion_single_value_with_result!(
1052 Decimal64Type,
1053 as_primitive,
1054 |v| VariantDecimal8::try_new(v, *s as u8),
1055 typed_value,
1056 index
1057 )
1058 }
1059 DataType::Decimal128(_, s) => {
1060 generic_conversion_single_value_with_result!(
1061 Decimal128Type,
1062 as_primitive,
1063 |v| VariantDecimal16::try_new(v, *s as u8),
1064 typed_value,
1065 index
1066 )
1067 }
1068 DataType::Date32 => {
1069 generic_conversion_single_value!(
1070 Date32Type,
1071 as_primitive,
1072 |v| Date32Type::to_naive_date_opt(v).unwrap(),
1073 typed_value,
1074 index
1075 )
1076 }
1077 DataType::Time64(TimeUnit::Microsecond) => {
1078 generic_conversion_single_value_with_result!(
1079 Time64MicrosecondType,
1080 as_primitive,
1081 |v| NaiveTime::from_num_seconds_from_midnight_opt(
1082 (v / 1_000_000) as u32,
1083 (v % 1_000_000) as u32 * 1000
1084 )
1085 .ok_or_else(|| format!("Invalid microsecond from midnight: {}", v)),
1086 typed_value,
1087 index
1088 )
1089 }
1090 DataType::Timestamp(TimeUnit::Microsecond, Some(_)) => {
1091 generic_conversion_single_value!(
1092 TimestampMicrosecondType,
1093 as_primitive,
1094 |v| DateTime::from_timestamp_micros(v).unwrap(),
1095 typed_value,
1096 index
1097 )
1098 }
1099 DataType::Timestamp(TimeUnit::Microsecond, None) => {
1100 generic_conversion_single_value!(
1101 TimestampMicrosecondType,
1102 as_primitive,
1103 |v| DateTime::from_timestamp_micros(v).unwrap().naive_utc(),
1104 typed_value,
1105 index
1106 )
1107 }
1108 DataType::Timestamp(TimeUnit::Nanosecond, Some(_)) => {
1109 generic_conversion_single_value!(
1110 TimestampNanosecondType,
1111 as_primitive,
1112 DateTime::from_timestamp_nanos,
1113 typed_value,
1114 index
1115 )
1116 }
1117 DataType::Timestamp(TimeUnit::Nanosecond, None) => {
1118 generic_conversion_single_value!(
1119 TimestampNanosecondType,
1120 as_primitive,
1121 |v| DateTime::from_timestamp_nanos(v).naive_utc(),
1122 typed_value,
1123 index
1124 )
1125 }
1126 _ => {
1129 debug_assert!(
1133 false,
1134 "Unsupported typed_value type: {}",
1135 typed_value.data_type()
1136 );
1137 Ok(Variant::Null)
1138 }
1139 }
1140}
1141
1142fn canonicalize_shredded_types(array: &dyn Array) -> Result<ArrayRef> {
1145 let new_type = canonicalize_and_verify_data_type(array.data_type())?;
1146 if let Cow::Borrowed(_) = new_type {
1147 if let Some(array) = array.as_struct_opt() {
1148 return Ok(Arc::new(array.clone())); }
1150 }
1151 cast(array, new_type.as_ref())
1152}
1153
1154fn canonicalize_and_verify_data_type(data_type: &DataType) -> Result<Cow<'_, DataType>> {
1158 use DataType::*;
1159
1160 macro_rules! fail {
1162 () => {
1163 return Err(ArrowError::InvalidArgumentError(format!(
1164 "Illegal shredded value type: {data_type}"
1165 )))
1166 };
1167 }
1168 macro_rules! borrow {
1169 () => {
1170 Cow::Borrowed(data_type)
1171 };
1172 }
1173
1174 let new_data_type = match data_type {
1175 Null | Boolean => borrow!(),
1177 Int8 | Int16 | Int32 | Int64 | Float32 | Float64 => borrow!(),
1178
1179 UInt8 | UInt16 | UInt32 | UInt64 | Float16 => fail!(),
1181
1182 Decimal64(p, s) | Decimal128(p, s)
1187 if VariantDecimal4::is_valid_precision_and_scale(p, s) =>
1188 {
1189 Cow::Owned(Decimal32(*p, *s))
1190 }
1191 Decimal128(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => {
1192 Cow::Owned(Decimal64(*p, *s))
1193 }
1194 Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => borrow!(),
1195 Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => borrow!(),
1196 Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => borrow!(),
1197 Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..) => fail!(),
1198
1199 Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, _) => borrow!(),
1201 Timestamp(TimeUnit::Millisecond | TimeUnit::Second, _) => fail!(),
1202
1203 Date32 | Time64(TimeUnit::Microsecond) => borrow!(),
1205 Date64 | Time32(_) | Time64(_) | Duration(_) | Interval(_) => fail!(),
1206
1207 Binary | LargeBinary | BinaryView | Utf8 | LargeUtf8 | Utf8View => borrow!(),
1209
1210 FixedSizeBinary(16) => borrow!(),
1212 FixedSizeBinary(_) | FixedSizeList(..) => fail!(),
1213
1214 List(field) => match canonicalize_and_verify_field(field)? {
1216 Cow::Borrowed(_) => borrow!(),
1217 Cow::Owned(new_field) => Cow::Owned(DataType::List(new_field)),
1218 },
1219 LargeList(field) => match canonicalize_and_verify_field(field)? {
1220 Cow::Borrowed(_) => borrow!(),
1221 Cow::Owned(new_field) => Cow::Owned(DataType::LargeList(new_field)),
1222 },
1223 ListView(field) => match canonicalize_and_verify_field(field)? {
1224 Cow::Borrowed(_) => borrow!(),
1225 Cow::Owned(new_field) => Cow::Owned(DataType::ListView(new_field)),
1226 },
1227 LargeListView(field) => match canonicalize_and_verify_field(field)? {
1228 Cow::Borrowed(_) => borrow!(),
1229 Cow::Owned(new_field) => Cow::Owned(DataType::LargeListView(new_field)),
1230 },
1231 Struct(fields) => {
1233 let mut new_fields = std::collections::HashMap::new();
1236 for (i, field) in fields.iter().enumerate() {
1237 if let Cow::Owned(new_field) = canonicalize_and_verify_field(field)? {
1238 new_fields.insert(i, new_field);
1239 }
1240 }
1241
1242 if new_fields.is_empty() {
1243 borrow!()
1244 } else {
1245 let new_fields = fields
1246 .iter()
1247 .enumerate()
1248 .map(|(i, field)| new_fields.remove(&i).unwrap_or_else(|| field.clone()));
1249 Cow::Owned(DataType::Struct(new_fields.collect()))
1250 }
1251 }
1252 Map(..) | Union(..) => fail!(),
1253
1254 Dictionary(..) | RunEndEncoded(..) => fail!(),
1256 };
1257 Ok(new_data_type)
1258}
1259
1260fn canonicalize_and_verify_field(field: &Arc<Field>) -> Result<Cow<'_, Arc<Field>>> {
1261 let new_data_type = canonicalize_and_verify_data_type(field.data_type())?;
1262
1263 if matches!(new_data_type.as_ref(), DataType::FixedSizeBinary(16))
1267 && !field.has_valid_extension_type::<UuidExtension>()
1268 {
1269 let new_field = field.as_ref().clone().with_extension_type(UuidExtension);
1270 return Ok(Cow::Owned(Arc::new(new_field)));
1271 }
1272
1273 let Cow::Owned(new_data_type) = new_data_type else {
1274 return Ok(Cow::Borrowed(field));
1275 };
1276 let new_field = field.as_ref().clone().with_data_type(new_data_type);
1277 Ok(Cow::Owned(Arc::new(new_field)))
1278}
1279
1280#[cfg(test)]
1283impl VariantArray {
1284 pub(crate) fn perfectly_shredded(
1285 metadata: ArrayRef,
1286 typed_value: ArrayRef,
1287 nulls: Option<NullBuffer>,
1288 ) -> Self {
1289 let value = all_null_value_column(typed_value.len());
1290 Self::from_parts(metadata, value, Some(typed_value), nulls)
1291 }
1292}
1293
1294#[cfg(test)]
1295impl ShreddedVariantFieldArray {
1296 pub(crate) fn perfectly_shredded(typed_value: ArrayRef) -> Self {
1297 let value = all_null_value_column(typed_value.len());
1298 Self::from_parts(value, Some(typed_value), None)
1299 }
1300}
1301
1302#[cfg(test)]
1303mod test {
1304 use crate::VariantArrayBuilder;
1305 use std::str::FromStr;
1306
1307 use super::*;
1308 use arrow::array::{
1309 BinaryArray, BinaryViewArray, Decimal32Array, Decimal64Array, Decimal128Array,
1310 FixedSizeBinaryArray, Int32Array, Int64Array, LargeBinaryArray, LargeListArray,
1311 LargeListViewArray, ListArray, ListViewArray, Time64MicrosecondArray,
1312 };
1313 use arrow::buffer::{OffsetBuffer, ScalarBuffer};
1314 use arrow_schema::{Field, Fields};
1315 use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, ShortString};
1316
1317 #[test]
1318 fn invalid_not_a_struct_array() {
1319 let array = make_binary_view_array();
1320 let err = VariantArray::try_new(&array);
1322 assert_eq!(
1323 err.unwrap_err().to_string(),
1324 "Invalid argument error: Invalid VariantArray: requires StructArray as input"
1325 );
1326 }
1327
1328 #[test]
1329 fn invalid_missing_metadata() {
1330 let fields = Fields::from(vec![Field::new("value", DataType::BinaryView, true)]);
1331 let array = StructArray::new(fields, vec![make_binary_view_array()], None);
1332 let err = VariantArray::try_new(&array);
1334 assert_eq!(
1335 err.unwrap_err().to_string(),
1336 "Invalid argument error: Invalid VariantArray: StructArray must contain a 'metadata' field"
1337 );
1338 }
1339
1340 #[test]
1341 fn read_missing_value_column() {
1342 let typed_value = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
1346 let metadata =
1347 BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 3));
1348 let struct_array = StructArrayBuilder::new()
1349 .with_field("metadata", Arc::new(metadata), false)
1350 .with_field("typed_value", typed_value, true)
1351 .build();
1352 assert!(struct_array.column_by_name("value").is_none());
1353
1354 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1355 assert_eq!(variant_array.value_column().len(), 3);
1356 assert_eq!(variant_array.value_column().null_count(), 3);
1357 assert!(variant_array.inner().column_by_name("value").is_some());
1358 assert!(variant_array.typed_value_column().is_some());
1359 assert_eq!(variant_array.value(0), Variant::from(1i64));
1360 assert_eq!(variant_array.value(1), Variant::Null);
1361 assert_eq!(variant_array.value(2), Variant::from(3i64));
1362
1363 let fields = Fields::from(vec![Field::new("metadata", DataType::BinaryView, false)]);
1365 let metadata_only = StructArray::new(fields, vec![make_binary_view_array()], None);
1366 let err = VariantArray::try_new(&metadata_only);
1367 assert_eq!(
1368 err.unwrap_err().to_string(),
1369 "Invalid argument error: Invalid VariantArray: StructArray must contain a 'value' field"
1370 );
1371 }
1372
1373 #[test]
1374 fn invalid_metadata_field_type() {
1375 let fields = Fields::from(vec![
1376 Field::new("metadata", DataType::Int32, true), Field::new("value", DataType::BinaryView, true),
1378 ]);
1379 let array = StructArray::new(
1380 fields,
1381 vec![make_int32_array(), make_binary_view_array()],
1382 None,
1383 );
1384 let err = VariantArray::try_new(&array);
1385 assert_eq!(
1386 err.unwrap_err().to_string(),
1387 "Invalid argument error: VariantArray 'metadata' field must be Binary, LargeBinary, or BinaryView, got Int32"
1388 );
1389 }
1390
1391 #[test]
1392 fn invalid_value_field_type() {
1393 let fields = Fields::from(vec![
1394 Field::new("metadata", DataType::BinaryView, true),
1395 Field::new("value", DataType::Int32, true),
1396 ]);
1397 let array = StructArray::new(
1398 fields,
1399 vec![make_binary_view_array(), make_int32_array()],
1400 None,
1401 );
1402 let err = VariantArray::try_new(&array);
1403 assert_eq!(
1404 err.unwrap_err().to_string(),
1405 "Invalid argument error: VariantArray 'value' field must be Binary, LargeBinary, or BinaryView, got Int32"
1406 );
1407 }
1408
1409 fn make_binary_view_array() -> ArrayRef {
1410 Arc::new(BinaryViewArray::from(vec![b"test" as &[u8]]))
1411 }
1412
1413 fn make_int32_array() -> ArrayRef {
1414 Arc::new(Int32Array::from(vec![1]))
1415 }
1416
1417 fn make_variant_struct_with_typed_value(typed_value: ArrayRef) -> StructArray {
1418 let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
1419 EMPTY_VARIANT_METADATA_BYTES,
1420 typed_value.len(),
1421 ));
1422 let value = new_null_array(&DataType::BinaryView, typed_value.len());
1423 StructArrayBuilder::new()
1424 .with_field("metadata", Arc::new(metadata), false)
1425 .with_field("value", value, true)
1426 .with_field("typed_value", typed_value, true)
1427 .build()
1428 }
1429
1430 #[test]
1431 fn try_new_tags_untagged_uuid_on_read() {
1432 let typed_value = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 2));
1435 let input = make_variant_struct_with_typed_value(Arc::new(typed_value.unwrap()));
1436
1437 let variant_array = VariantArray::try_new(&input).unwrap();
1439 let typed_value = variant_array.inner().field_by_name("typed_value").unwrap();
1440 assert_eq!(typed_value.data_type(), &DataType::FixedSizeBinary(16));
1441 assert!(typed_value.has_valid_extension_type::<UuidExtension>());
1442 }
1443
1444 #[test]
1445 fn try_new_tags_untagged_nested_uuid_on_read() {
1446 let leaf = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 1)).unwrap();
1449 let inner = StructArrayBuilder::new()
1450 .with_field("typed_value", Arc::new(leaf), true)
1451 .build();
1452 let object = StructArrayBuilder::new()
1453 .with_field("id", Arc::new(inner), false)
1454 .build();
1455 let input = make_variant_struct_with_typed_value(Arc::new(object));
1456
1457 let variant_array = VariantArray::try_new(&input).unwrap();
1459 let object = variant_array.typed_value_column().unwrap().as_struct();
1460 let id = object.column_by_name("id").unwrap().as_struct();
1461 let uuid_leaf = id.field_by_name("typed_value").unwrap();
1462 assert!(uuid_leaf.has_valid_extension_type::<UuidExtension>());
1463 }
1464
1465 #[test]
1466 fn all_null_value_column_is_valid_and_unshredded() {
1467 let metadata = BinaryViewArray::from(vec![b"test" as &[u8]; 3]);
1470 let value = new_null_array(&DataType::BinaryView, 3);
1471
1472 let fields = Fields::from(vec![
1473 Field::new("metadata", DataType::BinaryView, false),
1474 Field::new("value", DataType::BinaryView, true),
1475 ]);
1476 let struct_array = StructArray::new(fields, vec![Arc::new(metadata), value], None);
1477
1478 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1479 assert!(variant_array.typed_value_column().is_none());
1480
1481 for i in 0..variant_array.len() {
1484 assert!(variant_array.is_valid(i));
1485 assert_eq!(variant_array.value(i), Variant::Null);
1486 }
1487 }
1488
1489 #[test]
1490 fn canonicalize_and_verify_list_like_data_types() {
1491 let make_item_binary = || Arc::new(Field::new("item", DataType::Binary, true));
1495 let make_large_binary = || Arc::new(Field::new("item", DataType::LargeBinary, true));
1496 let make_item_binary_view = || Arc::new(Field::new("item", DataType::BinaryView, true));
1497
1498 let cases = vec![
1499 DataType::LargeList(make_item_binary()),
1501 DataType::ListView(make_item_binary()),
1502 DataType::LargeListView(make_item_binary()),
1503 DataType::LargeList(make_large_binary()),
1505 DataType::ListView(make_large_binary()),
1506 DataType::LargeListView(make_large_binary()),
1507 DataType::LargeList(make_item_binary_view()),
1509 DataType::ListView(make_item_binary_view()),
1510 DataType::LargeListView(make_item_binary_view()),
1511 ];
1512
1513 for input in cases {
1514 assert_eq!(
1515 canonicalize_and_verify_data_type(&input).unwrap().as_ref(),
1516 &input
1517 );
1518 }
1519 }
1520
1521 #[test]
1522 fn variant_array_try_new_supports_list_like_typed_value() {
1523 let item_field = Arc::new(Field::new("item", DataType::Int64, true));
1524 let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
1525
1526 let typed_values = vec![
1527 Arc::new(ListArray::new(
1528 item_field.clone(),
1529 OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
1530 values.clone(),
1531 None,
1532 )) as ArrayRef,
1533 Arc::new(LargeListArray::new(
1534 item_field.clone(),
1535 OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 3])),
1536 values.clone(),
1537 None,
1538 )) as ArrayRef,
1539 Arc::new(ListViewArray::new(
1540 item_field.clone(),
1541 ScalarBuffer::from(vec![0, 2]),
1542 ScalarBuffer::from(vec![2, 1]),
1543 values.clone(),
1544 None,
1545 )) as ArrayRef,
1546 Arc::new(LargeListViewArray::new(
1547 item_field,
1548 ScalarBuffer::from(vec![0_i64, 2]),
1549 ScalarBuffer::from(vec![2_i64, 1]),
1550 values,
1551 None,
1552 )) as ArrayRef,
1553 ];
1554
1555 for typed_value in typed_values {
1556 let input = make_variant_struct_with_typed_value(typed_value.clone());
1557 let variant_array = VariantArray::try_new(&input).unwrap();
1558 assert_eq!(
1559 variant_array.typed_value_column().unwrap().data_type(),
1560 typed_value.data_type(),
1561 );
1562 }
1563 }
1564
1565 #[test]
1566 fn test_variant_array_iterable() {
1567 let mut b = VariantArrayBuilder::new(6);
1568
1569 b.append_null();
1570 b.append_variant(Variant::from(1_i8));
1571 b.append_variant(Variant::Null);
1572 b.append_variant(Variant::from(2_i32));
1573 b.append_variant(Variant::from(3_i64));
1574 b.append_null();
1575
1576 let v = b.build();
1577
1578 let variants = v.iter().collect::<Vec<_>>();
1579
1580 assert_eq!(
1581 variants,
1582 vec![
1583 None,
1584 Some(Variant::Int8(1)),
1585 Some(Variant::Null),
1586 Some(Variant::Int32(2)),
1587 Some(Variant::Int64(3)),
1588 None,
1589 ]
1590 );
1591 }
1592
1593 #[test]
1594 fn test_variant_array_iter_double_ended() {
1595 let mut b = VariantArrayBuilder::new(5);
1596
1597 b.append_variant(Variant::from(0_i32));
1598 b.append_null();
1599 b.append_variant(Variant::from(2_i32));
1600 b.append_null();
1601 b.append_variant(Variant::from(4_i32));
1602
1603 let array = b.build();
1604 let mut iter = array.iter();
1605
1606 assert_eq!(iter.next(), Some(Some(Variant::from(0_i32))));
1607 assert_eq!(iter.next(), Some(None));
1608
1609 assert_eq!(iter.next_back(), Some(Some(Variant::from(4_i32))));
1610 assert_eq!(iter.next_back(), Some(None));
1611 assert_eq!(iter.next_back(), Some(Some(Variant::from(2_i32))));
1612
1613 assert_eq!(iter.next_back(), None);
1614 assert_eq!(iter.next(), None);
1615 }
1616
1617 #[test]
1618 fn test_variant_array_iter_reverse() {
1619 let mut b = VariantArrayBuilder::new(5);
1620
1621 b.append_variant(Variant::from("a"));
1622 b.append_null();
1623 b.append_variant(Variant::from("aaa"));
1624 b.append_null();
1625 b.append_variant(Variant::from("aaaaa"));
1626
1627 let array = b.build();
1628
1629 let result: Vec<_> = array.iter().rev().collect();
1630 assert_eq!(
1631 result,
1632 vec![
1633 Some(Variant::from("aaaaa")),
1634 None,
1635 Some(Variant::from("aaa")),
1636 None,
1637 Some(Variant::from("a")),
1638 ]
1639 );
1640 }
1641
1642 #[test]
1643 fn test_variant_array_iter_empty() {
1644 let v = VariantArrayBuilder::new(0).build();
1645 let mut i = v.iter();
1646 assert!(i.next().is_none());
1647 assert!(i.next_back().is_none());
1648 }
1649
1650 #[test]
1651 fn test_from_variant_opts_into_variant_array() {
1652 let v = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse), None];
1653
1654 let variant_array = VariantArray::from_iter(v);
1655
1656 assert_eq!(variant_array.len(), 4);
1657
1658 assert!(variant_array.is_null(0));
1659
1660 assert!(!variant_array.is_null(1));
1661 assert_eq!(variant_array.value(1), Variant::Null);
1662
1663 assert!(!variant_array.is_null(2));
1664 assert_eq!(variant_array.value(2), Variant::BooleanFalse);
1665
1666 assert!(variant_array.is_null(3));
1667 }
1668
1669 #[test]
1670 fn test_from_variants_into_variant_array() {
1671 let v = vec![
1672 Variant::Null,
1673 Variant::BooleanFalse,
1674 Variant::ShortString(ShortString::try_new("norm").unwrap()),
1675 ];
1676
1677 let variant_array = VariantArray::from_iter(v);
1678
1679 assert_eq!(variant_array.len(), 3);
1680
1681 assert!(!variant_array.is_null(0));
1682 assert_eq!(variant_array.value(0), Variant::Null);
1683
1684 assert!(!variant_array.is_null(1));
1685 assert_eq!(variant_array.value(1), Variant::BooleanFalse);
1686
1687 assert!(!variant_array.is_null(2));
1688 assert_eq!(
1689 variant_array.value(2),
1690 Variant::ShortString(ShortString::try_new("norm").unwrap())
1691 );
1692 }
1693
1694 #[test]
1695 fn test_variant_equality() {
1696 let v_iter = [None, Some(Variant::BooleanFalse), Some(Variant::Null), None];
1697 let v = VariantArray::from_iter(v_iter.clone());
1698
1699 {
1700 let v_copy = v.clone();
1701 assert_eq!(v, v_copy);
1702 }
1703
1704 {
1705 let v_iter_reversed = v_iter.iter().cloned().rev();
1706 let v_reversed = VariantArray::from_iter(v_iter_reversed);
1707
1708 assert_ne!(v, v_reversed);
1709 }
1710
1711 {
1712 let v_sliced = v.slice(0, 1);
1713 assert_ne!(v, v_sliced);
1714 }
1715 }
1716
1717 #[test]
1718 fn binary_typed_value_roundtrips() {
1719 let typed_value: ArrayRef = Arc::new(BinaryArray::from(vec![b"hello" as &[u8]]));
1721 let struct_array = make_variant_struct_with_typed_value(typed_value);
1722
1723 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1724 assert_eq!(variant_array.value(0), Variant::from(b"hello" as &[u8]));
1725 }
1726
1727 #[test]
1728 fn large_binary_typed_value_roundtrips() {
1729 let typed_value: ArrayRef = Arc::new(LargeBinaryArray::from(vec![b"world" as &[u8]]));
1731 let struct_array = make_variant_struct_with_typed_value(typed_value);
1732
1733 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1734 assert_eq!(variant_array.value(0), Variant::from(b"world" as &[u8]));
1735 }
1736
1737 macro_rules! invalid_variant_array_test {
1738 ($fn_name: ident, $invalid_typed_value: expr, $error_msg: literal) => {
1739 #[test]
1740 fn $fn_name() {
1741 let invalid_typed_value = $invalid_typed_value;
1742
1743 let struct_array =
1744 make_variant_struct_with_typed_value(Arc::new(invalid_typed_value));
1745
1746 let array: VariantArray = VariantArray::try_new(&struct_array)
1747 .expect("should create variant array")
1748 .into();
1749
1750 let result = array.try_value(0);
1751 assert!(result.is_err());
1752 let error = result.unwrap_err();
1753 assert!(matches!(error, ArrowError::CastError(_)));
1754
1755 let expected: &str = $error_msg;
1756 assert!(
1757 error.to_string().contains($error_msg),
1758 "error `{}` did not contain `{}`",
1759 error,
1760 expected
1761 )
1762 }
1763 };
1764 }
1765
1766 invalid_variant_array_test!(
1767 test_variant_array_invalide_time,
1768 Time64MicrosecondArray::from(vec![Some(86401000000)]),
1769 "Cast error: Cast failed at index 0 (array type: Time64(µs)): Invalid microsecond from midnight: 86401000000"
1770 );
1771
1772 invalid_variant_array_test!(
1773 test_variant_array_invalid_decimal32,
1774 Decimal32Array::from(vec![Some(1234567890)]),
1775 "Cast error: Cast failed at index 0 (array type: Decimal32(9, 2)): Invalid argument error: 1234567890 is wider than max precision 9"
1776 );
1777
1778 invalid_variant_array_test!(
1779 test_variant_array_invalid_decimal64,
1780 Decimal64Array::from(vec![Some(1234567890123456789)]),
1781 "Cast error: Cast failed at index 0 (array type: Decimal64(18, 6)): Invalid argument error: 1234567890123456789 is wider than max precision 18"
1782 );
1783
1784 invalid_variant_array_test!(
1785 test_variant_array_invalid_decimal128,
1786 Decimal128Array::from(vec![Some(
1787 i128::from_str("123456789012345678901234567890123456789").unwrap()
1788 ),]),
1789 "Cast error: Cast failed at index 0 (array type: Decimal128(38, 10)): Invalid argument error: 123456789012345678901234567890123456789 is wider than max precision 38"
1790 );
1791}