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::{
26 Array, ArrayRef, AsArray, StructArray, StructArrayAssembler, downcast_dictionary_array,
27 downcast_run_array, new_null_array,
28};
29use arrow::buffer::NullBuffer;
30use arrow::compute::cast;
31use arrow::datatypes::{
32 Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, Float16Type, Float32Type,
33 Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType,
34 TimestampMicrosecondType, TimestampNanosecondType,
35};
36use arrow::error::Result;
37use arrow_schema::extension::{ExtensionType, Uuid as UuidExtension};
38use arrow_schema::{ArrowError, DataType, Field, FieldRef, TimeUnit};
39use chrono::{DateTime, NaiveTime};
40use parquet_variant::{
41 Uuid, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16, VariantDecimalType as _,
42};
43
44use std::borrow::Cow;
45use std::sync::Arc;
46
47pub(crate) fn binary_array_value(array: &dyn Array, index: usize) -> Option<&[u8]> {
50 if array.is_null(index) {
51 return None;
52 }
53 match array.data_type() {
54 DataType::Binary => Some(array.as_binary::<i32>().value(index)),
55 DataType::LargeBinary => Some(array.as_binary::<i64>().value(index)),
56 DataType::BinaryView => Some(array.as_binary_view().value(index)),
57 DataType::Dictionary(..) => downcast_dictionary_array! {
58 array => {
59 let index = array.key(index)?;
60 binary_array_value(array.values().as_ref(), index)
61 },
62 _ => unreachable!(),
63 },
64 DataType::RunEndEncoded(..) => downcast_run_array! {
65 array => {
66 let index = array.get_physical_index(index);
67 binary_array_value(array.values().as_ref(), index)
68 },
69 _ => unreachable!(),
70 },
71 _ => None,
72 }
73}
74
75pub(crate) fn variant_from_arrays_at<'m, 'v>(
78 metadata: &'m dyn Array,
79 value: &'v dyn Array,
80 index: usize,
81) -> Option<Variant<'m, 'v>> {
82 let metadata = binary_array_value(metadata, index)?;
83 let value = binary_array_value(value, index)?;
84 Some(Variant::new(metadata, value))
85}
86
87pub(crate) fn all_null_value_column(len: usize) -> ArrayRef {
93 new_null_array(&DataType::BinaryView, len)
94}
95
96pub(crate) fn validate_binary_array(array: &dyn Array, field_name: &str) -> Result<()> {
98 match array.data_type() {
99 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()),
100 _ => Err(ArrowError::InvalidArgumentError(format!(
101 "VariantArray '{field_name}' field must be Binary, LargeBinary, or BinaryView, got {}",
102 array.data_type()
103 ))),
104 }
105}
106
107fn validate_metadata_array(array: &dyn Array) -> Result<()> {
109 let is_binary = |data_type: &DataType| {
110 matches!(
111 data_type,
112 DataType::Binary | DataType::LargeBinary | DataType::BinaryView
113 )
114 };
115 match array.data_type() {
116 data_type if is_binary(data_type) => Ok(()),
117 DataType::Dictionary(_, values) if is_binary(values) => Ok(()),
118 DataType::RunEndEncoded(_, values) if is_binary(values.data_type()) => Ok(()),
119 _ => Err(ArrowError::InvalidArgumentError(format!(
120 "VariantArray 'metadata' field must be Binary, LargeBinary, BinaryView, or a Dictionary or RunEndEncoded array of one of those types, got {}",
121 array.data_type()
122 ))),
123 }
124}
125
126pub struct VariantType;
131
132impl ExtensionType for VariantType {
133 const NAME: &'static str = "arrow.parquet.variant";
134
135 type Metadata = &'static str;
138
139 fn metadata(&self) -> &Self::Metadata {
140 &""
141 }
142
143 fn serialize_metadata(&self) -> Option<String> {
144 Some(String::new())
145 }
146
147 fn deserialize_metadata(_metadata: Option<&str>) -> Result<Self::Metadata> {
148 Ok("")
149 }
150
151 fn supports_data_type(&self, data_type: &DataType) -> Result<()> {
152 if matches!(data_type, DataType::Struct(_)) {
153 Ok(())
154 } else {
155 Err(ArrowError::InvalidArgumentError(format!(
156 "VariantType only supports StructArray, got {data_type}"
157 )))
158 }
159 }
160
161 fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result<Self> {
162 Self.supports_data_type(data_type)?;
163 Ok(Self)
164 }
165
166 fn validate(data_type: &DataType, _metadata: Self::Metadata) -> Result<()> {
167 Self.supports_data_type(data_type)
168 }
169}
170
171#[derive(Debug, Clone)]
313pub struct VariantArray {
314 inner: StructArray,
316
317 metadata: ArrayRef,
319
320 shredding_state: ShreddingState,
322}
323
324impl VariantArray {
325 pub fn try_new(inner: &dyn Array) -> Result<Self> {
349 let inner = canonicalize_shredded_types(inner)?;
351
352 let Some(inner) = inner.as_struct_opt() else {
353 return Err(ArrowError::InvalidArgumentError(
354 "Invalid VariantArray: requires StructArray as input".to_string(),
355 ));
356 };
357
358 let Some(metadata_col) = inner.column_by_name("metadata") else {
362 return Err(ArrowError::InvalidArgumentError(
363 "Invalid VariantArray: StructArray must contain a 'metadata' field".to_string(),
364 ));
365 };
366 validate_metadata_array(metadata_col.as_ref())?;
367
368 let shredding_state = ShreddingState::try_from(inner)?;
369
370 if inner.column_by_name("value").is_none() {
373 return Ok(Self::from_parts(
374 metadata_col.clone(),
375 shredding_state.value_column().clone(),
376 shredding_state.typed_value_column().cloned(),
377 inner.nulls().cloned(),
378 ));
379 }
380
381 Ok(Self {
383 inner: inner.clone(),
384 metadata: metadata_col.clone(),
385 shredding_state,
386 })
387 }
388
389 pub(crate) fn from_parts(
394 metadata: ArrayRef,
395 value: ArrayRef,
396 typed_value: Option<ArrayRef>,
397 nulls: Option<NullBuffer>,
398 ) -> Self {
399 Self::from_parts_with_nullable_value(metadata, value, typed_value, nulls, true)
400 }
401
402 pub(crate) fn from_parts_unshredded(
408 metadata: ArrayRef,
409 value: ArrayRef,
410 nulls: Option<NullBuffer>,
411 ) -> Self {
412 Self::from_parts_with_nullable_value(metadata, value, None, nulls, false)
413 }
414
415 fn from_parts_with_nullable_value(
416 metadata: ArrayRef,
417 value: ArrayRef,
418 typed_value: Option<ArrayRef>,
419 nulls: Option<NullBuffer>,
420 value_nullable: bool,
421 ) -> Self {
422 let mut builder = StructArrayAssembler::new()
423 .with_field("metadata", metadata.clone(), false)
424 .with_field("value", value.clone(), value_nullable);
425 if let Some(typed_value) = typed_value.clone() {
426 builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
427 }
428 if let Some(nulls) = nulls {
429 builder = builder.with_nulls(nulls);
430 }
431
432 Self {
433 inner: builder.build().unwrap(),
434 metadata,
435 shredding_state: ShreddingState::new(value, typed_value),
436 }
437 }
438
439 pub fn inner(&self) -> &StructArray {
441 &self.inner
442 }
443
444 pub fn into_inner(self) -> StructArray {
446 self.inner
447 }
448
449 pub fn shredding_state(&self) -> &ShreddingState {
451 &self.shredding_state
452 }
453
454 pub fn value(&self, index: usize) -> Variant<'_, '_> {
465 self.try_value(index)
466 .unwrap_or_else(|err| panic!("VariantArray::value({index}) failed: {err}"))
467 }
468
469 pub fn try_value(&self, index: usize) -> Result<Variant<'_, '_>> {
498 if self.len() <= index {
499 return Err(ArrowError::InvalidArgumentError(format!(
500 "Index {index} out of bounds for VariantArray of length {}",
501 self.len()
502 )));
503 }
504
505 let value = self.value_column();
506 match self.typed_value_column() {
507 Some(typed_value) if typed_value.is_valid(index) => {
509 if !matches!(typed_value.data_type(), DataType::Struct(_)) && value.is_valid(index) {
510 return Err(ArrowError::InvalidArgumentError(
512 "Invalid variant, conflicting value and typed_value".to_owned(),
513 ));
514 }
515 typed_value_to_variant(typed_value, index)
516 }
517 _ if value.is_valid(index) => variant_from_arrays_at(&self.metadata, value, index)
519 .ok_or_else(|| {
520 ArrowError::InvalidArgumentError(format!(
521 "metadata and value fields must be binary-like arrays, instead got {} and {}",
522 self.metadata.data_type(),
523 value.data_type()
524 ))
525 }),
526 _ => Ok(Variant::Null),
529 }
530 }
531
532 pub fn metadata_column(&self) -> &ArrayRef {
534 &self.metadata
535 }
536
537 pub fn value_column(&self) -> &ArrayRef {
539 self.shredding_state.value_column()
540 }
541
542 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
544 self.shredding_state.typed_value_column()
545 }
546
547 pub fn field(&self, name: impl Into<String>) -> Field {
550 Field::new(
551 name.into(),
552 self.data_type().clone(),
553 self.inner.is_nullable(),
554 )
555 .with_extension_type(VariantType)
556 }
557
558 pub fn data_type(&self) -> &DataType {
560 self.inner.data_type()
561 }
562
563 pub fn slice(&self, offset: usize, length: usize) -> Self {
564 let inner = self.inner.slice(offset, length);
565 let metadata = self.metadata.slice(offset, length);
566 let shredding_state = self.shredding_state.slice(offset, length);
567 Self {
568 inner,
569 metadata,
570 shredding_state,
571 }
572 }
573
574 pub fn len(&self) -> usize {
575 self.inner.len()
576 }
577
578 pub fn is_empty(&self) -> bool {
579 self.inner.is_empty()
580 }
581
582 pub fn nulls(&self) -> Option<&NullBuffer> {
583 self.inner.nulls()
584 }
585
586 pub fn is_null(&self, index: usize) -> bool {
588 self.nulls().is_some_and(|n| n.is_null(index))
589 }
590
591 pub fn is_valid(&self, index: usize) -> bool {
593 !self.is_null(index)
594 }
595
596 pub fn iter(&self) -> VariantArrayIter<'_> {
598 VariantArrayIter::new(self)
599 }
600}
601
602impl<'a> IntoIterator for &'a VariantArray {
603 type Item = Option<Variant<'a, 'a>>;
604 type IntoIter = VariantArrayIter<'a>;
605
606 fn into_iter(self) -> Self::IntoIter {
607 VariantArrayIter::new(self)
608 }
609}
610
611impl PartialEq for VariantArray {
612 fn eq(&self, other: &Self) -> bool {
613 self.inner == other.inner
614 }
615}
616
617impl From<VariantArray> for StructArray {
618 fn from(variant_array: VariantArray) -> Self {
619 variant_array.into_inner()
620 }
621}
622
623impl From<VariantArray> for ArrayRef {
624 fn from(variant_array: VariantArray) -> Self {
625 Arc::new(variant_array.into_inner())
626 }
627}
628
629impl<'m, 'v, V: Into<Variant<'m, 'v>>> FromIterator<Option<V>> for VariantArray {
630 fn from_iter<T: IntoIterator<Item = Option<V>>>(iter: T) -> Self {
631 let iter = iter.into_iter().map(|value| value.map(Into::into));
632
633 let mut b = VariantArrayBuilder::new(iter.size_hint().0);
634 b.extend(iter);
635 b.build()
636 }
637}
638
639impl<'m, 'v> FromIterator<Variant<'m, 'v>> for VariantArray {
640 fn from_iter<T: IntoIterator<Item = Variant<'m, 'v>>>(iter: T) -> Self {
641 Self::from_iter(iter.into_iter().map(Some))
642 }
643}
644
645#[derive(Debug)]
670pub struct VariantArrayIter<'a> {
671 array: &'a VariantArray,
672 head_i: usize,
673 tail_i: usize,
674}
675
676impl<'a> VariantArrayIter<'a> {
677 pub fn new(array: &'a VariantArray) -> Self {
679 Self {
680 array,
681 head_i: 0,
682 tail_i: array.len(),
683 }
684 }
685
686 fn value_opt(&self, i: usize) -> Option<Variant<'a, 'a>> {
687 self.array.is_valid(i).then(|| self.array.value(i))
688 }
689}
690
691impl<'a> Iterator for VariantArrayIter<'a> {
692 type Item = Option<Variant<'a, 'a>>;
693
694 #[inline]
695 fn next(&mut self) -> Option<Self::Item> {
696 if self.head_i == self.tail_i {
697 return None;
698 }
699
700 let out = self.value_opt(self.head_i);
701
702 self.head_i += 1;
703
704 Some(out)
705 }
706
707 fn size_hint(&self) -> (usize, Option<usize>) {
708 let remainder = self.tail_i - self.head_i;
709
710 (remainder, Some(remainder))
711 }
712}
713
714impl DoubleEndedIterator for VariantArrayIter<'_> {
715 fn next_back(&mut self) -> Option<Self::Item> {
716 if self.head_i == self.tail_i {
717 return None;
718 }
719
720 self.tail_i -= 1;
721
722 Some(self.value_opt(self.tail_i))
723 }
724}
725
726impl ExactSizeIterator for VariantArrayIter<'_> {}
727
728#[derive(Debug)]
763pub struct ShreddedVariantFieldArray {
764 inner: StructArray,
766 shredding_state: ShreddingState,
767}
768
769impl ShreddedVariantFieldArray {
770 pub fn try_new(inner: &dyn Array) -> Result<Self> {
790 let Some(inner_struct) = inner.as_struct_opt() else {
791 return Err(ArrowError::InvalidArgumentError(
792 "Invalid ShreddedVariantFieldArray: requires StructArray as input".to_string(),
793 ));
794 };
795
796 let shredding_state = ShreddingState::try_from(inner_struct)?;
797
798 if inner_struct.column_by_name("value").is_none() {
801 return Ok(Self::from_parts(
802 shredding_state.value_column().clone(),
803 shredding_state.typed_value_column().cloned(),
804 inner_struct.nulls().cloned(),
805 ));
806 }
807
808 Ok(Self {
810 inner: inner_struct.clone(),
811 shredding_state,
812 })
813 }
814
815 pub fn shredding_state(&self) -> &ShreddingState {
817 &self.shredding_state
818 }
819
820 pub fn value_column(&self) -> &ArrayRef {
822 self.shredding_state.value_column()
823 }
824
825 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
827 self.shredding_state.typed_value_column()
828 }
829
830 pub fn inner(&self) -> &StructArray {
832 &self.inner
833 }
834
835 pub(crate) fn from_parts(
836 value: ArrayRef,
837 typed_value: Option<ArrayRef>,
838 nulls: Option<NullBuffer>,
839 ) -> Self {
840 let mut builder = StructArrayAssembler::new().with_field("value", value.clone(), true);
841 if let Some(typed_value) = typed_value.clone() {
842 builder = builder.with_field_ref(typed_value_field(&typed_value), typed_value);
843 }
844 if let Some(nulls) = nulls {
845 builder = builder.with_nulls(nulls);
846 }
847
848 Self {
849 inner: builder.build().unwrap(),
850 shredding_state: ShreddingState::new(value, typed_value),
851 }
852 }
853
854 pub fn into_inner(self) -> StructArray {
856 self.inner
857 }
858
859 pub fn data_type(&self) -> &DataType {
860 self.inner.data_type()
861 }
862
863 pub fn len(&self) -> usize {
864 self.inner.len()
865 }
866
867 pub fn is_empty(&self) -> bool {
868 self.inner.is_empty()
869 }
870
871 pub fn offset(&self) -> usize {
872 self.inner.offset()
873 }
874
875 pub fn nulls(&self) -> Option<&NullBuffer> {
876 None
880 }
881 pub fn is_null(&self, index: usize) -> bool {
883 self.nulls().is_some_and(|n| n.is_null(index))
884 }
885
886 pub fn is_valid(&self, index: usize) -> bool {
888 !self.is_null(index)
889 }
890}
891
892impl From<ShreddedVariantFieldArray> for ArrayRef {
893 fn from(array: ShreddedVariantFieldArray) -> Self {
894 Arc::new(array.into_inner())
895 }
896}
897
898impl From<ShreddedVariantFieldArray> for StructArray {
899 fn from(array: ShreddedVariantFieldArray) -> Self {
900 array.into_inner()
901 }
902}
903
904#[derive(Debug, Clone)]
941pub struct ShreddingState {
942 value: ArrayRef,
943 typed_value: Option<ArrayRef>,
944}
945
946impl ShreddingState {
947 pub fn new(value: ArrayRef, typed_value: Option<ArrayRef>) -> Self {
962 Self { value, typed_value }
963 }
964
965 pub fn value_column(&self) -> &ArrayRef {
967 &self.value
968 }
969
970 pub fn typed_value_column(&self) -> Option<&ArrayRef> {
972 self.typed_value.as_ref()
973 }
974
975 pub fn slice(&self, offset: usize, length: usize) -> Self {
977 Self {
978 value: self.value.slice(offset, length),
979 typed_value: self.typed_value.as_ref().map(|tv| tv.slice(offset, length)),
980 }
981 }
982}
983
984impl TryFrom<&StructArray> for ShreddingState {
985 type Error = ArrowError;
986
987 fn try_from(inner_struct: &StructArray) -> Result<Self> {
988 let typed_value = inner_struct.column_by_name("typed_value").cloned();
989 let value = match inner_struct.column_by_name("value") {
990 Some(value) => {
991 validate_binary_array(value.as_ref(), "value")?;
992 value.clone()
993 }
994 None if typed_value.is_some() => all_null_value_column(inner_struct.len()),
997 None => {
998 return Err(ArrowError::InvalidArgumentError(
999 "Invalid VariantArray: StructArray must contain a 'value' field".to_string(),
1000 ));
1001 }
1002 };
1003 Ok(ShreddingState::new(value, typed_value))
1004 }
1005}
1006
1007fn typed_value_field(array: &ArrayRef) -> FieldRef {
1013 let mut field = Field::new("typed_value", array.data_type().clone(), true);
1014 if matches!(array.data_type(), DataType::FixedSizeBinary(16)) {
1015 field = field.with_extension_type(UuidExtension);
1016 }
1017 Arc::new(field)
1018}
1019
1020fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Result<Variant<'_, '_>> {
1022 let data_type = typed_value.data_type();
1023 match data_type {
1024 DataType::Null => Ok(Variant::Null),
1025 DataType::Boolean => {
1026 let boolean_array = typed_value.as_boolean();
1027 let value = boolean_array.value(index);
1028 Ok(Variant::from(value))
1029 }
1030 DataType::FixedSizeBinary(16) => {
1032 let array = typed_value.as_fixed_size_binary();
1033 let value = array.value(index);
1034 Ok(Uuid::from_slice(value).unwrap().into()) }
1036 DataType::Binary => {
1037 let array = typed_value.as_binary::<i32>();
1038 let value = array.value(index);
1039 Ok(Variant::from(value))
1040 }
1041 DataType::LargeBinary => {
1042 let array = typed_value.as_binary::<i64>();
1043 let value = array.value(index);
1044 Ok(Variant::from(value))
1045 }
1046 DataType::BinaryView => {
1047 let array = typed_value.as_binary_view();
1048 let value = array.value(index);
1049 Ok(Variant::from(value))
1050 }
1051 DataType::Utf8 => {
1052 let array = typed_value.as_string::<i32>();
1053 let value = array.value(index);
1054 Ok(Variant::from(value))
1055 }
1056 DataType::LargeUtf8 => {
1057 let array = typed_value.as_string::<i64>();
1058 let value = array.value(index);
1059 Ok(Variant::from(value))
1060 }
1061 DataType::Utf8View => {
1062 let array = typed_value.as_string_view();
1063 let value = array.value(index);
1064 Ok(Variant::from(value))
1065 }
1066 DataType::Int8 => {
1067 primitive_conversion_single_value!(Int8Type, typed_value, index)
1068 }
1069 DataType::Int16 => {
1070 primitive_conversion_single_value!(Int16Type, typed_value, index)
1071 }
1072 DataType::Int32 => {
1073 primitive_conversion_single_value!(Int32Type, typed_value, index)
1074 }
1075 DataType::Int64 => {
1076 primitive_conversion_single_value!(Int64Type, typed_value, index)
1077 }
1078 DataType::Float16 => {
1079 primitive_conversion_single_value!(Float16Type, typed_value, index)
1080 }
1081 DataType::Float32 => {
1082 primitive_conversion_single_value!(Float32Type, typed_value, index)
1083 }
1084 DataType::Float64 => {
1085 primitive_conversion_single_value!(Float64Type, typed_value, index)
1086 }
1087 DataType::Decimal32(_, s) => {
1088 generic_conversion_single_value_with_result!(
1089 Decimal32Type,
1090 as_primitive,
1091 |v| VariantDecimal4::try_new(v, *s as u8),
1092 typed_value,
1093 index
1094 )
1095 }
1096 DataType::Decimal64(_, s) => {
1097 generic_conversion_single_value_with_result!(
1098 Decimal64Type,
1099 as_primitive,
1100 |v| VariantDecimal8::try_new(v, *s as u8),
1101 typed_value,
1102 index
1103 )
1104 }
1105 DataType::Decimal128(_, s) => {
1106 generic_conversion_single_value_with_result!(
1107 Decimal128Type,
1108 as_primitive,
1109 |v| VariantDecimal16::try_new(v, *s as u8),
1110 typed_value,
1111 index
1112 )
1113 }
1114 DataType::Date32 => {
1115 generic_conversion_single_value!(
1116 Date32Type,
1117 as_primitive,
1118 |v| Date32Type::to_naive_date_opt(v).unwrap(),
1119 typed_value,
1120 index
1121 )
1122 }
1123 DataType::Time64(TimeUnit::Microsecond) => {
1124 generic_conversion_single_value_with_result!(
1125 Time64MicrosecondType,
1126 as_primitive,
1127 |v| NaiveTime::from_num_seconds_from_midnight_opt(
1128 (v / 1_000_000) as u32,
1129 (v % 1_000_000) as u32 * 1000
1130 )
1131 .ok_or_else(|| format!("Invalid microsecond from midnight: {v}")),
1132 typed_value,
1133 index
1134 )
1135 }
1136 DataType::Timestamp(TimeUnit::Microsecond, Some(_)) => {
1137 generic_conversion_single_value!(
1138 TimestampMicrosecondType,
1139 as_primitive,
1140 |v| DateTime::from_timestamp_micros(v).unwrap(),
1141 typed_value,
1142 index
1143 )
1144 }
1145 DataType::Timestamp(TimeUnit::Microsecond, None) => {
1146 generic_conversion_single_value!(
1147 TimestampMicrosecondType,
1148 as_primitive,
1149 |v| DateTime::from_timestamp_micros(v).unwrap().naive_utc(),
1150 typed_value,
1151 index
1152 )
1153 }
1154 DataType::Timestamp(TimeUnit::Nanosecond, Some(_)) => {
1155 generic_conversion_single_value!(
1156 TimestampNanosecondType,
1157 as_primitive,
1158 DateTime::from_timestamp_nanos,
1159 typed_value,
1160 index
1161 )
1162 }
1163 DataType::Timestamp(TimeUnit::Nanosecond, None) => {
1164 generic_conversion_single_value!(
1165 TimestampNanosecondType,
1166 as_primitive,
1167 |v| DateTime::from_timestamp_nanos(v).naive_utc(),
1168 typed_value,
1169 index
1170 )
1171 }
1172 _ => Err(ArrowError::NotYetImplemented(format!(
1179 "VariantArray::try_value cannot materialize typed_value of type {} \
1180 as a borrowed Variant; call unshred_variant first",
1181 typed_value.data_type()
1182 ))),
1183 }
1184}
1185
1186fn canonicalize_shredded_types(array: &dyn Array) -> Result<ArrayRef> {
1189 let new_type = canonicalize_and_verify_data_type_impl(array.data_type(), true)?;
1190 if let Cow::Borrowed(_) = new_type
1191 && let Some(array) = array.as_struct_opt()
1192 {
1193 return Ok(Arc::new(array.clone())); }
1195 cast(array, new_type.as_ref())
1196}
1197
1198fn canonicalize_and_verify_data_type(data_type: &DataType) -> Result<Cow<'_, DataType>> {
1202 canonicalize_and_verify_data_type_impl(data_type, false)
1203}
1204
1205fn canonicalize_and_verify_data_type_impl(
1206 data_type: &DataType,
1207 skip_top_level_metadata: bool,
1208) -> Result<Cow<'_, DataType>> {
1209 use DataType::*;
1210
1211 macro_rules! fail {
1213 () => {
1214 return Err(ArrowError::InvalidArgumentError(format!(
1215 "Illegal shredded value type: {data_type}"
1216 )))
1217 };
1218 }
1219 macro_rules! borrow {
1220 () => {
1221 Cow::Borrowed(data_type)
1222 };
1223 }
1224
1225 let new_data_type = match data_type {
1226 Null | Boolean => borrow!(),
1228 Int8 | Int16 | Int32 | Int64 | Float32 | Float64 => borrow!(),
1229
1230 UInt8 | UInt16 | UInt32 | UInt64 | Float16 => fail!(),
1232
1233 Decimal64(p, s) | Decimal128(p, s)
1238 if VariantDecimal4::is_valid_precision_and_scale(p, s) =>
1239 {
1240 Cow::Owned(Decimal32(*p, *s))
1241 }
1242 Decimal128(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => {
1243 Cow::Owned(Decimal64(*p, *s))
1244 }
1245 Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => borrow!(),
1246 Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => borrow!(),
1247 Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => borrow!(),
1248 Decimal32(..) | Decimal64(..) | Decimal128(..) | Decimal256(..) => fail!(),
1249
1250 Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, _) => borrow!(),
1252 Timestamp(TimeUnit::Millisecond | TimeUnit::Second, _) => fail!(),
1253
1254 Date32 | Time64(TimeUnit::Microsecond) => borrow!(),
1256 Date64 | Time32(_) | Time64(_) | Duration(_) | Interval(_) => fail!(),
1257
1258 Binary | LargeBinary | BinaryView | Utf8 | LargeUtf8 | Utf8View => borrow!(),
1260
1261 FixedSizeBinary(16) => borrow!(),
1263 FixedSizeBinary(_) => fail!(),
1264
1265 FixedSizeList(field, _) => match canonicalize_and_verify_field(field)? {
1269 Cow::Borrowed(_) => Cow::Owned(DataType::List(field.clone())),
1270 Cow::Owned(new_field) => Cow::Owned(DataType::List(new_field)),
1271 },
1272
1273 List(field) => match canonicalize_and_verify_field(field)? {
1275 Cow::Borrowed(_) => borrow!(),
1276 Cow::Owned(new_field) => Cow::Owned(DataType::List(new_field)),
1277 },
1278 LargeList(field) => match canonicalize_and_verify_field(field)? {
1279 Cow::Borrowed(_) => borrow!(),
1280 Cow::Owned(new_field) => Cow::Owned(DataType::LargeList(new_field)),
1281 },
1282 ListView(field) => match canonicalize_and_verify_field(field)? {
1283 Cow::Borrowed(_) => borrow!(),
1284 Cow::Owned(new_field) => Cow::Owned(DataType::ListView(new_field)),
1285 },
1286 LargeListView(field) => match canonicalize_and_verify_field(field)? {
1287 Cow::Borrowed(_) => borrow!(),
1288 Cow::Owned(new_field) => Cow::Owned(DataType::LargeListView(new_field)),
1289 },
1290 Struct(fields) => {
1292 let mut new_fields = std::collections::HashMap::new();
1295 for (i, field) in fields.iter().enumerate() {
1296 if skip_top_level_metadata && field.name() == "metadata" {
1297 continue;
1298 }
1299 if let Cow::Owned(new_field) = canonicalize_and_verify_field(field)? {
1300 new_fields.insert(i, new_field);
1301 }
1302 }
1303
1304 if new_fields.is_empty() {
1305 borrow!()
1306 } else {
1307 let new_fields = fields
1308 .iter()
1309 .enumerate()
1310 .map(|(i, field)| new_fields.remove(&i).unwrap_or_else(|| field.clone()));
1311 Cow::Owned(DataType::Struct(new_fields.collect()))
1312 }
1313 }
1314 Map(..) | Union(..) => fail!(),
1315
1316 Dictionary(..) | RunEndEncoded(..) => fail!(),
1318 };
1319 Ok(new_data_type)
1320}
1321
1322fn canonicalize_and_verify_field(field: &Arc<Field>) -> Result<Cow<'_, Arc<Field>>> {
1323 let new_data_type = canonicalize_and_verify_data_type(field.data_type())?;
1324
1325 if matches!(new_data_type.as_ref(), DataType::FixedSizeBinary(16))
1329 && !field.has_valid_extension_type::<UuidExtension>()
1330 {
1331 let new_field = field.as_ref().clone().with_extension_type(UuidExtension);
1332 return Ok(Cow::Owned(Arc::new(new_field)));
1333 }
1334
1335 let Cow::Owned(new_data_type) = new_data_type else {
1336 return Ok(Cow::Borrowed(field));
1337 };
1338 let new_field = field.as_ref().clone().with_data_type(new_data_type);
1339 Ok(Cow::Owned(Arc::new(new_field)))
1340}
1341
1342#[cfg(test)]
1345impl VariantArray {
1346 pub(crate) fn perfectly_shredded(
1347 metadata: ArrayRef,
1348 typed_value: ArrayRef,
1349 nulls: Option<NullBuffer>,
1350 ) -> Self {
1351 let value = all_null_value_column(typed_value.len());
1352 Self::from_parts(metadata, value, Some(typed_value), nulls)
1353 }
1354}
1355
1356#[cfg(test)]
1357impl ShreddedVariantFieldArray {
1358 pub(crate) fn perfectly_shredded(typed_value: ArrayRef) -> Self {
1359 let value = all_null_value_column(typed_value.len());
1360 Self::from_parts(value, Some(typed_value), None)
1361 }
1362}
1363
1364#[cfg(test)]
1365mod test {
1366 use crate::{GetOptions, VariantArrayBuilder, json_to_variant, variant_get, variant_to_json};
1367 use std::str::FromStr;
1368
1369 use super::*;
1370 use arrow::array::{
1371 BinaryArray, BinaryDictionaryBuilder, BinaryRunBuilder, BinaryViewArray, Decimal32Array,
1372 Decimal64Array, Decimal128Array, FixedSizeBinaryArray, FixedSizeListArray, Int8Array,
1373 Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray, ListArray,
1374 ListViewArray, StringArray, Time64MicrosecondArray,
1375 };
1376 use arrow::buffer::{OffsetBuffer, ScalarBuffer};
1377 use arrow_schema::{Field, Fields};
1378 use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, ShortString};
1379
1380 #[test]
1381 fn invalid_not_a_struct_array() {
1382 let array = make_binary_view_array();
1383 let err = VariantArray::try_new(&array);
1385 assert_eq!(
1386 err.unwrap_err().to_string(),
1387 "Invalid argument error: Invalid VariantArray: requires StructArray as input"
1388 );
1389 }
1390
1391 #[test]
1392 fn invalid_missing_metadata() {
1393 let fields = Fields::from(vec![Field::new("value", DataType::BinaryView, true)]);
1394 let array = StructArray::new(fields, vec![make_binary_view_array()], None);
1395 let err = VariantArray::try_new(&array);
1397 assert_eq!(
1398 err.unwrap_err().to_string(),
1399 "Invalid argument error: Invalid VariantArray: StructArray must contain a 'metadata' field"
1400 );
1401 }
1402
1403 #[test]
1404 fn read_missing_value_column() {
1405 let typed_value = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])) as ArrayRef;
1409 let metadata =
1410 BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 3));
1411 let struct_array = StructArrayAssembler::new()
1412 .with_field("metadata", Arc::new(metadata), false)
1413 .with_field("typed_value", typed_value, true)
1414 .build()
1415 .unwrap();
1416 assert!(struct_array.column_by_name("value").is_none());
1417
1418 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1419 assert_eq!(variant_array.value_column().len(), 3);
1420 assert_eq!(variant_array.value_column().null_count(), 3);
1421 assert!(variant_array.inner().column_by_name("value").is_some());
1422 assert!(variant_array.typed_value_column().is_some());
1423 assert_eq!(variant_array.value(0), Variant::from(1i64));
1424 assert_eq!(variant_array.value(1), Variant::Null);
1425 assert_eq!(variant_array.value(2), Variant::from(3i64));
1426
1427 let fields = Fields::from(vec![Field::new("metadata", DataType::BinaryView, false)]);
1429 let metadata_only = StructArray::new(fields, vec![make_binary_view_array()], None);
1430 let err = VariantArray::try_new(&metadata_only);
1431 assert_eq!(
1432 err.unwrap_err().to_string(),
1433 "Invalid argument error: Invalid VariantArray: StructArray must contain a 'value' field"
1434 );
1435 }
1436
1437 #[test]
1438 fn invalid_metadata_field_type() {
1439 let fields = Fields::from(vec![
1440 Field::new("metadata", DataType::Int32, true), Field::new("value", DataType::BinaryView, true),
1442 ]);
1443 let array = StructArray::new(
1444 fields,
1445 vec![make_int32_array(), make_binary_view_array()],
1446 None,
1447 );
1448 let err = VariantArray::try_new(&array);
1449 assert_eq!(
1450 err.unwrap_err().to_string(),
1451 "Invalid argument error: VariantArray 'metadata' field must be Binary, LargeBinary, BinaryView, or a Dictionary or RunEndEncoded array of one of those types, got Int32"
1452 );
1453 }
1454
1455 #[test]
1456 fn encoded_metadata_supports_nulls_slices_and_variant_get() {
1457 let json: ArrayRef = Arc::new(StringArray::from(vec![
1458 Some(r#"{"a":0}"#),
1459 Some(r#"{"a":1}"#),
1460 None,
1461 Some(r#"{"b":3}"#),
1462 Some(r#"{"b":4}"#),
1463 ]));
1464 let baseline = json_to_variant(&json).unwrap();
1465 let metadata = baseline.metadata_column().as_binary_view();
1466 let metadata_a = metadata.value(0);
1467 let metadata_b = metadata.value(3);
1468
1469 let logical_metadata = [
1470 Some(metadata_a),
1471 Some(metadata_a),
1472 None,
1473 Some(metadata_b),
1474 Some(metadata_b),
1475 ];
1476
1477 let mut dictionary = BinaryDictionaryBuilder::<Int8Type>::new();
1478 dictionary.extend(logical_metadata);
1479 let dictionary: ArrayRef = Arc::new(dictionary.finish());
1480
1481 let mut ree = BinaryRunBuilder::<Int16Type>::new();
1482 ree.extend(logical_metadata);
1483 let run_end_encoded: ArrayRef = Arc::new(ree.finish());
1484
1485 for metadata in [dictionary, run_end_encoded] {
1486 assert_eq!(binary_array_value(metadata.as_ref(), 2), None);
1487 let fields = Fields::from(vec![
1488 Field::new("metadata", metadata.data_type().clone(), false),
1489 Field::new("value", baseline.value_column().data_type().clone(), false),
1490 ]);
1491 let input = StructArray::try_new(
1492 fields,
1493 vec![metadata, baseline.value_column().clone()],
1494 baseline.nulls().cloned(),
1495 )
1496 .unwrap()
1497 .slice(1, 3);
1498
1499 let variant = VariantArray::try_new(&input).unwrap();
1500 assert_eq!(variant.value(0), baseline.value(1));
1501 assert!(variant.is_null(1));
1502 assert_eq!(variant.value(2), baseline.value(3));
1503
1504 let input: ArrayRef = Arc::new(input);
1505 let options = GetOptions::new_with_path("b".try_into().unwrap())
1506 .with_as_type(Some(Arc::new(Field::new("b", DataType::Int8, true))));
1507 let result = variant_get(&input, options).unwrap();
1508 assert_eq!(
1509 result.as_primitive::<Int8Type>(),
1510 &Int8Array::from(vec![None, None, Some(3)])
1511 );
1512 assert_eq!(
1513 variant_to_json(&input).unwrap(),
1514 StringArray::from(vec![Some(r#"{"a":1}"#), None, Some(r#"{"b":3}"#)])
1515 );
1516 }
1517 }
1518
1519 #[test]
1520 fn invalid_value_field_type() {
1521 let fields = Fields::from(vec![
1522 Field::new("metadata", DataType::BinaryView, true),
1523 Field::new("value", DataType::Int32, true),
1524 ]);
1525 let array = StructArray::new(
1526 fields,
1527 vec![make_binary_view_array(), make_int32_array()],
1528 None,
1529 );
1530 let err = VariantArray::try_new(&array);
1531 assert_eq!(
1532 err.unwrap_err().to_string(),
1533 "Invalid argument error: VariantArray 'value' field must be Binary, LargeBinary, or BinaryView, got Int32"
1534 );
1535 }
1536
1537 fn make_binary_view_array() -> ArrayRef {
1538 Arc::new(BinaryViewArray::from(vec![b"test" as &[u8]]))
1539 }
1540
1541 fn make_int32_array() -> ArrayRef {
1542 Arc::new(Int32Array::from(vec![1]))
1543 }
1544
1545 fn make_variant_struct_with_typed_value(typed_value: ArrayRef) -> StructArray {
1546 let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
1547 EMPTY_VARIANT_METADATA_BYTES,
1548 typed_value.len(),
1549 ));
1550 let value = new_null_array(&DataType::BinaryView, typed_value.len());
1551 StructArrayAssembler::new()
1552 .with_field("metadata", Arc::new(metadata), false)
1553 .with_field("value", value, true)
1554 .with_field("typed_value", typed_value, true)
1555 .build()
1556 .unwrap()
1557 }
1558
1559 #[test]
1560 fn try_new_tags_untagged_uuid_on_read() {
1561 let typed_value = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 2));
1564 let input = make_variant_struct_with_typed_value(Arc::new(typed_value.unwrap()));
1565
1566 let variant_array = VariantArray::try_new(&input).unwrap();
1568 let typed_value = variant_array.inner().field_by_name("typed_value").unwrap();
1569 assert_eq!(typed_value.data_type(), &DataType::FixedSizeBinary(16));
1570 assert!(typed_value.has_valid_extension_type::<UuidExtension>());
1571 }
1572
1573 #[test]
1574 fn try_new_tags_untagged_nested_uuid_on_read() {
1575 let leaf = FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n([0u8; 16], 1)).unwrap();
1578 let inner = StructArrayAssembler::new()
1579 .with_field("typed_value", Arc::new(leaf), true)
1580 .build()
1581 .unwrap();
1582 let object = StructArrayAssembler::new()
1583 .with_field("id", Arc::new(inner), false)
1584 .build()
1585 .unwrap();
1586 let input = make_variant_struct_with_typed_value(Arc::new(object));
1587
1588 let variant_array = VariantArray::try_new(&input).unwrap();
1590 let object = variant_array.typed_value_column().unwrap().as_struct();
1591 let id = object.column_by_name("id").unwrap().as_struct();
1592 let uuid_leaf = id.field_by_name("typed_value").unwrap();
1593 assert!(uuid_leaf.has_valid_extension_type::<UuidExtension>());
1594 }
1595
1596 #[test]
1597 fn all_null_value_column_is_valid_and_unshredded() {
1598 let metadata = BinaryViewArray::from(vec![b"test" as &[u8]; 3]);
1601 let value = new_null_array(&DataType::BinaryView, 3);
1602
1603 let fields = Fields::from(vec![
1604 Field::new("metadata", DataType::BinaryView, false),
1605 Field::new("value", DataType::BinaryView, true),
1606 ]);
1607 let struct_array = StructArray::new(fields, vec![Arc::new(metadata), value], None);
1608
1609 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1610 assert!(variant_array.typed_value_column().is_none());
1611
1612 for i in 0..variant_array.len() {
1615 assert!(variant_array.is_valid(i));
1616 assert_eq!(variant_array.value(i), Variant::Null);
1617 }
1618 }
1619
1620 #[test]
1621 fn canonicalize_and_verify_list_like_data_types() {
1622 let make_item_binary = || Arc::new(Field::new("item", DataType::Binary, true));
1626 let make_large_binary = || Arc::new(Field::new("item", DataType::LargeBinary, true));
1627 let make_item_binary_view = || Arc::new(Field::new("item", DataType::BinaryView, true));
1628
1629 let cases = vec![
1630 DataType::LargeList(make_item_binary()),
1632 DataType::ListView(make_item_binary()),
1633 DataType::LargeListView(make_item_binary()),
1634 DataType::LargeList(make_large_binary()),
1636 DataType::ListView(make_large_binary()),
1637 DataType::LargeListView(make_large_binary()),
1638 DataType::LargeList(make_item_binary_view()),
1640 DataType::ListView(make_item_binary_view()),
1641 DataType::LargeListView(make_item_binary_view()),
1642 ];
1643
1644 for input in cases {
1645 assert_eq!(
1646 canonicalize_and_verify_data_type(&input).unwrap().as_ref(),
1647 &input
1648 );
1649 }
1650 }
1651
1652 #[test]
1653 fn variant_array_try_new_supports_list_like_typed_value() {
1654 let item_field = Arc::new(Field::new("item", DataType::Int64, true));
1655 let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
1656
1657 let typed_values = vec![
1658 Arc::new(ListArray::new(
1659 item_field.clone(),
1660 OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
1661 values.clone(),
1662 None,
1663 )) as ArrayRef,
1664 Arc::new(LargeListArray::new(
1665 item_field.clone(),
1666 OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 3])),
1667 values.clone(),
1668 None,
1669 )) as ArrayRef,
1670 Arc::new(ListViewArray::new(
1671 item_field.clone(),
1672 ScalarBuffer::from(vec![0, 2]),
1673 ScalarBuffer::from(vec![2, 1]),
1674 values.clone(),
1675 None,
1676 )) as ArrayRef,
1677 Arc::new(LargeListViewArray::new(
1678 item_field,
1679 ScalarBuffer::from(vec![0_i64, 2]),
1680 ScalarBuffer::from(vec![2_i64, 1]),
1681 values,
1682 None,
1683 )) as ArrayRef,
1684 ];
1685
1686 for typed_value in typed_values {
1687 let input = make_variant_struct_with_typed_value(typed_value.clone());
1688 let variant_array = VariantArray::try_new(&input).unwrap();
1689 assert_eq!(
1690 variant_array.typed_value_column().unwrap().data_type(),
1691 typed_value.data_type(),
1692 );
1693 }
1694 }
1695
1696 #[test]
1697 fn variant_array_try_new_normalizes_fixed_size_list_typed_value() {
1698 let element_values: ArrayRef =
1699 ShreddedVariantFieldArray::perfectly_shredded(Arc::new(Int64Array::from(vec![
1700 1, 2, 3, 4,
1701 ])))
1702 .into();
1703 let item_field = Arc::new(Field::new("item", element_values.data_type().clone(), true));
1704 let typed_value: ArrayRef = Arc::new(FixedSizeListArray::new(
1705 item_field.clone(),
1706 2,
1707 element_values,
1708 None,
1709 ));
1710 let input = make_variant_struct_with_typed_value(typed_value);
1711
1712 let variant_array = VariantArray::try_new(&input).unwrap();
1713 assert_eq!(
1714 variant_array.typed_value_column().unwrap().data_type(),
1715 &DataType::List(item_field),
1716 );
1717
1718 let unshredded = crate::unshred_variant(&variant_array).unwrap();
1719 assert!(unshredded.typed_value_column().is_none());
1720 assert_eq!(unshredded.len(), 2);
1721 }
1722
1723 #[test]
1724 fn test_try_value_out_of_bounds() {
1725 let mut b = VariantArrayBuilder::new(2);
1726 b.append_variant(Variant::from(1_i8));
1727 b.append_variant(Variant::Null);
1728 let v = b.build();
1729
1730 assert_eq!(v.try_value(0).unwrap(), Variant::Int8(1));
1731 assert_eq!(v.try_value(1).unwrap(), Variant::Null);
1732
1733 let err = v.try_value(2).unwrap_err();
1734 assert_eq!(
1735 err.to_string(),
1736 "Invalid argument error: Index 2 out of bounds for VariantArray of length 2"
1737 );
1738 }
1739
1740 #[test]
1741 fn test_variant_array_iterable() {
1742 let mut b = VariantArrayBuilder::new(6);
1743
1744 b.append_null();
1745 b.append_variant(Variant::from(1_i8));
1746 b.append_variant(Variant::Null);
1747 b.append_variant(Variant::from(2_i32));
1748 b.append_variant(Variant::from(3_i64));
1749 b.append_null();
1750
1751 let v = b.build();
1752
1753 let variants = v.iter().collect::<Vec<_>>();
1754
1755 assert_eq!(
1756 variants,
1757 vec![
1758 None,
1759 Some(Variant::Int8(1)),
1760 Some(Variant::Null),
1761 Some(Variant::Int32(2)),
1762 Some(Variant::Int64(3)),
1763 None,
1764 ]
1765 );
1766 }
1767
1768 #[test]
1769 fn test_variant_array_iter_double_ended() {
1770 let array = VariantArray::from_iter([Some(0_i32), None, Some(2_i32), None, Some(4_i32)]);
1771 let mut iter = array.iter();
1772
1773 assert_eq!(iter.next(), Some(Some(Variant::from(0_i32))));
1774 assert_eq!(iter.next(), Some(None));
1775
1776 assert_eq!(iter.next_back(), Some(Some(Variant::from(4_i32))));
1777 assert_eq!(iter.next_back(), Some(None));
1778 assert_eq!(iter.next_back(), Some(Some(Variant::from(2_i32))));
1779
1780 assert_eq!(iter.next_back(), None);
1781 assert_eq!(iter.next(), None);
1782 }
1783
1784 #[test]
1785 fn test_variant_array_iter_reverse() {
1786 let array = VariantArray::from_iter([Some("a"), None, Some("aaa"), None, Some("aaaaa")]);
1787
1788 let result: Vec<_> = array.iter().rev().collect();
1789 assert_eq!(
1790 result,
1791 vec![
1792 Some(Variant::from("aaaaa")),
1793 None,
1794 Some(Variant::from("aaa")),
1795 None,
1796 Some(Variant::from("a")),
1797 ]
1798 );
1799 }
1800
1801 #[test]
1802 fn test_variant_array_iter_empty() {
1803 let v = VariantArrayBuilder::new(0).build();
1804 let mut i = v.iter();
1805 assert!(i.next().is_none());
1806 assert!(i.next_back().is_none());
1807 }
1808
1809 #[test]
1810 fn test_from_variant_opts_into_variant_array() {
1811 let v = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse), None];
1812
1813 let variant_array = VariantArray::from_iter(v);
1814
1815 assert_eq!(variant_array.len(), 4);
1816
1817 assert!(variant_array.is_null(0));
1818
1819 assert!(!variant_array.is_null(1));
1820 assert_eq!(variant_array.value(1), Variant::Null);
1821
1822 assert!(!variant_array.is_null(2));
1823 assert_eq!(variant_array.value(2), Variant::BooleanFalse);
1824
1825 assert!(variant_array.is_null(3));
1826 }
1827
1828 #[test]
1829 fn test_from_option_into_variants_into_variant_array() {
1830 let v = vec![Some(42_i64), None, Some(-1_i64)];
1833
1834 let variant_array = VariantArray::from_iter(v);
1835
1836 assert_eq!(variant_array.len(), 3);
1837
1838 assert!(!variant_array.is_null(0));
1839 assert_eq!(variant_array.value(0), Variant::Int64(42));
1840
1841 assert!(variant_array.is_null(1));
1842
1843 assert!(!variant_array.is_null(2));
1844 assert_eq!(variant_array.value(2), Variant::Int64(-1));
1845 }
1846
1847 #[test]
1848 fn test_from_option_str_into_variant_array() {
1849 let v = vec![
1850 Some("hello"),
1851 None,
1852 Some(
1853 "hello world this is much longer than the maximum short string length of sixty three bytes",
1854 ),
1855 ];
1856
1857 let variant_array: VariantArray = v.into_iter().collect();
1858
1859 assert_eq!(variant_array.len(), 3);
1860
1861 assert!(!variant_array.is_null(0));
1862 assert_eq!(
1863 variant_array.value(0),
1864 Variant::ShortString(ShortString::try_new("hello").unwrap())
1865 );
1866
1867 assert!(variant_array.is_null(1));
1868
1869 assert!(!variant_array.is_null(2));
1870 assert_eq!(
1871 variant_array.value(2),
1872 Variant::String(
1873 "hello world this is much longer than the maximum short string length of sixty three bytes"
1874 )
1875 );
1876 }
1877
1878 #[test]
1879 fn test_from_variants_into_variant_array() {
1880 let v = vec![
1881 Variant::Null,
1882 Variant::BooleanFalse,
1883 Variant::ShortString(ShortString::try_new("norm").unwrap()),
1884 ];
1885
1886 let variant_array = VariantArray::from_iter(v);
1887
1888 assert_eq!(variant_array.len(), 3);
1889
1890 assert!(!variant_array.is_null(0));
1891 assert_eq!(variant_array.value(0), Variant::Null);
1892
1893 assert!(!variant_array.is_null(1));
1894 assert_eq!(variant_array.value(1), Variant::BooleanFalse);
1895
1896 assert!(!variant_array.is_null(2));
1897 assert_eq!(
1898 variant_array.value(2),
1899 Variant::ShortString(ShortString::try_new("norm").unwrap())
1900 );
1901 }
1902
1903 #[test]
1904 fn test_variant_equality() {
1905 let v_iter = [None, Some(Variant::BooleanFalse), Some(Variant::Null), None];
1906 let v = VariantArray::from_iter(v_iter.clone());
1907
1908 {
1909 let v_copy = v.clone();
1910 assert_eq!(v, v_copy);
1911 }
1912
1913 {
1914 let v_iter_reversed = v_iter.iter().cloned().rev();
1915 let v_reversed = VariantArray::from_iter(v_iter_reversed);
1916
1917 assert_ne!(v, v_reversed);
1918 }
1919
1920 {
1921 let v_sliced = v.slice(0, 1);
1922 assert_ne!(v, v_sliced);
1923 }
1924 }
1925
1926 #[test]
1927 fn binary_typed_value_roundtrips() {
1928 let typed_value: ArrayRef = Arc::new(BinaryArray::from(vec![b"hello" as &[u8]]));
1930 let struct_array = make_variant_struct_with_typed_value(typed_value);
1931
1932 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1933 assert_eq!(variant_array.value(0), Variant::from(b"hello" as &[u8]));
1934 }
1935
1936 #[test]
1937 fn large_binary_typed_value_roundtrips() {
1938 let typed_value: ArrayRef = Arc::new(LargeBinaryArray::from(vec![b"world" as &[u8]]));
1940 let struct_array = make_variant_struct_with_typed_value(typed_value);
1941
1942 let variant_array = VariantArray::try_new(&struct_array).unwrap();
1943 assert_eq!(variant_array.value(0), Variant::from(b"world" as &[u8]));
1944 }
1945
1946 macro_rules! invalid_variant_array_test {
1947 ($fn_name: ident, $invalid_typed_value: expr, $error_msg: literal) => {
1948 #[test]
1949 fn $fn_name() {
1950 let invalid_typed_value = $invalid_typed_value;
1951
1952 let struct_array =
1953 make_variant_struct_with_typed_value(Arc::new(invalid_typed_value));
1954
1955 let array: VariantArray = VariantArray::try_new(&struct_array)
1956 .expect("should create variant array")
1957 .into();
1958
1959 let result = array.try_value(0);
1960 assert!(result.is_err());
1961 let error = result.unwrap_err();
1962 assert!(matches!(error, ArrowError::CastError(_)));
1963
1964 let expected: &str = $error_msg;
1965 assert!(
1966 error.to_string().contains($error_msg),
1967 "error `{}` did not contain `{}`",
1968 error,
1969 expected
1970 )
1971 }
1972 };
1973 }
1974
1975 invalid_variant_array_test!(
1976 test_variant_array_invalid_time,
1977 Time64MicrosecondArray::from(vec![Some(86401000000)]),
1978 "Cast error: Cast failed at index 0 (array type: Time64(µs)): Invalid microsecond from midnight: 86401000000"
1979 );
1980
1981 invalid_variant_array_test!(
1982 test_variant_array_invalid_decimal32,
1983 Decimal32Array::from(vec![Some(1234567890)]),
1984 "Cast error: Cast failed at index 0 (array type: Decimal32(9, 2)): Invalid argument error: 1234567890 is wider than max precision 9"
1985 );
1986
1987 invalid_variant_array_test!(
1988 test_variant_array_invalid_decimal64,
1989 Decimal64Array::from(vec![Some(1234567890123456789)]),
1990 "Cast error: Cast failed at index 0 (array type: Decimal64(18, 6)): Invalid argument error: 1234567890123456789 is wider than max precision 18"
1991 );
1992
1993 invalid_variant_array_test!(
1994 test_variant_array_invalid_decimal128,
1995 Decimal128Array::from(vec![Some(
1996 i128::from_str("123456789012345678901234567890123456789").unwrap()
1997 ),]),
1998 "Cast error: Cast failed at index 0 (array type: Decimal128(38, 10)): Invalid argument error: 123456789012345678901234567890123456789 is wider than max precision 38"
1999 );
2000 #[test]
2001 fn try_value_errors_on_unimplemented_typed_value_type() {
2002 use crate::{json_to_variant, shred_variant};
2003 use arrow::array::StringArray;
2004
2005 let json: ArrayRef = Arc::new(StringArray::from(vec![r#"{"qty": 3}"#]));
2006 let variant = json_to_variant(&json).unwrap();
2007 let shred_type = DataType::Struct(vec![Field::new("qty", DataType::Int64, true)].into());
2008 let shredded = shred_variant(&variant, &shred_type).unwrap();
2009 let err = shredded.try_value(0).unwrap_err();
2013 assert!(
2014 err.to_string().starts_with(
2015 "Not yet implemented: VariantArray::try_value cannot materialize typed_value"
2016 ),
2017 "unexpected error: {err}"
2018 );
2019 }
2020}