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<'a> VariantToShreddedVariantRowBuilder<'a> {
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 _ => {
643 *self = Self::Struct(BTreeMap::new());
644 match self {
645 Self::Struct(children) => children,
646 _ => 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 {:?}, got {:?}",
1507 data_type,
1508 err
1509 );
1510 }
1511 }
1512
1513 #[test]
1514 fn test_array_shredding_as_list() {
1515 let input = build_variant_array(vec![
1516 VariantRow::List(vec![
1518 VariantValue::from(1i64),
1519 VariantValue::from(2i64),
1520 VariantValue::from(3i64),
1521 ]),
1522 VariantRow::List(vec![
1524 VariantValue::from(1i64),
1525 VariantValue::from("two"),
1526 VariantValue::from(Variant::Null),
1527 ]),
1528 VariantRow::Value(VariantValue::from("not a list")),
1530 VariantRow::Null,
1532 VariantRow::List(vec![]),
1534 ]);
1535 let list_schema = DataType::List(Arc::new(Field::new("item", DataType::Int64, true)));
1536 let result = shred_variant(&input, &list_schema).unwrap();
1537 assert_eq!(result.len(), 5);
1538
1539 assert_list_structure_and_elements::<Int64Type, i32>(
1540 &result,
1541 5,
1542 &[0, 3, 6, 6, 6, 6],
1543 &[Some(3), Some(3), None, None, Some(0)],
1544 &[None, None, Some(Variant::from("not a list")), None, None],
1545 (
1546 &[Some(1), Some(2), Some(3), Some(1), None, None],
1547 &[
1548 None,
1549 None,
1550 None,
1551 None,
1552 Some(Variant::from("two")),
1553 Some(Variant::Null),
1554 ],
1555 ),
1556 );
1557 }
1558
1559 #[test]
1560 fn test_array_shredding_as_large_list() {
1561 let input = build_variant_array(vec![
1562 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1564 VariantRow::Value(VariantValue::from("not a list")),
1566 VariantRow::List(vec![]),
1568 ]);
1569 let list_schema = DataType::LargeList(Arc::new(Field::new("item", DataType::Int64, true)));
1570 let result = shred_variant(&input, &list_schema).unwrap();
1571 assert_eq!(result.len(), 3);
1572
1573 assert_list_structure_and_elements::<Int64Type, i64>(
1574 &result,
1575 3,
1576 &[0, 2, 2, 2],
1577 &[Some(2), None, Some(0)],
1578 &[None, Some(Variant::from("not a list")), None],
1579 (&[Some(1), Some(2)], &[None, None]),
1580 );
1581 }
1582
1583 #[test]
1584 fn test_array_shredding_as_list_view() {
1585 let input = build_variant_array(vec![
1586 VariantRow::List(vec![
1588 VariantValue::from(1i64),
1589 VariantValue::from(2i64),
1590 VariantValue::from(3i64),
1591 ]),
1592 VariantRow::List(vec![
1594 VariantValue::from(1i64),
1595 VariantValue::from("two"),
1596 VariantValue::from(Variant::Null),
1597 ]),
1598 VariantRow::Value(VariantValue::from("not a list")),
1600 VariantRow::Null,
1602 VariantRow::List(vec![]),
1604 ]);
1605 let list_schema = DataType::ListView(Arc::new(Field::new("item", DataType::Int64, true)));
1606 let result = shred_variant(&input, &list_schema).unwrap();
1607 assert_eq!(result.len(), 5);
1608
1609 assert_list_structure_and_elements::<Int64Type, i32>(
1610 &result,
1611 5,
1612 &[0, 3, 6, 6, 6],
1613 &[Some(3), Some(3), None, None, Some(0)],
1614 &[None, None, Some(Variant::from("not a list")), None, None],
1615 (
1616 &[Some(1), Some(2), Some(3), Some(1), None, None],
1617 &[
1618 None,
1619 None,
1620 None,
1621 None,
1622 Some(Variant::from("two")),
1623 Some(Variant::Null),
1624 ],
1625 ),
1626 );
1627 }
1628
1629 #[test]
1630 fn test_array_shredding_as_large_list_view() {
1631 let input = build_variant_array(vec![
1632 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1634 VariantRow::Value(VariantValue::from("fallback")),
1636 VariantRow::List(vec![]),
1638 ]);
1639 let list_schema =
1640 DataType::LargeListView(Arc::new(Field::new("item", DataType::Int64, true)));
1641 let result = shred_variant(&input, &list_schema).unwrap();
1642 assert_eq!(result.len(), 3);
1643
1644 assert_list_structure_and_elements::<Int64Type, i64>(
1645 &result,
1646 3,
1647 &[0, 2, 2],
1648 &[Some(2), None, Some(0)],
1649 &[None, Some(Variant::from("fallback")), None],
1650 (&[Some(1), Some(2)], &[None, None]),
1651 );
1652 }
1653
1654 #[test]
1655 fn test_array_shredding_as_fixed_size_list() {
1656 let input = build_variant_array(vec![
1657 VariantRow::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1658 VariantRow::Value(VariantValue::from("This should not be shredded")),
1659 VariantRow::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]),
1660 ]);
1661
1662 let list_schema =
1663 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
1664 let result = shred_variant(&input, &list_schema).unwrap();
1665 assert_eq!(result.len(), 3);
1666
1667 assert!(result.is_valid(0));
1670 assert!(result.value_column().is_null(0));
1671 assert!(result.typed_value_column().unwrap().is_valid(0));
1672
1673 assert!(result.is_valid(1));
1677 assert!(result.value_column().is_valid(1));
1678 assert!(result.typed_value_column().unwrap().is_null(1));
1679
1680 assert!(result.is_valid(2));
1683 assert!(result.value_column().is_null(2));
1684 assert!(result.typed_value_column().unwrap().is_valid(2));
1685
1686 let typed_value = result.typed_value_column().unwrap();
1687 let fixed_size_list = typed_value
1688 .as_any()
1689 .downcast_ref::<FixedSizeListArray>()
1690 .expect("Expected FixedSizeListArray");
1691
1692 assert_eq!(fixed_size_list.len(), 3);
1694 assert_eq!(fixed_size_list.value_length(), 2);
1695
1696 let val0 = fixed_size_list.value(0);
1698 let val0_struct = val0.as_any().downcast_ref::<StructArray>().unwrap();
1699 let val0_typed = val0_struct.column_by_name("typed_value").unwrap();
1700 let val0_ints = val0_typed.as_any().downcast_ref::<Int64Array>().unwrap();
1701 assert_eq!(val0_ints.values(), &[1i64, 2i64]);
1702
1703 assert!(fixed_size_list.is_null(1));
1706
1707 let val2 = fixed_size_list.value(2);
1709 let val2_struct = val2.as_any().downcast_ref::<StructArray>().unwrap();
1710 let val2_typed = val2_struct.column_by_name("typed_value").unwrap();
1711 let val2_ints = val2_typed.as_any().downcast_ref::<Int64Array>().unwrap();
1712 assert_eq!(val2_ints.values(), &[3i64, 4i64]);
1713 }
1714
1715 #[test]
1716 fn test_array_shredding_as_fixed_size_list_wrong_size() {
1717 let input = build_variant_array(vec![VariantRow::List(vec![
1718 VariantValue::from(1i64),
1719 VariantValue::from(2i64),
1720 VariantValue::from(3i64),
1721 ])]);
1722 let list_schema =
1723 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
1724
1725 let err = shred_variant(&input, &list_schema).unwrap_err();
1726 assert!(
1727 err.to_string()
1728 .contains("Expected fixed size list of size 2, got size 3"),
1729 "got: {err}",
1730 );
1731 }
1732
1733 #[test]
1734 fn test_array_shredding_with_array_elements() {
1735 let input = build_variant_array(vec![
1736 VariantRow::List(vec![
1738 VariantValue::List(vec![VariantValue::from(1i64), VariantValue::from(2i64)]),
1739 VariantValue::List(vec![VariantValue::from(3i64), VariantValue::from(4i64)]),
1740 VariantValue::List(vec![]),
1741 ]),
1742 VariantRow::List(vec![
1744 VariantValue::List(vec![
1745 VariantValue::from(5i64),
1746 VariantValue::from("bad"),
1747 VariantValue::from(Variant::Null),
1748 ]),
1749 VariantValue::from("not a list inner"),
1750 VariantValue::Null,
1751 ]),
1752 VariantRow::Value(VariantValue::from("not a list")),
1754 VariantRow::Null,
1756 ]);
1757 let inner_field = Arc::new(Field::new("item", DataType::Int64, true));
1758 let inner_list_schema = DataType::List(inner_field);
1759 let list_schema = DataType::List(Arc::new(Field::new(
1760 "item",
1761 inner_list_schema.clone(),
1762 true,
1763 )));
1764 let result = shred_variant(&input, &list_schema).unwrap();
1765 assert_eq!(result.len(), 4);
1766
1767 let typed_value = result
1768 .typed_value_column()
1769 .unwrap()
1770 .as_any()
1771 .downcast_ref::<ListArray>()
1772 .unwrap();
1773
1774 assert_list_structure::<i32>(
1775 &result,
1776 4,
1777 &[0, 3, 6, 6, 6],
1778 &[Some(3), Some(3), None, None],
1779 &[None, None, Some(Variant::from("not a list")), None],
1780 );
1781
1782 let outer_elements =
1783 ShreddedVariantFieldArray::try_new(typed_value.values().as_ref()).unwrap();
1784 assert_eq!(outer_elements.len(), 6);
1785 let outer_values = outer_elements
1786 .typed_value_column()
1787 .unwrap()
1788 .as_any()
1789 .downcast_ref::<ListArray>()
1790 .unwrap();
1791 let outer_fallbacks = outer_elements.value_column();
1792
1793 let outer_metadata = Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
1794 EMPTY_VARIANT_METADATA_BYTES,
1795 outer_elements.len(),
1796 )));
1797 let outer_variant = VariantArray::from_parts(
1798 outer_metadata,
1799 outer_fallbacks.clone(),
1800 Some(Arc::new(outer_values.clone())),
1801 None,
1802 );
1803
1804 assert_list_structure_and_elements::<Int64Type, i32>(
1805 &outer_variant,
1806 outer_elements.len(),
1807 &[0, 2, 4, 4, 7, 7, 7],
1808 &[Some(2), Some(2), Some(0), Some(3), None, None],
1809 &[
1810 None,
1811 None,
1812 None,
1813 None,
1814 Some(Variant::from("not a list inner")),
1815 Some(Variant::Null),
1816 ],
1817 (
1818 &[Some(1), Some(2), Some(3), Some(4), Some(5), None, None],
1819 &[
1820 None,
1821 None,
1822 None,
1823 None,
1824 None,
1825 Some(Variant::from("bad")),
1826 Some(Variant::Null),
1827 ],
1828 ),
1829 );
1830 }
1831
1832 #[test]
1833 fn test_array_shredding_with_object_elements() {
1834 let input = build_variant_array(vec![
1835 VariantRow::List(vec![
1837 VariantValue::Object(vec![
1838 ("id", VariantValue::from(1i64)),
1839 ("name", VariantValue::from("Alice")),
1840 ]),
1841 VariantValue::Object(vec![("id", VariantValue::from(Variant::Null))]),
1842 ]),
1843 VariantRow::Value(VariantValue::from("not a list")),
1845 VariantRow::Null,
1847 ]);
1848
1849 let object_fields = Fields::from(vec![
1851 Field::new("id", DataType::Int64, true),
1852 Field::new("name", DataType::Utf8, true),
1853 ]);
1854 let list_schema = DataType::List(Arc::new(Field::new(
1855 "item",
1856 DataType::Struct(object_fields),
1857 true,
1858 )));
1859 let result = shred_variant(&input, &list_schema).unwrap();
1860 assert_eq!(result.len(), 3);
1861
1862 assert_list_structure::<i32>(
1863 &result,
1864 3,
1865 &[0, 2, 2, 2],
1866 &[Some(2), None, None],
1867 &[None, Some(Variant::from("not a list")), None],
1868 );
1869
1870 let typed_value = result
1872 .typed_value_column()
1873 .unwrap()
1874 .as_any()
1875 .downcast_ref::<ListArray>()
1876 .unwrap();
1877 let element_array =
1878 ShreddedVariantFieldArray::try_new(typed_value.values().as_ref()).unwrap();
1879 assert_eq!(element_array.len(), 2);
1880 let element_objects = element_array
1881 .typed_value_column()
1882 .unwrap()
1883 .as_any()
1884 .downcast_ref::<arrow::array::StructArray>()
1885 .unwrap();
1886
1887 let id_field =
1889 ShreddedVariantFieldArray::try_new(element_objects.column_by_name("id").unwrap())
1890 .unwrap();
1891 let id_values = id_field.value_column();
1892 let id_typed_values = id_field
1893 .typed_value_column()
1894 .unwrap()
1895 .as_any()
1896 .downcast_ref::<Int64Array>()
1897 .unwrap();
1898 assert!(id_values.is_null(0));
1899 assert_eq!(id_typed_values.value(0), 1);
1900 assert!(id_values.is_valid(1));
1902 assert_eq!(
1903 Variant::new(
1904 EMPTY_VARIANT_METADATA_BYTES,
1905 binary_array_value(id_values.as_ref(), 1).unwrap()
1906 ),
1907 Variant::Null
1908 );
1909 assert!(id_typed_values.is_null(1));
1910
1911 let name_field =
1913 ShreddedVariantFieldArray::try_new(element_objects.column_by_name("name").unwrap())
1914 .unwrap();
1915 let name_values = name_field.value_column();
1916 let name_typed_values = name_field
1917 .typed_value_column()
1918 .unwrap()
1919 .as_any()
1920 .downcast_ref::<StringArray>()
1921 .unwrap();
1922 assert!(name_values.is_null(0));
1923 assert_eq!(name_typed_values.value(0), "Alice");
1924 assert!(name_values.is_null(1));
1926 assert!(name_typed_values.is_null(1));
1927 }
1928
1929 #[test]
1930 fn test_object_shredding_comprehensive() -> Result<()> {
1931 let input = build_variant_array(vec![
1932 VariantRow::Object(vec![
1934 ("score", VariantValue::from(95.5f64)),
1935 ("age", VariantValue::from(30i64)),
1936 ]),
1937 VariantRow::Object(vec![
1939 ("score", VariantValue::from(87.2f64)),
1940 ("age", VariantValue::from(25i64)),
1941 ("email", VariantValue::from("bob@example.com")),
1942 ]),
1943 VariantRow::Object(vec![("age", VariantValue::from(35i64))]),
1945 VariantRow::Object(vec![
1947 ("score", VariantValue::from("ninety-five")),
1948 ("age", VariantValue::from("thirty")),
1949 ]),
1950 VariantRow::Value(VariantValue::from("not an object")),
1952 VariantRow::Object(vec![]),
1954 VariantRow::Null,
1956 VariantRow::Object(vec![("foo", VariantValue::from(10))]),
1958 VariantRow::Object(vec![
1960 ("score", VariantValue::from(66.67f64)),
1961 ("foo", VariantValue::from(10)),
1962 ]),
1963 ]);
1964
1965 let target_schema = ShreddedSchemaBuilder::default()
1968 .with_path("score", &DataType::Float64)?
1969 .with_path("age", &DataType::Int64)?
1970 .build();
1971
1972 let result = shred_variant(&input, &target_schema).unwrap();
1973
1974 assert!(result.typed_value_column().is_some());
1976 assert_eq!(result.len(), 9);
1977
1978 let metadata = result.metadata_column();
1979 let value = result.value_column();
1980 let typed_value = result
1981 .typed_value_column()
1982 .unwrap()
1983 .as_any()
1984 .downcast_ref::<arrow::array::StructArray>()
1985 .unwrap();
1986
1987 let score_field =
1989 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("score").unwrap())
1990 .unwrap();
1991 let age_field =
1992 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("age").unwrap()).unwrap();
1993
1994 let score_value = score_field.value_column();
1995 let score_typed_value = score_field
1996 .typed_value_column()
1997 .unwrap()
1998 .as_any()
1999 .downcast_ref::<Float64Array>()
2000 .unwrap();
2001 let age_value = age_field.value_column();
2002 let age_typed_value = age_field
2003 .typed_value_column()
2004 .unwrap()
2005 .as_any()
2006 .downcast_ref::<Int64Array>()
2007 .unwrap();
2008
2009 struct ShreddedValue<'m, 'v, T> {
2011 value: Option<Variant<'m, 'v>>,
2012 typed_value: Option<T>,
2013 }
2014 struct ShreddedStruct<'m, 'v> {
2015 score: ShreddedValue<'m, 'v, f64>,
2016 age: ShreddedValue<'m, 'v, i64>,
2017 }
2018 fn get_value<'m, 'v>(
2019 i: usize,
2020 metadata: &'m dyn Array,
2021 value: &'v dyn Array,
2022 ) -> Variant<'m, 'v> {
2023 variant_from_arrays_at(metadata, value, i).unwrap()
2024 }
2025 let expect = |i, expected_result: Option<ShreddedValue<ShreddedStruct>>| {
2026 match expected_result {
2027 Some(ShreddedValue {
2028 value: expected_value,
2029 typed_value: expected_typed_value,
2030 }) => {
2031 assert!(result.is_valid(i));
2032 match expected_value {
2033 Some(expected_value) => {
2034 assert!(value.is_valid(i));
2035 assert_eq!(
2036 expected_value,
2037 get_value(i, metadata.as_ref(), value.as_ref())
2038 );
2039 }
2040 None => {
2041 assert!(value.is_null(i));
2042 }
2043 }
2044 match expected_typed_value {
2045 Some(ShreddedStruct {
2046 score: expected_score,
2047 age: expected_age,
2048 }) => {
2049 assert!(typed_value.is_valid(i));
2050 assert!(score_field.is_valid(i)); assert!(age_field.is_valid(i)); match expected_score.value {
2053 Some(expected_score_value) => {
2054 assert!(score_value.is_valid(i));
2055 assert_eq!(
2056 expected_score_value,
2057 get_value(i, metadata.as_ref(), score_value.as_ref())
2058 );
2059 }
2060 None => {
2061 assert!(score_value.is_null(i));
2062 }
2063 }
2064 match expected_score.typed_value {
2065 Some(expected_score) => {
2066 assert!(score_typed_value.is_valid(i));
2067 assert_eq!(expected_score, score_typed_value.value(i));
2068 }
2069 None => {
2070 assert!(score_typed_value.is_null(i));
2071 }
2072 }
2073 match expected_age.value {
2074 Some(expected_age_value) => {
2075 assert!(age_value.is_valid(i));
2076 assert_eq!(
2077 expected_age_value,
2078 get_value(i, metadata.as_ref(), age_value.as_ref())
2079 );
2080 }
2081 None => {
2082 assert!(age_value.is_null(i));
2083 }
2084 }
2085 match expected_age.typed_value {
2086 Some(expected_age) => {
2087 assert!(age_typed_value.is_valid(i));
2088 assert_eq!(expected_age, age_typed_value.value(i));
2089 }
2090 None => {
2091 assert!(age_typed_value.is_null(i));
2092 }
2093 }
2094 }
2095 None => {
2096 assert!(typed_value.is_null(i));
2097 }
2098 }
2099 }
2100 None => {
2101 assert!(result.is_null(i));
2102 }
2103 };
2104 };
2105
2106 expect(
2108 0,
2109 Some(ShreddedValue {
2110 value: None,
2111 typed_value: Some(ShreddedStruct {
2112 score: ShreddedValue {
2113 value: None,
2114 typed_value: Some(95.5),
2115 },
2116 age: ShreddedValue {
2117 value: None,
2118 typed_value: Some(30),
2119 },
2120 }),
2121 }),
2122 );
2123
2124 let mut builder = VariantBuilder::new();
2126 builder
2127 .new_object()
2128 .with_field("email", "bob@example.com")
2129 .finish();
2130 let (m, v) = builder.finish();
2131 let expected_value = Variant::new(&m, &v);
2132
2133 expect(
2134 1,
2135 Some(ShreddedValue {
2136 value: Some(expected_value),
2137 typed_value: Some(ShreddedStruct {
2138 score: ShreddedValue {
2139 value: None,
2140 typed_value: Some(87.2),
2141 },
2142 age: ShreddedValue {
2143 value: None,
2144 typed_value: Some(25),
2145 },
2146 }),
2147 }),
2148 );
2149
2150 expect(
2152 2,
2153 Some(ShreddedValue {
2154 value: None,
2155 typed_value: Some(ShreddedStruct {
2156 score: ShreddedValue {
2157 value: None,
2158 typed_value: None,
2159 },
2160 age: ShreddedValue {
2161 value: None,
2162 typed_value: Some(35),
2163 },
2164 }),
2165 }),
2166 );
2167
2168 expect(
2170 3,
2171 Some(ShreddedValue {
2172 value: None,
2173 typed_value: Some(ShreddedStruct {
2174 score: ShreddedValue {
2175 value: Some(Variant::from("ninety-five")),
2176 typed_value: None,
2177 },
2178 age: ShreddedValue {
2179 value: Some(Variant::from("thirty")),
2180 typed_value: None,
2181 },
2182 }),
2183 }),
2184 );
2185
2186 expect(
2188 4,
2189 Some(ShreddedValue {
2190 value: Some(Variant::from("not an object")),
2191 typed_value: None,
2192 }),
2193 );
2194
2195 expect(
2197 5,
2198 Some(ShreddedValue {
2199 value: None,
2200 typed_value: Some(ShreddedStruct {
2201 score: ShreddedValue {
2202 value: None,
2203 typed_value: None,
2204 },
2205 age: ShreddedValue {
2206 value: None,
2207 typed_value: None,
2208 },
2209 }),
2210 }),
2211 );
2212
2213 expect(6, None);
2215
2216 let object_with_foo_field = |i| {
2218 use parquet_variant::{ParentState, ValueBuilder, VariantMetadata};
2219 let metadata = VariantMetadata::new(binary_array_value(metadata.as_ref(), i).unwrap());
2220 let mut metadata_builder = ReadOnlyMetadataBuilder::new(&metadata);
2221 let mut value_builder = ValueBuilder::new();
2222 let state = ParentState::variant(&mut value_builder, &mut metadata_builder);
2223 ObjectBuilder::new(state, false)
2224 .with_field("foo", 10)
2225 .finish();
2226 (metadata, value_builder.into_inner())
2227 };
2228
2229 let (m, v) = object_with_foo_field(7);
2231 expect(
2232 7,
2233 Some(ShreddedValue {
2234 value: Some(Variant::new_with_metadata(m, &v)),
2235 typed_value: Some(ShreddedStruct {
2236 score: ShreddedValue {
2237 value: None,
2238 typed_value: None,
2239 },
2240 age: ShreddedValue {
2241 value: None,
2242 typed_value: None,
2243 },
2244 }),
2245 }),
2246 );
2247
2248 let (m, v) = object_with_foo_field(8);
2250 expect(
2251 8,
2252 Some(ShreddedValue {
2253 value: Some(Variant::new_with_metadata(m, &v)),
2254 typed_value: Some(ShreddedStruct {
2255 score: ShreddedValue {
2256 value: None,
2257 typed_value: Some(66.67),
2258 },
2259 age: ShreddedValue {
2260 value: None,
2261 typed_value: None,
2262 },
2263 }),
2264 }),
2265 );
2266 Ok(())
2267 }
2268
2269 #[test]
2270 fn test_object_shredding_with_array_field() {
2271 let input = build_variant_array(vec![
2272 VariantRow::Object(vec![(
2274 "scores",
2275 VariantValue::List(vec![VariantValue::from(10i64), VariantValue::from(20i64)]),
2276 )]),
2277 VariantRow::Object(vec![(
2279 "scores",
2280 VariantValue::List(vec![
2281 VariantValue::from("oops"),
2282 VariantValue::from(Variant::Null),
2283 ]),
2284 )]),
2285 VariantRow::Object(vec![]),
2287 VariantRow::Value(VariantValue::from("not an object")),
2289 VariantRow::Null,
2291 ]);
2292 let list_field = Arc::new(Field::new("item", DataType::Int64, true));
2293 let inner_list_schema = DataType::List(list_field);
2294 let schema = DataType::Struct(Fields::from(vec![Field::new(
2295 "scores",
2296 inner_list_schema.clone(),
2297 true,
2298 )]));
2299
2300 let result = shred_variant(&input, &schema).unwrap();
2301 assert_eq!(result.len(), 5);
2302
2303 let value_field = result.value_column();
2305 let typed_struct = result
2306 .typed_value_column()
2307 .unwrap()
2308 .as_any()
2309 .downcast_ref::<arrow::array::StructArray>()
2310 .unwrap();
2311
2312 assert!(value_field.is_null(0));
2314 assert!(value_field.is_null(1));
2315 assert!(value_field.is_null(2));
2316 assert!(value_field.is_valid(3));
2317 assert_eq!(
2318 variant_from_arrays_at(result.metadata_column(), value_field, 3).unwrap(),
2319 Variant::from("not an object")
2320 );
2321 assert!(value_field.is_null(4));
2322
2323 assert!(typed_struct.is_valid(0));
2325 assert!(typed_struct.is_valid(1));
2326 assert!(typed_struct.is_valid(2));
2327 assert!(typed_struct.is_null(3));
2328 assert!(typed_struct.is_null(4));
2329
2330 let scores_field =
2332 ShreddedVariantFieldArray::try_new(typed_struct.column_by_name("scores").unwrap())
2333 .unwrap();
2334 assert_list_structure_and_elements::<Int64Type, i32>(
2335 &VariantArray::from_parts(
2336 Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
2337 EMPTY_VARIANT_METADATA_BYTES,
2338 scores_field.len(),
2339 ))),
2340 scores_field.value_column().clone(),
2341 Some(scores_field.typed_value_column().unwrap().clone()),
2342 None,
2343 ),
2344 scores_field.len(),
2345 &[0i32, 2, 4, 4, 4, 4],
2346 &[Some(2), Some(2), None, None, None],
2347 &[None, None, None, None, None],
2348 (
2349 &[Some(10), Some(20), None, None],
2350 &[None, None, Some(Variant::from("oops")), Some(Variant::Null)],
2351 ),
2352 );
2353 }
2354
2355 #[test]
2356 fn test_object_different_schemas() -> Result<()> {
2357 let input = build_variant_array(vec![VariantRow::Object(vec![
2359 ("id", VariantValue::from(123i32)),
2360 ("age", VariantValue::from(25i64)),
2361 ("score", VariantValue::from(95.5f64)),
2362 ])]);
2363
2364 let schema1 = ShreddedSchemaBuilder::default()
2366 .with_path("id", &DataType::Int32)?
2367 .build();
2368 let result1 = shred_variant(&input, &schema1).unwrap();
2369 let value_field1 = result1.value_column();
2370 assert!(!value_field1.is_null(0)); let schema2 = ShreddedSchemaBuilder::default()
2374 .with_path("id", &DataType::Int32)?
2375 .with_path("age", &DataType::Int64)?
2376 .build();
2377 let result2 = shred_variant(&input, &schema2).unwrap();
2378 let value_field2 = result2.value_column();
2379 assert!(!value_field2.is_null(0)); let schema3 = ShreddedSchemaBuilder::default()
2383 .with_path("id", &DataType::Int32)?
2384 .with_path("age", &DataType::Int64)?
2385 .with_path("score", &DataType::Float64)?
2386 .build();
2387 let result3 = shred_variant(&input, &schema3).unwrap();
2388 let value_field3 = result3.value_column();
2389 assert!(value_field3.is_null(0)); Ok(())
2392 }
2393
2394 #[test]
2395 fn test_uuid_shredding_in_objects() -> Result<()> {
2396 let mock_uuid_1 = Uuid::new_v4();
2397 let mock_uuid_2 = Uuid::new_v4();
2398 let mock_uuid_3 = Uuid::new_v4();
2399
2400 let input = build_variant_array(vec![
2401 VariantRow::Object(vec![
2403 ("id", VariantValue::from(mock_uuid_1)),
2404 ("session_id", VariantValue::from(mock_uuid_2)),
2405 ]),
2406 VariantRow::Object(vec![
2408 ("id", VariantValue::from(mock_uuid_2)),
2409 ("session_id", VariantValue::from(mock_uuid_3)),
2410 ("name", VariantValue::from("test_user")),
2411 ]),
2412 VariantRow::Object(vec![("id", VariantValue::from(mock_uuid_1))]),
2414 VariantRow::Object(vec![
2416 ("id", VariantValue::from(mock_uuid_3)),
2417 ("session_id", VariantValue::from("not-a-uuid")),
2418 ]),
2419 VariantRow::Object(vec![
2421 ("id", VariantValue::from(12345i64)),
2422 ("session_id", VariantValue::from(mock_uuid_1)),
2423 ]),
2424 VariantRow::Null,
2426 ]);
2427
2428 let target_schema = ShreddedSchemaBuilder::default()
2429 .with_path("id", DataType::FixedSizeBinary(16))?
2430 .with_path("session_id", DataType::FixedSizeBinary(16))?
2431 .build();
2432
2433 let result = shred_variant(&input, &target_schema).unwrap();
2434
2435 assert!(result.typed_value_column().is_some());
2436 assert_eq!(result.len(), 6);
2437
2438 let metadata = result.metadata_column();
2439 let value = result.value_column();
2440 let typed_value = result
2441 .typed_value_column()
2442 .unwrap()
2443 .as_any()
2444 .downcast_ref::<arrow::array::StructArray>()
2445 .unwrap();
2446
2447 let id_field =
2449 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("id").unwrap()).unwrap();
2450 let session_id_field =
2451 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("session_id").unwrap())
2452 .unwrap();
2453
2454 let id_value = id_field.value_column();
2455 let id_typed_value = id_field
2456 .typed_value_column()
2457 .unwrap()
2458 .as_any()
2459 .downcast_ref::<FixedSizeBinaryArray>()
2460 .unwrap();
2461 let session_id_value = session_id_field.value_column();
2462 let session_id_typed_value = session_id_field
2463 .typed_value_column()
2464 .unwrap()
2465 .as_any()
2466 .downcast_ref::<FixedSizeBinaryArray>()
2467 .unwrap();
2468
2469 assert!(result.is_valid(0));
2471
2472 assert!(value.is_null(0)); assert!(id_value.is_null(0));
2474 assert!(session_id_value.is_null(0));
2475
2476 assert!(typed_value.is_valid(0));
2477 assert!(id_typed_value.is_valid(0));
2478 assert!(session_id_typed_value.is_valid(0));
2479
2480 assert_eq!(id_typed_value.value(0), mock_uuid_1.as_bytes());
2481 assert_eq!(session_id_typed_value.value(0), mock_uuid_2.as_bytes());
2482
2483 assert!(result.is_valid(1));
2485
2486 assert!(value.is_valid(1)); assert!(typed_value.is_valid(1));
2488
2489 assert!(id_value.is_null(1));
2490 assert!(id_typed_value.is_valid(1));
2491 assert_eq!(id_typed_value.value(1), mock_uuid_2.as_bytes());
2492
2493 assert!(session_id_value.is_null(1));
2494 assert!(session_id_typed_value.is_valid(1));
2495 assert_eq!(session_id_typed_value.value(1), mock_uuid_3.as_bytes());
2496
2497 let row_1_variant = variant_from_arrays_at(metadata, value, 1).unwrap();
2499 let Variant::Object(obj) = row_1_variant else {
2500 panic!("Expected object");
2501 };
2502
2503 assert_eq!(obj.get("name"), Some(Variant::from("test_user")));
2504
2505 assert!(result.is_valid(2));
2507
2508 assert!(value.is_null(2)); assert!(typed_value.is_valid(2));
2510
2511 assert!(id_value.is_null(2));
2512 assert!(id_typed_value.is_valid(2));
2513 assert_eq!(id_typed_value.value(2), mock_uuid_1.as_bytes());
2514
2515 assert!(session_id_value.is_null(2));
2516 assert!(session_id_typed_value.is_null(2)); assert!(result.is_valid(3));
2520
2521 assert!(value.is_null(3)); assert!(typed_value.is_valid(3));
2523
2524 assert!(id_value.is_null(3));
2525 assert!(id_typed_value.is_valid(3));
2526 assert_eq!(id_typed_value.value(3), mock_uuid_3.as_bytes());
2527
2528 assert!(session_id_value.is_valid(3)); assert!(session_id_typed_value.is_null(3));
2530 let session_id_variant = variant_from_arrays_at(metadata, session_id_value, 3).unwrap();
2531 assert_eq!(session_id_variant, Variant::from("not-a-uuid"));
2532
2533 assert!(result.is_valid(4));
2535
2536 assert!(value.is_null(4)); assert!(typed_value.is_valid(4));
2538
2539 assert!(id_value.is_valid(4)); assert!(id_typed_value.is_null(4));
2541 let id_variant = variant_from_arrays_at(metadata, id_value, 4).unwrap();
2542 assert_eq!(id_variant, Variant::from(12345i64));
2543
2544 assert!(session_id_value.is_null(4));
2545 assert!(session_id_typed_value.is_valid(4));
2546 assert_eq!(session_id_typed_value.value(4), mock_uuid_1.as_bytes());
2547
2548 assert!(result.is_null(5));
2550
2551 Ok(())
2552 }
2553
2554 macro_rules! validate_decimal_shredding {
2555 ($shred_type: expr, $array_type: ty, $expected_typed_value: ident $(, $expected_precision: literal, $expected_scale:literal)? $(,)?) => {{
2556 let input = VariantArray::from_iter(vec![
2557 Variant::from(12i8),
2558 Variant::from(234i16),
2559 Variant::from(456i32),
2560 Variant::from(456i64),
2561 Variant::from(VariantDecimal4::try_new(1200, 2).unwrap()),
2562 Variant::from(VariantDecimal8::try_new(1230, 2).unwrap()),
2563 Variant::from(VariantDecimal16::try_new(1234, 2).unwrap()),
2564 ]);
2565
2566 let result = shred_variant(&input, &$shred_type).unwrap();
2567
2568 assert!(result.typed_value_column().is_some());
2569 assert_eq!(result.len(), input.len());
2570
2571 let value = result.value_column();
2572 let typed_value = result
2573 .typed_value_column()
2574 .unwrap()
2575 .as_any()
2576 .downcast_ref::<$array_type>()
2577 .unwrap();
2578
2579 $(assert_eq!(typed_value.precision(), $expected_precision);)?
2580 $(assert_eq!(typed_value.scale(), $expected_scale);)?
2581
2582 for i in 0..$expected_typed_value.len() {
2583 assert_eq!(value.is_valid(i), $expected_typed_value.is_null(i));
2584 assert_eq!(typed_value.is_valid(i), $expected_typed_value.is_valid(i));
2585 assert_eq!(typed_value.value(i), $expected_typed_value.value(i));
2586 }
2587 }};
2588 }
2589
2590 #[test]
2591 fn test_shredding_decimal32_with_same_scale() {
2592 let expected_array = Decimal32Array::from(vec![
2593 Some(1200),
2594 None, None, None, Some(1200),
2598 Some(1230),
2599 Some(1234),
2600 ])
2601 .with_precision_and_scale(4, 2)
2602 .unwrap();
2603 validate_decimal_shredding!(
2604 DataType::Decimal32(4, 2),
2605 arrow::array::Decimal32Array,
2606 expected_array,
2607 4,
2608 2,
2609 );
2610 }
2611
2612 #[test]
2613 fn test_shredding_decimal32_with_bigger_scale() {
2614 let expected_array = Decimal32Array::from(vec![
2615 Some(12000),
2616 Some(234000),
2617 Some(456000),
2618 Some(456000),
2619 Some(12000),
2620 Some(12300),
2621 Some(12340),
2622 ])
2623 .with_precision_and_scale(6, 3)
2624 .unwrap();
2625
2626 validate_decimal_shredding!(
2627 DataType::Decimal32(6, 3),
2628 arrow::array::Decimal32Array,
2629 expected_array,
2630 6,
2631 3,
2632 );
2633 }
2634
2635 #[test]
2636 fn test_shredding_decimal32_with_smaller_scale() {
2637 let expected_array = Decimal32Array::from(vec![
2638 Some(12),
2639 Some(234),
2640 Some(456),
2641 Some(456),
2642 Some(12),
2643 None, None, ])
2646 .with_precision_and_scale(6, 0)
2647 .unwrap();
2648 validate_decimal_shredding!(
2649 DataType::Decimal32(6, 0),
2650 arrow::array::Decimal32Array,
2651 expected_array,
2652 6,
2653 0
2654 );
2655 }
2656
2657 #[test]
2658 fn test_shredding_decimal64_with_same_scale() {
2659 let expected_array_decimal64_same_scale = Decimal64Array::from(vec![
2660 Some(1200),
2661 None, None, None, Some(1200),
2665 Some(1230),
2666 Some(1234),
2667 ])
2668 .with_precision_and_scale(4, 2)
2669 .unwrap();
2670 validate_decimal_shredding!(
2671 DataType::Decimal64(4, 2),
2672 arrow::array::Decimal64Array,
2673 expected_array_decimal64_same_scale,
2674 4,
2675 2
2676 );
2677 }
2678
2679 #[test]
2680 fn test_shredding_decimal64_with_big_scale() {
2681 let expected_array = Decimal64Array::from(vec![
2682 Some(12000),
2683 Some(234000),
2684 Some(456000),
2685 Some(456000),
2686 Some(12000),
2687 Some(12300),
2688 Some(12340),
2689 ])
2690 .with_precision_and_scale(6, 3)
2691 .unwrap();
2692 validate_decimal_shredding!(
2693 DataType::Decimal64(6, 3),
2694 arrow::array::Decimal64Array,
2695 expected_array,
2696 6,
2697 3,
2698 );
2699 }
2700
2701 #[test]
2702 fn test_shredding_decimal64_with_smaller_scale() {
2703 let expected_array = Decimal64Array::from(vec![
2704 Some(12),
2705 Some(234),
2706 Some(456),
2707 Some(456),
2708 Some(12),
2709 None, None, ])
2712 .with_precision_and_scale(6, 0)
2713 .unwrap();
2714 validate_decimal_shredding!(
2715 DataType::Decimal64(6, 0),
2716 arrow::array::Decimal64Array,
2717 expected_array,
2718 6,
2719 0
2720 );
2721 }
2722
2723 #[test]
2724 fn test_shredding_decimal128_with_same_scale() {
2725 let expected_array = Decimal128Array::from(vec![
2726 Some(1200),
2727 None, None, None, Some(1200),
2731 Some(1230),
2732 Some(1234),
2733 ])
2734 .with_precision_and_scale(4, 2)
2735 .unwrap();
2736
2737 validate_decimal_shredding!(
2738 DataType::Decimal128(4, 2),
2739 arrow::array::Decimal128Array,
2740 expected_array,
2741 4,
2742 2,
2743 );
2744 }
2745
2746 #[test]
2747 fn test_shredding_decimal128_with_big_scale() {
2748 let expected_array = Decimal128Array::from(vec![
2749 Some(12000),
2750 Some(234000),
2751 Some(456000),
2752 Some(456000),
2753 Some(12000),
2754 Some(12300),
2755 Some(12340),
2756 ])
2757 .with_precision_and_scale(6, 3)
2758 .unwrap();
2759 validate_decimal_shredding!(
2760 DataType::Decimal128(6, 3),
2761 arrow::array::Decimal128Array,
2762 expected_array,
2763 6,
2764 3
2765 );
2766 }
2767
2768 #[test]
2769 fn test_shredding_decimal128_with_smaller_scale() {
2770 let expected_array = Decimal128Array::from(vec![
2771 Some(12),
2772 Some(234),
2773 Some(456),
2774 Some(456),
2775 Some(12),
2776 None, None, ])
2779 .with_precision_and_scale(6, 0)
2780 .unwrap();
2781 validate_decimal_shredding!(
2782 DataType::Decimal128(6, 0),
2783 arrow::array::Decimal128Array,
2784 expected_array,
2785 6,
2786 0
2787 );
2788 }
2789
2790 #[test]
2791 fn test_shredding_decimal128_to_integer() {
2792 let expected_array = Int64Array::from(vec![
2793 Some(12),
2794 Some(234),
2795 Some(456),
2796 Some(456),
2797 Some(12),
2798 None, None, ]);
2801
2802 validate_decimal_shredding!(DataType::Int64, arrow::array::Int64Array, expected_array);
2803 }
2804
2805 #[test]
2806 fn test_spec_compliance() {
2807 let input = VariantArray::from_iter(vec![Variant::from(42i64), Variant::from("hello")]);
2808
2809 let result = shred_variant(&input, &DataType::Int64).unwrap();
2810
2811 let inner_struct = result.inner();
2813 assert!(inner_struct.column_by_name("metadata").is_some());
2814 assert!(inner_struct.column_by_name("value").is_some());
2815 assert!(inner_struct.column_by_name("typed_value").is_some());
2816
2817 assert_eq!(
2819 result.metadata_column().len(),
2820 input.metadata_column().len()
2821 );
2822 assert_eq!(
2825 result.metadata_column().len(),
2826 input.metadata_column().len()
2827 );
2828
2829 assert_eq!(result.len(), input.len());
2831 assert!(result.typed_value_column().is_some());
2832
2833 let value_field = result.value_column();
2836 let typed_value_field = result
2837 .typed_value_column()
2838 .unwrap()
2839 .as_any()
2840 .downcast_ref::<Int64Array>()
2841 .unwrap();
2842
2843 for i in 0..result.len() {
2844 if !result.is_null(i) {
2845 let value_is_null = value_field.is_null(i);
2846 let typed_value_is_null = typed_value_field.is_null(i);
2847 assert!(
2849 value_is_null || typed_value_is_null,
2850 "Row {}: both value and typed_value are non-null for primitive shredding",
2851 i
2852 );
2853 }
2854 }
2855 }
2856
2857 #[test]
2858 fn test_variant_schema_builder_simple() -> Result<()> {
2859 let shredding_type = ShreddedSchemaBuilder::default()
2860 .with_path("a", &DataType::Int64)?
2861 .with_path("b", &DataType::Float64)?
2862 .build();
2863
2864 assert_eq!(
2865 shredding_type,
2866 DataType::Struct(Fields::from(vec![
2867 Field::new("a", DataType::Int64, true),
2868 Field::new("b", DataType::Float64, true),
2869 ]))
2870 );
2871
2872 Ok(())
2873 }
2874
2875 #[test]
2876 fn test_variant_schema_builder_nested() -> Result<()> {
2877 let shredding_type = ShreddedSchemaBuilder::default()
2878 .with_path("a", &DataType::Int64)?
2879 .with_path("b.c", &DataType::Utf8)?
2880 .with_path("b.d", &DataType::Float64)?
2881 .build();
2882
2883 assert_eq!(
2884 shredding_type,
2885 DataType::Struct(Fields::from(vec![
2886 Field::new("a", DataType::Int64, true),
2887 Field::new(
2888 "b",
2889 DataType::Struct(Fields::from(vec![
2890 Field::new("c", DataType::Utf8, true),
2891 Field::new("d", DataType::Float64, true),
2892 ])),
2893 true
2894 ),
2895 ]))
2896 );
2897
2898 Ok(())
2899 }
2900
2901 #[test]
2902 fn test_variant_schema_builder_with_path_variant_path_arg() -> Result<()> {
2903 let path = VariantPath::from_iter([VariantPathElement::from("a.b")]);
2904 let shredding_type = ShreddedSchemaBuilder::default()
2905 .with_path(path, &DataType::Int64)?
2906 .build();
2907
2908 match shredding_type {
2909 DataType::Struct(fields) => {
2910 assert_eq!(fields.len(), 1);
2911 assert_eq!(fields[0].name(), "a.b");
2912 assert_eq!(fields[0].data_type(), &DataType::Int64);
2913 }
2914 _ => panic!("expected struct data type"),
2915 }
2916
2917 Ok(())
2918 }
2919
2920 #[test]
2921 fn test_variant_schema_builder_custom_nullability() -> Result<()> {
2922 let shredding_type = ShreddedSchemaBuilder::default()
2923 .with_path(
2924 "foo",
2925 Arc::new(Field::new("should_be_renamed", DataType::Utf8, false)),
2926 )?
2927 .with_path("bar", (&DataType::Int64, false))?
2928 .build();
2929
2930 let DataType::Struct(fields) = shredding_type else {
2931 panic!("expected struct data type");
2932 };
2933
2934 let foo = fields.iter().find(|f| f.name() == "foo").unwrap();
2935 assert_eq!(foo.data_type(), &DataType::Utf8);
2936 assert!(!foo.is_nullable());
2937
2938 let bar = fields.iter().find(|f| f.name() == "bar").unwrap();
2939 assert_eq!(bar.data_type(), &DataType::Int64);
2940 assert!(!bar.is_nullable());
2941
2942 Ok(())
2943 }
2944
2945 #[test]
2946 fn test_variant_schema_builder_with_shred_variant() -> Result<()> {
2947 let input = build_variant_array(vec![
2948 VariantRow::Object(vec![
2949 ("time", VariantValue::from(1234567890i64)),
2950 ("hostname", VariantValue::from("server1")),
2951 ("extra", VariantValue::from(42)),
2952 ]),
2953 VariantRow::Object(vec![
2954 ("time", VariantValue::from(9876543210i64)),
2955 ("hostname", VariantValue::from("server2")),
2956 ]),
2957 VariantRow::Null,
2958 ]);
2959
2960 let shredding_type = ShreddedSchemaBuilder::default()
2961 .with_path("time", &DataType::Int64)?
2962 .with_path("hostname", &DataType::Utf8)?
2963 .build();
2964
2965 let result = shred_variant(&input, &shredding_type).unwrap();
2966
2967 assert_eq!(
2968 result.data_type(),
2969 &DataType::Struct(Fields::from(vec![
2970 Field::new("metadata", DataType::BinaryView, false),
2971 Field::new("value", DataType::BinaryView, true),
2972 Field::new(
2973 "typed_value",
2974 DataType::Struct(Fields::from(vec![
2975 Field::new(
2976 "hostname",
2977 DataType::Struct(Fields::from(vec![
2978 Field::new("value", DataType::BinaryView, true),
2979 Field::new("typed_value", DataType::Utf8, true),
2980 ])),
2981 false,
2982 ),
2983 Field::new(
2984 "time",
2985 DataType::Struct(Fields::from(vec![
2986 Field::new("value", DataType::BinaryView, true),
2987 Field::new("typed_value", DataType::Int64, true),
2988 ])),
2989 false,
2990 ),
2991 ])),
2992 true,
2993 ),
2994 ]))
2995 );
2996
2997 assert_eq!(result.len(), 3);
2998 assert!(result.typed_value_column().is_some());
2999
3000 let typed_value = result
3001 .typed_value_column()
3002 .unwrap()
3003 .as_any()
3004 .downcast_ref::<arrow::array::StructArray>()
3005 .unwrap();
3006
3007 let time_field =
3008 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("time").unwrap())
3009 .unwrap();
3010 let hostname_field =
3011 ShreddedVariantFieldArray::try_new(typed_value.column_by_name("hostname").unwrap())
3012 .unwrap();
3013
3014 let time_typed = time_field
3015 .typed_value_column()
3016 .unwrap()
3017 .as_any()
3018 .downcast_ref::<Int64Array>()
3019 .unwrap();
3020 let hostname_typed = hostname_field
3021 .typed_value_column()
3022 .unwrap()
3023 .as_any()
3024 .downcast_ref::<arrow::array::StringArray>()
3025 .unwrap();
3026
3027 assert!(!result.is_null(0));
3029 assert_eq!(time_typed.value(0), 1234567890);
3030 assert_eq!(hostname_typed.value(0), "server1");
3031
3032 assert!(!result.is_null(1));
3034 assert_eq!(time_typed.value(1), 9876543210);
3035 assert_eq!(hostname_typed.value(1), "server2");
3036
3037 assert!(result.is_null(2));
3039
3040 Ok(())
3041 }
3042
3043 #[test]
3044 fn test_variant_schema_builder_conflicting_path() -> Result<()> {
3045 let shredding_type = ShreddedSchemaBuilder::default()
3046 .with_path("a", &DataType::Int64)?
3047 .with_path("a", &DataType::Float64)?
3048 .build();
3049
3050 assert_eq!(
3051 shredding_type,
3052 DataType::Struct(Fields::from(
3053 vec![Field::new("a", DataType::Float64, true),]
3054 ))
3055 );
3056
3057 Ok(())
3058 }
3059
3060 #[test]
3061 fn test_variant_schema_builder_root_path() -> Result<()> {
3062 let path = VariantPath::new(vec![]);
3063 let shredding_type = ShreddedSchemaBuilder::default()
3064 .with_path(path, &DataType::Int64)?
3065 .build();
3066
3067 assert_eq!(shredding_type, DataType::Int64);
3068
3069 Ok(())
3070 }
3071
3072 #[test]
3073 fn test_variant_schema_builder_empty_path() -> Result<()> {
3074 let shredding_type = ShreddedSchemaBuilder::default()
3075 .with_path("", &DataType::Int64)?
3076 .build();
3077
3078 assert_eq!(shredding_type, DataType::Int64);
3079 Ok(())
3080 }
3081
3082 #[test]
3083 fn test_variant_schema_builder_default() {
3084 let shredding_type = ShreddedSchemaBuilder::default().build();
3085 assert_eq!(shredding_type, DataType::Null);
3086 }
3087
3088 #[test]
3090 fn test_variant_type_shredded_correctly() {
3091 let mut array_builder = VariantArrayBuilder::new(30);
3093 array_builder.append_value(Variant::Null);
3094 array_builder.append_value(Variant::Int8(1));
3095 array_builder.append_value(Variant::Int16(2));
3096 array_builder.append_value(Variant::Int32(3));
3097 array_builder.append_value(Variant::Int64(4));
3098 array_builder.append_value(Variant::Date(NaiveDate::from_epoch_days(12345).unwrap()));
3099 array_builder.append_value(Variant::TimestampMicros(
3100 DateTime::from_timestamp_micros(123456789).unwrap(),
3101 ));
3102 array_builder.append_value(Variant::TimestampNtzMicros(
3103 DateTime::from_timestamp_micros(123456789)
3104 .unwrap()
3105 .naive_utc(),
3106 ));
3107 array_builder.append_value(Variant::TimestampNanos(DateTime::from_timestamp_nanos(
3108 1234567890000,
3109 )));
3110 array_builder.append_value(Variant::TimestampNtzNanos(
3111 DateTime::from_timestamp_nanos(1234567890000).naive_utc(),
3112 ));
3113 array_builder.append_value(VariantDecimal4::try_new(123, 0).unwrap());
3114 array_builder.append_value(VariantDecimal8::try_new(123, 0).unwrap());
3115 array_builder.append_value(VariantDecimal16::try_new(123, 0).unwrap());
3116 array_builder.append_value(Variant::Float(5.0));
3117 array_builder.append_value(Variant::Double(6f64));
3118 array_builder.append_value(Variant::BooleanTrue);
3119 array_builder.append_value(Variant::BooleanFalse);
3120 array_builder.append_value(Variant::Binary("helow".as_bytes()));
3121 array_builder.append_value(Variant::String("hello"));
3122 array_builder.append_value(Variant::ShortString(
3123 ShortString::try_from("world").unwrap(),
3124 ));
3125 array_builder.append_value(Variant::Time(
3126 NaiveTime::from_num_seconds_from_midnight_opt(12345, 123).unwrap(),
3127 ));
3128
3129 let array = array_builder.build();
3130
3131 fn can_shred_to(v: &Variant, dt: &DataType) -> bool {
3132 matches!(
3133 (v, dt),
3134 (
3135 Variant::Int8(_)
3136 | Variant::Int16(_)
3137 | Variant::Int32(_)
3138 | Variant::Int64(_)
3139 | Variant::Decimal4(_)
3140 | Variant::Decimal8(_)
3141 | Variant::Decimal16(_),
3142 DataType::Int8
3143 | DataType::Int16
3144 | DataType::Int32
3145 | DataType::Int64
3146 | DataType::Decimal32(_, _)
3147 | DataType::Decimal64(_, _)
3148 | DataType::Decimal128(_, _)
3149 ) | (Variant::Date(_), DataType::Date32)
3150 | (
3151 Variant::TimestampMicros(_) | Variant::TimestampNanos(_),
3152 DataType::Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, Some(_))
3153 )
3154 | (
3155 Variant::TimestampNtzMicros(_) | Variant::TimestampNtzNanos(_),
3156 DataType::Timestamp(TimeUnit::Microsecond | TimeUnit::Nanosecond, None)
3157 )
3158 | (Variant::Float(_), DataType::Float32)
3159 | (Variant::Double(_), DataType::Float64)
3160 | (
3161 Variant::BooleanFalse | Variant::BooleanTrue,
3162 DataType::Boolean
3163 )
3164 | (
3165 Variant::Binary(_),
3166 DataType::Binary | DataType::BinaryView | DataType::LargeBinary
3167 )
3168 | (
3169 Variant::ShortString(_) | Variant::String(_),
3170 DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8
3171 )
3172 | (Variant::Time(_), DataType::Time64(_))
3173 )
3174 }
3175
3176 macro_rules! assert_shred_type {
3177 ($shred_type:expr, $expected_value_valid_bits:expr) => {
3178 let shredded_array_result = shred_variant(&array, &$shred_type);
3179 match shredded_array_result {
3180 Ok(shredded_array) => {
3181 let value_column = shredded_array.inner().column_by_name("value").unwrap();
3182 for (idx, valid) in $expected_value_valid_bits.iter().enumerate() {
3183 match valid {
3184 true => assert!(
3185 value_column.is_null(idx),
3186 "{:?} should be shredded to {}",
3187 array.value(idx),
3188 $shred_type
3189 ),
3190 false => assert!(
3191 value_column.is_valid(idx),
3192 "{:?} should not be shredded to {}",
3193 array.value(idx),
3194 $shred_type
3195 ),
3196 }
3197 }
3198 }
3199 Err(e) => {
3200 let error_msg = format!("is not a valid variant shredding type");
3201 assert!(
3202 e.to_string().contains(error_msg.as_str()),
3203 "{} => {}",
3204 $shred_type,
3205 e.to_string()
3206 );
3207 }
3208 }
3209 };
3210 }
3211
3212 let types = [
3213 DataType::Null,
3214 DataType::Boolean,
3215 DataType::Int8,
3216 DataType::Int16,
3217 DataType::Int32,
3218 DataType::Int64,
3219 DataType::UInt8,
3220 DataType::UInt16,
3221 DataType::UInt32,
3222 DataType::UInt64,
3223 DataType::Float32,
3224 DataType::Float64,
3225 DataType::Timestamp(TimeUnit::Second, Some("+00:00".into())),
3226 DataType::Timestamp(TimeUnit::Second, None),
3227 DataType::Timestamp(TimeUnit::Millisecond, Some("-00:00".into())),
3228 DataType::Timestamp(TimeUnit::Millisecond, None),
3229 DataType::Timestamp(TimeUnit::Microsecond, Some("-00:00".into())),
3230 DataType::Timestamp(TimeUnit::Microsecond, None),
3231 DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
3232 DataType::Timestamp(TimeUnit::Nanosecond, None),
3233 DataType::Date32,
3234 DataType::Date64,
3235 DataType::Time32(TimeUnit::Second),
3236 DataType::Time32(TimeUnit::Millisecond),
3237 DataType::Time64(TimeUnit::Microsecond),
3238 DataType::Time64(TimeUnit::Nanosecond),
3239 DataType::Duration(TimeUnit::Nanosecond),
3240 DataType::Interval(IntervalUnit::DayTime),
3241 DataType::Binary,
3242 DataType::FixedSizeBinary(16), DataType::FixedSizeBinary(32),
3244 DataType::LargeBinary,
3245 DataType::BinaryView,
3246 DataType::Utf8,
3247 DataType::LargeUtf8,
3248 DataType::Utf8View,
3249 DataType::Decimal32(7, 4),
3250 DataType::Decimal64(7, 4),
3251 DataType::Decimal128(7, 4),
3252 DataType::Decimal256(7, 4),
3253 ];
3254
3255 for data_type in types {
3256 let expected_bits = array
3257 .iter()
3258 .map(|v| can_shred_to(&v.unwrap(), &data_type))
3259 .collect::<Vec<bool>>();
3260 assert_shred_type!(data_type, expected_bits);
3261 }
3262 }
3263}