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