1use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
21use crate::variant_to_arrow::{
22 ArrayVariantToArrowRowBuilder, PrimitiveVariantToArrowRowBuilder,
23 make_primitive_variant_to_arrow_row_builder,
24};
25use crate::{VariantArray, VariantValueArrayBuilder};
26use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
27use arrow::buffer::NullBuffer;
28use arrow::compute::CastOptions;
29use arrow::datatypes::{DataType, Field, FieldRef, Fields, TimeUnit};
30use arrow::error::{ArrowError, Result};
31use indexmap::IndexMap;
32use parquet_variant::{Variant, VariantBuilderExt, VariantPath, VariantPathElement};
33use std::collections::BTreeMap;
34use std::sync::Arc;
35
36pub fn shred_variant(array: &VariantArray, as_type: &DataType) -> Result<VariantArray> {
71 shred_variant_with_options(array, as_type, &CastOptions::default())
72}
73
74pub(crate) fn shred_variant_with_options(
75 array: &VariantArray,
76 as_type: &DataType,
77 cast_options: &CastOptions,
78) -> Result<VariantArray> {
79 if array.typed_value_column().is_some() {
80 return Err(ArrowError::InvalidArgumentError(
81 "Input is already shredded".to_string(),
82 ));
83 }
84
85 let mut builder = make_variant_to_shredded_variant_arrow_row_builder(
86 as_type,
87 cast_options,
88 array.len(),
89 NullValue::TopLevelVariant,
90 true,
91 )?;
92 for i in 0..array.len() {
93 if array.is_null(i) {
94 builder.append_null()?;
95 } else {
96 builder.append_value(array.value(i))?;
97 }
98 }
99 let (value, typed_value, nulls) = builder.finish()?;
100 Ok(VariantArray::from_parts(
101 array.metadata_column().clone(),
102 Arc::new(value),
103 Some(typed_value),
104 nulls,
105 ))
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub(crate) enum NullValue {
117 TopLevelVariant,
118 ObjectField,
119 ArrayElement,
120}
121
122impl NullValue {
123 fn append_to(
124 self,
125 nulls: &mut NullBufferBuilder,
126 value_builder: &mut VariantValueArrayBuilder,
127 ) {
128 match self {
129 Self::TopLevelVariant => nulls.append_null(),
130 Self::ObjectField | Self::ArrayElement => nulls.append_non_null(),
131 }
132 match self {
133 Self::TopLevelVariant | Self::ObjectField => value_builder.append_null(),
134 Self::ArrayElement => value_builder.append_value(Variant::Null),
135 }
136 }
137}
138
139pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>(
140 data_type: &'a DataType,
141 cast_options: &'a CastOptions,
142 capacity: usize,
143 null_value: NullValue,
144 shred: bool,
145) -> Result<VariantToShreddedVariantRowBuilder<'a>> {
146 let builder = match data_type {
147 DataType::Struct(fields) => {
148 let typed_value_builder = VariantToShreddedObjectVariantRowBuilder::try_new(
149 fields,
150 cast_options,
151 capacity,
152 null_value,
153 shred,
154 )?;
155 VariantToShreddedVariantRowBuilder::Object(typed_value_builder)
156 }
157 DataType::List(_)
158 | DataType::LargeList(_)
159 | DataType::ListView(_)
160 | DataType::LargeListView(_)
161 | DataType::FixedSizeList(..) => {
162 let typed_value_builder = VariantToShreddedArrayVariantRowBuilder::try_new(
163 data_type,
164 cast_options,
165 capacity,
166 null_value,
167 )?;
168 VariantToShreddedVariantRowBuilder::Array(typed_value_builder)
169 }
170 DataType::Boolean
173 | DataType::Int8
174 | DataType::Int16
175 | DataType::Int32
176 | DataType::Int64
177 | DataType::Float32
178 | DataType::Float64
179 | DataType::Decimal32(..)
180 | DataType::Decimal64(..)
181 | DataType::Decimal128(..)
182 | DataType::Date32
183 | DataType::Time64(TimeUnit::Microsecond)
184 | DataType::Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, _)
185 | DataType::Binary
186 | DataType::BinaryView
187 | DataType::LargeBinary
188 | DataType::Utf8
189 | DataType::Utf8View
190 | DataType::LargeUtf8
191 | DataType::FixedSizeBinary(16) => {
193 let builder =
194 make_primitive_variant_to_arrow_row_builder(data_type, cast_options, capacity, shred)?;
195 let typed_value_builder =
196 VariantToShreddedPrimitiveVariantRowBuilder::new(builder, capacity, null_value);
197 VariantToShreddedVariantRowBuilder::Primitive(typed_value_builder)
198 }
199 DataType::FixedSizeBinary(_) => {
200 return Err(ArrowError::InvalidArgumentError(format!("{data_type} is not a valid variant shredding type. Only FixedSizeBinary(16) for UUID is supported.")))
201 }
202 _ => {
203 return Err(ArrowError::InvalidArgumentError(format!("{data_type} is not a valid variant shredding type")))
204 }
205 };
206 Ok(builder)
207}
208
209pub(crate) enum VariantToShreddedVariantRowBuilder<'a> {
210 Primitive(VariantToShreddedPrimitiveVariantRowBuilder<'a>),
211 Array(VariantToShreddedArrayVariantRowBuilder<'a>),
212 Object(VariantToShreddedObjectVariantRowBuilder<'a>),
213}
214
215impl VariantToShreddedVariantRowBuilder<'_> {
216 pub fn append_null(&mut self) -> Result<()> {
217 use VariantToShreddedVariantRowBuilder::*;
218 match self {
219 Primitive(b) => b.append_null(),
220 Array(b) => b.append_null(),
221 Object(b) => b.append_null(),
222 }
223 }
224
225 pub fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
226 use VariantToShreddedVariantRowBuilder::*;
227 match self {
228 Primitive(b) => b.append_value(value),
229 Array(b) => b.append_value(value),
230 Object(b) => b.append_value(value),
231 }
232 }
233
234 pub fn finish(self) -> Result<(BinaryViewArray, ArrayRef, Option<NullBuffer>)> {
235 use VariantToShreddedVariantRowBuilder::*;
236 match self {
237 Primitive(b) => b.finish(),
238 Array(b) => b.finish(),
239 Object(b) => b.finish(),
240 }
241 }
242}
243
244pub(crate) struct VariantToShreddedPrimitiveVariantRowBuilder<'a> {
246 value_builder: VariantValueArrayBuilder,
247 typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
248 nulls: NullBufferBuilder,
249 null_value: NullValue,
250}
251
252impl<'a> VariantToShreddedPrimitiveVariantRowBuilder<'a> {
253 pub(crate) fn new(
254 typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
255 capacity: usize,
256 null_value: NullValue,
257 ) -> Self {
258 Self {
259 value_builder: VariantValueArrayBuilder::new(capacity),
260 typed_value_builder,
261 nulls: NullBufferBuilder::new(capacity),
262 null_value,
263 }
264 }
265
266 fn append_null(&mut self) -> Result<()> {
267 self.null_value
268 .append_to(&mut self.nulls, &mut self.value_builder);
269 self.typed_value_builder.append_null()
270 }
271
272 fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
273 self.nulls.append_non_null();
274 if self.typed_value_builder.append_value(&value)? {
275 self.value_builder.append_null();
276 } else {
277 self.value_builder.append_value(value);
278 }
279 Ok(true)
280 }
281
282 fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef, Option<NullBuffer>)> {
283 Ok((
284 self.value_builder.build()?,
285 self.typed_value_builder.finish()?,
286 self.nulls.finish(),
287 ))
288 }
289}
290
291pub(crate) struct VariantToShreddedArrayVariantRowBuilder<'a> {
292 value_builder: VariantValueArrayBuilder,
293 typed_value_builder: ArrayVariantToArrowRowBuilder<'a>,
294 nulls: NullBufferBuilder,
295 null_value: NullValue,
296}
297
298impl<'a> VariantToShreddedArrayVariantRowBuilder<'a> {
299 fn try_new(
300 data_type: &'a DataType,
301 cast_options: &'a CastOptions,
302 capacity: usize,
303 null_value: NullValue,
304 ) -> Result<Self> {
305 Ok(Self {
306 value_builder: VariantValueArrayBuilder::new(capacity),
307 typed_value_builder: ArrayVariantToArrowRowBuilder::try_new(
308 data_type,
309 cast_options,
310 capacity,
311 true,
312 )?,
313 nulls: NullBufferBuilder::new(capacity),
314 null_value,
315 })
316 }
317
318 fn append_null(&mut self) -> Result<()> {
319 self.null_value
320 .append_to(&mut self.nulls, &mut self.value_builder);
321 self.typed_value_builder.append_null()?;
322 Ok(())
323 }
324
325 fn append_value(&mut self, variant: Variant<'_, '_>) -> Result<bool> {
326 match variant {
329 Variant::List(list) => {
330 self.nulls.append_non_null();
331 self.value_builder.append_null();
332
333 self.typed_value_builder
335 .append_value(&Variant::List(list))?;
336 Ok(true)
337 }
338 other => {
339 self.nulls.append_non_null();
340 self.value_builder.append_value(other);
341 self.typed_value_builder.append_null()?;
342 Ok(false)
343 }
344 }
345 }
346
347 fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef, Option<NullBuffer>)> {
348 Ok((
349 self.value_builder.build()?,
350 self.typed_value_builder.finish()?,
351 self.nulls.finish(),
352 ))
353 }
354}
355
356pub(crate) struct VariantToShreddedObjectVariantRowBuilder<'a> {
357 value_builder: VariantValueArrayBuilder,
358 typed_value_builders: IndexMap<&'a str, VariantToShreddedVariantRowBuilder<'a>>,
359 typed_value_nulls: NullBufferBuilder,
360 nulls: NullBufferBuilder,
361 null_value: NullValue,
362}
363
364impl<'a> VariantToShreddedObjectVariantRowBuilder<'a> {
365 fn try_new(
366 fields: &'a Fields,
367 cast_options: &'a CastOptions,
368 capacity: usize,
369 null_value: NullValue,
370 shred: bool,
371 ) -> Result<Self> {
372 let typed_value_builders = fields.iter().map(|field| {
373 let builder = make_variant_to_shredded_variant_arrow_row_builder(
374 field.data_type(),
375 cast_options,
376 capacity,
377 NullValue::ObjectField,
378 shred,
379 )?;
380 Ok((field.name().as_str(), builder))
381 });
382 Ok(Self {
383 value_builder: VariantValueArrayBuilder::new(capacity),
384 typed_value_builders: typed_value_builders.collect::<Result<_>>()?,
385 typed_value_nulls: NullBufferBuilder::new(capacity),
386 nulls: NullBufferBuilder::new(capacity),
387 null_value,
388 })
389 }
390
391 fn append_null(&mut self) -> Result<()> {
392 self.null_value
393 .append_to(&mut self.nulls, &mut self.value_builder);
394 self.typed_value_nulls.append_null();
395 for (_, typed_value_builder) in &mut self.typed_value_builders {
396 typed_value_builder.append_null()?;
397 }
398 Ok(())
399 }
400
401 fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
402 let Variant::Object(ref obj) = value else {
403 self.nulls.append_non_null();
405 self.value_builder.append_value(value);
406 self.typed_value_nulls.append_null();
407 for (_, typed_value_builder) in &mut self.typed_value_builders {
408 typed_value_builder.append_null()?;
409 }
410 return Ok(false);
411 };
412
413 let mut builder = self.value_builder.builder_ext(value.metadata());
415 let mut object_builder = builder.try_new_object()?;
416 let mut seen = std::collections::HashSet::new();
417 let mut partially_shredded = false;
418 for (field_name, value) in obj.iter() {
419 match self.typed_value_builders.get_mut(field_name) {
420 Some(typed_value_builder) => {
421 typed_value_builder.append_value(value)?;
422 seen.insert(field_name);
423 }
424 None => {
425 object_builder.insert_bytes(field_name, value);
426 partially_shredded = true;
427 }
428 }
429 }
430
431 for (field_name, typed_value_builder) in &mut self.typed_value_builders {
433 if !seen.contains(field_name) {
434 typed_value_builder.append_null()?;
435 }
436 }
437
438 if partially_shredded {
440 object_builder.finish();
441 } else {
442 drop(object_builder);
443 self.value_builder.append_null();
444 }
445
446 self.typed_value_nulls.append_non_null();
447 self.nulls.append_non_null();
448 Ok(true)
449 }
450
451 fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef, Option<NullBuffer>)> {
452 let mut builder = StructArrayBuilder::new();
453 for (field_name, typed_value_builder) in self.typed_value_builders {
454 let (value, typed_value, nulls) = typed_value_builder.finish()?;
455 let array =
456 ShreddedVariantFieldArray::from_parts(Arc::new(value), Some(typed_value), nulls);
457 builder = builder.with_field(field_name, ArrayRef::from(array), false);
458 }
459 if let Some(nulls) = self.typed_value_nulls.finish() {
460 builder = builder.with_nulls(nulls);
461 }
462 Ok((
463 self.value_builder.build()?,
464 Arc::new(builder.build()),
465 self.nulls.finish(),
466 ))
467 }
468}
469
470#[derive(Clone)]
472pub struct ShreddingField {
473 data_type: DataType,
474 nullable: bool,
475}
476
477impl ShreddingField {
478 fn new(data_type: DataType, nullable: bool) -> Self {
479 Self {
480 data_type,
481 nullable,
482 }
483 }
484
485 fn null() -> Self {
486 Self::new(DataType::Null, true)
487 }
488}
489
490pub trait IntoShreddingField {
492 fn into_shredding_field(self) -> ShreddingField;
493}
494
495impl IntoShreddingField for FieldRef {
496 fn into_shredding_field(self) -> ShreddingField {
497 ShreddingField::new(self.data_type().clone(), self.is_nullable())
498 }
499}
500
501impl IntoShreddingField for &DataType {
502 fn into_shredding_field(self) -> ShreddingField {
503 ShreddingField::new(self.clone(), true)
504 }
505}
506
507impl IntoShreddingField for DataType {
508 fn into_shredding_field(self) -> ShreddingField {
509 ShreddingField::new(self, true)
510 }
511}
512
513impl IntoShreddingField for (&DataType, bool) {
514 fn into_shredding_field(self) -> ShreddingField {
515 ShreddingField::new(self.0.clone(), self.1)
516 }
517}
518
519impl IntoShreddingField for (DataType, bool) {
520 fn into_shredding_field(self) -> ShreddingField {
521 ShreddingField::new(self.0, self.1)
522 }
523}
524
525#[derive(Default, Clone)]
566pub struct ShreddedSchemaBuilder {
567 root: VariantSchemaNode,
568}
569
570impl ShreddedSchemaBuilder {
571 pub fn new() -> Self {
573 Self::default()
574 }
575
576 pub fn with_path<'a, P, F>(mut self, path: P, field: F) -> Result<Self>
588 where
589 P: TryInto<VariantPath<'a>>,
590 P::Error: std::fmt::Debug,
591 F: IntoShreddingField,
592 {
593 let path: VariantPath<'a> = path
594 .try_into()
595 .map_err(|e| ArrowError::InvalidArgumentError(format!("{e:?}")))?;
596 self.root.insert_path(&path, field.into_shredding_field());
597 Ok(self)
598 }
599
600 pub fn build(self) -> DataType {
602 let shredding_type = self.root.to_shredding_type();
603 match shredding_type {
604 Some(shredding_type) => shredding_type,
605 None => DataType::Null,
606 }
607 }
608}
609
610#[derive(Clone)]
612enum VariantSchemaNode {
613 Leaf(ShreddingField),
615 Struct(BTreeMap<String, VariantSchemaNode>),
617}
618
619impl Default for VariantSchemaNode {
620 fn default() -> Self {
621 Self::Leaf(ShreddingField::null())
622 }
623}
624
625impl VariantSchemaNode {
626 fn insert_path(&mut self, path: &VariantPath<'_>, field: ShreddingField) {
628 self.insert_path_elements(path, field);
629 }
630
631 fn insert_path_elements(&mut self, segments: &[VariantPathElement<'_>], field: ShreddingField) {
632 let Some((head, tail)) = segments.split_first() else {
633 *self = Self::Leaf(field);
634 return;
635 };
636
637 match head {
638 VariantPathElement::Field { name } => {
639 let children = match self {
641 Self::Struct(children) => children,
642 Self::Leaf(_) => {
643 *self = Self::Struct(BTreeMap::new());
644 match self {
645 Self::Struct(children) => children,
646 Self::Leaf(_) => unreachable!(),
647 }
648 }
649 };
650
651 children
652 .entry(name.to_string())
653 .or_default()
654 .insert_path_elements(tail, field);
655 }
656 VariantPathElement::Index { .. } => {
657 unreachable!("List paths are not supported yet");
659 }
660 }
661 }
662
663 fn to_shredding_type(&self) -> Option<DataType> {
667 match self {
668 Self::Leaf(field) => Some(field.data_type.clone()),
669 Self::Struct(children) => {
670 let child_fields: Vec<_> = children
671 .iter()
672 .filter_map(|(name, child)| child.to_shredding_field(name))
673 .collect();
674 if child_fields.is_empty() {
675 None
676 } else {
677 Some(DataType::Struct(Fields::from(child_fields)))
678 }
679 }
680 }
681 }
682
683 fn to_shredding_field(&self, name: &str) -> Option<FieldRef> {
684 match self {
685 Self::Leaf(field) => Some(Arc::new(Field::new(
686 name,
687 field.data_type.clone(),
688 field.nullable,
689 ))),
690 Self::Struct(_) => self
691 .to_shredding_type()
692 .map(|data_type| Arc::new(Field::new(name, data_type, true))),
693 }
694 }
695}
696
697#[cfg(test)]
698mod tests {
699 use super::*;
700 use crate::VariantArrayBuilder;
701 use crate::variant_array::{all_null_value_column, binary_array_value, variant_from_arrays_at};
702 use arrow::array::{
703 Array, BinaryViewArray, Decimal32Array, Decimal64Array, Decimal128Array,
704 FixedSizeBinaryArray, FixedSizeListArray, Float64Array, GenericListArray,
705 GenericListViewArray, Int64Array, LargeBinaryArray, LargeStringArray, ListArray,
706 ListLikeArray, OffsetSizeTrait, PrimitiveArray, StringArray, StructArray,
707 };
708 use arrow::datatypes::{
709 ArrowPrimitiveType, DataType, Field, Fields, Int64Type, TimeUnit, UnionFields, UnionMode,
710 };
711 use arrow_schema::IntervalUnit;
712 use chrono::{DateTime, NaiveDate, NaiveTime};
713 use parquet_variant::{
714 BuilderSpecificState, EMPTY_VARIANT_METADATA_BYTES, ObjectBuilder, ReadOnlyMetadataBuilder,
715 ShortString, Variant, VariantBuilder, VariantDecimal4, VariantDecimal8, VariantDecimal16,
716 VariantPath, VariantPathElement,
717 };
718 use std::sync::Arc;
719 use uuid::Uuid;
720
721 const NULL_VALUES: [NullValue; 3] = [
722 NullValue::TopLevelVariant,
723 NullValue::ObjectField,
724 NullValue::ArrayElement,
725 ];
726
727 #[derive(Clone)]
728 enum VariantValue<'a> {
729 Value(Variant<'a, 'a>),
730 List(Vec<VariantValue<'a>>),
731 Object(Vec<(&'a str, VariantValue<'a>)>),
732 Null,
733 }
734
735 impl<'a, T> From<T> for VariantValue<'a>
736 where
737 T: Into<Variant<'a, 'a>>,
738 {
739 fn from(value: T) -> Self {
740 Self::Value(value.into())
741 }
742 }
743
744 #[derive(Clone)]
745 enum VariantRow<'a> {
746 Value(VariantValue<'a>),
747 List(Vec<VariantValue<'a>>),
748 Object(Vec<(&'a str, VariantValue<'a>)>),
749 Null,
750 }
751
752 fn build_variant_array(rows: Vec<VariantRow<'static>>) -> VariantArray {
753 let mut builder = VariantArrayBuilder::new(rows.len());
754
755 fn append_variant_value<B: VariantBuilderExt>(builder: &mut B, value: VariantValue) {
756 match value {
757 VariantValue::Value(v) => builder.append_value(v),
758 VariantValue::List(values) => {
759 let mut list = builder.new_list();
760 for v in values {
761 append_variant_value(&mut list, v);
762 }
763 list.finish();
764 }
765 VariantValue::Object(fields) => {
766 let mut object = builder.new_object();
767 for (name, value) in fields {
768 append_variant_field(&mut object, name, value);
769 }
770 object.finish();
771 }
772 VariantValue::Null => builder.append_null(),
773 }
774 }
775
776 fn append_variant_field<'a, S: BuilderSpecificState>(
777 object: &mut ObjectBuilder<'_, S>,
778 name: &'a str,
779 value: VariantValue<'a>,
780 ) {
781 match value {
782 VariantValue::Value(v) => {
783 object.insert(name, v);
784 }
785 VariantValue::List(values) => {
786 let mut list = object.new_list(name);
787 for v in values {
788 append_variant_value(&mut list, v);
789 }
790 list.finish();
791 }
792 VariantValue::Object(fields) => {
793 let mut nested = object.new_object(name);
794 for (field_name, v) in fields {
795 append_variant_field(&mut nested, field_name, v);
796 }
797 nested.finish();
798 }
799 VariantValue::Null => {
800 object.insert(name, Variant::Null);
801 }
802 }
803 }
804
805 rows.into_iter().for_each(|row| match row {
806 VariantRow::Value(value) => append_variant_value(&mut builder, value),
807 VariantRow::List(values) => {
808 let mut list = builder.new_list();
809 for value in values {
810 append_variant_value(&mut list, value);
811 }
812 list.finish();
813 }
814 VariantRow::Object(fields) => {
815 let mut object = builder.new_object();
816 for (name, value) in fields {
817 append_variant_field(&mut object, name, value);
818 }
819 object.finish();
820 }
821 VariantRow::Null => builder.append_null(),
822 });
823 builder.build()
824 }
825
826 trait TestListLikeArray: ListLikeArray {
827 type OffsetSize: OffsetSizeTrait;
828 fn value_offsets(&self) -> Option<&[Self::OffsetSize]>;
829 fn value_size(&self, index: usize) -> Self::OffsetSize;
830 }
831
832 impl<O: OffsetSizeTrait> TestListLikeArray for GenericListArray<O> {
833 type OffsetSize = O;
834
835 fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
836 Some(GenericListArray::value_offsets(self))
837 }
838
839 fn value_size(&self, index: usize) -> Self::OffsetSize {
840 GenericListArray::value_length(self, index)
841 }
842 }
843
844 impl<O: OffsetSizeTrait> TestListLikeArray for GenericListViewArray<O> {
845 type OffsetSize = O;
846
847 fn value_offsets(&self) -> Option<&[Self::OffsetSize]> {
848 Some(GenericListViewArray::value_offsets(self))
849 }
850
851 fn value_size(&self, index: usize) -> Self::OffsetSize {
852 GenericListViewArray::value_size(self, index)
853 }
854 }
855
856 fn downcast_list_like_array<O: OffsetSizeTrait>(
857 array: &VariantArray,
858 ) -> &dyn TestListLikeArray<OffsetSize = O> {
859 let typed_value = array.typed_value_column().unwrap();
860 if let Some(list) = typed_value.as_any().downcast_ref::<GenericListArray<O>>() {
861 list
862 } else if let Some(list_view) = typed_value
863 .as_any()
864 .downcast_ref::<GenericListViewArray<O>>()
865 {
866 list_view
867 } else {
868 panic!(
869 "Expected list-like typed_value with matching offset type, got {}",
870 typed_value.data_type()
871 );
872 }
873 }
874
875 fn assert_list_structure<O: OffsetSizeTrait>(
876 array: &VariantArray,
877 expected_len: usize,
878 expected_offsets: &[O],
879 expected_sizes: &[Option<O>],
880 expected_fallbacks: &[Option<Variant<'static, 'static>>],
881 ) {
882 assert_eq!(array.len(), expected_len);
883
884 let fallback_value = array.value_column();
885 let fallback_metadata = array.metadata_column();
886 let array = downcast_list_like_array::<O>(array);
887
888 assert_eq!(
889 array.value_offsets().unwrap(),
890 expected_offsets,
891 "list offsets mismatch"
892 );
893 assert_eq!(
894 array.len(),
895 expected_sizes.len(),
896 "expected_sizes should match array length"
897 );
898 assert_eq!(
899 array.len(),
900 expected_fallbacks.len(),
901 "expected_fallbacks should match array length"
902 );
903 assert_eq!(
904 array.len(),
905 fallback_value.len(),
906 "fallbacks value field should match array length"
907 );
908
909 for (idx, (expected_size, expected_fallback)) in expected_sizes
911 .iter()
912 .zip(expected_fallbacks.iter())
913 .enumerate()
914 {
915 match expected_size {
916 Some(len) => {
917 assert!(array.is_valid(idx));
919 assert_eq!(array.value_size(idx), *len);
920 assert!(fallback_value.is_null(idx));
921 }
922 None => {
923 assert!(array.is_null(idx));
925 assert_eq!(array.value_size(idx), O::zero());
926 match expected_fallback {
927 Some(expected_variant) => {
928 assert!(fallback_value.is_valid(idx));
929 let metadata_bytes =
930 binary_array_value(fallback_metadata.as_ref(), idx).unwrap();
931 let metadata_bytes =
932 if fallback_metadata.is_valid(idx) && !metadata_bytes.is_empty() {
933 metadata_bytes
934 } else {
935 EMPTY_VARIANT_METADATA_BYTES
936 };
937 assert_eq!(
938 Variant::new(
939 metadata_bytes,
940 binary_array_value(fallback_value.as_ref(), idx).unwrap()
941 ),
942 expected_variant.clone()
943 );
944 }
945 None => {
946 assert!(fallback_value.is_null(idx));
947 }
948 }
949 }
950 }
951 }
952 }
953
954 fn assert_list_structure_and_elements<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
955 array: &VariantArray,
956 expected_len: usize,
957 expected_offsets: &[O],
958 expected_sizes: &[Option<O>],
959 expected_fallbacks: &[Option<Variant<'static, 'static>>],
960 expected_shredded_elements: (&[Option<T::Native>], &[Option<Variant<'static, 'static>>]),
961 ) {
962 assert_list_structure(
963 array,
964 expected_len,
965 expected_offsets,
966 expected_sizes,
967 expected_fallbacks,
968 );
969 let array = downcast_list_like_array::<O>(array);
970
971 let (expected_values, expected_fallbacks) = expected_shredded_elements;
973 assert_eq!(
974 expected_values.len(),
975 expected_fallbacks.len(),
976 "expected_values and expected_fallbacks should be aligned"
977 );
978
979 let element_array = ShreddedVariantFieldArray::try_new(array.values().as_ref()).unwrap();
981 let element_values = element_array
982 .typed_value_column()
983 .unwrap()
984 .as_any()
985 .downcast_ref::<PrimitiveArray<T>>()
986 .unwrap();
987 assert_eq!(element_values.len(), expected_values.len());
988 for (idx, expected_value) in expected_values.iter().enumerate() {
989 match expected_value {
990 Some(value) => {
991 assert!(element_values.is_valid(idx));
992 assert_eq!(element_values.value(idx), *value);
993 }
994 None => assert!(element_values.is_null(idx)),
995 }
996 }
997
998 let element_fallbacks = element_array.value_column();
1000 assert_eq!(element_fallbacks.len(), expected_fallbacks.len());
1001 for (idx, expected_fallback) in expected_fallbacks.iter().enumerate() {
1002 match expected_fallback {
1003 Some(expected_variant) => {
1004 assert!(element_fallbacks.is_valid(idx));
1005 assert_eq!(
1006 Variant::new(
1007 EMPTY_VARIANT_METADATA_BYTES,
1008 binary_array_value(element_fallbacks.as_ref(), idx).unwrap()
1009 ),
1010 expected_variant.clone()
1011 );
1012 }
1013 None => assert!(element_fallbacks.is_null(idx)),
1014 }
1015 }
1016 }
1017
1018 fn assert_append_null_mode_value_and_struct_nulls(
1019 mode: NullValue,
1020 value: &BinaryViewArray,
1021 nulls: Option<&arrow::buffer::NullBuffer>,
1022 ) {
1023 if mode == NullValue::TopLevelVariant {
1024 assert!(nulls.is_some_and(|n| n.is_null(0)));
1025 } else {
1026 assert!(nulls.is_none());
1027 }
1028
1029 if mode == NullValue::ArrayElement {
1030 assert!(value.is_valid(0));
1031 assert_eq!(
1032 Variant::new(EMPTY_VARIANT_METADATA_BYTES, value.value(0)),
1033 Variant::Null
1034 );
1035 } else {
1036 assert!(value.is_null(0));
1037 }
1038 }
1039
1040 #[test]
1041 fn test_append_null_mode_semantics_primitive_builder() {
1042 let cast_options = arrow::compute::CastOptions::default();
1043
1044 for mode in NULL_VALUES {
1045 let mut primitive_builder = make_variant_to_shredded_variant_arrow_row_builder(
1046 &DataType::Int64,
1047 &cast_options,
1048 1,
1049 mode,
1050 true,
1051 )
1052 .unwrap();
1053 primitive_builder.append_null().unwrap();
1054 let (primitive_value, primitive_typed_value, primitive_nulls) =
1055 primitive_builder.finish().unwrap();
1056 let primitive_typed_value = primitive_typed_value
1057 .as_any()
1058 .downcast_ref::<Int64Array>()
1059 .unwrap();
1060
1061 assert!(primitive_typed_value.is_null(0));
1062 assert_append_null_mode_value_and_struct_nulls(
1063 mode,
1064 &primitive_value,
1065 primitive_nulls.as_ref(),
1066 );
1067 }
1068 }
1069
1070 #[test]
1071 fn test_append_null_mode_semantics_array_builder() {
1072 let cast_options = arrow::compute::CastOptions::default();
1073 let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
1074
1075 for mode in NULL_VALUES {
1076 let mut array_builder = make_variant_to_shredded_variant_arrow_row_builder(
1077 &list_type,
1078 &cast_options,
1079 1,
1080 mode,
1081 true,
1082 )
1083 .unwrap();
1084 array_builder.append_null().unwrap();
1085 let (value, typed_value, nulls) = array_builder.finish().unwrap();
1086
1087 assert_append_null_mode_value_and_struct_nulls(mode, &value, nulls.as_ref());
1088
1089 let typed_value = typed_value.as_any().downcast_ref::<ListArray>().unwrap();
1090 assert_eq!(typed_value.len(), 1);
1091 assert!(typed_value.is_null(0));
1092 assert_eq!(typed_value.values().len(), 0);
1093 }
1094 }
1095
1096 #[test]
1097 fn test_append_null_mode_semantics_object_builder() {
1098 let cast_options = arrow::compute::CastOptions::default();
1099 let object_type = DataType::Struct(Fields::from(vec![
1100 Field::new("id", DataType::Int64, true),
1101 Field::new("name", DataType::Utf8, true),
1102 ]));
1103
1104 for mode in NULL_VALUES {
1105 let mut object_builder = make_variant_to_shredded_variant_arrow_row_builder(
1106 &object_type,
1107 &cast_options,
1108 1,
1109 mode,
1110 true,
1111 )
1112 .unwrap();
1113 object_builder.append_null().unwrap();
1114 let (value, typed_value, nulls) = object_builder.finish().unwrap();
1115
1116 assert_append_null_mode_value_and_struct_nulls(mode, &value, nulls.as_ref());
1117
1118 let typed_struct = typed_value
1119 .as_any()
1120 .downcast_ref::<arrow::array::StructArray>()
1121 .unwrap();
1122 assert_eq!(typed_struct.len(), 1);
1123 assert!(typed_struct.is_null(0));
1124
1125 for field_name in ["id", "name"] {
1126 let field = ShreddedVariantFieldArray::try_new(
1127 typed_struct.column_by_name(field_name).unwrap(),
1128 )
1129 .unwrap();
1130 assert!(field.value_column().is_null(0));
1131 assert!(field.typed_value_column().unwrap().is_null(0));
1132 }
1133 }
1134 }
1135
1136 #[test]
1137 fn test_already_shredded_input_error() {
1138 let temp_array = VariantArray::from_iter(vec![Some(Variant::from("test"))]);
1141 let metadata = temp_array.metadata_column().clone();
1142 let value = temp_array.value_column().clone();
1143 let typed_value = Arc::new(Int64Array::from(vec![42])) as ArrayRef;
1144
1145 let shredded_array = VariantArray::from_parts(metadata, value, Some(typed_value), None);
1146
1147 let result = shred_variant(&shredded_array, &DataType::Int64);
1148 assert!(matches!(
1149 result.unwrap_err(),
1150 ArrowError::InvalidArgumentError(_)
1151 ));
1152 }
1153
1154 #[test]
1155 fn test_all_null_input() {
1156 let metadata = Arc::new(BinaryViewArray::from_iter_values([
1158 EMPTY_VARIANT_METADATA_BYTES,
1159 ]));
1160 let all_null_array =
1161 VariantArray::from_parts(metadata, all_null_value_column(1), None, None);
1162 let result = shred_variant(&all_null_array, &DataType::Int64).unwrap();
1163
1164 assert!(result.typed_value_column().unwrap().is_null(0));
1167 assert_eq!(result.value(0), Variant::Null);
1168 }
1169
1170 #[test]
1171 fn test_invalid_fixed_size_binary_shredding() {
1172 let mock_uuid_1 = Uuid::new_v4();
1173
1174 let input = VariantArray::from_iter([Some(Variant::from(mock_uuid_1)), None]);
1175
1176 let err = shred_variant(&input, &DataType::FixedSizeBinary(17)).unwrap_err();
1178
1179 assert_eq!(
1180 err.to_string(),
1181 "Invalid argument error: FixedSizeBinary(17) is not a valid variant shredding type. Only FixedSizeBinary(16) for UUID is supported."
1182 );
1183 }
1184
1185 #[test]
1186 fn test_uuid_shredding() {
1187 let mock_uuid_1 = Uuid::new_v4();
1188 let mock_uuid_2 = Uuid::new_v4();
1189
1190 let input = VariantArray::from_iter([
1191 Some(Variant::from(mock_uuid_1)),
1192 None,
1193 Some(Variant::from(false)),
1194 Some(Variant::from(mock_uuid_2)),
1195 ]);
1196
1197 let variant_array = shred_variant(&input, &DataType::FixedSizeBinary(16)).unwrap();
1198
1199 let typed_value_field = variant_array.inner().field_by_name("typed_value").unwrap();
1200
1201 assert!(typed_value_field.has_valid_extension_type::<arrow_schema::extension::Uuid>());
1202
1203 let uuids = variant_array
1205 .typed_value_column()
1206 .unwrap()
1207 .as_any()
1208 .downcast_ref::<FixedSizeBinaryArray>()
1209 .unwrap();
1210
1211 assert_eq!(uuids.len(), 4);
1212
1213 assert!(!uuids.is_null(0));
1214
1215 let got_uuid_1: &[u8] = uuids.value(0);
1216 assert_eq!(got_uuid_1, mock_uuid_1.as_bytes());
1217
1218 assert!(uuids.is_null(1));
1219 assert!(uuids.is_null(2));
1220
1221 assert!(!uuids.is_null(3));
1222
1223 let got_uuid_2: &[u8] = uuids.value(3);
1224 assert_eq!(got_uuid_2, mock_uuid_2.as_bytes());
1225 }
1226
1227 #[test]
1228 fn test_uuid_nested_shredding() {
1229 let mock_uuid = Uuid::new_v4();
1230 let input = build_variant_array(vec![VariantRow::Object(vec![(
1231 "id",
1232 VariantValue::from(mock_uuid),
1233 )])]);
1234 let target = ShreddedSchemaBuilder::default()
1235 .with_path("id", DataType::FixedSizeBinary(16))
1236 .unwrap()
1237 .build();
1238
1239 let result = shred_variant(&input, &target).unwrap();
1240
1241 let typed_value = result.typed_value_column().unwrap();
1242 let typed_struct = typed_value.as_any().downcast_ref::<StructArray>().unwrap();
1243 let id =
1244 ShreddedVariantFieldArray::try_new(typed_struct.column_by_name("id").unwrap()).unwrap();
1245
1246 let leaf = id.inner().field_by_name("typed_value").unwrap();
1248
1249 assert_eq!(leaf.data_type(), &DataType::FixedSizeBinary(16));
1250 assert!(leaf.has_valid_extension_type::<arrow_schema::extension::Uuid>());
1251 }
1252
1253 #[test]
1254 fn test_primitive_shredding_comprehensive() {
1255 let input = VariantArray::from_iter(vec![
1257 Some(Variant::from(42i64)), Some(Variant::from("hello")), Some(Variant::from(100i64)), None, Some(Variant::Null), Some(Variant::from(3i8)), ]);
1264
1265 let result = shred_variant(&input, &DataType::Int64).unwrap();
1266
1267 let metadata_field = result.metadata_column();
1269 let value_field = result.value_column();
1270 let typed_value_field = result
1271 .typed_value_column()
1272 .unwrap()
1273 .as_any()
1274 .downcast_ref::<Int64Array>()
1275 .unwrap();
1276
1277 assert_eq!(result.len(), 6);
1279
1280 assert!(!result.is_null(0));
1282 assert!(value_field.is_null(0)); assert!(!typed_value_field.is_null(0));
1284 assert_eq!(typed_value_field.value(0), 42);
1285
1286 assert!(!result.is_null(1));
1288 assert!(!value_field.is_null(1)); assert!(typed_value_field.is_null(1)); assert_eq!(
1291 variant_from_arrays_at(metadata_field, value_field, 1).unwrap(),
1292 Variant::from("hello")
1293 );
1294
1295 assert!(!result.is_null(2));
1297 assert!(value_field.is_null(2));
1298 assert_eq!(typed_value_field.value(2), 100);
1299
1300 assert!(result.is_null(3));
1302
1303 assert!(!result.is_null(4));
1305 assert!(!value_field.is_null(4)); assert_eq!(
1307 variant_from_arrays_at(metadata_field, value_field, 4).unwrap(),
1308 Variant::Null
1309 );
1310 assert!(typed_value_field.is_null(4));
1311
1312 assert!(!result.is_null(5));
1314 assert!(value_field.is_null(5)); assert!(!typed_value_field.is_null(5));
1316 assert_eq!(typed_value_field.value(5), 3);
1317 }
1318
1319 #[test]
1320 fn test_primitive_different_target_types() {
1321 let input = VariantArray::from_iter(vec![
1322 Variant::from(42i32),
1323 Variant::from(3.15f64),
1324 Variant::from("not_a_number"),
1325 ]);
1326
1327 let result_int32 = shred_variant(&input, &DataType::Int32).unwrap();
1329 let typed_value_int32 = result_int32
1330 .typed_value_column()
1331 .unwrap()
1332 .as_any()
1333 .downcast_ref::<arrow::array::Int32Array>()
1334 .unwrap();
1335 assert_eq!(typed_value_int32.value(0), 42);
1336 assert!(typed_value_int32.is_null(1)); assert!(typed_value_int32.is_null(2)); let result_float64 = shred_variant(&input, &DataType::Float64).unwrap();
1341 let typed_value_float64 = result_float64
1342 .typed_value_column()
1343 .unwrap()
1344 .as_any()
1345 .downcast_ref::<Float64Array>()
1346 .unwrap();
1347 assert!(typed_value_float64.is_null(0)); assert_eq!(typed_value_float64.value(1), 3.15);
1349 assert!(typed_value_float64.is_null(2)); }
1351
1352 #[test]
1353 fn test_largeutf8_shredding() {
1354 let input = VariantArray::from_iter(vec![
1355 Some(Variant::from("hello")),
1356 Some(Variant::from(42i64)),
1357 None,
1358 Some(Variant::Null),
1359 Some(Variant::from("world")),
1360 ]);
1361
1362 let result = shred_variant(&input, &DataType::LargeUtf8).unwrap();
1363 let metadata = result.metadata_column();
1364 let value = result.value_column();
1365 let typed_value = result
1366 .typed_value_column()
1367 .unwrap()
1368 .as_any()
1369 .downcast_ref::<LargeStringArray>()
1370 .unwrap();
1371
1372 assert_eq!(result.len(), 5);
1373
1374 assert!(result.is_valid(0));
1376 assert!(value.is_null(0));
1377 assert_eq!(typed_value.value(0), "hello");
1378
1379 assert!(result.is_valid(1));
1381 assert!(value.is_valid(1));
1382 assert!(typed_value.is_null(1));
1383 assert_eq!(
1384 variant_from_arrays_at(metadata, value, 1).unwrap(),
1385 Variant::from(42i64)
1386 );
1387
1388 assert!(result.is_null(2));
1390 assert!(value.is_null(2));
1391 assert!(typed_value.is_null(2));
1392
1393 assert!(result.is_valid(3));
1395 assert!(value.is_valid(3));
1396 assert!(typed_value.is_null(3));
1397 assert_eq!(
1398 variant_from_arrays_at(metadata, value, 3).unwrap(),
1399 Variant::Null
1400 );
1401
1402 assert!(result.is_valid(4));
1404 assert!(value.is_null(4));
1405 assert_eq!(typed_value.value(4), "world");
1406 }
1407
1408 #[test]
1409 fn test_largebinary_shredding() {
1410 let input = VariantArray::from_iter(vec![
1411 Some(Variant::from(&b"\x00\x01\x02"[..])),
1412 Some(Variant::from("not_binary")),
1413 None,
1414 Some(Variant::Null),
1415 Some(Variant::from(&b"\xff\xaa"[..])),
1416 ]);
1417
1418 let result = shred_variant(&input, &DataType::LargeBinary).unwrap();
1419 let metadata = result.metadata_column();
1420 let value = result.value_column();
1421 let typed_value = result
1422 .typed_value_column()
1423 .unwrap()
1424 .as_any()
1425 .downcast_ref::<LargeBinaryArray>()
1426 .unwrap();
1427
1428 assert_eq!(result.len(), 5);
1429
1430 assert!(result.is_valid(0));
1432 assert!(value.is_null(0));
1433 assert_eq!(typed_value.value(0), &[0x00, 0x01, 0x02]);
1434
1435 assert!(result.is_valid(1));
1437 assert!(value.is_valid(1));
1438 assert!(typed_value.is_null(1));
1439 assert_eq!(
1440 variant_from_arrays_at(metadata, value, 1).unwrap(),
1441 Variant::from("not_binary")
1442 );
1443
1444 assert!(result.is_null(2));
1446 assert!(value.is_null(2));
1447 assert!(typed_value.is_null(2));
1448
1449 assert!(result.is_valid(3));
1451 assert!(value.is_valid(3));
1452 assert!(typed_value.is_null(3));
1453 assert_eq!(
1454 variant_from_arrays_at(metadata, value, 3).unwrap(),
1455 Variant::Null
1456 );
1457
1458 assert!(result.is_valid(4));
1460 assert!(value.is_null(4));
1461 assert_eq!(typed_value.value(4), &[0xff, 0xaa]);
1462 }
1463
1464 #[test]
1465 fn test_invalid_shredded_types_rejected() {
1466 let input = VariantArray::from_iter([Variant::from(42)]);
1467
1468 let invalid_types = vec![
1469 DataType::UInt8,
1470 DataType::Float16,
1471 DataType::Decimal256(38, 10),
1472 DataType::Date64,
1473 DataType::Time32(TimeUnit::Second),
1474 DataType::Time64(TimeUnit::Nanosecond),
1475 DataType::Timestamp(TimeUnit::Millisecond, None),
1476 DataType::FixedSizeBinary(17),
1477 DataType::Union(
1478 UnionFields::from_fields(vec![
1479 Field::new("int_field", DataType::Int32, false),
1480 Field::new("str_field", DataType::Utf8, true),
1481 ]),
1482 UnionMode::Dense,
1483 ),
1484 DataType::Map(
1485 Arc::new(Field::new(
1486 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1487 DataType::Struct(Fields::from(vec![
1488 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1489 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
1490 ])),
1491 false,
1492 )),
1493 false,
1494 ),
1495 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1496 DataType::RunEndEncoded(
1497 Arc::new(Field::new("run_ends", DataType::Int32, false)),
1498 Arc::new(Field::new("values", DataType::Utf8, true)),
1499 ),
1500 ];
1501
1502 for data_type in invalid_types {
1503 let err = shred_variant(&input, &data_type).unwrap_err();
1504 assert!(
1505 matches!(err, ArrowError::InvalidArgumentError(_)),
1506 "expected InvalidArgumentError for {data_type:?}, got {err:?}"
1507 );
1508 }
1509 }
1510
1511 #[test]
1512 fn test_array_shredding_as_list() {
1513 let input = build_variant_array(vec![
1514 VariantRow::List(vec![
1516 VariantValue::from(1i64),
1517 VariantValue::from(2i64),
1518 VariantValue::from(3i64),
1519 ]),
1520 VariantRow::List(vec![
1522 VariantValue::from(1i64),
1523 VariantValue::from("two"),
1524 VariantValue::from(Variant::Null),
1525 ]),
1526 VariantRow::Value(VariantValue::from("not a list")),
1528 VariantRow::Null,
1530 VariantRow::List(vec![]),
1532 ]);
1533 let list_schema = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
1534 let result = shred_variant(&input, &list_schema).unwrap();
1535 assert_eq!(result.len(), 5);
1536
1537 assert_list_structure_and_elements::<Int64Type, i32>(
1538 &result,
1539 5,
1540 &[0, 3, 6, 6, 6, 6],
1541 &[Some(3), Some(3), None, None, Some(0)],
1542 &[None, None, Some(Variant::from("not a list")), None, None],
1543 (
1544 &[Some(1), Some(2), Some(3), Some(1), None, None],
1545 &[
1546 None,
1547 None,
1548 None,
1549 None,
1550 Some(Variant::from("two")),
1551 Some(Variant::Null),
1552 ],
1553 ),
1554 );
1555 }
1556
1557 #[test]
1558 fn test_array_shredding_as_large_list() {
1559 let input = build_variant_array(vec![
1560 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1562 VariantRow::Value(VariantValue::from("not a list")),
1564 VariantRow::List(vec![]),
1566 ]);
1567 let list_schema = DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true)));
1568 let result = shred_variant(&input, &list_schema).unwrap();
1569 assert_eq!(result.len(), 3);
1570
1571 assert_list_structure_and_elements::<Int64Type, i64>(
1572 &result,
1573 3,
1574 &[0, 2, 2, 2],
1575 &[Some(2), None, Some(0)],
1576 &[None, Some(Variant::from("not a list")), None],
1577 (&[Some(1), Some(2)], &[None, None]),
1578 );
1579 }
1580
1581 #[test]
1582 fn test_array_shredding_as_list_view() {
1583 let input = build_variant_array(vec![
1584 VariantRow::List(vec![
1586 VariantValue::from(1i64),
1587 VariantValue::from(2i64),
1588 VariantValue::from(3i64),
1589 ]),
1590 VariantRow::List(vec![
1592 VariantValue::from(1i64),
1593 VariantValue::from("two"),
1594 VariantValue::from(Variant::Null),
1595 ]),
1596 VariantRow::Value(VariantValue::from("not a list")),
1598 VariantRow::Null,
1600 VariantRow::List(vec![]),
1602 ]);
1603 let list_schema = DataType::ListView(Arc::new(Field::new("item", DataType::Int64, true)));
1604 let result = shred_variant(&input, &list_schema).unwrap();
1605 assert_eq!(result.len(), 5);
1606
1607 assert_list_structure_and_elements::<Int64Type, i32>(
1608 &result,
1609 5,
1610 &[0, 3, 6, 6, 6],
1611 &[Some(3), Some(3), None, None, Some(0)],
1612 &[None, None, Some(Variant::from("not a list")), None, None],
1613 (
1614 &[Some(1), Some(2), Some(3), Some(1), None, None],
1615 &[
1616 None,
1617 None,
1618 None,
1619 None,
1620 Some(Variant::from("two")),
1621 Some(Variant::Null),
1622 ],
1623 ),
1624 );
1625 }
1626
1627 #[test]
1628 fn test_array_shredding_as_large_list_view() {
1629 let input = build_variant_array(vec![
1630 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1632 VariantRow::Value(VariantValue::from("fallback")),
1634 VariantRow::List(vec![]),
1636 ]);
1637 let list_schema =
1638 DataType::LargeListView(Arc::new(Field::new("item", DataType::Int64, true)));
1639 let result = shred_variant(&input, &list_schema).unwrap();
1640 assert_eq!(result.len(), 3);
1641
1642 assert_list_structure_and_elements::<Int64Type, i64>(
1643 &result,
1644 3,
1645 &[0, 2, 2],
1646 &[Some(2), None, Some(0)],
1647 &[None, Some(Variant::from("fallback")), None],
1648 (&[Some(1), Some(2)], &[None, None]),
1649 );
1650 }
1651
1652 #[test]
1653 fn test_array_shredding_as_fixed_size_list() {
1654 let input = build_variant_array(vec![
1655 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1656 VariantRow::Value(VariantValue::from("This should not be shredded")),
1657 VariantRow::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]),
1658 ]);
1659
1660 let list_schema =
1661 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
1662 let result = shred_variant(&input, &list_schema).unwrap();
1663 assert_eq!(result.len(), 3);
1664
1665 assert!(result.is_valid(0));
1668 assert!(result.value_column().is_null(0));
1669 assert!(result.typed_value_column().unwrap().is_valid(0));
1670
1671 assert!(result.is_valid(1));
1675 assert!(result.value_column().is_valid(1));
1676 assert!(result.typed_value_column().unwrap().is_null(1));
1677
1678 assert!(result.is_valid(2));
1681 assert!(result.value_column().is_null(2));
1682 assert!(result.typed_value_column().unwrap().is_valid(2));
1683
1684 let typed_value = result.typed_value_column().unwrap();
1685 let fixed_size_list = typed_value
1686 .as_any()
1687 .downcast_ref::<FixedSizeListArray>()
1688 .expect("Expected FixedSizeListArray");
1689
1690 assert_eq!(fixed_size_list.len(), 3);
1692 assert_eq!(fixed_size_list.value_length(), 2);
1693
1694 let val0 = fixed_size_list.value(0);
1696 let val0_struct = val0.as_any().downcast_ref::<StructArray>().unwrap();
1697 let val0_typed = val0_struct.column_by_name("typed_value").unwrap();
1698 let val0_ints = val0_typed.as_any().downcast_ref::<Int64Array>().unwrap();
1699 assert_eq!(val0_ints.values(), &[1i64, 2i64]);
1700
1701 assert!(fixed_size_list.is_null(1));
1704
1705 let val2 = fixed_size_list.value(2);
1707 let val2_struct = val2.as_any().downcast_ref::<StructArray>().unwrap();
1708 let val2_typed = val2_struct.column_by_name("typed_value").unwrap();
1709 let val2_ints = val2_typed.as_any().downcast_ref::<Int64Array>().unwrap();
1710 assert_eq!(val2_ints.values(), &[3i64, 4i64]);
1711 }
1712
1713 #[test]
1714 fn test_array_shredding_as_fixed_size_list_wrong_size() {
1715 let input = build_variant_array(vec![VariantRow::List(vec![
1716 VariantValue::from(1i64),
1717 VariantValue::from(2i64),
1718 VariantValue::from(3i64),
1719 ])]);
1720 let list_schema =
1721 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
1722
1723 let err = shred_variant(&input, &list_schema).unwrap_err();
1724 assert!(
1725 err.to_string()
1726 .contains("Expected fixed size list of size 2, got size 3"),
1727 "got: {err}",
1728 );
1729 }
1730
1731 #[test]
1732 fn test_array_shredding_with_array_elements() {
1733 let input = build_variant_array(vec![
1734 VariantRow::List(vec![
1736 VariantValue::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1737 VariantValue::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]),
1738 VariantValue::List(vec![]),
1739 ]),
1740 VariantRow::List(vec![
1742 VariantValue::List(vec![
1743 VariantValue::from(5i64),
1744 VariantValue::from("bad"),
1745 VariantValue::from(Variant::Null),
1746 ]),
1747 VariantValue::from("not a list inner"),
1748 VariantValue::Null,
1749 ]),
1750 VariantRow::Value(VariantValue::from("not a list")),
1752 VariantRow::Null,
1754 ]);
1755 let inner_field = Arc::new(Field::new("item", DataType::Int64, true));
1756 let inner_list_schema = DataType::List(inner_field);
1757 let list_schema = DataType::List(Arc::new(Field::new(
1758 "item",
1759 inner_list_schema.clone(),
1760 true,
1761 )));
1762 let result = shred_variant(&input, &list_schema).unwrap();
1763 assert_eq!(result.len(), 4);
1764
1765 let typed_value = result
1766 .typed_value_column()
1767 .unwrap()
1768 .as_any()
1769 .downcast_ref::<ListArray>()
1770 .unwrap();
1771
1772 assert_list_structure::<i32>(
1773 &result,
1774 4,
1775 &[0, 3, 6, 6, 6],
1776 &[Some(3), Some(3), None, None],
1777 &[None, None, Some(Variant::from("not a list")), None],
1778 );
1779
1780 let outer_elements =
1781 ShreddedVariantFieldArray::try_new(typed_value.values().as_ref()).unwrap();
1782 assert_eq!(outer_elements.len(), 6);
1783 let outer_values = outer_elements
1784 .typed_value_column()
1785 .unwrap()
1786 .as_any()
1787 .downcast_ref::<ListArray>()
1788 .unwrap();
1789 let outer_fallbacks = outer_elements.value_column();
1790
1791 let outer_metadata = Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
1792 EMPTY_VARIANT_METADATA_BYTES,
1793 outer_elements.len(),
1794 )));
1795 let outer_variant = VariantArray::from_parts(
1796 outer_metadata,
1797 outer_fallbacks.clone(),
1798 Some(Arc::new(outer_values.clone())),
1799 None,
1800 );
1801
1802 assert_list_structure_and_elements::<Int64Type, i32>(
1803 &outer_variant,
1804 outer_elements.len(),
1805 &[0, 2, 4, 4, 7, 7, 7],
1806 &[Some(2), Some(2), Some(0), Some(3), None, None],
1807 &[
1808 None,
1809 None,
1810 None,
1811 None,
1812 Some(Variant::from("not a list inner")),
1813 Some(Variant::Null),
1814 ],
1815 (
1816 &[Some(1), Some(2), Some(3), Some(4), Some(5), None, None],
1817 &[
1818 None,
1819 None,
1820 None,
1821 None,
1822 None,
1823 Some(Variant::from("bad")),
1824 Some(Variant::Null),
1825 ],
1826 ),
1827 );
1828 }
1829
1830 #[test]
1831 fn test_array_shredding_with_object_elements() {
1832 let input = build_variant_array(vec![
1833 VariantRow::List(vec![
1835 VariantValue::Object(vec![
1836 ("id", VariantValue::from(1i64)),
1837 ("name", VariantValue::from("Alice")),
1838 ]),
1839 VariantValue::Object(vec![("id", VariantValue::from(Variant::Null))]),
1840 ]),
1841 VariantRow::Value(VariantValue::from("not a list")),
1843 VariantRow::Null,
1845 ]);
1846
1847 let object_fields = Fields::from(vec![
1849 Field::new("id", DataType::Int64, true),
1850 Field::new("name", DataType::Utf8, true),
1851 ]);
1852 let list_schema = DataType::List(Arc::new(Field::new(
1853 "item",
1854 DataType::Struct(object_fields),
1855 true,
1856 )));
1857 let result = shred_variant(&input, &list_schema).unwrap();
1858 assert_eq!(result.len(), 3);
1859
1860 assert_list_structure::<i32>(
1861 &result,
1862 3,
1863 &[0, 2, 2, 2],
1864 &[Some(2), None, None],
1865 &[None, Some(Variant::from("not a list")), None],
1866 );
1867
1868 let typed_value = result
1870 .typed_value_column()
1871 .unwrap()
1872 .as_any()
1873 .downcast_ref::<ListArray>()
1874 .unwrap();
1875 let element_array =
1876 ShreddedVariantFieldArray::try_new(typed_value.values().as_ref()).unwrap();
1877 assert_eq!(element_array.len(), 2);
1878 let element_objects = element_array
1879 .typed_value_column()
1880 .unwrap()
1881 .as_any()
1882 .downcast_ref::<arrow::array::StructArray>()
1883 .unwrap();
1884
1885 let id_field =
1887 ShreddedVariantFieldArray::try_new(element_objects.column_by_name("id").unwrap())
1888 .unwrap();
1889 let id_values = id_field.value_column();
1890 let id_typed_values = id_field
1891 .typed_value_column()
1892 .unwrap()
1893 .as_any()
1894 .downcast_ref::<Int64Array>()
1895 .unwrap();
1896 assert!(id_values.is_null(0));
1897 assert_eq!(id_typed_values.value(0), 1);
1898 assert!(id_values.is_valid(1));
1900 assert_eq!(
1901 Variant::new(
1902 EMPTY_VARIANT_METADATA_BYTES,
1903 binary_array_value(id_values.as_ref(), 1).unwrap()
1904 ),
1905 Variant::Null
1906 );
1907 assert!(id_typed_values.is_null(1));
1908
1909 let name_field =
1911 ShreddedVariantFieldArray::try_new(element_objects.column_by_name("name").unwrap())
1912 .unwrap();
1913 let name_values = name_field.value_column();
1914 let name_typed_values = name_field
1915 .typed_value_column()
1916 .unwrap()
1917 .as_any()
1918 .downcast_ref::<StringArray>()
1919 .unwrap();
1920 assert!(name_values.is_null(0));
1921 assert_eq!(name_typed_values.value(0), "Alice");
1922 assert!(name_values.is_null(1));
1924 assert!(name_typed_values.is_null(1));
1925 }
1926
1927 #[test]
1928 fn test_object_shredding_comprehensive() -> Result<()> {
1929 let input = build_variant_array(vec![
1930 VariantRow::Object(vec![
1932 ("score", VariantValue::from(95.5f64)),
1933 ("age", VariantValue::from(30i64)),
1934 ]),
1935 VariantRow::Object(vec![
1937 ("score", VariantValue::from(87.2f64)),
1938 ("age", VariantValue::from(25i64)),
1939 ("email", VariantValue::from("bob@example.com")),
1940 ]),
1941 VariantRow::Object(vec![("age", VariantValue::from(35i64))]),
1943 VariantRow::Object(vec![
1945 ("score", VariantValue::from("ninety-five")),
1946 ("age", VariantValue::from("thirty")),
1947 ]),
1948 VariantRow::Value(VariantValue::from("not an object")),
1950 VariantRow::Object(vec![]),
1952 VariantRow::Null,
1954 VariantRow::Object(vec![("foo", VariantValue::from(10))]),
1956 VariantRow::Object(vec![
1958 ("score", VariantValue::from(66.67f64)),
1959 ("foo", VariantValue::from(10)),
1960 ]),
1961 ]);
1962
1963 let target_schema = ShreddedSchemaBuilder::default()
1966 .with_path("score", &DataType::Float64)?
1967 .with_path("age", &DataType::Int64)?
1968 .build();
1969
1970 let result = shred_variant(&input, &target_schema).unwrap();
1971
1972 assert!(result.typed_value_column().is_some());
1974 assert_eq!(result.len(), 9);
1975
1976 let metadata = result.metadata_column();
1977 let value = result.value_column();
1978 let typed_value = result
1979 .typed_value_column()
1980 .unwrap()
1981 .as_any()
1982 .downcast_ref::<arrow::array::StructArray>()
1983 .unwrap();
1984
1985 let score_field =
1987 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("score").unwrap())
1988 .unwrap();
1989 let age_field =
1990 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("age").unwrap()).unwrap();
1991
1992 let score_value = score_field.value_column();
1993 let score_typed_value = score_field
1994 .typed_value_column()
1995 .unwrap()
1996 .as_any()
1997 .downcast_ref::<Float64Array>()
1998 .unwrap();
1999 let age_value = age_field.value_column();
2000 let age_typed_value = age_field
2001 .typed_value_column()
2002 .unwrap()
2003 .as_any()
2004 .downcast_ref::<Int64Array>()
2005 .unwrap();
2006
2007 struct ShreddedValue<'m, 'v, T> {
2009 value: Option<Variant<'m, 'v>>,
2010 typed_value: Option<T>,
2011 }
2012 struct ShreddedStruct<'m, 'v> {
2013 score: ShreddedValue<'m, 'v, f64>,
2014 age: ShreddedValue<'m, 'v, i64>,
2015 }
2016 fn get_value<'m, 'v>(
2017 i: usize,
2018 metadata: &'m dyn Array,
2019 value: &'v dyn Array,
2020 ) -> Variant<'m, 'v> {
2021 variant_from_arrays_at(metadata, value, i).unwrap()
2022 }
2023 let expect = |i, expected_result: Option<ShreddedValue<ShreddedStruct>>| {
2024 match expected_result {
2025 Some(ShreddedValue {
2026 value: expected_value,
2027 typed_value: expected_typed_value,
2028 }) => {
2029 assert!(result.is_valid(i));
2030 match expected_value {
2031 Some(expected_value) => {
2032 assert!(value.is_valid(i));
2033 assert_eq!(
2034 expected_value,
2035 get_value(i, metadata.as_ref(), value.as_ref())
2036 );
2037 }
2038 None => {
2039 assert!(value.is_null(i));
2040 }
2041 }
2042 match expected_typed_value {
2043 Some(ShreddedStruct {
2044 score: expected_score,
2045 age: expected_age,
2046 }) => {
2047 assert!(typed_value.is_valid(i));
2048 assert!(score_field.is_valid(i)); assert!(age_field.is_valid(i)); match expected_score.value {
2051 Some(expected_score_value) => {
2052 assert!(score_value.is_valid(i));
2053 assert_eq!(
2054 expected_score_value,
2055 get_value(i, metadata.as_ref(), score_value.as_ref())
2056 );
2057 }
2058 None => {
2059 assert!(score_value.is_null(i));
2060 }
2061 }
2062 match expected_score.typed_value {
2063 Some(expected_score) => {
2064 assert!(score_typed_value.is_valid(i));
2065 assert_eq!(expected_score, score_typed_value.value(i));
2066 }
2067 None => {
2068 assert!(score_typed_value.is_null(i));
2069 }
2070 }
2071 match expected_age.value {
2072 Some(expected_age_value) => {
2073 assert!(age_value.is_valid(i));
2074 assert_eq!(
2075 expected_age_value,
2076 get_value(i, metadata.as_ref(), age_value.as_ref())
2077 );
2078 }
2079 None => {
2080 assert!(age_value.is_null(i));
2081 }
2082 }
2083 match expected_age.typed_value {
2084 Some(expected_age) => {
2085 assert!(age_typed_value.is_valid(i));
2086 assert_eq!(expected_age, age_typed_value.value(i));
2087 }
2088 None => {
2089 assert!(age_typed_value.is_null(i));
2090 }
2091 }
2092 }
2093 None => {
2094 assert!(typed_value.is_null(i));
2095 }
2096 }
2097 }
2098 None => {
2099 assert!(result.is_null(i));
2100 }
2101 }
2102 };
2103
2104 expect(
2106 0,
2107 Some(ShreddedValue {
2108 value: None,
2109 typed_value: Some(ShreddedStruct {
2110 score: ShreddedValue {
2111 value: None,
2112 typed_value: Some(95.5),
2113 },
2114 age: ShreddedValue {
2115 value: None,
2116 typed_value: Some(30),
2117 },
2118 }),
2119 }),
2120 );
2121
2122 let mut builder = VariantBuilder::new();
2124 builder
2125 .new_object()
2126 .with_field("email", "bob@example.com")
2127 .finish();
2128 let (m, v) = builder.finish();
2129 let expected_value = Variant::new(&m, &v);
2130
2131 expect(
2132 1,
2133 Some(ShreddedValue {
2134 value: Some(expected_value),
2135 typed_value: Some(ShreddedStruct {
2136 score: ShreddedValue {
2137 value: None,
2138 typed_value: Some(87.2),
2139 },
2140 age: ShreddedValue {
2141 value: None,
2142 typed_value: Some(25),
2143 },
2144 }),
2145 }),
2146 );
2147
2148 expect(
2150 2,
2151 Some(ShreddedValue {
2152 value: None,
2153 typed_value: Some(ShreddedStruct {
2154 score: ShreddedValue {
2155 value: None,
2156 typed_value: None,
2157 },
2158 age: ShreddedValue {
2159 value: None,
2160 typed_value: Some(35),
2161 },
2162 }),
2163 }),
2164 );
2165
2166 expect(
2168 3,
2169 Some(ShreddedValue {
2170 value: None,
2171 typed_value: Some(ShreddedStruct {
2172 score: ShreddedValue {
2173 value: Some(Variant::from("ninety-five")),
2174 typed_value: None,
2175 },
2176 age: ShreddedValue {
2177 value: Some(Variant::from("thirty")),
2178 typed_value: None,
2179 },
2180 }),
2181 }),
2182 );
2183
2184 expect(
2186 4,
2187 Some(ShreddedValue {
2188 value: Some(Variant::from("not an object")),
2189 typed_value: None,
2190 }),
2191 );
2192
2193 expect(
2195 5,
2196 Some(ShreddedValue {
2197 value: None,
2198 typed_value: Some(ShreddedStruct {
2199 score: ShreddedValue {
2200 value: None,
2201 typed_value: None,
2202 },
2203 age: ShreddedValue {
2204 value: None,
2205 typed_value: None,
2206 },
2207 }),
2208 }),
2209 );
2210
2211 expect(6, None);
2213
2214 let object_with_foo_field = |i| {
2216 use parquet_variant::{ParentState, ValueBuilder, VariantMetadata};
2217 let metadata = VariantMetadata::new(binary_array_value(metadata.as_ref(), i).unwrap());
2218 let mut metadata_builder = ReadOnlyMetadataBuilder::new(&metadata);
2219 let mut value_builder = ValueBuilder::new();
2220 let state = ParentState::variant(&mut value_builder, &mut metadata_builder);
2221 ObjectBuilder::new(state, false)
2222 .with_field("foo", 10)
2223 .finish();
2224 (metadata, value_builder.into_inner())
2225 };
2226
2227 let (m, v) = object_with_foo_field(7);
2229 expect(
2230 7,
2231 Some(ShreddedValue {
2232 value: Some(Variant::new_with_metadata(m, &v)),
2233 typed_value: Some(ShreddedStruct {
2234 score: ShreddedValue {
2235 value: None,
2236 typed_value: None,
2237 },
2238 age: ShreddedValue {
2239 value: None,
2240 typed_value: None,
2241 },
2242 }),
2243 }),
2244 );
2245
2246 let (m, v) = object_with_foo_field(8);
2248 expect(
2249 8,
2250 Some(ShreddedValue {
2251 value: Some(Variant::new_with_metadata(m, &v)),
2252 typed_value: Some(ShreddedStruct {
2253 score: ShreddedValue {
2254 value: None,
2255 typed_value: Some(66.67),
2256 },
2257 age: ShreddedValue {
2258 value: None,
2259 typed_value: None,
2260 },
2261 }),
2262 }),
2263 );
2264 Ok(())
2265 }
2266
2267 #[test]
2268 fn test_object_shredding_with_array_field() {
2269 let input = build_variant_array(vec![
2270 VariantRow::Object(vec![(
2272 "scores",
2273 VariantValue::List(vec![VariantValue::from(10i64), VariantValue::from(20i64)]),
2274 )]),
2275 VariantRow::Object(vec![(
2277 "scores",
2278 VariantValue::List(vec![
2279 VariantValue::from("oops"),
2280 VariantValue::from(Variant::Null),
2281 ]),
2282 )]),
2283 VariantRow::Object(vec![]),
2285 VariantRow::Value(VariantValue::from("not an object")),
2287 VariantRow::Null,
2289 ]);
2290 let list_field = Arc::new(Field::new("item", DataType::Int64, true));
2291 let inner_list_schema = DataType::List(list_field);
2292 let schema = DataType::Struct(Fields::from(vec![Field::new(
2293 "scores",
2294 inner_list_schema.clone(),
2295 true,
2296 )]));
2297
2298 let result = shred_variant(&input, &schema).unwrap();
2299 assert_eq!(result.len(), 5);
2300
2301 let value_field = result.value_column();
2303 let typed_struct = result
2304 .typed_value_column()
2305 .unwrap()
2306 .as_any()
2307 .downcast_ref::<arrow::array::StructArray>()
2308 .unwrap();
2309
2310 assert!(value_field.is_null(0));
2312 assert!(value_field.is_null(1));
2313 assert!(value_field.is_null(2));
2314 assert!(value_field.is_valid(3));
2315 assert_eq!(
2316 variant_from_arrays_at(result.metadata_column(), value_field, 3).unwrap(),
2317 Variant::from("not an object")
2318 );
2319 assert!(value_field.is_null(4));
2320
2321 assert!(typed_struct.is_valid(0));
2323 assert!(typed_struct.is_valid(1));
2324 assert!(typed_struct.is_valid(2));
2325 assert!(typed_struct.is_null(3));
2326 assert!(typed_struct.is_null(4));
2327
2328 let scores_field =
2330 ShreddedVariantFieldArray::try_new(typed_struct.column_by_name("scores").unwrap())
2331 .unwrap();
2332 assert_list_structure_and_elements::<Int64Type, i32>(
2333 &VariantArray::from_parts(
2334 Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
2335 EMPTY_VARIANT_METADATA_BYTES,
2336 scores_field.len(),
2337 ))),
2338 scores_field.value_column().clone(),
2339 Some(scores_field.typed_value_column().unwrap().clone()),
2340 None,
2341 ),
2342 scores_field.len(),
2343 &[0i32, 2, 4, 4, 4, 4],
2344 &[Some(2), Some(2), None, None, None],
2345 &[None, None, None, None, None],
2346 (
2347 &[Some(10), Some(20), None, None],
2348 &[None, None, Some(Variant::from("oops")), Some(Variant::Null)],
2349 ),
2350 );
2351 }
2352
2353 #[test]
2354 fn test_object_different_schemas() -> Result<()> {
2355 let input = build_variant_array(vec![VariantRow::Object(vec![
2357 ("id", VariantValue::from(123i32)),
2358 ("age", VariantValue::from(25i64)),
2359 ("score", VariantValue::from(95.5f64)),
2360 ])]);
2361
2362 let schema1 = ShreddedSchemaBuilder::default()
2364 .with_path("id", &DataType::Int32)?
2365 .build();
2366 let result1 = shred_variant(&input, &schema1).unwrap();
2367 let value_field1 = result1.value_column();
2368 assert!(!value_field1.is_null(0)); let schema2 = ShreddedSchemaBuilder::default()
2372 .with_path("id", &DataType::Int32)?
2373 .with_path("age", &DataType::Int64)?
2374 .build();
2375 let result2 = shred_variant(&input, &schema2).unwrap();
2376 let value_field2 = result2.value_column();
2377 assert!(!value_field2.is_null(0)); let schema3 = ShreddedSchemaBuilder::default()
2381 .with_path("id", &DataType::Int32)?
2382 .with_path("age", &DataType::Int64)?
2383 .with_path("score", &DataType::Float64)?
2384 .build();
2385 let result3 = shred_variant(&input, &schema3).unwrap();
2386 let value_field3 = result3.value_column();
2387 assert!(value_field3.is_null(0)); Ok(())
2390 }
2391
2392 #[test]
2393 fn test_uuid_shredding_in_objects() -> Result<()> {
2394 let mock_uuid_1 = Uuid::new_v4();
2395 let mock_uuid_2 = Uuid::new_v4();
2396 let mock_uuid_3 = Uuid::new_v4();
2397
2398 let input = build_variant_array(vec![
2399 VariantRow::Object(vec![
2401 ("id", VariantValue::from(mock_uuid_1)),
2402 ("session_id", VariantValue::from(mock_uuid_2)),
2403 ]),
2404 VariantRow::Object(vec![
2406 ("id", VariantValue::from(mock_uuid_2)),
2407 ("session_id", VariantValue::from(mock_uuid_3)),
2408 ("name", VariantValue::from("test_user")),
2409 ]),
2410 VariantRow::Object(vec![("id", VariantValue::from(mock_uuid_1))]),
2412 VariantRow::Object(vec![
2414 ("id", VariantValue::from(mock_uuid_3)),
2415 ("session_id", VariantValue::from("not-a-uuid")),
2416 ]),
2417 VariantRow::Object(vec![
2419 ("id", VariantValue::from(12345i64)),
2420 ("session_id", VariantValue::from(mock_uuid_1)),
2421 ]),
2422 VariantRow::Null,
2424 ]);
2425
2426 let target_schema = ShreddedSchemaBuilder::default()
2427 .with_path("id", DataType::FixedSizeBinary(16))?
2428 .with_path("session_id", DataType::FixedSizeBinary(16))?
2429 .build();
2430
2431 let result = shred_variant(&input, &target_schema).unwrap();
2432
2433 assert!(result.typed_value_column().is_some());
2434 assert_eq!(result.len(), 6);
2435
2436 let metadata = result.metadata_column();
2437 let value = result.value_column();
2438 let typed_value = result
2439 .typed_value_column()
2440 .unwrap()
2441 .as_any()
2442 .downcast_ref::<arrow::array::StructArray>()
2443 .unwrap();
2444
2445 let id_field =
2447 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("id").unwrap()).unwrap();
2448 let session_id_field =
2449 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("session_id").unwrap())
2450 .unwrap();
2451
2452 let id_value = id_field.value_column();
2453 let id_typed_value = id_field
2454 .typed_value_column()
2455 .unwrap()
2456 .as_any()
2457 .downcast_ref::<FixedSizeBinaryArray>()
2458 .unwrap();
2459 let session_id_value = session_id_field.value_column();
2460 let session_id_typed_value = session_id_field
2461 .typed_value_column()
2462 .unwrap()
2463 .as_any()
2464 .downcast_ref::<FixedSizeBinaryArray>()
2465 .unwrap();
2466
2467 assert!(result.is_valid(0));
2469
2470 assert!(value.is_null(0)); assert!(id_value.is_null(0));
2472 assert!(session_id_value.is_null(0));
2473
2474 assert!(typed_value.is_valid(0));
2475 assert!(id_typed_value.is_valid(0));
2476 assert!(session_id_typed_value.is_valid(0));
2477
2478 assert_eq!(id_typed_value.value(0), mock_uuid_1.as_bytes());
2479 assert_eq!(session_id_typed_value.value(0), mock_uuid_2.as_bytes());
2480
2481 assert!(result.is_valid(1));
2483
2484 assert!(value.is_valid(1)); assert!(typed_value.is_valid(1));
2486
2487 assert!(id_value.is_null(1));
2488 assert!(id_typed_value.is_valid(1));
2489 assert_eq!(id_typed_value.value(1), mock_uuid_2.as_bytes());
2490
2491 assert!(session_id_value.is_null(1));
2492 assert!(session_id_typed_value.is_valid(1));
2493 assert_eq!(session_id_typed_value.value(1), mock_uuid_3.as_bytes());
2494
2495 let row_1_variant = variant_from_arrays_at(metadata, value, 1).unwrap();
2497 let Variant::Object(obj) = row_1_variant else {
2498 panic!("Expected object");
2499 };
2500
2501 assert_eq!(obj.get("name"), Some(Variant::from("test_user")));
2502
2503 assert!(result.is_valid(2));
2505
2506 assert!(value.is_null(2)); assert!(typed_value.is_valid(2));
2508
2509 assert!(id_value.is_null(2));
2510 assert!(id_typed_value.is_valid(2));
2511 assert_eq!(id_typed_value.value(2), mock_uuid_1.as_bytes());
2512
2513 assert!(session_id_value.is_null(2));
2514 assert!(session_id_typed_value.is_null(2)); assert!(result.is_valid(3));
2518
2519 assert!(value.is_null(3)); assert!(typed_value.is_valid(3));
2521
2522 assert!(id_value.is_null(3));
2523 assert!(id_typed_value.is_valid(3));
2524 assert_eq!(id_typed_value.value(3), mock_uuid_3.as_bytes());
2525
2526 assert!(session_id_value.is_valid(3)); assert!(session_id_typed_value.is_null(3));
2528 let session_id_variant = variant_from_arrays_at(metadata, session_id_value, 3).unwrap();
2529 assert_eq!(session_id_variant, Variant::from("not-a-uuid"));
2530
2531 assert!(result.is_valid(4));
2533
2534 assert!(value.is_null(4)); assert!(typed_value.is_valid(4));
2536
2537 assert!(id_value.is_valid(4)); assert!(id_typed_value.is_null(4));
2539 let id_variant = variant_from_arrays_at(metadata, id_value, 4).unwrap();
2540 assert_eq!(id_variant, Variant::from(12345i64));
2541
2542 assert!(session_id_value.is_null(4));
2543 assert!(session_id_typed_value.is_valid(4));
2544 assert_eq!(session_id_typed_value.value(4), mock_uuid_1.as_bytes());
2545
2546 assert!(result.is_null(5));
2548
2549 Ok(())
2550 }
2551
2552 macro_rules! validate_decimal_shredding {
2553 ($shred_type: expr, $array_type: ty, $expected_typed_value: ident $(, $expected_precision: literal, $expected_scale:literal)? $(,)?) => {{
2554 let input = VariantArray::from_iter(vec![
2555 Variant::from(12i8),
2556 Variant::from(234i16),
2557 Variant::from(456i32),
2558 Variant::from(456i64),
2559 Variant::from(VariantDecimal4::try_new(1200, 2).unwrap()),
2560 Variant::from(VariantDecimal8::try_new(1230, 2).unwrap()),
2561 Variant::from(VariantDecimal16::try_new(1234, 2).unwrap()),
2562 ]);
2563
2564 let result = shred_variant(&input, &$shred_type).unwrap();
2565
2566 assert!(result.typed_value_column().is_some());
2567 assert_eq!(result.len(), input.len());
2568
2569 let value = result.value_column();
2570 let typed_value = result
2571 .typed_value_column()
2572 .unwrap()
2573 .as_any()
2574 .downcast_ref::<$array_type>()
2575 .unwrap();
2576
2577 $(assert_eq!(typed_value.precision(), $expected_precision);)?
2578 $(assert_eq!(typed_value.scale(), $expected_scale);)?
2579
2580 for i in 0..$expected_typed_value.len() {
2581 assert_eq!(value.is_valid(i), $expected_typed_value.is_null(i));
2582 assert_eq!(typed_value.is_valid(i), $expected_typed_value.is_valid(i));
2583 assert_eq!(typed_value.value(i), $expected_typed_value.value(i));
2584 }
2585 }};
2586 }
2587
2588 #[test]
2589 fn test_shredding_decimal32_with_same_scale() {
2590 let expected_array = Decimal32Array::from(vec![
2591 Some(1200),
2592 None, None, None, Some(1200),
2596 Some(1230),
2597 Some(1234),
2598 ])
2599 .with_precision_and_scale(4, 2)
2600 .unwrap();
2601 validate_decimal_shredding!(
2602 DataType::Decimal32(4, 2),
2603 arrow::array::Decimal32Array,
2604 expected_array,
2605 4,
2606 2,
2607 );
2608 }
2609
2610 #[test]
2611 fn test_shredding_decimal32_with_bigger_scale() {
2612 let expected_array = Decimal32Array::from(vec![
2613 Some(12000),
2614 Some(234000),
2615 Some(456000),
2616 Some(456000),
2617 Some(12000),
2618 Some(12300),
2619 Some(12340),
2620 ])
2621 .with_precision_and_scale(6, 3)
2622 .unwrap();
2623
2624 validate_decimal_shredding!(
2625 DataType::Decimal32(6, 3),
2626 arrow::array::Decimal32Array,
2627 expected_array,
2628 6,
2629 3,
2630 );
2631 }
2632
2633 #[test]
2634 fn test_shredding_decimal32_with_smaller_scale() {
2635 let expected_array = Decimal32Array::from(vec![
2636 Some(12),
2637 Some(234),
2638 Some(456),
2639 Some(456),
2640 Some(12),
2641 None, None, ])
2644 .with_precision_and_scale(6, 0)
2645 .unwrap();
2646 validate_decimal_shredding!(
2647 DataType::Decimal32(6, 0),
2648 arrow::array::Decimal32Array,
2649 expected_array,
2650 6,
2651 0
2652 );
2653 }
2654
2655 #[test]
2656 fn test_shredding_decimal64_with_same_scale() {
2657 let expected_array_decimal64_same_scale = Decimal64Array::from(vec![
2658 Some(1200),
2659 None, None, None, Some(1200),
2663 Some(1230),
2664 Some(1234),
2665 ])
2666 .with_precision_and_scale(4, 2)
2667 .unwrap();
2668 validate_decimal_shredding!(
2669 DataType::Decimal64(4, 2),
2670 arrow::array::Decimal64Array,
2671 expected_array_decimal64_same_scale,
2672 4,
2673 2
2674 );
2675 }
2676
2677 #[test]
2678 fn test_shredding_decimal64_with_big_scale() {
2679 let expected_array = Decimal64Array::from(vec![
2680 Some(12000),
2681 Some(234000),
2682 Some(456000),
2683 Some(456000),
2684 Some(12000),
2685 Some(12300),
2686 Some(12340),
2687 ])
2688 .with_precision_and_scale(6, 3)
2689 .unwrap();
2690 validate_decimal_shredding!(
2691 DataType::Decimal64(6, 3),
2692 arrow::array::Decimal64Array,
2693 expected_array,
2694 6,
2695 3,
2696 );
2697 }
2698
2699 #[test]
2700 fn test_shredding_decimal64_with_smaller_scale() {
2701 let expected_array = Decimal64Array::from(vec![
2702 Some(12),
2703 Some(234),
2704 Some(456),
2705 Some(456),
2706 Some(12),
2707 None, None, ])
2710 .with_precision_and_scale(6, 0)
2711 .unwrap();
2712 validate_decimal_shredding!(
2713 DataType::Decimal64(6, 0),
2714 arrow::array::Decimal64Array,
2715 expected_array,
2716 6,
2717 0
2718 );
2719 }
2720
2721 #[test]
2722 fn test_shredding_decimal128_with_same_scale() {
2723 let expected_array = Decimal128Array::from(vec![
2724 Some(1200),
2725 None, None, None, Some(1200),
2729 Some(1230),
2730 Some(1234),
2731 ])
2732 .with_precision_and_scale(4, 2)
2733 .unwrap();
2734
2735 validate_decimal_shredding!(
2736 DataType::Decimal128(4, 2),
2737 arrow::array::Decimal128Array,
2738 expected_array,
2739 4,
2740 2,
2741 );
2742 }
2743
2744 #[test]
2745 fn test_shredding_decimal128_with_big_scale() {
2746 let expected_array = Decimal128Array::from(vec![
2747 Some(12000),
2748 Some(234000),
2749 Some(456000),
2750 Some(456000),
2751 Some(12000),
2752 Some(12300),
2753 Some(12340),
2754 ])
2755 .with_precision_and_scale(6, 3)
2756 .unwrap();
2757 validate_decimal_shredding!(
2758 DataType::Decimal128(6, 3),
2759 arrow::array::Decimal128Array,
2760 expected_array,
2761 6,
2762 3
2763 );
2764 }
2765
2766 #[test]
2767 fn test_shredding_decimal128_with_smaller_scale() {
2768 let expected_array = Decimal128Array::from(vec![
2769 Some(12),
2770 Some(234),
2771 Some(456),
2772 Some(456),
2773 Some(12),
2774 None, None, ])
2777 .with_precision_and_scale(6, 0)
2778 .unwrap();
2779 validate_decimal_shredding!(
2780 DataType::Decimal128(6, 0),
2781 arrow::array::Decimal128Array,
2782 expected_array,
2783 6,
2784 0
2785 );
2786 }
2787
2788 #[test]
2789 fn test_shredding_decimal128_to_integer() {
2790 let expected_array = Int64Array::from(vec![
2791 Some(12),
2792 Some(234),
2793 Some(456),
2794 Some(456),
2795 Some(12),
2796 None, None, ]);
2799
2800 validate_decimal_shredding!(DataType::Int64, arrow::array::Int64Array, expected_array);
2801 }
2802
2803 #[test]
2804 fn test_spec_compliance() {
2805 let input = VariantArray::from_iter(vec![Variant::from(42i64), Variant::from("hello")]);
2806
2807 let result = shred_variant(&input, &DataType::Int64).unwrap();
2808
2809 let inner_struct = result.inner();
2811 assert!(inner_struct.column_by_name("metadata").is_some());
2812 assert!(inner_struct.column_by_name("value").is_some());
2813 assert!(inner_struct.column_by_name("typed_value").is_some());
2814
2815 assert_eq!(
2817 result.metadata_column().len(),
2818 input.metadata_column().len()
2819 );
2820 assert_eq!(
2823 result.metadata_column().len(),
2824 input.metadata_column().len()
2825 );
2826
2827 assert_eq!(result.len(), input.len());
2829 assert!(result.typed_value_column().is_some());
2830
2831 let value_field = result.value_column();
2834 let typed_value_field = result
2835 .typed_value_column()
2836 .unwrap()
2837 .as_any()
2838 .downcast_ref::<Int64Array>()
2839 .unwrap();
2840
2841 for i in 0..result.len() {
2842 if !result.is_null(i) {
2843 let value_is_null = value_field.is_null(i);
2844 let typed_value_is_null = typed_value_field.is_null(i);
2845 assert!(
2847 value_is_null || typed_value_is_null,
2848 "Row {i}: both value and typed_value are non-null for primitive shredding"
2849 );
2850 }
2851 }
2852 }
2853
2854 #[test]
2855 fn test_variant_schema_builder_simple() -> Result<()> {
2856 let shredding_type = ShreddedSchemaBuilder::default()
2857 .with_path("a", &DataType::Int64)?
2858 .with_path("b", &DataType::Float64)?
2859 .build();
2860
2861 assert_eq!(
2862 shredding_type,
2863 DataType::Struct(Fields::from(vec![
2864 Field::new("a", DataType::Int64, true),
2865 Field::new("b", DataType::Float64, true),
2866 ]))
2867 );
2868
2869 Ok(())
2870 }
2871
2872 #[test]
2873 fn test_variant_schema_builder_nested() -> Result<()> {
2874 let shredding_type = ShreddedSchemaBuilder::default()
2875 .with_path("a", &DataType::Int64)?
2876 .with_path("b.c", &DataType::Utf8)?
2877 .with_path("b.d", &DataType::Float64)?
2878 .build();
2879
2880 assert_eq!(
2881 shredding_type,
2882 DataType::Struct(Fields::from(vec![
2883 Field::new("a", DataType::Int64, true),
2884 Field::new(
2885 "b",
2886 DataType::Struct(Fields::from(vec![
2887 Field::new("c", DataType::Utf8, true),
2888 Field::new("d", DataType::Float64, true),
2889 ])),
2890 true
2891 ),
2892 ]))
2893 );
2894
2895 Ok(())
2896 }
2897
2898 #[test]
2899 fn test_variant_schema_builder_with_path_variant_path_arg() -> Result<()> {
2900 let path = VariantPath::from_iter([VariantPathElement::from("a.b")]);
2901 let shredding_type = ShreddedSchemaBuilder::default()
2902 .with_path(path, &DataType::Int64)?
2903 .build();
2904
2905 match shredding_type {
2906 DataType::Struct(fields) => {
2907 assert_eq!(fields.len(), 1);
2908 assert_eq!(fields[0].name(), "a.b");
2909 assert_eq!(fields[0].data_type(), &DataType::Int64);
2910 }
2911 _ => panic!("expected struct data type"),
2912 }
2913
2914 Ok(())
2915 }
2916
2917 #[test]
2918 fn test_variant_schema_builder_custom_nullability() -> Result<()> {
2919 let shredding_type = ShreddedSchemaBuilder::default()
2920 .with_path(
2921 "foo",
2922 Arc::new(Field::new("should_be_renamed", DataType::Utf8, false)),
2923 )?
2924 .with_path("bar", (&DataType::Int64, false))?
2925 .build();
2926
2927 let DataType::Struct(fields) = shredding_type else {
2928 panic!("expected struct data type");
2929 };
2930
2931 let foo = fields.iter().find(|f| f.name() == "foo").unwrap();
2932 assert_eq!(foo.data_type(), &DataType::Utf8);
2933 assert!(!foo.is_nullable());
2934
2935 let bar = fields.iter().find(|f| f.name() == "bar").unwrap();
2936 assert_eq!(bar.data_type(), &DataType::Int64);
2937 assert!(!bar.is_nullable());
2938
2939 Ok(())
2940 }
2941
2942 #[test]
2943 fn test_variant_schema_builder_with_shred_variant() -> Result<()> {
2944 let input = build_variant_array(vec![
2945 VariantRow::Object(vec![
2946 ("time", VariantValue::from(1234567890i64)),
2947 ("hostname", VariantValue::from("server1")),
2948 ("extra", VariantValue::from(42)),
2949 ]),
2950 VariantRow::Object(vec![
2951 ("time", VariantValue::from(9876543210i64)),
2952 ("hostname", VariantValue::from("server2")),
2953 ]),
2954 VariantRow::Null,
2955 ]);
2956
2957 let shredding_type = ShreddedSchemaBuilder::default()
2958 .with_path("time", &DataType::Int64)?
2959 .with_path("hostname", &DataType::Utf8)?
2960 .build();
2961
2962 let result = shred_variant(&input, &shredding_type).unwrap();
2963
2964 assert_eq!(
2965 result.data_type(),
2966 &DataType::Struct(Fields::from(vec![
2967 Field::new("metadata", DataType::BinaryView, false),
2968 Field::new("value", DataType::BinaryView, true),
2969 Field::new(
2970 "typed_value",
2971 DataType::Struct(Fields::from(vec![
2972 Field::new(
2973 "hostname",
2974 DataType::Struct(Fields::from(vec![
2975 Field::new("value", DataType::BinaryView, true),
2976 Field::new("typed_value", DataType::Utf8, true),
2977 ])),
2978 false,
2979 ),
2980 Field::new(
2981 "time",
2982 DataType::Struct(Fields::from(vec![
2983 Field::new("value", DataType::BinaryView, true),
2984 Field::new("typed_value", DataType::Int64, true),
2985 ])),
2986 false,
2987 ),
2988 ])),
2989 true,
2990 ),
2991 ]))
2992 );
2993
2994 assert_eq!(result.len(), 3);
2995 assert!(result.typed_value_column().is_some());
2996
2997 let typed_value = result
2998 .typed_value_column()
2999 .unwrap()
3000 .as_any()
3001 .downcast_ref::<arrow::array::StructArray>()
3002 .unwrap();
3003
3004 let time_field =
3005 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("time").unwrap())
3006 .unwrap();
3007 let hostname_field =
3008 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("hostname").unwrap())
3009 .unwrap();
3010
3011 let time_typed = time_field
3012 .typed_value_column()
3013 .unwrap()
3014 .as_any()
3015 .downcast_ref::<Int64Array>()
3016 .unwrap();
3017 let hostname_typed = hostname_field
3018 .typed_value_column()
3019 .unwrap()
3020 .as_any()
3021 .downcast_ref::<arrow::array::StringArray>()
3022 .unwrap();
3023
3024 assert!(!result.is_null(0));
3026 assert_eq!(time_typed.value(0), 1234567890);
3027 assert_eq!(hostname_typed.value(0), "server1");
3028
3029 assert!(!result.is_null(1));
3031 assert_eq!(time_typed.value(1), 9876543210);
3032 assert_eq!(hostname_typed.value(1), "server2");
3033
3034 assert!(result.is_null(2));
3036
3037 Ok(())
3038 }
3039
3040 #[test]
3041 fn test_variant_schema_builder_conflicting_path() -> Result<()> {
3042 let shredding_type = ShreddedSchemaBuilder::default()
3043 .with_path("a", &DataType::Int64)?
3044 .with_path("a", &DataType::Float64)?
3045 .build();
3046
3047 assert_eq!(
3048 shredding_type,
3049 DataType::Struct(Fields::from(
3050 vec![Field::new("a", DataType::Float64, true),]
3051 ))
3052 );
3053
3054 Ok(())
3055 }
3056
3057 #[test]
3058 fn test_variant_schema_builder_root_path() -> Result<()> {
3059 let path = VariantPath::new(vec![]);
3060 let shredding_type = ShreddedSchemaBuilder::default()
3061 .with_path(path, &DataType::Int64)?
3062 .build();
3063
3064 assert_eq!(shredding_type, DataType::Int64);
3065
3066 Ok(())
3067 }
3068
3069 #[test]
3070 fn test_variant_schema_builder_empty_path() -> Result<()> {
3071 let shredding_type = ShreddedSchemaBuilder::default()
3072 .with_path("", &DataType::Int64)?
3073 .build();
3074
3075 assert_eq!(shredding_type, DataType::Int64);
3076 Ok(())
3077 }
3078
3079 #[test]
3080 fn test_variant_schema_builder_default() {
3081 let shredding_type = ShreddedSchemaBuilder::default().build();
3082 assert_eq!(shredding_type, DataType::Null);
3083 }
3084
3085 #[test]
3087 fn test_variant_type_shredded_correctly() {
3088 let mut array_builder = VariantArrayBuilder::new(30);
3090 array_builder.append_value(Variant::Null);
3091 array_builder.append_value(Variant::Int8(1));
3092 array_builder.append_value(Variant::Int16(2));
3093 array_builder.append_value(Variant::Int32(3));
3094 array_builder.append_value(Variant::Int64(4));
3095 array_builder.append_value(Variant::Date(NaiveDate::from_epoch_days(12345).unwrap()));
3096 array_builder.append_value(Variant::TimestampMicros(
3097 DateTime::from_timestamp_micros(123456789).unwrap(),
3098 ));
3099 array_builder.append_value(Variant::TimestampNtzMicros(
3100 DateTime::from_timestamp_micros(123456789)
3101 .unwrap()
3102 .naive_utc(),
3103 ));
3104 array_builder.append_value(Variant::TimestampNanos(DateTime::from_timestamp_nanos(
3105 1234567890000,
3106 )));
3107 array_builder.append_value(Variant::TimestampNtzNanos(
3108 DateTime::from_timestamp_nanos(1234567890000).naive_utc(),
3109 ));
3110 array_builder.append_value(VariantDecimal4::try_new(123, 0).unwrap());
3111 array_builder.append_value(VariantDecimal8::try_new(123, 0).unwrap());
3112 array_builder.append_value(VariantDecimal16::try_new(123, 0).unwrap());
3113 array_builder.append_value(Variant::Float(5.0));
3114 array_builder.append_value(Variant::Double(6f64));
3115 array_builder.append_value(Variant::BooleanTrue);
3116 array_builder.append_value(Variant::BooleanFalse);
3117 array_builder.append_value(Variant::Binary(b"helow"));
3118 array_builder.append_value(Variant::String("hello"));
3119 array_builder.append_value(Variant::ShortString(
3120 ShortString::try_from("world").unwrap(),
3121 ));
3122 array_builder.append_value(Variant::Time(
3123 NaiveTime::from_num_seconds_from_midnight_opt(12345, 123).unwrap(),
3124 ));
3125
3126 let array = array_builder.build();
3127
3128 fn can_shred_to(v: &Variant, dt: &DataType) -> bool {
3129 matches!(
3130 (v, dt),
3131 (
3132 Variant::Int8(_)
3133 | Variant::Int16(_)
3134 | Variant::Int32(_)
3135 | Variant::Int64(_)
3136 | Variant::Decimal4(_)
3137 | Variant::Decimal8(_)
3138 | Variant::Decimal16(_),
3139 DataType::Int8
3140 | DataType::Int16
3141 | DataType::Int32
3142 | DataType::Int64
3143 | DataType::Decimal32(_, _)
3144 | DataType::Decimal64(_, _)
3145 | DataType::Decimal128(_, _)
3146 ) | (Variant::Date(_), DataType::Date32)
3147 | (
3148 Variant::TimestampMicros(_) | Variant::TimestampNanos(_),
3149 DataType::Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, Some(_))
3150 )
3151 | (
3152 Variant::TimestampNtzMicros(_) | Variant::TimestampNtzNanos(_),
3153 DataType::Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, None)
3154 )
3155 | (Variant::Float(_), DataType::Float32)
3156 | (Variant::Double(_), DataType::Float64)
3157 | (
3158 Variant::BooleanFalse | Variant::BooleanTrue,
3159 DataType::Boolean
3160 )
3161 | (
3162 Variant::Binary(_),
3163 DataType::Binary | DataType::BinaryView | DataType::LargeBinary
3164 )
3165 | (
3166 Variant::ShortString(_) | Variant::String(_),
3167 DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8
3168 )
3169 | (Variant::Time(_), DataType::Time64(_))
3170 )
3171 }
3172
3173 macro_rules! assert_shred_type {
3174 ($shred_type:expr, $expected_value_valid_bits:expr) => {
3175 let shredded_array_result = shred_variant(&array, &$shred_type);
3176 match shredded_array_result {
3177 Ok(shredded_array) => {
3178 let value_column = shredded_array.inner().column_by_name("value").unwrap();
3179 for (idx, valid) in $expected_value_valid_bits.iter().enumerate() {
3180 match valid {
3181 true => assert!(
3182 value_column.is_null(idx),
3183 "{:?} should be shredded to {}",
3184 array.value(idx),
3185 $shred_type
3186 ),
3187 false => assert!(
3188 value_column.is_valid(idx),
3189 "{:?} should not be shredded to {}",
3190 array.value(idx),
3191 $shred_type
3192 ),
3193 }
3194 }
3195 }
3196 Err(e) => {
3197 let error_msg = format!("is not a valid variant shredding type");
3198 assert!(
3199 e.to_string().contains(error_msg.as_str()),
3200 "{} => {}",
3201 $shred_type,
3202 e.to_string()
3203 );
3204 }
3205 }
3206 };
3207 }
3208
3209 let types = [
3210 DataType::Null,
3211 DataType::Boolean,
3212 DataType::Int8,
3213 DataType::Int16,
3214 DataType::Int32,
3215 DataType::Int64,
3216 DataType::UInt8,
3217 DataType::UInt16,
3218 DataType::UInt32,
3219 DataType::UInt64,
3220 DataType::Float32,
3221 DataType::Float64,
3222 DataType::Timestamp(TimeUnit::Second, Some("+00:00".into())),
3223 DataType::Timestamp(TimeUnit::Second, None),
3224 DataType::Timestamp(TimeUnit::Millisecond, Some("-00:00".into())),
3225 DataType::Timestamp(TimeUnit::Millisecond, None),
3226 DataType::Timestamp(TimeUnit::Microsecond, Some("-00:00".into())),
3227 DataType::Timestamp(TimeUnit::Microsecond, None),
3228 DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
3229 DataType::Timestamp(TimeUnit::Nanosecond, None),
3230 DataType::Date32,
3231 DataType::Date64,
3232 DataType::Time32(TimeUnit::Second),
3233 DataType::Time32(TimeUnit::Millisecond),
3234 DataType::Time64(TimeUnit::Microsecond),
3235 DataType::Time64(TimeUnit::Nanosecond),
3236 DataType::Duration(TimeUnit::Nanosecond),
3237 DataType::Interval(IntervalUnit::DayTime),
3238 DataType::Binary,
3239 DataType::FixedSizeBinary(16), DataType::FixedSizeBinary(32),
3241 DataType::LargeBinary,
3242 DataType::BinaryView,
3243 DataType::Utf8,
3244 DataType::LargeUtf8,
3245 DataType::Utf8View,
3246 DataType::Decimal32(7, 4),
3247 DataType::Decimal64(7, 4),
3248 DataType::Decimal128(7, 4),
3249 DataType::Decimal256(7, 4),
3250 ];
3251
3252 for data_type in types {
3253 let expected_bits = array
3254 .iter()
3255 .map(|v| can_shred_to(&v.unwrap(), &data_type))
3256 .collect::<Vec<bool>>();
3257 assert_shred_type!(data_type, expected_bits);
3258 }
3259 }
3260}