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