1use arrow::{
18 array::{
19 self, Array, ArrayRef, GenericListArray, GenericListViewArray, ListLikeArray, StructArray,
20 UInt64Array, make_array,
21 },
22 buffer::NullBuffer,
23 compute::{CastOptions, take},
24 datatypes::Field,
25 error::Result,
26};
27use arrow_schema::{ArrowError, DataType, FieldRef};
28use parquet_variant::{VariantPath, VariantPathElement};
29
30use crate::ShreddingState;
31use crate::variant_array::all_null_value_column;
32use crate::variant_to_arrow::make_variant_to_arrow_row_builder;
33use crate::{VariantArray, VariantType, unshred_variant};
34
35use arrow::array::AsArray;
36use std::sync::Arc;
37
38pub(crate) enum ShreddedPathStep {
39 Success(ShreddingState),
41 Missing,
44 NotShredded,
48}
49
50fn take_list_like_index_as_shredding_state<L: ListLikeArray + 'static>(
53 typed_value: &dyn Array,
54 index: usize,
55) -> Result<Option<ShreddingState>> {
56 let list_array = typed_value.as_any().downcast_ref::<L>().ok_or_else(|| {
57 ArrowError::ComputeError(format!(
58 "Expected array type '{}' while handling list-like path step, got '{}'",
59 std::any::type_name::<L>(),
60 typed_value.data_type()
61 ))
62 })?;
63
64 let values = list_array.values();
65
66 let Some(struct_array) = values.as_struct_opt() else {
67 return Ok(None);
68 };
69 let shredding_state = ShreddingState::try_from(struct_array)?;
70
71 let value_array = shredding_state.value_column();
72 let typed_array = shredding_state.typed_value_column();
73
74 if typed_array.is_none() && value_array.null_count() == value_array.len() {
76 return Ok(None);
77 }
78
79 let mut take_indices = Vec::with_capacity(list_array.len());
80 for row in 0..list_array.len() {
81 let row_range = list_array.element_range(row);
82 let take_index = (index < row_range.len()).then(|| (row_range.start + index) as u64);
83 take_indices.push(take_index);
84 }
85
86 let index_array = UInt64Array::from(take_indices);
87
88 let taken_value = take(value_array, &index_array, None)?;
90 let taken_typed = typed_array
91 .map(|typed| take(typed, &index_array, None))
92 .transpose()?;
93
94 Ok(Some(ShreddingState::new(taken_value, taken_typed)))
95}
96
97pub(crate) fn follow_shredded_path_element(
113 shredding_state: &ShreddingState,
114 path_element: &VariantPathElement<'_>,
115 _cast_options: &CastOptions,
116) -> Result<ShreddedPathStep> {
117 let missing_path_step = || {
120 let value = shredding_state.value_column();
121 if value.null_count() == value.len() {
122 ShreddedPathStep::Missing
123 } else {
124 ShreddedPathStep::NotShredded
125 }
126 };
127
128 let Some(typed_value) = shredding_state.typed_value_column() else {
129 return Ok(missing_path_step());
130 };
131
132 match path_element {
133 VariantPathElement::Field { name } => {
134 let Some(struct_array) = typed_value.as_struct_opt() else {
137 return Ok(missing_path_step());
139 };
140
141 let Some(field) = struct_array.column_by_name(name) else {
143 return Ok(missing_path_step());
145 };
146
147 let struct_array = field.as_struct_opt().ok_or_else(|| {
148 ArrowError::InvalidArgumentError(format!(
151 "Shredded object field '{name}' must be a Struct containing 'value' and/or \
152 'typed_value', got {}",
153 field.data_type(),
154 ))
155 })?;
156
157 let state = ShreddingState::try_from(struct_array)?;
158 Ok(ShreddedPathStep::Success(state))
159 }
160 VariantPathElement::Index { index } => {
161 let state = match typed_value.data_type() {
162 DataType::List(_) => take_list_like_index_as_shredding_state::<
163 GenericListArray<i32>,
164 >(typed_value.as_ref(), *index)?,
165 DataType::LargeList(_) => take_list_like_index_as_shredding_state::<
166 GenericListArray<i64>,
167 >(typed_value.as_ref(), *index)?,
168 DataType::ListView(_) => take_list_like_index_as_shredding_state::<
169 GenericListViewArray<i32>,
170 >(typed_value.as_ref(), *index)?,
171 DataType::LargeListView(_) => take_list_like_index_as_shredding_state::<
172 GenericListViewArray<i64>,
173 >(typed_value.as_ref(), *index)?,
174 _ => {
175 return Ok(missing_path_step());
177 }
178 };
179
180 match state {
181 Some(state) => Ok(ShreddedPathStep::Success(state)),
182 None => Ok(missing_path_step()),
183 }
184 }
185 VariantPathElement::ListElement => Err(ArrowError::InvalidArgumentError(
186 "variant_get does not support [*] path elements".to_string(),
187 )),
188 }
189}
190
191fn shredded_get_path(
195 input: &VariantArray,
196 path: &[VariantPathElement<'_>],
197 as_field: Option<&Field>,
198 cast_options: &CastOptions,
199) -> Result<ArrayRef> {
200 let make_target_variant =
203 |value: ArrayRef, typed_value: Option<ArrayRef>, accumulated_nulls: Option<NullBuffer>| {
204 let metadata = input.metadata_column().clone();
205 VariantArray::from_parts(metadata, value, typed_value, accumulated_nulls)
206 };
207
208 let shred_basic_variant =
211 |target: VariantArray, path: VariantPath<'_>, as_field: Option<&Field>| {
212 let requested_variant =
215 as_field.is_some_and(Field::has_valid_extension_type::<VariantType>);
216
217 if requested_variant && requested_field_is_shredded(as_field) {
222 return Err(ArrowError::NotYetImplemented(
223 "variant_get with shredded `Variant` output is not yet supported".to_string(),
224 ));
225 }
226
227 let target = if requested_variant {
233 unshred_variant(&target)?
234 } else {
235 target
236 };
237
238 if requested_variant && path.is_empty() {
240 return Ok(ArrayRef::from(target));
241 }
242
243 let as_type = if requested_variant {
244 None
245 } else {
246 as_field.map(|f| f.data_type())
247 };
248 let mut builder = make_variant_to_arrow_row_builder(
249 target.metadata_column(),
250 path,
251 as_type,
252 cast_options,
253 target.len(),
254 )?;
255 for i in 0..target.len() {
256 if target.is_null(i) {
257 builder.append_null()?;
258 } else if !cast_options.safe {
259 let value = target.try_value(i)?;
260 builder.append_value(value)?;
261 } else {
262 let _ = match target.try_value(i) {
263 Ok(v) => builder.append_value(v)?,
264 Err(_) => {
265 builder.append_null()?;
266 false }
268 };
269 }
270 }
271 builder.finish()
272 };
273
274 let mut shredding_state = input.shredding_state().clone();
277 let mut accumulated_nulls = input.inner().nulls().cloned();
278 let mut path_index = 0;
279 for path_element in path {
280 match follow_shredded_path_element(&shredding_state, path_element, cast_options)? {
281 ShreddedPathStep::Success(state) => {
282 if let Some(typed_value) = shredding_state.typed_value_column() {
284 accumulated_nulls =
285 NullBuffer::union(accumulated_nulls.as_ref(), typed_value.nulls());
286 }
287 shredding_state = state;
288 path_index += 1;
289 }
290 ShreddedPathStep::Missing => {
291 let num_rows = input.len();
292 if as_field.is_some_and(Field::has_valid_extension_type::<VariantType>) {
293 let all_nulls = Some(arrow::buffer::NullBuffer::from(vec![false; num_rows]));
294 let metadata = input.metadata_column().clone();
297 let arr = VariantArray::from_parts_unshredded(
298 metadata,
299 all_null_value_column(num_rows),
300 all_nulls,
301 );
302 return Ok(ArrayRef::from(arr));
303 }
304 let arr = match as_field.map(|f| f.data_type()) {
305 Some(data_type) => array::new_null_array(data_type, num_rows),
306 None => Arc::new(array::NullArray::new(num_rows)) as _,
307 };
308 return Ok(arr);
309 }
310 ShreddedPathStep::NotShredded => {
311 let target = make_target_variant(
312 shredding_state.value_column().clone(),
313 None,
314 accumulated_nulls,
315 );
316 return shred_basic_variant(target, path[path_index..].into(), as_field);
317 }
318 }
319 }
320
321 let target = make_target_variant(
323 shredding_state.value_column().clone(),
324 shredding_state.typed_value_column().cloned(),
325 accumulated_nulls,
326 );
327
328 let Some(as_field) = as_field else {
330 return Ok(ArrayRef::from(target));
331 };
332
333 if let Some(shredded) = try_perfect_shredding(&target, as_field) {
335 return Ok(shredded);
336 }
337
338 if !as_field.has_valid_extension_type::<VariantType>()
348 && let DataType::Struct(fields) = as_field.data_type()
349 {
350 if target.typed_value_column().is_none() {
351 return shred_basic_variant(target, VariantPath::default(), Some(as_field));
352 }
353
354 let children = fields
355 .iter()
356 .map(|field| {
357 let path = &[VariantPathElement::from(field.name().as_str())];
358 shredded_get_path(&target, path, Some(field), cast_options)
359 })
360 .collect::<Result<Vec<_>>>()?;
361
362 return Ok(Arc::new(StructArray::try_new(
363 fields.clone(),
364 children,
365 target.nulls().cloned(),
366 )?));
367 }
368
369 shred_basic_variant(target, VariantPath::default(), Some(as_field))
371}
372
373fn requested_field_is_shredded(as_field: Option<&Field>) -> bool {
378 as_field.is_some_and(|f| match f.data_type() {
379 DataType::Struct(fields) => fields.iter().any(|field| field.name() == "typed_value"),
380 _ => false,
381 })
382}
383
384fn try_perfect_shredding(variant_array: &VariantArray, as_field: &Field) -> Option<ArrayRef> {
385 if matches!(as_field.data_type(), DataType::Struct(_)) {
387 return None;
388 }
389 let typed_value = variant_array.typed_value_column()?;
390
391 let value = variant_array.value_column();
392 if typed_value.data_type() == as_field.data_type() && value.null_count() == value.len() {
393 let parent_nulls = variant_array.nulls();
400
401 let target_array = if parent_nulls.is_none() || typed_value.data_type().is_null() {
403 typed_value.clone()
404 } else {
405 let merged_nulls = NullBuffer::union(parent_nulls, typed_value.nulls());
406 let data = typed_value
407 .to_data()
408 .into_builder()
409 .nulls(merged_nulls)
410 .build()
411 .ok()?;
412 make_array(data)
413 };
414
415 return Some(target_array);
416 }
417
418 None
419}
420
421pub fn variant_get(input: &ArrayRef, options: GetOptions) -> Result<ArrayRef> {
453 let variant_array = VariantArray::try_new(input)?;
454
455 let GetOptions {
456 as_type,
457 path,
458 cast_options,
459 } = options;
460
461 if path
462 .iter()
463 .any(|element| matches!(element, VariantPathElement::ListElement))
464 {
465 return Err(ArrowError::InvalidArgumentError(
466 "variant_get does not support [*] path elements".to_string(),
467 ));
468 }
469
470 shredded_get_path(&variant_array, &path, as_type.as_deref(), &cast_options)
471}
472
473#[derive(Debug, Clone, Default)]
475pub struct GetOptions<'a> {
476 pub path: VariantPath<'a>,
478 pub as_type: Option<FieldRef>,
482 pub cast_options: CastOptions<'a>,
484}
485
486impl<'a> GetOptions<'a> {
487 pub fn new() -> Self {
489 Default::default()
490 }
491
492 pub fn new_with_path(path: VariantPath<'a>) -> Self {
494 Self {
495 path,
496 as_type: None,
497 cast_options: Default::default(),
498 }
499 }
500
501 pub fn with_as_type(mut self, as_type: Option<FieldRef>) -> Self {
503 self.as_type = as_type;
504 self
505 }
506
507 pub fn with_cast_options(mut self, cast_options: CastOptions<'a>) -> Self {
509 self.cast_options = cast_options;
510 self
511 }
512}
513
514#[cfg(test)]
515mod test {
516 use std::str::FromStr;
517 use std::sync::Arc;
518
519 use super::{GetOptions, requested_field_is_shredded, variant_get};
520 use crate::variant_array::{
521 ShreddedVariantFieldArray, StructArrayBuilder, all_null_value_column,
522 };
523 use crate::{
524 ShreddedSchemaBuilder, VariantArray, VariantArrayBuilder, cast_to_variant, json_to_variant,
525 shred_variant,
526 };
527 use arrow::array::{
528 Array, ArrayRef, AsArray, BinaryArray, BinaryViewArray, BooleanArray, Date32Array,
529 Date64Array, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array,
530 FixedSizeListArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
531 Int64Array, Int64Builder, LargeBinaryArray, LargeListArray, LargeListViewArray,
532 LargeStringArray, ListArray, ListBuilder, ListViewArray, MapBuilder, NullArray,
533 NullBuilder, StringArray, StringBuilder, StringViewArray, StructArray,
534 Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray,
535 UnionArray,
536 };
537 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
538 use arrow::compute::{CastOptions, cast};
539 use arrow::datatypes::DataType::{Int16, Int32, Int64};
540 use arrow::datatypes::i256;
541 use arrow::util::display::FormatOptions;
542 use arrow_schema::ArrowError;
543 use arrow_schema::DataType::{Boolean, Float32, Float64, Int8};
544 use arrow_schema::{
545 DataType, Field, FieldRef, Fields, IntervalUnit, TimeUnit, UnionFields, UnionMode,
546 };
547 use chrono::DateTime;
548 use parquet_variant::{
549 EMPTY_VARIANT_METADATA_BYTES, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16,
550 VariantDecimalType, VariantPath,
551 };
552
553 fn single_variant_get_test(input_json: &str, path: VariantPath, expected_json: &str) {
554 let input_array_ref: ArrayRef = Arc::new(StringArray::from(vec![Some(input_json)]));
556 let input_variant_array_ref = ArrayRef::from(json_to_variant(&input_array_ref).unwrap());
557
558 let result =
559 variant_get(&input_variant_array_ref, GetOptions::new_with_path(path)).unwrap();
560
561 let expected_array_ref: ArrayRef = Arc::new(StringArray::from(vec![Some(expected_json)]));
563 let expected_variant_array = json_to_variant(&expected_array_ref).unwrap();
564
565 let result_array = VariantArray::try_new(&result).unwrap();
566 assert_eq!(
567 result_array.len(),
568 1,
569 "Expected result array to have length 1"
570 );
571 assert!(
572 result_array.nulls().is_none(),
573 "Expected no nulls in result array"
574 );
575 let result_variant = result_array.value(0);
576 let expected_variant = expected_variant_array.value(0);
577 assert_eq!(
578 result_variant, expected_variant,
579 "Result variant does not match expected variant"
580 );
581 }
582
583 #[test]
584 fn get_primitive_variant_field() {
585 single_variant_get_test(
586 r#"{"some_field": 1234}"#,
587 VariantPath::try_from("some_field").unwrap(),
588 "1234",
589 );
590 }
591
592 #[test]
593 fn get_primitive_variant_list_index() {
594 single_variant_get_test("[1234, 5678]", VariantPath::from(0), "1234");
595 }
596
597 #[test]
598 fn get_primitive_variant_inside_object_of_object() {
599 single_variant_get_test(
600 r#"{"top_level_field": {"inner_field": 1234}}"#,
601 VariantPath::try_from("top_level_field")
602 .unwrap()
603 .join("inner_field"),
604 "1234",
605 );
606 }
607
608 #[test]
609 fn get_primitive_variant_inside_list_of_object() {
610 single_variant_get_test(
611 r#"[{"some_field": 1234}]"#,
612 VariantPath::from(0).join("some_field"),
613 "1234",
614 );
615 }
616
617 #[test]
618 fn get_primitive_variant_inside_object_of_list() {
619 single_variant_get_test(
620 r#"{"some_field": [1234]}"#,
621 VariantPath::try_from("some_field[0]").unwrap(),
622 "1234",
623 );
624 }
625
626 #[test]
627 fn get_complex_variant() {
628 single_variant_get_test(
629 r#"{"top_level_field": {"inner_field": 1234}}"#,
630 VariantPath::try_from("top_level_field").unwrap(),
631 r#"{"inner_field": 1234}"#,
632 );
633 }
634
635 macro_rules! numeric_partially_shredded_test {
637 ($primitive_type:ty, $data_fn:ident) => {
638 let array = $data_fn();
639 let options = GetOptions::new();
640 let result = variant_get(&array, options).unwrap();
641
642 let result = VariantArray::try_new(&result).unwrap();
644 assert_eq!(result.len(), 4);
645
646 assert_eq!(
648 result.value(0),
649 Variant::from(<$primitive_type>::try_from(34u8).unwrap())
650 );
651 assert!(!result.is_valid(1));
652 assert_eq!(result.value(2), Variant::from("n/a"));
653 assert_eq!(
654 result.value(3),
655 Variant::from(<$primitive_type>::try_from(100u8).unwrap())
656 );
657 };
658 }
659
660 macro_rules! partially_shredded_variant_array_gen {
663 ($func_name:ident, $typed_value_array_gen: expr) => {
664 partially_shredded_variant_array_gen!(
665 $func_name,
666 $typed_value_array_gen,
667 Variant::from("n/a")
668 );
669 };
670 ($func_name:ident, $typed_value_array_gen: expr, $fallback_variant:expr) => {
671 fn $func_name() -> ArrayRef {
672 let typed_value: ArrayRef = Arc::new($typed_value_array_gen());
673 let typed_as_variant = cast_to_variant(typed_value.as_ref())
674 .expect("should cast typed array to variant");
675 let mut input_builder = VariantArrayBuilder::new(typed_as_variant.len());
676 input_builder.append_variant(typed_as_variant.value(0));
677 input_builder.append_null();
678 input_builder.append_variant($fallback_variant);
679 input_builder.append_variant(typed_as_variant.value(3));
680
681 let variant_array = shred_variant(&input_builder.build(), typed_value.data_type())
682 .expect("should shred variant array");
683 ArrayRef::from(variant_array)
684 }
685 };
686 }
687
688 macro_rules! numeric_partially_shredded_variant_array_fn {
690 ($func:ident, $array_type:ident, $primitive_type:ty) => {
691 partially_shredded_variant_array_gen!($func, || $array_type::from(vec![
692 Some(<$primitive_type>::try_from(34u8).unwrap()),
693 None,
694 None,
695 Some(<$primitive_type>::try_from(100u8).unwrap()),
696 ]));
697 };
698 }
699
700 numeric_partially_shredded_variant_array_fn!(
701 partially_shredded_int8_variant_array,
702 Int8Array,
703 i8
704 );
705 numeric_partially_shredded_variant_array_fn!(
706 partially_shredded_int16_variant_array,
707 Int16Array,
708 i16
709 );
710 numeric_partially_shredded_variant_array_fn!(
711 partially_shredded_int32_variant_array,
712 Int32Array,
713 i32
714 );
715 numeric_partially_shredded_variant_array_fn!(
716 partially_shredded_int64_variant_array,
717 Int64Array,
718 i64
719 );
720 numeric_partially_shredded_variant_array_fn!(
721 partially_shredded_float32_variant_array,
722 Float32Array,
723 f32
724 );
725 numeric_partially_shredded_variant_array_fn!(
726 partially_shredded_float64_variant_array,
727 Float64Array,
728 f64
729 );
730
731 partially_shredded_variant_array_gen!(partially_shredded_bool_variant_array, || {
732 arrow::array::BooleanArray::from(vec![Some(true), None, None, Some(false)])
733 });
734
735 partially_shredded_variant_array_gen!(
736 partially_shredded_utf8_variant_array,
737 || { StringArray::from(vec![Some("hello"), None, None, Some("world")]) },
738 Variant::from(42i32)
739 );
740
741 partially_shredded_variant_array_gen!(partially_shredded_date32_variant_array, || {
742 Date32Array::from(vec![
743 Some(20348), None,
745 None,
746 Some(20340), ])
748 });
749
750 #[test]
751 fn get_variant_partially_shredded_int8_as_variant() {
752 numeric_partially_shredded_test!(i8, partially_shredded_int8_variant_array);
753 }
754
755 #[test]
756 fn get_variant_partially_shredded_int16_as_variant() {
757 numeric_partially_shredded_test!(i16, partially_shredded_int16_variant_array);
758 }
759
760 #[test]
761 fn get_variant_partially_shredded_int32_as_variant() {
762 numeric_partially_shredded_test!(i32, partially_shredded_int32_variant_array);
763 }
764
765 #[test]
766 fn get_variant_partially_shredded_int64_as_variant() {
767 numeric_partially_shredded_test!(i64, partially_shredded_int64_variant_array);
768 }
769
770 #[test]
771 fn get_variant_partially_shredded_float32_as_variant() {
772 numeric_partially_shredded_test!(f32, partially_shredded_float32_variant_array);
773 }
774
775 #[test]
776 fn get_variant_partially_shredded_float64_as_variant() {
777 numeric_partially_shredded_test!(f64, partially_shredded_float64_variant_array);
778 }
779
780 #[test]
781 fn get_variant_partially_shredded_bool_as_variant() {
782 let array = partially_shredded_bool_variant_array();
783 let options = GetOptions::new();
784 let result = variant_get(&array, options).unwrap();
785
786 let result = VariantArray::try_new(&result).unwrap();
788 assert_eq!(result.len(), 4);
789
790 assert_eq!(result.value(0), Variant::from(true));
792 assert!(!result.is_valid(1));
793 assert_eq!(result.value(2), Variant::from("n/a"));
794 assert_eq!(result.value(3), Variant::from(false));
795 }
796
797 #[test]
798 fn get_variant_partially_shredded_utf8_as_variant() {
799 let array = partially_shredded_utf8_variant_array();
800 let options = GetOptions::new();
801 let result = variant_get(&array, options).unwrap();
802
803 let result = VariantArray::try_new(&result).unwrap();
805 assert_eq!(result.len(), 4);
806
807 assert_eq!(result.value(0), Variant::from("hello"));
809 assert!(!result.is_valid(1));
810 assert_eq!(result.value(2), Variant::from(42i32));
811 assert_eq!(result.value(3), Variant::from("world"));
812 }
813
814 partially_shredded_variant_array_gen!(partially_shredded_binary_view_variant_array, || {
815 BinaryViewArray::from(vec![
816 Some(&[1u8, 2u8, 3u8][..]), None, None, Some(&[4u8, 5u8, 6u8][..]), ])
821 });
822
823 #[test]
824 fn get_variant_partially_shredded_date32_as_variant() {
825 let array = partially_shredded_date32_variant_array();
826 let options = GetOptions::new();
827 let result = variant_get(&array, options).unwrap();
828
829 let result = VariantArray::try_new(&result).unwrap();
831 assert_eq!(result.len(), 4);
832
833 use chrono::NaiveDate;
835 let date1 = NaiveDate::from_ymd_opt(2025, 9, 17).unwrap();
836 let date2 = NaiveDate::from_ymd_opt(2025, 9, 9).unwrap();
837 assert_eq!(result.value(0), Variant::from(date1));
838 assert!(!result.is_valid(1));
839 assert_eq!(result.value(2), Variant::from("n/a"));
840 assert_eq!(result.value(3), Variant::from(date2));
841 }
842
843 #[test]
844 fn get_variant_partially_shredded_binary_view_as_variant() {
845 let array = partially_shredded_binary_view_variant_array();
846 let options = GetOptions::new();
847 let result = variant_get(&array, options).unwrap();
848
849 let result = VariantArray::try_new(&result).unwrap();
851 assert_eq!(result.len(), 4);
852
853 assert_eq!(result.value(0), Variant::from(&[1u8, 2u8, 3u8][..]));
855 assert!(!result.is_valid(1));
856 assert_eq!(result.value(2), Variant::from("n/a"));
857 assert_eq!(result.value(3), Variant::from(&[4u8, 5u8, 6u8][..]));
858 }
859
860 macro_rules! assert_variant_get_as_variant_array_with_default_option {
862 ($variant_array: expr, $array_expected: expr) => {{
863 let options = GetOptions::new();
864 let array = $variant_array;
865 let result = variant_get(&array, options).unwrap();
866 let result = VariantArray::try_new(&result).unwrap();
867
868 assert_eq!(result.len(), $array_expected.len());
869
870 for (idx, item) in $array_expected.into_iter().enumerate() {
871 match item {
872 Some(item) => assert_eq!(result.value(idx), item),
873 None => assert!(result.is_null(idx)),
874 }
875 }
876 }};
877 }
878
879 partially_shredded_variant_array_gen!(
880 partially_shredded_timestamp_micro_ntz_variant_array,
881 || {
882 arrow::array::TimestampMicrosecondArray::from(vec![
883 Some(-456000),
884 None,
885 None,
886 Some(1758602096000000),
887 ])
888 }
889 );
890
891 #[test]
892 fn get_variant_partial_shredded_timestamp_micro_ntz_as_variant() {
893 let array = partially_shredded_timestamp_micro_ntz_variant_array();
894 assert_variant_get_as_variant_array_with_default_option!(
895 array,
896 vec![
897 Some(Variant::from(
898 DateTime::from_timestamp_micros(-456000i64)
899 .unwrap()
900 .naive_utc(),
901 )),
902 None,
903 Some(Variant::from("n/a")),
904 Some(Variant::from(
905 DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
906 .unwrap()
907 .naive_utc(),
908 )),
909 ]
910 )
911 }
912
913 partially_shredded_variant_array_gen!(partially_shredded_timestamp_micro_variant_array, || {
914 arrow::array::TimestampMicrosecondArray::from(vec![
915 Some(-456000),
916 None,
917 None,
918 Some(1758602096000000),
919 ])
920 .with_timezone("+00:00")
921 });
922
923 #[test]
924 fn get_variant_partial_shredded_timestamp_micro_as_variant() {
925 let array = partially_shredded_timestamp_micro_variant_array();
926 assert_variant_get_as_variant_array_with_default_option!(
927 array,
928 vec![
929 Some(Variant::from(
930 DateTime::from_timestamp_micros(-456000i64)
931 .unwrap()
932 .to_utc(),
933 )),
934 None,
935 Some(Variant::from("n/a")),
936 Some(Variant::from(
937 DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
938 .unwrap()
939 .to_utc(),
940 )),
941 ]
942 )
943 }
944
945 partially_shredded_variant_array_gen!(
946 partially_shredded_timestamp_nano_ntz_variant_array,
947 || {
948 arrow::array::TimestampNanosecondArray::from(vec![
949 Some(-4999999561),
950 None,
951 None,
952 Some(1758602096000000000),
953 ])
954 }
955 );
956
957 #[test]
958 fn get_variant_partial_shredded_timestamp_nano_ntz_as_variant() {
959 let array = partially_shredded_timestamp_nano_ntz_variant_array();
960 assert_variant_get_as_variant_array_with_default_option!(
961 array,
962 vec![
963 Some(Variant::from(
964 DateTime::from_timestamp(-5, 439).unwrap().naive_utc()
965 )),
966 None,
967 Some(Variant::from("n/a")),
968 Some(Variant::from(
969 DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
970 .unwrap()
971 .naive_utc()
972 )),
973 ]
974 )
975 }
976
977 partially_shredded_variant_array_gen!(partially_shredded_timestamp_nano_variant_array, || {
978 arrow::array::TimestampNanosecondArray::from(vec![
979 Some(-4999999561),
980 None,
981 None,
982 Some(1758602096000000000),
983 ])
984 .with_timezone("+00:00")
985 });
986
987 #[test]
988 fn get_variant_partial_shredded_timestamp_nano_as_variant() {
989 let array = partially_shredded_timestamp_nano_variant_array();
990 assert_variant_get_as_variant_array_with_default_option!(
991 array,
992 vec![
993 Some(Variant::from(
994 DateTime::from_timestamp(-5, 439).unwrap().to_utc()
995 )),
996 None,
997 Some(Variant::from("n/a")),
998 Some(Variant::from(
999 DateTime::parse_from_rfc3339("2025-09-23T12:34:56+08:00")
1000 .unwrap()
1001 .to_utc()
1002 )),
1003 ]
1004 )
1005 }
1006
1007 #[test]
1009 fn get_variant_shredded_int32_as_int32_safe_cast() {
1010 let array = partially_shredded_int32_variant_array();
1012 let field = Field::new("typed_value", DataType::Int32, true);
1014 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
1015 let result = variant_get(&array, options).unwrap();
1016 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
1017 Some(34),
1018 None,
1019 None, Some(100),
1021 ]));
1022 assert_eq!(&result, &expected)
1023 }
1024
1025 #[test]
1027 fn get_variant_shredded_int32_as_int32_unsafe_cast() {
1028 let array = partially_shredded_int32_variant_array();
1030 let field = Field::new("typed_value", DataType::Int32, true);
1031 let cast_options = CastOptions {
1032 safe: false, ..Default::default()
1034 };
1035 let options = GetOptions::new()
1036 .with_as_type(Some(FieldRef::from(field)))
1037 .with_cast_options(cast_options);
1038
1039 let err = variant_get(&array, options).unwrap_err();
1040 assert_eq!(
1042 err.to_string(),
1043 "Cast error: Failed to extract primitive of type Int32 from variant ShortString(ShortString(\"n/a\")) at path VariantPath([])"
1044 );
1045 }
1046
1047 macro_rules! numeric_perfectly_shredded_test {
1049 ($primitive_type:ty, $data_fn:ident) => {
1050 let array = $data_fn();
1051 let options = GetOptions::new();
1052 let result = variant_get(&array, options).unwrap();
1053
1054 let result = VariantArray::try_new(&result).unwrap();
1056 assert_eq!(result.len(), 3);
1057
1058 assert_eq!(
1060 result.value(0),
1061 Variant::from(<$primitive_type>::try_from(1u8).unwrap())
1062 );
1063 assert_eq!(
1064 result.value(1),
1065 Variant::from(<$primitive_type>::try_from(2u8).unwrap())
1066 );
1067 assert_eq!(
1068 result.value(2),
1069 Variant::from(<$primitive_type>::try_from(3u8).unwrap())
1070 );
1071 };
1072 }
1073
1074 #[test]
1075 fn get_variant_perfectly_shredded_int8_as_variant() {
1076 numeric_perfectly_shredded_test!(i8, perfectly_shredded_int8_variant_array);
1077 }
1078
1079 #[test]
1080 fn get_variant_perfectly_shredded_int16_as_variant() {
1081 numeric_perfectly_shredded_test!(i16, perfectly_shredded_int16_variant_array);
1082 }
1083
1084 #[test]
1085 fn get_variant_perfectly_shredded_int32_as_variant() {
1086 numeric_perfectly_shredded_test!(i32, perfectly_shredded_int32_variant_array);
1087 }
1088
1089 #[test]
1090 fn get_variant_perfectly_shredded_int64_as_variant() {
1091 numeric_perfectly_shredded_test!(i64, perfectly_shredded_int64_variant_array);
1092 }
1093
1094 #[test]
1095 fn get_variant_perfectly_shredded_float32_as_variant() {
1096 numeric_perfectly_shredded_test!(f32, perfectly_shredded_float32_variant_array);
1097 }
1098
1099 #[test]
1100 fn get_variant_perfectly_shredded_float64_as_variant() {
1101 numeric_perfectly_shredded_test!(f64, perfectly_shredded_float64_variant_array);
1102 }
1103
1104 #[test]
1106 fn get_variant_all_null_as_variant() {
1107 let array = all_null_variant_array();
1108 let options = GetOptions::new();
1109 let result = variant_get(&array, options).unwrap();
1110
1111 let result = VariantArray::try_new(&result).unwrap();
1113 assert_eq!(result.len(), 3);
1114
1115 assert!(!result.is_valid(0));
1117 assert!(!result.is_valid(1));
1118 assert!(!result.is_valid(2));
1119 }
1120
1121 #[test]
1123 fn get_variant_all_null_as_int32() {
1124 let array = all_null_variant_array();
1125 let field = Field::new("typed_value", DataType::Int32, true);
1127 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
1128 let result = variant_get(&array, options).unwrap();
1129
1130 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
1131 Option::<i32>::None,
1132 Option::<i32>::None,
1133 Option::<i32>::None,
1134 ]));
1135 assert_eq!(&result, &expected)
1136 }
1137
1138 macro_rules! perfectly_shredded_to_arrow_primitive_test {
1139 ($name:ident, $primitive_type:expr, $perfectly_shredded_array_gen_fun:ident, $expected_array:expr) => {
1140 #[test]
1141 fn $name() {
1142 let array = $perfectly_shredded_array_gen_fun();
1143 let field = Field::new("typed_value", $primitive_type, true);
1144 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
1145 let result = variant_get(&array, options).unwrap();
1146 let expected_array: ArrayRef = Arc::new($expected_array);
1147 assert_eq!(&result, &expected_array);
1148 }
1149 };
1150 }
1151
1152 perfectly_shredded_to_arrow_primitive_test!(
1153 get_variant_perfectly_shredded_int18_as_int8,
1154 Int8,
1155 perfectly_shredded_int8_variant_array,
1156 Int8Array::from(vec![Some(1), Some(2), Some(3)])
1157 );
1158
1159 perfectly_shredded_to_arrow_primitive_test!(
1160 get_variant_perfectly_shredded_int16_as_int16,
1161 Int16,
1162 perfectly_shredded_int16_variant_array,
1163 Int16Array::from(vec![Some(1), Some(2), Some(3)])
1164 );
1165
1166 perfectly_shredded_to_arrow_primitive_test!(
1167 get_variant_perfectly_shredded_int32_as_int32,
1168 Int32,
1169 perfectly_shredded_int32_variant_array,
1170 Int32Array::from(vec![Some(1), Some(2), Some(3)])
1171 );
1172
1173 perfectly_shredded_to_arrow_primitive_test!(
1174 get_variant_perfectly_shredded_int64_as_int64,
1175 Int64,
1176 perfectly_shredded_int64_variant_array,
1177 Int64Array::from(vec![Some(1), Some(2), Some(3)])
1178 );
1179
1180 perfectly_shredded_to_arrow_primitive_test!(
1181 get_variant_perfectly_shredded_float32_as_float32,
1182 Float32,
1183 perfectly_shredded_float32_variant_array,
1184 Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])
1185 );
1186
1187 perfectly_shredded_to_arrow_primitive_test!(
1188 get_variant_perfectly_shredded_float64_as_float64,
1189 Float64,
1190 perfectly_shredded_float64_variant_array,
1191 Float64Array::from(vec![Some(1.0), Some(2.0), Some(3.0)])
1192 );
1193
1194 perfectly_shredded_to_arrow_primitive_test!(
1195 get_variant_perfectly_shredded_boolean_as_boolean,
1196 Boolean,
1197 perfectly_shredded_bool_variant_array,
1198 BooleanArray::from(vec![Some(true), Some(false), Some(true)])
1199 );
1200
1201 perfectly_shredded_to_arrow_primitive_test!(
1202 get_variant_perfectly_shredded_utf8_as_utf8,
1203 DataType::Utf8,
1204 perfectly_shredded_utf8_variant_array,
1205 StringArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1206 );
1207
1208 perfectly_shredded_to_arrow_primitive_test!(
1209 get_variant_perfectly_shredded_large_utf8_as_utf8,
1210 DataType::Utf8,
1211 perfectly_shredded_large_utf8_variant_array,
1212 StringArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1213 );
1214
1215 perfectly_shredded_to_arrow_primitive_test!(
1216 get_variant_perfectly_shredded_utf8_view_as_utf8,
1217 DataType::Utf8,
1218 perfectly_shredded_utf8_view_variant_array,
1219 StringArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1220 );
1221
1222 macro_rules! perfectly_shredded_variant_array_fn {
1223 ($func:ident, $typed_value_gen:expr) => {
1224 fn $func() -> ArrayRef {
1225 let typed_value: ArrayRef = Arc::new($typed_value_gen());
1228 if let Some(shredded) = cast_to_variant(typed_value.as_ref())
1229 .ok()
1230 .and_then(|unshredded| shred_variant(&unshredded, typed_value.data_type()).ok())
1231 {
1232 return shredded.into();
1233 }
1234
1235 let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
1236 EMPTY_VARIANT_METADATA_BYTES,
1237 typed_value.len(),
1238 ));
1239 VariantArray::perfectly_shredded(Arc::new(metadata), typed_value, None).into()
1240 }
1241 };
1242 }
1243
1244 perfectly_shredded_variant_array_fn!(perfectly_shredded_utf8_variant_array, || {
1245 StringArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1246 });
1247
1248 perfectly_shredded_variant_array_fn!(perfectly_shredded_large_utf8_variant_array, || {
1249 LargeStringArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1250 });
1251
1252 perfectly_shredded_variant_array_fn!(perfectly_shredded_utf8_view_variant_array, || {
1253 StringViewArray::from(vec![Some("foo"), Some("bar"), Some("baz")])
1254 });
1255
1256 perfectly_shredded_variant_array_fn!(perfectly_shredded_bool_variant_array, || {
1257 BooleanArray::from(vec![Some(true), Some(false), Some(true)])
1258 });
1259
1260 macro_rules! numeric_perfectly_shredded_variant_array_fn {
1272 ($func:ident, $array_type:ident, $primitive_type:ty) => {
1273 perfectly_shredded_variant_array_fn!($func, || {
1274 $array_type::from(vec![
1275 Some(<$primitive_type>::try_from(1u8).unwrap()),
1276 Some(<$primitive_type>::try_from(2u8).unwrap()),
1277 Some(<$primitive_type>::try_from(3u8).unwrap()),
1278 ])
1279 });
1280 };
1281 }
1282
1283 numeric_perfectly_shredded_variant_array_fn!(
1284 perfectly_shredded_int8_variant_array,
1285 Int8Array,
1286 i8
1287 );
1288 numeric_perfectly_shredded_variant_array_fn!(
1289 perfectly_shredded_int16_variant_array,
1290 Int16Array,
1291 i16
1292 );
1293 numeric_perfectly_shredded_variant_array_fn!(
1294 perfectly_shredded_int32_variant_array,
1295 Int32Array,
1296 i32
1297 );
1298 numeric_perfectly_shredded_variant_array_fn!(
1299 perfectly_shredded_int64_variant_array,
1300 Int64Array,
1301 i64
1302 );
1303 numeric_perfectly_shredded_variant_array_fn!(
1304 perfectly_shredded_float32_variant_array,
1305 Float32Array,
1306 f32
1307 );
1308 numeric_perfectly_shredded_variant_array_fn!(
1309 perfectly_shredded_float64_variant_array,
1310 Float64Array,
1311 f64
1312 );
1313
1314 perfectly_shredded_variant_array_fn!(
1315 perfectly_shredded_timestamp_micro_ntz_variant_array,
1316 || {
1317 arrow::array::TimestampMicrosecondArray::from(vec![
1318 Some(-456000),
1319 Some(1758602096000001),
1320 Some(1758602096000002),
1321 ])
1322 }
1323 );
1324
1325 perfectly_shredded_to_arrow_primitive_test!(
1326 get_variant_perfectly_shredded_timestamp_micro_ntz_as_timestamp_micro_ntz,
1327 DataType::Timestamp(TimeUnit::Microsecond, None),
1328 perfectly_shredded_timestamp_micro_ntz_variant_array,
1329 arrow::array::TimestampMicrosecondArray::from(vec![
1330 Some(-456000),
1331 Some(1758602096000001),
1332 Some(1758602096000002),
1333 ])
1334 );
1335
1336 perfectly_shredded_to_arrow_primitive_test!(
1338 get_variant_perfectly_shredded_timestamp_micro_ntz_as_nano_ntz,
1339 DataType::Timestamp(TimeUnit::Nanosecond, None),
1340 perfectly_shredded_timestamp_micro_ntz_variant_array,
1341 arrow::array::TimestampNanosecondArray::from(vec![
1342 Some(-456000000),
1343 Some(1758602096000001000),
1344 Some(1758602096000002000)
1345 ])
1346 );
1347
1348 perfectly_shredded_variant_array_fn!(perfectly_shredded_timestamp_micro_variant_array, || {
1349 arrow::array::TimestampMicrosecondArray::from(vec![
1350 Some(-456000),
1351 Some(1758602096000001),
1352 Some(1758602096000002),
1353 ])
1354 .with_timezone("+00:00")
1355 });
1356
1357 perfectly_shredded_to_arrow_primitive_test!(
1358 get_variant_perfectly_shredded_timestamp_micro_as_timestamp_micro,
1359 DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("+00:00"))),
1360 perfectly_shredded_timestamp_micro_variant_array,
1361 arrow::array::TimestampMicrosecondArray::from(vec![
1362 Some(-456000),
1363 Some(1758602096000001),
1364 Some(1758602096000002),
1365 ])
1366 .with_timezone("+00:00")
1367 );
1368
1369 perfectly_shredded_to_arrow_primitive_test!(
1371 get_variant_perfectly_shredded_timestamp_micro_as_nano,
1372 DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+00:00"))),
1373 perfectly_shredded_timestamp_micro_variant_array,
1374 arrow::array::TimestampNanosecondArray::from(vec![
1375 Some(-456000000),
1376 Some(1758602096000001000),
1377 Some(1758602096000002000)
1378 ])
1379 .with_timezone("+00:00")
1380 );
1381
1382 perfectly_shredded_variant_array_fn!(
1383 perfectly_shredded_timestamp_nano_ntz_variant_array,
1384 || {
1385 arrow::array::TimestampNanosecondArray::from(vec![
1386 Some(-4999999561),
1387 Some(1758602096000000001),
1388 Some(1758602096000000002),
1389 ])
1390 }
1391 );
1392
1393 perfectly_shredded_variant_array_fn!(
1394 perfectly_shredded_timestamp_micro_variant_array_for_second_and_milli_second,
1395 || {
1396 arrow::array::TimestampMicrosecondArray::from(vec![
1397 Some(1234), Some(1234000), Some(1234000000), ])
1401 .with_timezone("+00:00")
1402 }
1403 );
1404
1405 perfectly_shredded_to_arrow_primitive_test!(
1408 get_variant_perfectly_shredded_timestamp_micro_as_timestamp_second,
1409 DataType::Timestamp(TimeUnit::Second, Some(Arc::from("+00:00"))),
1410 perfectly_shredded_timestamp_micro_variant_array_for_second_and_milli_second,
1411 arrow::array::TimestampSecondArray::from(vec![
1412 None,
1413 None, Some(1234)
1415 ])
1416 .with_timezone("+00:00")
1417 );
1418
1419 perfectly_shredded_to_arrow_primitive_test!(
1420 get_variant_perfectly_shredded_timestamp_micro_as_timestamp_milli,
1421 DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+00:00"))),
1422 perfectly_shredded_timestamp_micro_variant_array_for_second_and_milli_second,
1423 arrow::array::TimestampMillisecondArray::from(vec![
1424 None, Some(1234),
1426 Some(1234000)
1427 ])
1428 .with_timezone("+00:00")
1429 );
1430
1431 perfectly_shredded_variant_array_fn!(
1432 perfectly_shredded_timestamp_micro_ntz_variant_array_for_second_and_milli_second,
1433 || {
1434 arrow::array::TimestampMicrosecondArray::from(vec![
1435 Some(1234), Some(1234000), Some(1234000000), ])
1439 }
1440 );
1441
1442 perfectly_shredded_to_arrow_primitive_test!(
1445 get_variant_perfectly_shredded_timestamp_micro_ntz_as_timestamp_second,
1446 DataType::Timestamp(TimeUnit::Second, None),
1447 perfectly_shredded_timestamp_micro_ntz_variant_array_for_second_and_milli_second,
1448 arrow::array::TimestampSecondArray::from(vec![
1449 None,
1450 None, Some(1234)
1452 ])
1453 );
1454
1455 perfectly_shredded_to_arrow_primitive_test!(
1456 get_variant_perfectly_shredded_timestamp_micro_ntz_as_timestamp_milli,
1457 DataType::Timestamp(TimeUnit::Millisecond, None),
1458 perfectly_shredded_timestamp_micro_ntz_variant_array_for_second_and_milli_second,
1459 arrow::array::TimestampMillisecondArray::from(vec![
1460 None, Some(1234),
1462 Some(1234000)
1463 ])
1464 );
1465
1466 perfectly_shredded_variant_array_fn!(
1467 perfectly_shredded_timestamp_nano_variant_array_for_second_and_milli_second,
1468 || {
1469 arrow::array::TimestampNanosecondArray::from(vec![
1470 Some(1234000), Some(1234000000), Some(1234000000000), ])
1474 .with_timezone("+00:00")
1475 }
1476 );
1477
1478 perfectly_shredded_to_arrow_primitive_test!(
1481 get_variant_perfectly_shredded_timestamp_nano_as_timestamp_second,
1482 DataType::Timestamp(TimeUnit::Second, Some(Arc::from("+00:00"))),
1483 perfectly_shredded_timestamp_nano_variant_array_for_second_and_milli_second,
1484 arrow::array::TimestampSecondArray::from(vec![
1485 None,
1486 None, Some(1234)
1488 ])
1489 .with_timezone("+00:00")
1490 );
1491
1492 perfectly_shredded_to_arrow_primitive_test!(
1493 get_variant_perfectly_shredded_timestamp_nano_as_timestamp_milli,
1494 DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+00:00"))),
1495 perfectly_shredded_timestamp_nano_variant_array_for_second_and_milli_second,
1496 arrow::array::TimestampMillisecondArray::from(vec![
1497 None, Some(1234),
1499 Some(1234000)
1500 ])
1501 .with_timezone("+00:00")
1502 );
1503
1504 perfectly_shredded_variant_array_fn!(
1505 perfectly_shredded_timestamp_nano_ntz_variant_array_for_second_and_milli_second,
1506 || {
1507 arrow::array::TimestampNanosecondArray::from(vec![
1508 Some(1234000), Some(1234000000), Some(1234000000000), ])
1512 }
1513 );
1514
1515 perfectly_shredded_to_arrow_primitive_test!(
1518 get_variant_perfectly_shredded_timestamp_nano_ntz_as_timestamp_second,
1519 DataType::Timestamp(TimeUnit::Second, None),
1520 perfectly_shredded_timestamp_nano_ntz_variant_array_for_second_and_milli_second,
1521 arrow::array::TimestampSecondArray::from(vec![
1522 None,
1523 None, Some(1234)
1525 ])
1526 );
1527
1528 perfectly_shredded_to_arrow_primitive_test!(
1529 get_variant_perfectly_shredded_timestamp_nano_ntz_as_timestamp_milli,
1530 DataType::Timestamp(TimeUnit::Millisecond, None),
1531 perfectly_shredded_timestamp_nano_ntz_variant_array_for_second_and_milli_second,
1532 arrow::array::TimestampMillisecondArray::from(vec![
1533 None, Some(1234),
1535 Some(1234000)
1536 ])
1537 );
1538
1539 perfectly_shredded_to_arrow_primitive_test!(
1540 get_variant_perfectly_shredded_timestamp_nano_ntz_as_timestamp_nano_ntz,
1541 DataType::Timestamp(TimeUnit::Nanosecond, None),
1542 perfectly_shredded_timestamp_nano_ntz_variant_array,
1543 arrow::array::TimestampNanosecondArray::from(vec![
1544 Some(-4999999561),
1545 Some(1758602096000000001),
1546 Some(1758602096000000002),
1547 ])
1548 );
1549
1550 perfectly_shredded_variant_array_fn!(perfectly_shredded_timestamp_nano_variant_array, || {
1551 arrow::array::TimestampNanosecondArray::from(vec![
1552 Some(-4999999561),
1553 Some(1758602096000000001),
1554 Some(1758602096000000002),
1555 ])
1556 .with_timezone("+00:00")
1557 });
1558
1559 perfectly_shredded_to_arrow_primitive_test!(
1560 get_variant_perfectly_shredded_timestamp_nano_as_timestamp_nano,
1561 DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+00:00"))),
1562 perfectly_shredded_timestamp_nano_variant_array,
1563 arrow::array::TimestampNanosecondArray::from(vec![
1564 Some(-4999999561),
1565 Some(1758602096000000001),
1566 Some(1758602096000000002),
1567 ])
1568 .with_timezone("+00:00")
1569 );
1570
1571 perfectly_shredded_variant_array_fn!(perfectly_shredded_date_variant_array, || {
1572 Date32Array::from(vec![Some(-12345), Some(17586), Some(20000)])
1573 });
1574
1575 perfectly_shredded_to_arrow_primitive_test!(
1576 get_variant_perfectly_shredded_date_as_date,
1577 DataType::Date32,
1578 perfectly_shredded_date_variant_array,
1579 Date32Array::from(vec![Some(-12345), Some(17586), Some(20000)])
1580 );
1581
1582 perfectly_shredded_to_arrow_primitive_test!(
1583 get_variant_perfectly_shredded_date_as_date64,
1584 DataType::Date64,
1585 perfectly_shredded_date_variant_array,
1586 Date64Array::from(vec![
1587 Some(-1066608000000),
1588 Some(1519430400000),
1589 Some(1728000000000)
1590 ])
1591 );
1592
1593 perfectly_shredded_variant_array_fn!(perfectly_shredded_time_variant_array, || {
1594 Time64MicrosecondArray::from(vec![Some(12345000), Some(87654000), Some(135792000)])
1595 });
1596
1597 perfectly_shredded_to_arrow_primitive_test!(
1598 get_variant_perfectly_shredded_time_as_time,
1599 DataType::Time64(TimeUnit::Microsecond),
1600 perfectly_shredded_time_variant_array,
1601 Time64MicrosecondArray::from(vec![Some(12345000), Some(87654000), Some(135792000)])
1602 );
1603
1604 perfectly_shredded_to_arrow_primitive_test!(
1605 get_variant_perfectly_shredded_time_as_time64_nano,
1606 DataType::Time64(TimeUnit::Nanosecond),
1607 perfectly_shredded_time_variant_array,
1608 Time64NanosecondArray::from(vec![
1609 Some(12345000000),
1610 Some(87654000000),
1611 Some(135792000000)
1612 ])
1613 );
1614
1615 perfectly_shredded_variant_array_fn!(perfectly_shredded_time_variant_array_for_time32, || {
1616 Time64MicrosecondArray::from(vec![
1617 Some(1234), Some(7654000), Some(35792000000), ])
1621 });
1622
1623 perfectly_shredded_to_arrow_primitive_test!(
1624 get_variant_perfectly_shredded_time_as_time32_second,
1625 DataType::Time32(TimeUnit::Second),
1626 perfectly_shredded_time_variant_array_for_time32,
1627 Time32SecondArray::from(vec![
1628 None,
1629 None, Some(35792)
1631 ])
1632 );
1633
1634 perfectly_shredded_to_arrow_primitive_test!(
1635 get_variant_perfectly_shredded_time_as_time32_milli,
1636 DataType::Time32(TimeUnit::Millisecond),
1637 perfectly_shredded_time_variant_array_for_time32,
1638 Time32MillisecondArray::from(vec![
1639 None, Some(7654),
1641 Some(35792000)
1642 ])
1643 );
1644
1645 perfectly_shredded_variant_array_fn!(perfectly_shredded_null_variant_array, || {
1646 let mut builder = NullBuilder::new();
1647 builder.append_nulls(3);
1648 builder.finish()
1649 });
1650
1651 perfectly_shredded_to_arrow_primitive_test!(
1652 get_variant_perfectly_shredded_null_as_null,
1653 DataType::Null,
1654 perfectly_shredded_null_variant_array,
1655 arrow::array::NullArray::new(3)
1656 );
1657
1658 perfectly_shredded_variant_array_fn!(perfectly_shredded_null_variant_array_with_int, || {
1659 Int32Array::from(vec![Some(32), Some(64), Some(48)])
1660 });
1661
1662 perfectly_shredded_to_arrow_primitive_test!(
1664 get_variant_perfectly_shredded_null_with_type_mismatch_in_safe_mode,
1665 DataType::Null,
1666 perfectly_shredded_null_variant_array_with_int,
1667 arrow::array::NullArray::new(3)
1668 );
1669
1670 #[test]
1672 fn get_variant_perfectly_shredded_null_as_null_with_type_mismatch_in_strict_mode() {
1673 let array = perfectly_shredded_null_variant_array_with_int();
1674 let field = Field::new("typed_value", DataType::Null, true);
1675 let options = GetOptions::new()
1676 .with_as_type(Some(FieldRef::from(field)))
1677 .with_cast_options(CastOptions {
1678 safe: false,
1679 format_options: FormatOptions::default(),
1680 });
1681
1682 let result = variant_get(&array, options);
1683
1684 assert!(result.is_err());
1685 let error_msg = format!("{}", result.unwrap_err());
1686 assert!(
1687 error_msg
1688 .contains("Cast error: Failed to extract primitive of type Null from variant Int32(32) at path VariantPath([])"),
1689 "Expected=[Cast error: Failed to extract primitive of type Null from variant Int32(32) at path VariantPath([])],\
1690 Got error message=[{error_msg}]"
1691 );
1692 }
1693
1694 perfectly_shredded_variant_array_fn!(perfectly_shredded_decimal4_variant_array, || {
1695 Decimal32Array::from(vec![Some(12345), Some(23400), Some(-12342)])
1696 .with_precision_and_scale(5, 2)
1697 .unwrap()
1698 });
1699
1700 perfectly_shredded_to_arrow_primitive_test!(
1701 get_variant_perfectly_shredded_decimal4_as_decimal4,
1702 DataType::Decimal32(5, 2),
1703 perfectly_shredded_decimal4_variant_array,
1704 Decimal32Array::from(vec![Some(12345), Some(23400), Some(-12342)])
1705 .with_precision_and_scale(5, 2)
1706 .unwrap()
1707 );
1708
1709 perfectly_shredded_variant_array_fn!(
1710 perfectly_shredded_decimal8_variant_array_cast2decimal32,
1711 || {
1712 Decimal64Array::from(vec![Some(123456), Some(145678), Some(-123456)])
1713 .with_precision_and_scale(6, 1)
1714 .unwrap()
1715 }
1716 );
1717
1718 perfectly_shredded_to_arrow_primitive_test!(
1722 get_variant_perfectly_shredded_decimal8_through_decimal32_as_decimal8,
1723 DataType::Decimal64(6, 1),
1724 perfectly_shredded_decimal8_variant_array_cast2decimal32,
1725 Decimal64Array::from(vec![Some(123456), Some(145678), Some(-123456)])
1726 .with_precision_and_scale(6, 1)
1727 .unwrap()
1728 );
1729
1730 perfectly_shredded_variant_array_fn!(perfectly_shredded_decimal8_variant_array, || {
1733 Decimal64Array::from(vec![Some(1234567809), Some(1456787000), Some(-1234561203)])
1734 .with_precision_and_scale(10, 1)
1735 .unwrap()
1736 });
1737
1738 perfectly_shredded_to_arrow_primitive_test!(
1739 get_variant_perfectly_shredded_decimal8_as_decimal8,
1740 DataType::Decimal64(10, 1),
1741 perfectly_shredded_decimal8_variant_array,
1742 Decimal64Array::from(vec![Some(1234567809), Some(1456787000), Some(-1234561203)])
1743 .with_precision_and_scale(10, 1)
1744 .unwrap()
1745 );
1746
1747 perfectly_shredded_variant_array_fn!(
1750 perfectly_shredded_decimal16_within_decimal4_variant_array,
1751 || {
1752 Decimal128Array::from(vec![
1753 Some(i128::from(1234589)),
1754 Some(i128::from(2344444)),
1755 Some(i128::from(-1234789)),
1756 ])
1757 .with_precision_and_scale(7, 3)
1758 .unwrap()
1759 }
1760 );
1761
1762 perfectly_shredded_to_arrow_primitive_test!(
1765 get_variant_perfectly_shredded_decimal16_within_decimal4_as_decimal16,
1766 DataType::Decimal128(7, 3),
1767 perfectly_shredded_decimal16_within_decimal4_variant_array,
1768 Decimal128Array::from(vec![
1769 Some(i128::from(1234589)),
1770 Some(i128::from(2344444)),
1771 Some(i128::from(-1234789)),
1772 ])
1773 .with_precision_and_scale(7, 3)
1774 .unwrap()
1775 );
1776
1777 perfectly_shredded_variant_array_fn!(
1778 perfectly_shredded_decimal16_within_decimal8_variant_array,
1779 || {
1780 Decimal128Array::from(vec![Some(1234567809), Some(1456787000), Some(-1234561203)])
1781 .with_precision_and_scale(10, 1)
1782 .unwrap()
1783 }
1784 );
1785
1786 perfectly_shredded_to_arrow_primitive_test!(
1789 get_variant_perfectly_shredded_decimal16_within8_as_decimal16,
1790 DataType::Decimal128(10, 1),
1791 perfectly_shredded_decimal16_within_decimal8_variant_array,
1792 Decimal128Array::from(vec![Some(1234567809), Some(1456787000), Some(-1234561203)])
1793 .with_precision_and_scale(10, 1)
1794 .unwrap()
1795 );
1796
1797 perfectly_shredded_variant_array_fn!(perfectly_shredded_decimal16_variant_array, || {
1798 Decimal128Array::from(vec![
1799 Some(i128::from_str("12345678901234567899").unwrap()),
1800 Some(i128::from_str("23445677483748324300").unwrap()),
1801 Some(i128::from_str("-12345678901234567899").unwrap()),
1802 ])
1803 .with_precision_and_scale(20, 3)
1804 .unwrap()
1805 });
1806
1807 perfectly_shredded_to_arrow_primitive_test!(
1810 get_variant_perfectly_shredded_decimal16_as_decimal16,
1811 DataType::Decimal128(20, 3),
1812 perfectly_shredded_decimal16_variant_array,
1813 Decimal128Array::from(vec![
1814 Some(i128::from_str("12345678901234567899").unwrap()),
1815 Some(i128::from_str("23445677483748324300").unwrap()),
1816 Some(i128::from_str("-12345678901234567899").unwrap())
1817 ])
1818 .with_precision_and_scale(20, 3)
1819 .unwrap()
1820 );
1821
1822 perfectly_shredded_variant_array_fn!(perfectly_shredded_binary_variant_array, || {
1823 BinaryArray::from(vec![
1824 Some(b"Apache" as &[u8]),
1825 Some(b"Arrow-rs" as &[u8]),
1826 Some(b"Parquet-variant" as &[u8]),
1827 ])
1828 });
1829
1830 perfectly_shredded_to_arrow_primitive_test!(
1831 get_variant_perfectly_shredded_binary_as_binary,
1832 DataType::Binary,
1833 perfectly_shredded_binary_variant_array,
1834 BinaryArray::from(vec![
1835 Some(b"Apache" as &[u8]),
1836 Some(b"Arrow-rs" as &[u8]),
1837 Some(b"Parquet-variant" as &[u8]),
1838 ])
1839 );
1840
1841 perfectly_shredded_variant_array_fn!(perfectly_shredded_large_binary_variant_array, || {
1842 LargeBinaryArray::from(vec![
1843 Some(b"Apache" as &[u8]),
1844 Some(b"Arrow-rs" as &[u8]),
1845 Some(b"Parquet-variant" as &[u8]),
1846 ])
1847 });
1848
1849 perfectly_shredded_to_arrow_primitive_test!(
1850 get_variant_perfectly_shredded_large_binary_as_large_binary,
1851 DataType::LargeBinary,
1852 perfectly_shredded_large_binary_variant_array,
1853 LargeBinaryArray::from(vec![
1854 Some(b"Apache" as &[u8]),
1855 Some(b"Arrow-rs" as &[u8]),
1856 Some(b"Parquet-variant" as &[u8]),
1857 ])
1858 );
1859
1860 perfectly_shredded_variant_array_fn!(perfectly_shredded_binary_view_variant_array, || {
1861 BinaryViewArray::from(vec![
1862 Some(b"Apache" as &[u8]),
1863 Some(b"Arrow-rs" as &[u8]),
1864 Some(b"Parquet-variant" as &[u8]),
1865 ])
1866 });
1867
1868 perfectly_shredded_to_arrow_primitive_test!(
1869 get_variant_perfectly_shredded_binary_view_as_binary_view,
1870 DataType::BinaryView,
1871 perfectly_shredded_binary_view_variant_array,
1872 BinaryViewArray::from(vec![
1873 Some(b"Apache" as &[u8]),
1874 Some(b"Arrow-rs" as &[u8]),
1875 Some(b"Parquet-variant" as &[u8]),
1876 ])
1877 );
1878
1879 fn all_null_variant_array() -> ArrayRef {
1896 let nulls = NullBuffer::from(vec![
1897 false, false, false, ]);
1901
1902 let metadata =
1904 BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 3));
1905
1906 ArrayRef::from(VariantArray::from_parts(
1907 Arc::new(metadata),
1908 all_null_value_column(3),
1909 None,
1910 Some(nulls),
1911 ))
1912 }
1913
1914 #[test]
1918 fn test_shredded_object_field_access() {
1919 let array = shredded_object_with_x_field_variant_array();
1920
1921 let options = GetOptions::new_with_path(VariantPath::try_from("x").unwrap());
1923 let result = variant_get(&array, options).unwrap();
1924
1925 let result_variant = VariantArray::try_new(&result).unwrap();
1926 assert_eq!(result_variant.len(), 2);
1927
1928 assert_eq!(result_variant.value(0), Variant::Int32(1));
1930 assert_eq!(result_variant.value(1), Variant::Int32(42));
1932 }
1933
1934 #[test]
1935 fn test_malformed_shredded_object_field_reports_field_and_type() {
1936 let metadata =
1937 BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 2));
1938 let typed_value = StructArray::try_new(
1939 Fields::from(vec![Field::new("x", DataType::Int32, true)]),
1940 vec![Arc::new(Int32Array::from(vec![Some(1), Some(42)]))],
1941 None,
1942 )
1943 .unwrap();
1944 let array = ArrayRef::from(VariantArray::from_parts(
1945 Arc::new(metadata),
1946 all_null_value_column(2),
1947 Some(Arc::new(typed_value)),
1948 None,
1949 ));
1950
1951 let options = GetOptions::new_with_path(VariantPath::try_from("x").unwrap());
1952 let err = variant_get(&array, options).unwrap_err();
1953
1954 assert_eq!(
1955 err.to_string(),
1956 "Invalid argument error: Shredded object field 'x' must be a Struct containing \
1957 'value' and/or 'typed_value', got Int32"
1958 );
1959 }
1960
1961 #[test]
1963 fn test_shredded_object_field_as_int32() {
1964 let array = shredded_object_with_x_field_variant_array();
1965
1966 let field = Field::new("x", DataType::Int32, false);
1968 let options = GetOptions::new_with_path(VariantPath::try_from("x").unwrap())
1969 .with_as_type(Some(FieldRef::from(field)));
1970 let result = variant_get(&array, options).unwrap();
1971
1972 let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(42)]));
1974 assert_eq!(&result, &expected);
1975 }
1976
1977 type ShreddedListLikeArrayGen = fn() -> ArrayRef;
1978 type ShreddedListLikeCase = (&'static str, ShreddedListLikeArrayGen);
1979
1980 fn shredded_list_like_cases() -> [ShreddedListLikeCase; 4] {
1981 [
1982 ("list", shredded_list_variant_array),
1983 ("large_list", shredded_large_list_variant_array),
1984 ("list_view", shredded_list_view_variant_array),
1985 ("large_list_view", shredded_large_list_view_variant_array),
1986 ]
1987 }
1988
1989 #[test]
1990 fn test_shredded_list_like_index_access_from_value_field() {
1991 let options = GetOptions::new_with_path(VariantPath::from(1));
1992
1993 for (case, array_gen) in shredded_list_like_cases() {
1994 let array = array_gen();
1995 let result = variant_get(&array, options.clone()).unwrap();
1996 let result_variant = VariantArray::try_new(&result).unwrap();
1997
1998 assert_eq!(result_variant.value(0), Variant::from("drama"), "{case}");
1999 assert_eq!(result_variant.value(1).as_int64(), Some(123), "{case}");
2000 }
2001 }
2002
2003 #[test]
2004 fn test_shredded_list_like_index_out_of_bounds_unsafe_cast_returns_null() {
2005 let options =
2006 GetOptions::new_with_path(VariantPath::from(10)).with_cast_options(CastOptions {
2007 safe: false,
2008 ..Default::default()
2009 });
2010
2011 for (case, array_gen) in shredded_list_like_cases() {
2012 let result = variant_get(&array_gen(), options.clone()).unwrap();
2013 let result_variant = VariantArray::try_new(&result).unwrap();
2014 assert_eq!(result_variant.value(0), Variant::Null, "{case}");
2015 assert_eq!(result_variant.value(1), Variant::Null, "{case}");
2016 }
2017 }
2018
2019 #[test]
2021 fn test_shredded_list_like_as_string() {
2022 let field = Field::new("typed_value", DataType::Utf8, false);
2023 let options = GetOptions::new_with_path(VariantPath::from(0))
2024 .with_as_type(Some(FieldRef::from(field)));
2025 let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("comedy"), Some("horror")]));
2026
2027 for (case, array_gen) in shredded_list_like_cases() {
2028 let result = variant_get(&array_gen(), options.clone()).unwrap();
2029 assert_eq!(&result, &expected, "{case}");
2030 }
2031 }
2032
2033 #[test]
2034 fn test_shredded_list_like_index_access_from_value_field_as_int64() {
2035 let field = Field::new("typed_value", DataType::Int64, true);
2036 let options = GetOptions::new_with_path(VariantPath::from(1))
2037 .with_as_type(Some(FieldRef::from(field)));
2038 let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)]));
2039
2040 for (case, array_gen) in shredded_list_like_cases() {
2041 let result = variant_get(&array_gen(), options.clone()).unwrap();
2042 assert_eq!(&result, &expected, "{case}");
2044 }
2045 }
2046
2047 #[test]
2048 fn test_shredded_list_in_struct_index_access() {
2049 let array = shredded_struct_with_list_variant_array();
2050 let options = GetOptions::new_with_path(VariantPath::try_from("a[1]").unwrap());
2051 let result = variant_get(&array, options).unwrap();
2052 let result_variant = VariantArray::try_new(&result).unwrap();
2053
2054 assert_eq!(result_variant.value(0), Variant::from("drama"));
2055 assert_eq!(result_variant.value(1).as_int64(), Some(123));
2056 }
2057
2058 #[test]
2059 fn test_shredded_struct_in_list_field_access() {
2060 let array = shredded_list_of_struct_variant_array();
2061 let field = Field::new("x", DataType::Int32, true);
2062 let path = VariantPath::from(0).join("x");
2063 let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2064 let result = variant_get(&array, options).unwrap();
2065
2066 let expected: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(3)]));
2067 assert_eq!(&result, &expected);
2068 }
2069
2070 #[test]
2071 fn test_shredded_list_of_lists_index_access() {
2072 let array = shredded_list_of_lists_variant_array();
2073 let path = VariantPath::from(0).join(1);
2074
2075 let result = variant_get(&array, GetOptions::new_with_path(path.clone())).unwrap();
2076 let result_variant = VariantArray::try_new(&result).unwrap();
2077 assert_eq!(result_variant.value(0), Variant::from("b"));
2078 assert_eq!(result_variant.value(1).as_int64(), Some(123));
2079
2080 let field = Field::new("typed_value", DataType::Int64, true);
2081 let casted = variant_get(
2082 &array,
2083 GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field))),
2084 )
2085 .unwrap();
2086 let expected: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(123)]));
2087 assert_eq!(&casted, &expected);
2088 }
2089
2090 fn shredded_list_like_variant_array(list_schema: DataType) -> ArrayRef {
2096 let json_rows: ArrayRef = Arc::new(StringArray::from(vec![
2097 Some(r#"["comedy", "drama"]"#),
2098 Some(r#"["horror", 123]"#),
2099 ]));
2100 let input = json_to_variant(&json_rows).unwrap();
2101
2102 let shredded = shred_variant(&input, &list_schema).unwrap();
2103 ArrayRef::from(shredded)
2104 }
2105
2106 fn shredded_list_of_lists_variant_array() -> ArrayRef {
2107 let json_rows: ArrayRef = Arc::new(StringArray::from(vec![
2108 Some(r#"[["a", "b"], ["c", "d"]]"#),
2109 Some(r#"[["x", 123], ["y", "z"]]"#),
2110 ]));
2111 let input = json_to_variant(&json_rows).unwrap();
2112
2113 let inner_list = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2114 let outer_list = DataType::List(Arc::new(Field::new("item", inner_list, true)));
2115 let shredded = shred_variant(&input, &outer_list).unwrap();
2116 ArrayRef::from(shredded)
2117 }
2118
2119 fn shredded_list_variant_array() -> ArrayRef {
2120 shredded_list_like_variant_array(DataType::List(Arc::new(Field::new(
2121 "item",
2122 DataType::Utf8,
2123 true,
2124 ))))
2125 }
2126
2127 fn shredded_large_list_variant_array() -> ArrayRef {
2128 shredded_list_like_variant_array(DataType::LargeList(Arc::new(Field::new(
2129 "item",
2130 DataType::Utf8,
2131 true,
2132 ))))
2133 }
2134
2135 fn shredded_list_view_variant_array() -> ArrayRef {
2136 shredded_list_like_variant_array(DataType::ListView(Arc::new(Field::new(
2137 "item",
2138 DataType::Utf8,
2139 true,
2140 ))))
2141 }
2142
2143 fn shredded_large_list_view_variant_array() -> ArrayRef {
2144 shredded_list_like_variant_array(DataType::LargeListView(Arc::new(Field::new(
2145 "item",
2146 DataType::Utf8,
2147 true,
2148 ))))
2149 }
2150
2151 fn shredded_struct_with_list_variant_array() -> ArrayRef {
2152 let json_rows: ArrayRef = Arc::new(StringArray::from(vec![
2153 Some(r#"{"a": ["comedy", "drama"]}"#),
2154 Some(r#"{"a": ["horror", 123]}"#),
2155 ]));
2156 let input = json_to_variant(&json_rows).unwrap();
2157
2158 let list_schema = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2159 let shredding_schema = ShreddedSchemaBuilder::default()
2160 .with_path("a", &list_schema)
2161 .unwrap()
2162 .build();
2163 let shredded = shred_variant(&input, &shredding_schema).unwrap();
2164 ArrayRef::from(shredded)
2165 }
2166
2167 fn shredded_list_of_struct_variant_array() -> ArrayRef {
2168 let json_rows: ArrayRef = Arc::new(StringArray::from(vec![
2169 Some(r#"[{"x": 1}, {"x": 2}]"#),
2170 Some(r#"[{"x": 3}, {"y": 4}]"#),
2171 ]));
2172 let input = json_to_variant(&json_rows).unwrap();
2173
2174 let struct_type =
2175 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, true)]));
2176 let list_schema = DataType::List(Arc::new(Field::new("item", struct_type, true)));
2177 let shredded = shred_variant(&input, &list_schema).unwrap();
2178 ArrayRef::from(shredded)
2179 }
2180
2181 fn shredded_object_with_x_field_variant_array() -> ArrayRef {
2193 let (metadata, y_field_value) = {
2195 let mut builder = parquet_variant::VariantBuilder::new();
2196 let mut obj = builder.new_object();
2197 obj.insert("x", Variant::Int32(42));
2198 obj.insert("y", Variant::from("foo"));
2199 obj.finish();
2200 builder.finish()
2201 };
2202
2203 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 2));
2205
2206 let empty_object_value = {
2211 let mut builder = parquet_variant::VariantBuilder::new();
2212 let obj = builder.new_object();
2213 obj.finish();
2214 let (_, value) = builder.finish();
2215 value
2216 };
2217
2218 let value_array = BinaryViewArray::from(vec![
2219 Some(y_field_value.as_slice()), Some(empty_object_value.as_slice()), ]);
2222
2223 let x_field_typed_value = Int32Array::from(vec![Some(1), Some(42)]);
2226
2227 let x_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
2229 x_field_typed_value,
2230 ) as ArrayRef);
2231
2232 let typed_value_fields = Fields::from(vec![Field::new(
2234 "x",
2235 x_field_shredded.data_type().clone(),
2236 true,
2237 )]);
2238 let typed_value_struct = StructArray::try_new(
2239 typed_value_fields,
2240 vec![ArrayRef::from(x_field_shredded)],
2241 None, )
2243 .unwrap();
2244
2245 ArrayRef::from(VariantArray::from_parts(
2247 Arc::new(metadata_array),
2248 Arc::new(value_array),
2249 Some(Arc::new(typed_value_struct)),
2250 None,
2251 ))
2252 }
2253
2254 #[test]
2256 fn test_simple_nested_path_support() {
2257 println!("Testing path parsing:");
2259
2260 let path_x = VariantPath::try_from("x").unwrap();
2261 let elements_x: Vec<_> = path_x.iter().collect();
2262 println!(" 'x' -> {} elements: {:?}", elements_x.len(), elements_x);
2263
2264 let path_ax = VariantPath::try_from("a.x").unwrap();
2265 let elements_ax: Vec<_> = path_ax.iter().collect();
2266 println!(
2267 " 'a.x' -> {} elements: {:?}",
2268 elements_ax.len(),
2269 elements_ax
2270 );
2271
2272 let path_ax_alt = VariantPath::try_from("$.a.x").unwrap();
2273 let elements_ax_alt: Vec<_> = path_ax_alt.iter().collect();
2274 println!(
2275 " '$.a.x' -> {} elements: {:?}",
2276 elements_ax_alt.len(),
2277 elements_ax_alt
2278 );
2279
2280 let path_nested = VariantPath::try_from("a").unwrap().join("x");
2281 let elements_nested: Vec<_> = path_nested.iter().collect();
2282 println!(
2283 " VariantPath::try_from('a').unwrap().join('x') -> {} elements: {:?}",
2284 elements_nested.len(),
2285 elements_nested
2286 );
2287
2288 let array = shredded_object_with_x_field_variant_array();
2290
2291 let real_nested_path = VariantPath::try_from("a").unwrap().join("x");
2293 let options = GetOptions::new_with_path(real_nested_path);
2294 let result = variant_get(&array, options);
2295
2296 match result {
2297 Ok(_) => {
2298 println!("Nested path 'a.x' works unexpectedly!");
2299 }
2300 Err(e) => {
2301 println!("Nested path 'a.x' error: {e}");
2302 if e.to_string().contains("Not yet implemented")
2303 || e.to_string().contains("NotYetImplemented")
2304 {
2305 println!("This is expected - nested paths are not implemented");
2306 return;
2307 }
2308 println!("This shows nested paths need implementation");
2310 }
2311 }
2312 }
2313
2314 #[test]
2318 fn test_depth_0_int32_conversion() {
2319 println!("=== Testing Depth 0: Direct field access ===");
2320
2321 let unshredded_array = create_depth_0_test_data();
2323
2324 let field = Field::new("result", DataType::Int32, true);
2325 let path = VariantPath::try_from("x").unwrap();
2326 let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2327 let result = variant_get(&unshredded_array, options).unwrap();
2328
2329 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2330 Some(42), None, None, ]));
2334 assert_eq!(&result, &expected);
2335 println!("Depth 0 (unshredded) passed");
2336
2337 let shredded_array = create_depth_0_shredded_test_data_simple();
2339
2340 let field = Field::new("result", DataType::Int32, true);
2341 let path = VariantPath::try_from("x").unwrap();
2342 let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2343 let result = variant_get(&shredded_array, options).unwrap();
2344
2345 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2346 Some(42), None, ]));
2349 assert_eq!(&result, &expected);
2350 println!("Depth 0 (shredded) passed");
2351 }
2352
2353 #[test]
2356 fn test_depth_1_int32_conversion() {
2357 println!("=== Testing Depth 1: Single nested field access ===");
2358
2359 let unshredded_array = create_nested_path_test_data();
2361
2362 let field = Field::new("result", DataType::Int32, true);
2363 let path = VariantPath::try_from("a.x").unwrap(); let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2365 let result = variant_get(&unshredded_array, options).unwrap();
2366
2367 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2368 Some(55), None, ]));
2371 assert_eq!(&result, &expected);
2372 println!("Depth 1 (unshredded) passed");
2373
2374 let shredded_array = create_depth_1_shredded_test_data_working();
2376
2377 let field = Field::new("result", DataType::Int32, true);
2378 let path = VariantPath::try_from("a.x").unwrap(); let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2380 let result = variant_get(&shredded_array, options).unwrap();
2381
2382 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2383 Some(55), None, ]));
2386 assert_eq!(&result, &expected);
2387 println!("Depth 1 (shredded) passed");
2388 }
2389
2390 #[test]
2393 fn test_depth_2_int32_conversion() {
2394 println!("=== Testing Depth 2: Double nested field access ===");
2395
2396 let unshredded_array = create_depth_2_test_data();
2398
2399 let field = Field::new("result", DataType::Int32, true);
2400 let path = VariantPath::try_from("a.b.x").unwrap(); let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2402 let result = variant_get(&unshredded_array, options).unwrap();
2403
2404 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2405 Some(100), None, None, ]));
2409 assert_eq!(&result, &expected);
2410 println!("Depth 2 (unshredded) passed");
2411
2412 let shredded_array = create_depth_2_shredded_test_data_working();
2414
2415 let field = Field::new("result", DataType::Int32, true);
2416 let path = VariantPath::try_from("a.b.x").unwrap(); let options = GetOptions::new_with_path(path).with_as_type(Some(FieldRef::from(field)));
2418 let result = variant_get(&shredded_array, options).unwrap();
2419
2420 let expected: ArrayRef = Arc::new(Int32Array::from(vec![
2421 Some(100), None, None, ]));
2425 assert_eq!(&result, &expected);
2426 println!("Depth 2 (shredded) passed");
2427 }
2428
2429 #[test]
2434 fn test_current_nested_path_functionality() {
2435 let array = shredded_object_with_x_field_variant_array();
2436
2437 let single_path = VariantPath::try_from("x").unwrap();
2439 let field = Field::new("result", DataType::Int32, true);
2440 let options =
2441 GetOptions::new_with_path(single_path).with_as_type(Some(FieldRef::from(field)));
2442 let result = variant_get(&array, options).unwrap();
2443
2444 println!("Single path 'x' works - result: {result:?}");
2445
2446 let nested_path = VariantPath::try_from("a").unwrap().join("x");
2448 let field = Field::new("result", DataType::Int32, true);
2449 let options =
2450 GetOptions::new_with_path(nested_path).with_as_type(Some(FieldRef::from(field)));
2451 let result = variant_get(&array, options).unwrap();
2452
2453 println!("Nested path 'a.x' result: {result:?}");
2454 }
2455
2456 #[test]
2457 fn test_variant_get_as_variant_from_unshredded_input() {
2458 let (unshredded, _) = create_variant_get_as_variant_test_data();
2459 let unshredded_field = VariantArray::try_new(&unshredded).unwrap().field("result");
2460 assert_variant_field_extraction_returns_unshredded_variant(&unshredded, &unshredded_field);
2461 }
2462
2463 #[test]
2464 fn test_variant_get_as_variant_from_shredded_input() {
2465 let (unshredded, shredded) = create_variant_get_as_variant_test_data();
2466 let unshredded_field = VariantArray::try_new(&unshredded).unwrap().field("result");
2467 assert_variant_field_extraction_returns_unshredded_variant(&shredded, &unshredded_field);
2468 }
2469
2470 #[test]
2471 fn test_variant_get_as_shredded_variant_is_not_yet_supported() {
2472 let (_, shredded) = create_variant_get_as_variant_test_data();
2473 let shredded_field = VariantArray::try_new(&shredded).unwrap().field("result");
2478 assert!(requested_field_is_shredded(Some(&shredded_field)));
2479
2480 let options = GetOptions::new_with_path(VariantPath::try_from("field_name").unwrap())
2481 .with_as_type(Some(FieldRef::from(shredded_field)));
2482 let err = variant_get(&shredded, options).unwrap_err();
2483 assert!(
2484 matches!(err, ArrowError::NotYetImplemented(_)),
2485 "expected NotYetImplemented, got {err:?}"
2486 );
2487 }
2488
2489 #[test]
2490 fn test_variant_get_list_element_wildcard_is_invalid_argument() {
2491 let (unshredded, _) = create_variant_get_as_variant_test_data();
2492 let options = GetOptions::new_with_path(VariantPath::try_from("field_name[*]").unwrap());
2493 let err = variant_get(&unshredded, options).unwrap_err();
2494 assert!(
2495 matches!(err, ArrowError::InvalidArgumentError(_)),
2496 "expected InvalidArgumentError, got {err:?}"
2497 );
2498 }
2499
2500 #[test]
2501 fn test_variant_get_missing_path_as_variant_annotates_value_non_nullable() {
2502 let (unshredded, shredded) = create_variant_get_as_variant_test_data();
2503 let variant_field = VariantArray::try_new(&unshredded).unwrap().field("result");
2504
2505 let options = GetOptions::new_with_path(VariantPath::try_from("field_name[0]").unwrap())
2507 .with_as_type(Some(FieldRef::from(variant_field)));
2508 let result = variant_get(&shredded, options).unwrap();
2509 let result_variant = VariantArray::try_new(&result).unwrap();
2510
2511 assert_eq!(result_variant.inner().null_count(), result_variant.len());
2512 let value_field = result_variant.inner().field_by_name("value").unwrap();
2513 assert!(!value_field.is_nullable());
2514 }
2515
2516 fn create_variant_get_as_variant_test_data() -> (ArrayRef, ArrayRef) {
2517 let input_json: ArrayRef = Arc::new(StringArray::from(vec![
2518 Some(r#"{"field_name": {"k": 100000}}"#),
2519 Some(r#"{"field_name": {"k": "s"}}"#),
2520 ]));
2521
2522 let unshredded = ArrayRef::from(json_to_variant(&input_json).unwrap());
2523 let unshredded_variant = VariantArray::try_new(&unshredded).unwrap();
2524
2525 let as_type = DataType::Struct(Fields::from(vec![Field::new(
2526 "field_name",
2527 DataType::Struct(Fields::from(vec![Field::new("k", DataType::Int32, true)])),
2528 true,
2529 )]));
2530 let shredded = ArrayRef::from(shred_variant(&unshredded_variant, &as_type).unwrap());
2531
2532 (unshredded, shredded)
2533 }
2534
2535 fn assert_variant_field_extraction_returns_unshredded_variant(
2536 input: &ArrayRef,
2537 variant_field: &Field,
2538 ) {
2539 let options = GetOptions::new_with_path(VariantPath::try_from("field_name").unwrap())
2540 .with_as_type(Some(FieldRef::from(variant_field.clone())));
2541
2542 let result = variant_get(input, options).unwrap();
2543 let result_variant = VariantArray::try_new(&result).unwrap();
2544
2545 assert!(result_variant.typed_value_column().is_none());
2546 assert!(result_variant.value_column().null_count() < result_variant.len());
2547 let value_field = result_variant.inner().field_by_name("value").unwrap();
2548 assert!(!value_field.is_nullable());
2549
2550 let expected_json: ArrayRef = Arc::new(StringArray::from(vec![
2551 Some(r#"{"k":100000}"#),
2552 Some(r#"{"k":"s"}"#),
2553 ]));
2554 let expected = json_to_variant(&expected_json).unwrap();
2555
2556 assert_eq!(result_variant.len(), expected.len());
2557 for i in 0..result_variant.len() {
2558 assert_eq!(result_variant.is_null(i), expected.is_null(i));
2559 if !result_variant.is_null(i) {
2560 assert_eq!(result_variant.value(i), expected.value(i));
2561 }
2562 }
2563 }
2564
2565 fn create_depth_0_test_data() -> ArrayRef {
2568 let mut builder = crate::VariantArrayBuilder::new(3);
2569
2570 {
2572 let json_str = r#"{"x": 42}"#;
2573 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2574 if let Ok(variant_array) = json_to_variant(&string_array) {
2575 builder.append_variant(variant_array.value(0));
2576 } else {
2577 builder.append_null();
2578 }
2579 }
2580
2581 {
2583 let json_str = r#"{"x": "foo"}"#;
2584 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2585 if let Ok(variant_array) = json_to_variant(&string_array) {
2586 builder.append_variant(variant_array.value(0));
2587 } else {
2588 builder.append_null();
2589 }
2590 }
2591
2592 {
2594 let json_str = r#"{"y": 10}"#;
2595 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2596 if let Ok(variant_array) = json_to_variant(&string_array) {
2597 builder.append_variant(variant_array.value(0));
2598 } else {
2599 builder.append_null();
2600 }
2601 }
2602
2603 ArrayRef::from(builder.build())
2604 }
2605
2606 fn create_nested_path_test_data() -> ArrayRef {
2609 let mut builder = crate::VariantArrayBuilder::new(2);
2610
2611 {
2613 let json_str = r#"{"a": {"x": 55}, "b": 42}"#;
2614 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2615 if let Ok(variant_array) = json_to_variant(&string_array) {
2616 builder.append_variant(variant_array.value(0));
2617 } else {
2618 builder.append_null();
2619 }
2620 }
2621
2622 {
2624 let json_str = r#"{"a": {"x": "foo"}, "b": 42}"#;
2625 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2626 if let Ok(variant_array) = json_to_variant(&string_array) {
2627 builder.append_variant(variant_array.value(0));
2628 } else {
2629 builder.append_null();
2630 }
2631 }
2632
2633 ArrayRef::from(builder.build())
2634 }
2635
2636 fn create_depth_2_test_data() -> ArrayRef {
2639 let mut builder = crate::VariantArrayBuilder::new(3);
2640
2641 {
2643 let json_str = r#"{"a": {"b": {"x": 100}}}"#;
2644 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2645 if let Ok(variant_array) = json_to_variant(&string_array) {
2646 builder.append_variant(variant_array.value(0));
2647 } else {
2648 builder.append_null();
2649 }
2650 }
2651
2652 {
2654 let json_str = r#"{"a": {"b": {"x": "bar"}}}"#;
2655 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2656 if let Ok(variant_array) = json_to_variant(&string_array) {
2657 builder.append_variant(variant_array.value(0));
2658 } else {
2659 builder.append_null();
2660 }
2661 }
2662
2663 {
2665 let json_str = r#"{"a": {"b": {"y": 200}}}"#;
2666 let string_array: ArrayRef = Arc::new(StringArray::from(vec![json_str]));
2667 if let Ok(variant_array) = json_to_variant(&string_array) {
2668 builder.append_variant(variant_array.value(0));
2669 } else {
2670 builder.append_null();
2671 }
2672 }
2673
2674 ArrayRef::from(builder.build())
2675 }
2676
2677 fn create_depth_0_shredded_test_data_simple() -> ArrayRef {
2680 let (metadata, string_x_value) = {
2682 let mut builder = parquet_variant::VariantBuilder::new();
2683 let mut obj = builder.new_object();
2684 obj.insert("x", Variant::from("foo"));
2685 obj.finish();
2686 builder.finish()
2687 };
2688
2689 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 2));
2691
2692 let empty_object_value = {
2696 let mut builder = parquet_variant::VariantBuilder::new();
2697 let obj = builder.new_object();
2698 obj.finish();
2699 let (_, value) = builder.finish();
2700 value
2701 };
2702
2703 let value_array = BinaryViewArray::from(vec![
2704 Some(empty_object_value.as_slice()), Some(string_x_value.as_slice()), ]);
2707
2708 let x_field_typed_value = Int32Array::from(vec![Some(42), None]);
2710
2711 let x_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
2713 x_field_typed_value,
2714 ) as ArrayRef);
2715
2716 let typed_value_fields = Fields::from(vec![Field::new(
2718 "x",
2719 x_field_shredded.data_type().clone(),
2720 true,
2721 )]);
2722 let typed_value_struct = StructArray::try_new(
2723 typed_value_fields,
2724 vec![ArrayRef::from(x_field_shredded)],
2725 None,
2726 )
2727 .unwrap();
2728
2729 ArrayRef::from(VariantArray::from_parts(
2731 Arc::new(metadata_array),
2732 Arc::new(value_array),
2733 Some(Arc::new(typed_value_struct)),
2734 None,
2735 ))
2736 }
2737
2738 fn create_depth_1_shredded_test_data_working() -> ArrayRef {
2743 let (metadata, _) = {
2745 let mut builder = parquet_variant::VariantBuilder::new();
2747 let mut obj = builder.new_object();
2748
2749 let mut a_obj = obj.new_object("a");
2751 a_obj.insert("x", Variant::Int32(55));
2752 a_obj.finish();
2753
2754 obj.insert("b", Variant::Int32(42));
2755 obj.finish();
2756 builder.finish()
2757 };
2758
2759 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 2));
2760
2761 let empty_object_value = {
2764 let mut builder = parquet_variant::VariantBuilder::new();
2765 let obj = builder.new_object();
2766 obj.finish();
2767 let (_, value) = builder.finish();
2768 value
2769 };
2770
2771 let row1_fallback = {
2774 let mut builder = parquet_variant::VariantBuilder::new();
2775 let mut obj = builder.new_object();
2776 obj.insert("fallback", Variant::from("data"));
2777 obj.finish();
2778 let (_, value) = builder.finish();
2779 value
2780 };
2781
2782 let value_array = BinaryViewArray::from(vec![
2783 Some(empty_object_value.as_slice()), Some(row1_fallback.as_slice()), ]);
2786
2787 let x_typed_value = Int32Array::from(vec![Some(55), None]);
2790 let x_field_shredded =
2791 ShreddedVariantFieldArray::perfectly_shredded(Arc::new(x_typed_value) as ArrayRef);
2792
2793 let a_value_data = {
2798 let mut builder = parquet_variant::VariantBuilder::new();
2799 let obj = builder.new_object();
2800 obj.finish();
2801 let (_, value) = builder.finish();
2802 value
2803 };
2804 let a_value_array = BinaryViewArray::from(vec![
2805 None, Some(a_value_data.as_slice()), ]);
2808
2809 let a_inner_fields = Fields::from(vec![Field::new(
2810 "x",
2811 x_field_shredded.data_type().clone(),
2812 true,
2813 )]);
2814 let a_inner_typed_value = Arc::new(
2815 StructArray::try_new(a_inner_fields, vec![ArrayRef::from(x_field_shredded)], None)
2816 .unwrap(),
2817 ) as ArrayRef;
2818 let a_field_shredded = ShreddedVariantFieldArray::from_parts(
2819 Arc::new(a_value_array),
2820 Some(a_inner_typed_value),
2821 None,
2822 );
2823
2824 let typed_value_fields = Fields::from(vec![Field::new(
2826 "a",
2827 a_field_shredded.data_type().clone(),
2828 true,
2829 )]);
2830 let typed_value_struct = StructArray::try_new(
2831 typed_value_fields,
2832 vec![ArrayRef::from(a_field_shredded)],
2833 None,
2834 )
2835 .unwrap();
2836
2837 ArrayRef::from(VariantArray::from_parts(
2839 Arc::new(metadata_array),
2840 Arc::new(value_array),
2841 Some(Arc::new(typed_value_struct)),
2842 None,
2843 ))
2844 }
2845
2846 fn create_depth_2_shredded_test_data_working() -> ArrayRef {
2852 let (metadata, _) = {
2854 let mut builder = parquet_variant::VariantBuilder::new();
2856 let mut obj = builder.new_object();
2857
2858 let mut a_obj = obj.new_object("a");
2860 let mut b_obj = a_obj.new_object("b");
2861 b_obj.insert("x", Variant::Int32(100));
2862 b_obj.finish();
2863 a_obj.finish();
2864
2865 obj.finish();
2866 builder.finish()
2867 };
2868
2869 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 3));
2870
2871 let empty_object_value = {
2873 let mut builder = parquet_variant::VariantBuilder::new();
2874 let obj = builder.new_object();
2875 obj.finish();
2876 let (_, value) = builder.finish();
2877 value
2878 };
2879
2880 let value_array = BinaryViewArray::from(vec![
2882 Some(empty_object_value.as_slice()), Some(empty_object_value.as_slice()), Some(empty_object_value.as_slice()), ]);
2886
2887 let x_typed_value = Int32Array::from(vec![Some(100), None, None]);
2891 let x_field_shredded =
2892 ShreddedVariantFieldArray::perfectly_shredded(Arc::new(x_typed_value) as ArrayRef);
2893
2894 let b_value_data = {
2896 let mut builder = parquet_variant::VariantBuilder::new();
2897 let obj = builder.new_object();
2898 obj.finish();
2899 let (_, value) = builder.finish();
2900 value
2901 };
2902 let b_value_array = BinaryViewArray::from(vec![
2903 None, Some(b_value_data.as_slice()), Some(b_value_data.as_slice()), ]);
2907
2908 let b_inner_fields = Fields::from(vec![Field::new(
2909 "x",
2910 x_field_shredded.data_type().clone(),
2911 true,
2912 )]);
2913 let b_inner_typed_value = Arc::new(
2914 StructArray::try_new(b_inner_fields, vec![ArrayRef::from(x_field_shredded)], None)
2915 .unwrap(),
2916 ) as ArrayRef;
2917 let b_field_shredded = ShreddedVariantFieldArray::from_parts(
2918 Arc::new(b_value_array),
2919 Some(b_inner_typed_value),
2920 None,
2921 );
2922
2923 let a_value_data = {
2925 let mut builder = parquet_variant::VariantBuilder::new();
2926 let obj = builder.new_object();
2927 obj.finish();
2928 let (_, value) = builder.finish();
2929 value
2930 };
2931 let a_value_array = BinaryViewArray::from(vec![
2932 None, Some(a_value_data.as_slice()), Some(a_value_data.as_slice()), ]);
2936
2937 let a_inner_fields = Fields::from(vec![Field::new(
2938 "b",
2939 b_field_shredded.data_type().clone(),
2940 true,
2941 )]);
2942 let a_inner_typed_value = Arc::new(
2943 StructArray::try_new(a_inner_fields, vec![ArrayRef::from(b_field_shredded)], None)
2944 .unwrap(),
2945 ) as ArrayRef;
2946 let a_field_shredded = ShreddedVariantFieldArray::from_parts(
2947 Arc::new(a_value_array),
2948 Some(a_inner_typed_value),
2949 None,
2950 );
2951
2952 let typed_value_fields = Fields::from(vec![Field::new(
2954 "a",
2955 a_field_shredded.data_type().clone(),
2956 true,
2957 )]);
2958 let typed_value_struct = StructArray::try_new(
2959 typed_value_fields,
2960 vec![ArrayRef::from(a_field_shredded)],
2961 None,
2962 )
2963 .unwrap();
2964
2965 ArrayRef::from(VariantArray::from_parts(
2967 Arc::new(metadata_array),
2968 Arc::new(value_array),
2969 Some(Arc::new(typed_value_struct)),
2970 None,
2971 ))
2972 }
2973
2974 #[test]
2975 fn test_field_path_non_struct_returns_missing_path_step() {
2976 let variant_array = perfectly_shredded_int32_variant_array();
2978
2979 for safe in [true, false] {
2980 let options = GetOptions {
2981 path: VariantPath::try_from("nonexistent_field").unwrap(),
2982 as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))),
2983 cast_options: CastOptions {
2984 safe,
2985 ..Default::default()
2986 },
2987 };
2988
2989 let result_array = variant_get(&variant_array, options).unwrap();
2990 assert_eq!(result_array.len(), 3);
2991 assert!(result_array.is_null(0));
2992 assert!(result_array.is_null(1));
2993 assert!(result_array.is_null(2));
2994 }
2995 }
2996
2997 #[test]
2998 fn test_strict_cast_options_index_on_non_list_returns_null() {
2999 use arrow::compute::CastOptions;
3000 use arrow::datatypes::{DataType, Field};
3001 use parquet_variant::VariantPath;
3002 use std::sync::Arc;
3003
3004 let variant_array = perfectly_shredded_int32_variant_array();
3006 let options = GetOptions {
3007 path: VariantPath::from(0),
3008 as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))),
3009 cast_options: CastOptions {
3010 safe: false,
3011 ..Default::default()
3012 },
3013 };
3014
3015 let variant_array_ref: Arc<dyn Array> = variant_array.clone();
3016 let result = variant_get(&variant_array_ref, options).unwrap();
3017
3018 assert_eq!(result.len(), 3);
3019 assert!(result.is_null(0));
3020 assert!(result.is_null(1));
3021 assert!(result.is_null(2));
3022 }
3023
3024 #[test]
3025 fn test_error_message_boolean_type_display() {
3026 let mut builder = VariantArrayBuilder::new(1);
3027 builder.append_variant(Variant::from("abcd"));
3028 let variant_array: ArrayRef = ArrayRef::from(builder.build());
3029
3030 let options = GetOptions {
3032 path: VariantPath::default(),
3033 as_type: Some(Arc::new(Field::new("result", DataType::Boolean, true))),
3034 cast_options: CastOptions {
3035 safe: false,
3036 ..Default::default()
3037 },
3038 };
3039
3040 let err = variant_get(&variant_array, options).unwrap_err();
3041 let msg = err.to_string();
3042 assert!(msg.contains("Failed to extract primitive of type Boolean"));
3043 }
3044
3045 #[test]
3046 fn test_error_message_numeric_type_display() {
3047 let mut builder = VariantArrayBuilder::new(1);
3048 builder.append_variant(Variant::from("abcd"));
3049 let variant_array: ArrayRef = ArrayRef::from(builder.build());
3050
3051 let options = GetOptions {
3053 path: VariantPath::default(),
3054 as_type: Some(Arc::new(Field::new("result", DataType::Float32, true))),
3055 cast_options: CastOptions {
3056 safe: false,
3057 ..Default::default()
3058 },
3059 };
3060
3061 let err = variant_get(&variant_array, options).unwrap_err();
3062 let msg = err.to_string();
3063 assert!(msg.contains("Failed to extract primitive of type Float32"));
3064 }
3065
3066 #[test]
3067 fn test_error_message_temporal_type_display() {
3068 let mut builder = VariantArrayBuilder::new(1);
3069 builder.append_variant(Variant::BooleanFalse);
3070 let variant_array: ArrayRef = ArrayRef::from(builder.build());
3071
3072 let options = GetOptions {
3074 path: VariantPath::default(),
3075 as_type: Some(Arc::new(Field::new(
3076 "result",
3077 DataType::Timestamp(TimeUnit::Nanosecond, None),
3078 true,
3079 ))),
3080 cast_options: CastOptions {
3081 safe: false,
3082 ..Default::default()
3083 },
3084 };
3085
3086 let err = variant_get(&variant_array, options).unwrap_err();
3087 let msg = err.to_string();
3088 assert!(msg.contains("Failed to extract primitive of type Timestamp(ns)"));
3089 }
3090
3091 #[test]
3092 fn test_null_buffer_union_for_shredded_paths() {
3093 let variant_array = create_depth_1_shredded_test_data_working();
3102
3103 let options = GetOptions {
3108 path: VariantPath::try_from("a.x").unwrap(),
3109 as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))),
3110 cast_options: CastOptions::default(),
3111 };
3112
3113 let result = variant_get(&variant_array, options).unwrap();
3114
3115 assert_eq!(result.len(), variant_array.len());
3117
3118 assert!(!result.is_null(0), "Row 0 should have valid Int32 data");
3122 assert!(
3123 result.is_null(1),
3124 "Row 1 should be null due to type casting failure"
3125 );
3126
3127 let int32_result = result.as_any().downcast_ref::<Int32Array>().unwrap();
3129 assert_eq!(int32_result.value(0), 55); }
3131
3132 #[test]
3133 fn test_struct_null_mask_union_from_children() {
3134 let json_strings = vec![
3139 r#"{"a": 42, "b": "hello"}"#, r#"{"a": "world", "b": 100}"#, r#"{"a": 55, "b": 77}"#, ];
3143
3144 let string_array: Arc<dyn arrow::array::Array> = Arc::new(StringArray::from(json_strings));
3145 let variant_array = json_to_variant(&string_array).unwrap();
3146
3147 let struct_fields = Fields::from(vec![
3150 Field::new("a", DataType::Int32, true),
3151 Field::new("b", DataType::Int32, true),
3152 ]);
3153 let struct_type = DataType::Struct(struct_fields);
3154
3155 let options = GetOptions {
3156 path: VariantPath::default(), as_type: Some(Arc::new(Field::new("result", struct_type, true))),
3158 cast_options: CastOptions::default(),
3159 };
3160
3161 let variant_array_ref = ArrayRef::from(variant_array);
3162 let result = variant_get(&variant_array_ref, options).unwrap();
3163
3164 let struct_result = result.as_struct();
3166 assert_eq!(struct_result.len(), 3);
3167
3168 let field_a = struct_result
3170 .column(0)
3171 .as_any()
3172 .downcast_ref::<Int32Array>()
3173 .unwrap();
3174 let field_b = struct_result
3175 .column(1)
3176 .as_any()
3177 .downcast_ref::<Int32Array>()
3178 .unwrap();
3179
3180 assert!(!field_a.is_null(0));
3183 assert_eq!(field_a.value(0), 42);
3184 assert!(field_b.is_null(0)); assert!(field_a.is_null(1)); assert!(!field_b.is_null(1));
3189 assert_eq!(field_b.value(1), 100);
3190
3191 assert!(!field_a.is_null(2));
3193 assert_eq!(field_a.value(2), 55);
3194 assert!(!field_b.is_null(2));
3195 assert_eq!(field_b.value(2), 77);
3196
3197 assert!(!struct_result.is_null(0)); assert!(!struct_result.is_null(1)); assert!(!struct_result.is_null(2)); }
3204
3205 #[test]
3206 fn test_field_nullability_preservation() {
3207 let json_strings = vec![
3210 r#"{"x": 42}"#, r#"{"x": "not_a_number"}"#, r#"{"x": null}"#, r#"{"x": "hello"}"#, r#"{"y": 100}"#, r#"{"x": 127}"#, r#"{"x": 32767}"#, r#"{"x": 2147483647}"#, r#"{"x": 9223372036854775807}"#, ];
3220
3221 let string_array: Arc<dyn arrow::array::Array> = Arc::new(StringArray::from(json_strings));
3222 let variant_array = json_to_variant(&string_array).unwrap();
3223
3224 let nullable_field = Arc::new(Field::new("result", DataType::Int32, true));
3226 let options_nullable = GetOptions {
3227 path: VariantPath::try_from("x").unwrap(),
3228 as_type: Some(nullable_field.clone()),
3229 cast_options: CastOptions::default(),
3230 };
3231
3232 let variant_array_ref = ArrayRef::from(variant_array);
3233 let result_nullable = variant_get(&variant_array_ref, options_nullable).unwrap();
3234
3235 let int32_result = result_nullable
3237 .as_any()
3238 .downcast_ref::<Int32Array>()
3239 .unwrap();
3240 assert_eq!(int32_result.len(), 9);
3241
3242 assert!(!int32_result.is_null(0));
3244 assert_eq!(int32_result.value(0), 42);
3245
3246 assert!(int32_result.is_null(1));
3248
3249 assert!(int32_result.is_null(2));
3251
3252 assert!(int32_result.is_null(3));
3254
3255 assert!(int32_result.is_null(4));
3257
3258 assert!(!int32_result.is_null(5));
3261 assert_eq!(int32_result.value(5), 127);
3262
3263 assert!(!int32_result.is_null(6));
3266 assert_eq!(int32_result.value(6), 32767);
3267
3268 assert!(!int32_result.is_null(7));
3271 assert_eq!(int32_result.value(7), 2147483647);
3272
3273 assert!(int32_result.is_null(8));
3276
3277 let non_nullable_field = Arc::new(Field::new("result", DataType::Int32, false));
3279 let options_non_nullable = GetOptions {
3280 path: VariantPath::try_from("x").unwrap(),
3281 as_type: Some(non_nullable_field.clone()),
3282 cast_options: CastOptions::default(), };
3284
3285 let variant_array_2 = json_to_variant(&string_array).unwrap();
3287 let variant_array_ref_2 = ArrayRef::from(variant_array_2);
3288 let result_non_nullable = variant_get(&variant_array_ref_2, options_non_nullable).unwrap();
3289 let int32_result_2 = result_non_nullable
3290 .as_any()
3291 .downcast_ref::<Int32Array>()
3292 .unwrap();
3293
3294 assert_eq!(int32_result_2.len(), 9);
3296
3297 assert!(!int32_result_2.is_null(0));
3299 assert_eq!(int32_result_2.value(0), 42);
3300
3301 assert!(int32_result_2.is_null(1)); assert!(int32_result_2.is_null(2)); assert!(int32_result_2.is_null(3)); assert!(int32_result_2.is_null(4)); assert!(!int32_result_2.is_null(5)); assert_eq!(int32_result_2.value(5), 127);
3311 assert!(!int32_result_2.is_null(6)); assert_eq!(int32_result_2.value(6), 32767);
3313 assert!(!int32_result_2.is_null(7)); assert_eq!(int32_result_2.value(7), 2147483647);
3315
3316 assert!(int32_result_2.is_null(8)); }
3319
3320 #[test]
3321 fn test_struct_extraction_subset_superset_schema_perfectly_shredded() {
3322 let variant_array = create_comprehensive_shredded_variant();
3324
3325 let struct_fields = Fields::from(vec![
3327 Field::new("a", DataType::Int32, true),
3328 Field::new("b", DataType::Int32, true),
3329 Field::new("d", DataType::Int32, true),
3330 ]);
3331 let struct_type = DataType::Struct(struct_fields);
3332
3333 let options = GetOptions {
3334 path: VariantPath::default(),
3335 as_type: Some(Arc::new(Field::new("result", struct_type, true))),
3336 cast_options: CastOptions::default(),
3337 };
3338
3339 let result = variant_get(&variant_array, options).unwrap();
3340
3341 let struct_result = result.as_any().downcast_ref::<StructArray>().unwrap();
3343 assert_eq!(struct_result.len(), 5);
3344 assert_eq!(struct_result.num_columns(), 3);
3345
3346 let field_a = struct_result
3347 .column(0)
3348 .as_any()
3349 .downcast_ref::<Int32Array>()
3350 .unwrap();
3351 let field_b = struct_result
3352 .column(1)
3353 .as_any()
3354 .downcast_ref::<Int32Array>()
3355 .unwrap();
3356 let field_d = struct_result
3357 .column(2)
3358 .as_any()
3359 .downcast_ref::<Int32Array>()
3360 .unwrap();
3361
3362 assert!(!struct_result.is_null(0));
3364 assert_eq!(field_a.value(0), 1);
3365 assert_eq!(field_b.value(0), 2);
3366 assert!(field_d.is_null(0)); assert!(struct_result.is_null(1));
3370
3371 assert!(!struct_result.is_null(2));
3373 assert!(field_a.is_null(2)); assert_eq!(field_b.value(2), 2);
3375 assert!(field_d.is_null(2)); assert!(!struct_result.is_null(3));
3379 assert_eq!(field_a.value(3), 1);
3380 assert!(field_b.is_null(3)); assert!(field_d.is_null(3)); assert!(!struct_result.is_null(4));
3385 assert!(field_a.is_null(4)); assert!(field_b.is_null(4)); assert!(field_d.is_null(4)); }
3389
3390 #[test]
3391 fn test_nested_struct_extraction_perfectly_shredded() {
3392 let variant_array = create_comprehensive_nested_shredded_variant();
3394 println!("variant_array: {variant_array:?}");
3395
3396 let inner_field = Field::new("inner", DataType::Int32, true);
3398 let inner_type = DataType::Struct(Fields::from(vec![inner_field]));
3399 let outer_field = Field::new("outer", inner_type, true);
3400 let result_type = DataType::Struct(Fields::from(vec![outer_field]));
3401
3402 let options = GetOptions {
3403 path: VariantPath::default(),
3404 as_type: Some(Arc::new(Field::new("result", result_type, true))),
3405 cast_options: CastOptions::default(),
3406 };
3407
3408 let result = variant_get(&variant_array, options).unwrap();
3409 println!("result: {result:?}");
3410
3411 let outer_struct = result.as_any().downcast_ref::<StructArray>().unwrap();
3413 assert_eq!(outer_struct.len(), 4);
3414 assert_eq!(outer_struct.num_columns(), 1);
3415
3416 let inner_struct = outer_struct
3418 .column(0)
3419 .as_any()
3420 .downcast_ref::<StructArray>()
3421 .unwrap();
3422 assert_eq!(inner_struct.num_columns(), 1);
3423
3424 let leaf_field = inner_struct
3426 .column(0)
3427 .as_any()
3428 .downcast_ref::<Int32Array>()
3429 .unwrap();
3430
3431 assert!(!outer_struct.is_null(0));
3433 assert!(!inner_struct.is_null(0));
3434 assert_eq!(leaf_field.value(0), 42);
3435
3436 assert!(!outer_struct.is_null(1));
3438 assert!(!inner_struct.is_null(1)); assert!(leaf_field.is_null(1)); assert!(!outer_struct.is_null(2));
3443 assert!(inner_struct.is_null(2)); assert!(outer_struct.is_null(3));
3447 }
3448
3449 #[test]
3450 fn test_path_based_null_masks_one_step() {
3451 let variant_array = create_comprehensive_nested_shredded_variant();
3453
3454 let path = VariantPath::try_from("outer").unwrap();
3456 let inner_field = Field::new("inner", DataType::Int32, true);
3457 let result_type = DataType::Struct(Fields::from(vec![inner_field]));
3458
3459 let options = GetOptions {
3460 path,
3461 as_type: Some(Arc::new(Field::new("result", result_type, true))),
3462 cast_options: CastOptions::default(),
3463 };
3464
3465 let result = variant_get(&variant_array, options).unwrap();
3466
3467 let outer_result = result.as_any().downcast_ref::<StructArray>().unwrap();
3469 assert_eq!(outer_result.len(), 4);
3470 assert_eq!(outer_result.num_columns(), 1);
3471
3472 let inner_field = outer_result
3474 .column(0)
3475 .as_any()
3476 .downcast_ref::<Int32Array>()
3477 .unwrap();
3478
3479 assert!(!outer_result.is_null(0));
3481 assert_eq!(inner_field.value(0), 42);
3482
3483 assert!(!outer_result.is_null(1));
3485 assert!(inner_field.is_null(1));
3486
3487 assert!(outer_result.is_null(2));
3489
3490 assert!(outer_result.is_null(3));
3492 }
3493
3494 #[test]
3495 fn test_path_based_null_masks_two_steps() {
3496 let variant_array = create_comprehensive_nested_shredded_variant();
3498
3499 let path = VariantPath::try_from("outer").unwrap().join("inner");
3501
3502 let options = GetOptions {
3503 path,
3504 as_type: Some(Arc::new(Field::new("result", DataType::Int32, true))),
3505 cast_options: CastOptions::default(),
3506 };
3507
3508 let result = variant_get(&variant_array, options).unwrap();
3509
3510 let int_result = result.as_any().downcast_ref::<Int32Array>().unwrap();
3512 assert_eq!(int_result.len(), 4);
3513
3514 assert!(!int_result.is_null(0));
3516 assert_eq!(int_result.value(0), 42);
3517
3518 assert!(int_result.is_null(1));
3520
3521 assert!(int_result.is_null(2));
3523
3524 assert!(int_result.is_null(3));
3526 }
3527
3528 #[test]
3529 fn test_struct_extraction_mixed_and_unshredded() {
3530 let variant_array = create_mixed_and_unshredded_variant();
3532
3533 let struct_fields = Fields::from(vec![
3535 Field::new("x", DataType::Int32, true),
3536 Field::new("y", DataType::Int32, true),
3537 ]);
3538 let struct_type = DataType::Struct(struct_fields);
3539
3540 let options = GetOptions {
3541 path: VariantPath::default(),
3542 as_type: Some(Arc::new(Field::new("result", struct_type, true))),
3543 cast_options: CastOptions::default(),
3544 };
3545
3546 let result = variant_get(&variant_array, options).unwrap();
3547
3548 let struct_result = result.as_any().downcast_ref::<StructArray>().unwrap();
3550 assert_eq!(struct_result.len(), 4);
3551 assert_eq!(struct_result.num_columns(), 2);
3552
3553 let field_x = struct_result
3554 .column(0)
3555 .as_any()
3556 .downcast_ref::<Int32Array>()
3557 .unwrap();
3558 let field_y = struct_result
3559 .column(1)
3560 .as_any()
3561 .downcast_ref::<Int32Array>()
3562 .unwrap();
3563
3564 assert_eq!(field_x.value(0), 1);
3566 assert_eq!(field_y.value(0), 42);
3567
3568 assert_eq!(field_x.value(1), 2);
3570 assert!(field_y.is_null(1));
3571
3572 assert_eq!(field_x.value(2), 3);
3574 assert!(field_y.is_null(2));
3575
3576 assert!(struct_result.is_null(3));
3578 }
3579
3580 #[test]
3581 fn test_struct_row_builder_handles_unshredded_nested_structs() {
3582 let json_strings = vec![
3584 r#"{"outer": {"inner": 42}}"#,
3585 r#"{"outer": {"inner": 100}}"#,
3586 ];
3587 let string_array: Arc<dyn Array> = Arc::new(StringArray::from(json_strings));
3588 let variant_array = json_to_variant(&string_array).unwrap();
3589
3590 let inner_fields = Fields::from(vec![Field::new("inner", DataType::Int32, true)]);
3592 let inner_struct_type = DataType::Struct(inner_fields);
3593 let outer_fields = Fields::from(vec![Field::new("outer", inner_struct_type, true)]);
3594 let outer_struct_type = DataType::Struct(outer_fields);
3595
3596 let options = GetOptions {
3597 path: VariantPath::default(),
3598 as_type: Some(Arc::new(Field::new("result", outer_struct_type, true))),
3599 cast_options: CastOptions::default(),
3600 };
3601
3602 let variant_array_ref = ArrayRef::from(variant_array);
3603 let result = variant_get(&variant_array_ref, options).unwrap();
3604
3605 let outer_struct = result.as_struct();
3606 assert_eq!(outer_struct.len(), 2);
3607 assert_eq!(outer_struct.num_columns(), 1);
3608
3609 let inner_struct = outer_struct.column(0).as_struct();
3610 assert_eq!(inner_struct.num_columns(), 1);
3611
3612 let inner_values = inner_struct
3613 .column(0)
3614 .as_any()
3615 .downcast_ref::<Int32Array>()
3616 .unwrap();
3617 assert_eq!(inner_values.value(0), 42);
3618 assert_eq!(inner_values.value(1), 100);
3619 }
3620
3621 #[test]
3622 fn test_unshredded_struct_safe_cast_and_field_mismatches() {
3623 let json_strings = vec![r#"{"a": 1, "b": 2, "extra": 3}"#, "123", "{}"];
3624 let string_array: Arc<dyn Array> = Arc::new(StringArray::from(json_strings));
3625 let variant_array_ref = ArrayRef::from(json_to_variant(&string_array).unwrap());
3626
3627 let struct_fields = Fields::from(vec![
3628 Field::new("a", DataType::Int32, true),
3629 Field::new("b", DataType::Int32, true),
3630 ]);
3631 let options = GetOptions {
3632 path: VariantPath::default(),
3633 as_type: Some(Arc::new(Field::new(
3634 "result",
3635 DataType::Struct(struct_fields),
3636 true,
3637 ))),
3638 cast_options: CastOptions::default(),
3639 };
3640
3641 let result = variant_get(&variant_array_ref, options).unwrap();
3642 let struct_result = result.as_struct();
3643 let field_a = struct_result
3644 .column(0)
3645 .as_primitive::<arrow::datatypes::Int32Type>();
3646 let field_b = struct_result
3647 .column(1)
3648 .as_primitive::<arrow::datatypes::Int32Type>();
3649
3650 assert!(!struct_result.is_null(0));
3653 assert_eq!(field_a.value(0), 1);
3654 assert_eq!(field_b.value(0), 2);
3655
3656 assert!(struct_result.is_null(1));
3658 assert!(field_a.is_null(1));
3659 assert!(field_b.is_null(1));
3660
3661 assert!(!struct_result.is_null(2));
3663 assert!(field_a.is_null(2));
3664 assert!(field_b.is_null(2));
3665 }
3666
3667 #[test]
3668 fn test_unshredded_struct_missing_non_nullable_field_errors() {
3669 let string_array: Arc<dyn Array> = Arc::new(StringArray::from(vec![r#"{"a": 1}"#]));
3670 let variant_array_ref = ArrayRef::from(json_to_variant(&string_array).unwrap());
3671
3672 let struct_fields = Fields::from(vec![
3673 Field::new("a", DataType::Int32, false),
3674 Field::new("missing", DataType::Int32, false),
3675 ]);
3676 let options = GetOptions {
3677 path: VariantPath::default(),
3678 as_type: Some(Arc::new(Field::new(
3679 "result",
3680 DataType::Struct(struct_fields),
3681 true,
3682 ))),
3683 cast_options: CastOptions::default(),
3684 };
3685
3686 let err = variant_get(&variant_array_ref, options).unwrap_err();
3687 assert!(
3688 err.to_string()
3689 .contains("unmasked nulls for non-nullable StructArray field \"missing\""),
3690 "unexpected error: {err}"
3691 );
3692 }
3693
3694 #[test]
3695 fn test_unshredded_struct_strict_cast_non_object_errors() {
3696 let json_strings = vec![r#"{"a": 1, "b": 2}"#, "123"];
3697 let string_array: Arc<dyn Array> = Arc::new(StringArray::from(json_strings));
3698 let variant_array_ref = ArrayRef::from(json_to_variant(&string_array).unwrap());
3699
3700 let struct_fields = Fields::from(vec![
3701 Field::new("a", DataType::Int32, true),
3702 Field::new("b", DataType::Int32, true),
3703 ]);
3704 let options = GetOptions {
3705 path: VariantPath::default(),
3706 as_type: Some(Arc::new(Field::new(
3707 "result",
3708 DataType::Struct(struct_fields),
3709 true,
3710 ))),
3711 cast_options: CastOptions {
3712 safe: false,
3713 ..Default::default()
3714 },
3715 };
3716
3717 let err = variant_get(&variant_array_ref, options).unwrap_err();
3718 assert!(
3719 err.to_string()
3720 .contains("Failed to extract struct from variant")
3721 );
3722 }
3723
3724 fn create_comprehensive_shredded_variant() -> ArrayRef {
3727 let (metadata, _) = {
3728 let mut builder = parquet_variant::VariantBuilder::new();
3729 let obj = builder.new_object();
3730 obj.finish();
3731 builder.finish()
3732 };
3733
3734 let nulls = NullBuffer::from(vec![
3736 true, false, true, true, true, ]);
3742
3743 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 5));
3744
3745 let a_field_typed_value = Int32Array::from(vec![Some(1), None, None, Some(1), None]);
3748 let a_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
3749 a_field_typed_value,
3750 ) as ArrayRef);
3751
3752 let b_field_typed_value = Int32Array::from(vec![Some(2), None, Some(2), None, None]);
3754 let b_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
3755 b_field_typed_value,
3756 ) as ArrayRef);
3757
3758 let c_field_typed_value = Int32Array::from(vec![Some(3), None, None, None, None]);
3760 let c_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
3761 c_field_typed_value,
3762 ) as ArrayRef);
3763
3764 let typed_value_fields = Fields::from(vec![
3766 Field::new("a", a_field_shredded.data_type().clone(), true),
3767 Field::new("b", b_field_shredded.data_type().clone(), true),
3768 Field::new("c", c_field_shredded.data_type().clone(), true),
3769 ]);
3770 let typed_value_struct = StructArray::try_new(
3771 typed_value_fields,
3772 vec![
3773 ArrayRef::from(a_field_shredded),
3774 ArrayRef::from(b_field_shredded),
3775 ArrayRef::from(c_field_shredded),
3776 ],
3777 None,
3778 )
3779 .unwrap();
3780
3781 ArrayRef::from(VariantArray::perfectly_shredded(
3783 Arc::new(metadata_array),
3784 Arc::new(typed_value_struct),
3785 Some(nulls),
3786 ))
3787 }
3788
3789 fn create_comprehensive_nested_shredded_variant() -> ArrayRef {
3794 let inner_typed_value = Int32Array::from(vec![Some(42), None, None, None]); let inner =
3798 ShreddedVariantFieldArray::perfectly_shredded(Arc::new(inner_typed_value) as ArrayRef);
3799
3800 let outer_typed_value_nulls = NullBuffer::from(vec![
3801 true, false, false, false, ]);
3806 let outer_typed_value = StructArrayBuilder::new()
3807 .with_field("inner", ArrayRef::from(inner), false)
3808 .with_nulls(outer_typed_value_nulls)
3809 .build();
3810
3811 let outer =
3812 ShreddedVariantFieldArray::perfectly_shredded(Arc::new(outer_typed_value) as ArrayRef);
3813
3814 let typed_value_nulls = NullBuffer::from(vec![
3815 true, true, false, false, ]);
3820 let typed_value = StructArrayBuilder::new()
3821 .with_field("outer", ArrayRef::from(outer), false)
3822 .with_nulls(typed_value_nulls)
3823 .build();
3824
3825 let metadata_array =
3827 BinaryViewArray::from_iter_values(std::iter::repeat_n(EMPTY_VARIANT_METADATA_BYTES, 4));
3828 let nulls = NullBuffer::from(vec![
3829 true, true, true, false, ]);
3834 ArrayRef::from(VariantArray::perfectly_shredded(
3835 Arc::new(metadata_array),
3836 Arc::new(typed_value),
3837 Some(nulls),
3838 ))
3839 }
3840
3841 fn create_mixed_and_unshredded_variant() -> ArrayRef {
3844 let (metadata, y_field_value) = {
3849 let mut builder = parquet_variant::VariantBuilder::new();
3850 let mut obj = builder.new_object();
3851 obj.insert("y", Variant::from(42));
3852 obj.finish();
3853 builder.finish()
3854 };
3855
3856 let metadata_array = BinaryViewArray::from_iter_values(std::iter::repeat_n(&metadata, 4));
3857
3858 let empty_object_value = {
3865 let mut builder = parquet_variant::VariantBuilder::new();
3866 builder.new_object().finish();
3867 let (_, value) = builder.finish();
3868 value
3869 };
3870
3871 let y_null_value = {
3872 let mut builder = parquet_variant::VariantBuilder::new();
3873 builder.new_object().with_field("y", Variant::Null).finish();
3874 let (_, value) = builder.finish();
3875 value
3876 };
3877
3878 let value_array = BinaryViewArray::from(vec![
3879 Some(y_field_value.as_slice()), Some(empty_object_value.as_slice()), Some(y_null_value.as_slice()), Some(empty_object_value.as_slice()), ]);
3884
3885 let x_field_typed_value = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(0)]);
3888 let x_field_shredded = ShreddedVariantFieldArray::perfectly_shredded(Arc::new(
3889 x_field_typed_value,
3890 ) as ArrayRef);
3891
3892 let typed_value_struct = StructArrayBuilder::new()
3894 .with_field("x", ArrayRef::from(x_field_shredded), false)
3895 .build();
3896
3897 let variant_nulls = NullBuffer::from(vec![true, true, true, false]); ArrayRef::from(VariantArray::from_parts(
3901 Arc::new(metadata_array),
3902 Arc::new(value_array),
3903 Some(Arc::new(typed_value_struct)),
3904 Some(variant_nulls),
3905 ))
3906 }
3907
3908 #[test]
3909 fn get_decimal32_rescaled_to_scale2() {
3910 let mut builder = crate::VariantArrayBuilder::new(5);
3912 builder.append_variant(VariantDecimal4::try_new(1234, 2).unwrap().into()); builder.append_variant(VariantDecimal4::try_new(1234, 3).unwrap().into()); builder.append_variant(VariantDecimal4::try_new(1234, 0).unwrap().into()); builder.append_null();
3916 builder.append_variant(
3917 VariantDecimal8::try_new((VariantDecimal4::MAX_UNSCALED_VALUE as i64) + 1, 3)
3918 .unwrap()
3919 .into(),
3920 ); let variant_array: ArrayRef = ArrayRef::from(builder.build());
3922
3923 let field = Field::new("result", DataType::Decimal32(9, 2), true);
3924 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
3925 let result = variant_get(&variant_array, options).unwrap();
3926 let result = result.as_any().downcast_ref::<Decimal32Array>().unwrap();
3927
3928 assert_eq!(result.precision(), 9);
3929 assert_eq!(result.scale(), 2);
3930 assert_eq!(result.value(0), 1234);
3931 assert_eq!(result.value(1), 123);
3932 assert_eq!(result.value(2), 123400);
3933 assert!(result.is_null(3));
3934 assert_eq!(
3935 result.value(4),
3936 VariantDecimal4::MAX_UNSCALED_VALUE / 10 + 1
3937 ); }
3939
3940 #[test]
3941 fn get_decimal32_scale_down_rounding() {
3942 let mut builder = crate::VariantArrayBuilder::new(7);
3943 builder.append_variant(VariantDecimal4::try_new(1235, 0).unwrap().into());
3944 builder.append_variant(VariantDecimal4::try_new(1245, 0).unwrap().into());
3945 builder.append_variant(VariantDecimal4::try_new(-1235, 0).unwrap().into());
3946 builder.append_variant(VariantDecimal4::try_new(-1245, 0).unwrap().into());
3947 builder.append_variant(VariantDecimal4::try_new(1235, 2).unwrap().into()); builder.append_variant(VariantDecimal4::try_new(1235, 3).unwrap().into()); builder.append_variant(VariantDecimal4::try_new(5235, 3).unwrap().into()); let variant_array: ArrayRef = ArrayRef::from(builder.build());
3951
3952 let field = Field::new("result", DataType::Decimal32(9, -1), true);
3953 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
3954 let result = variant_get(&variant_array, options).unwrap();
3955 let result = result.as_any().downcast_ref::<Decimal32Array>().unwrap();
3956
3957 assert_eq!(result.precision(), 9);
3958 assert_eq!(result.scale(), -1);
3959 assert_eq!(result.value(0), 124);
3960 assert_eq!(result.value(1), 125);
3961 assert_eq!(result.value(2), -124);
3962 assert_eq!(result.value(3), -125);
3963 assert_eq!(result.value(4), 1);
3964 assert!(result.is_valid(5));
3965 assert_eq!(result.value(5), 0);
3966 assert_eq!(result.value(6), 1);
3967 }
3968
3969 #[test]
3970 fn get_decimal32_large_scale_reduction() {
3971 let mut builder = crate::VariantArrayBuilder::new(2);
3972 builder.append_variant(
3973 VariantDecimal4::try_new(-VariantDecimal4::MAX_UNSCALED_VALUE, 0)
3974 .unwrap()
3975 .into(),
3976 );
3977 builder.append_variant(
3978 VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0)
3979 .unwrap()
3980 .into(),
3981 );
3982 let variant_array: ArrayRef = ArrayRef::from(builder.build());
3983
3984 let field = Field::new("result", DataType::Decimal32(9, -9), true);
3985 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
3986 let result = variant_get(&variant_array, options).unwrap();
3987 let result = result.as_any().downcast_ref::<Decimal32Array>().unwrap();
3988
3989 assert_eq!(result.precision(), 9);
3990 assert_eq!(result.scale(), -9);
3991 assert_eq!(result.value(0), -1);
3992 assert_eq!(result.value(1), 1);
3993
3994 let field = Field::new("result", DataType::Decimal32(9, -10), true);
3995 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
3996 let result = variant_get(&variant_array, options).unwrap();
3997 let result = result.as_any().downcast_ref::<Decimal32Array>().unwrap();
3998
3999 assert_eq!(result.precision(), 9);
4000 assert_eq!(result.scale(), -10);
4001 assert!(result.is_valid(0));
4002 assert_eq!(result.value(0), 0);
4003 assert!(result.is_valid(1));
4004 assert_eq!(result.value(1), 0);
4005 }
4006
4007 #[test]
4008 fn get_decimal32_precision_overflow_safe() {
4009 let mut builder = crate::VariantArrayBuilder::new(2);
4011 builder.append_variant(
4012 VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0)
4013 .unwrap()
4014 .into(),
4015 );
4016 builder.append_variant(
4017 VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 9)
4018 .unwrap()
4019 .into(),
4020 ); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4022
4023 let field = Field::new("result", DataType::Decimal32(2, 2), true);
4024 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4025 let result = variant_get(&variant_array, options).unwrap();
4026 let result = result.as_any().downcast_ref::<Decimal32Array>().unwrap();
4027
4028 assert!(result.is_null(0));
4029 assert!(result.is_null(1)); }
4031
4032 #[test]
4033 fn get_decimal32_precision_overflow_unsafe_errors() {
4034 let mut builder = crate::VariantArrayBuilder::new(1);
4035 builder.append_variant(
4036 VariantDecimal4::try_new(VariantDecimal4::MAX_UNSCALED_VALUE, 0)
4037 .unwrap()
4038 .into(),
4039 );
4040 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4041
4042 let field = Field::new("result", DataType::Decimal32(9, 2), true);
4043 let cast_options = CastOptions {
4044 safe: false,
4045 ..Default::default()
4046 };
4047 let options = GetOptions::new()
4048 .with_as_type(Some(FieldRef::from(field)))
4049 .with_cast_options(cast_options);
4050 let err = variant_get(&variant_array, options).unwrap_err();
4051
4052 assert!(
4053 err.to_string().contains(
4054 "Failed to cast to Decimal32(precision=9, scale=2) from variant Decimal4"
4055 )
4056 );
4057 }
4058
4059 #[test]
4060 fn get_decimal64_rescaled_to_scale2() {
4061 let mut builder = crate::VariantArrayBuilder::new(5);
4062 builder.append_variant(VariantDecimal8::try_new(1234, 2).unwrap().into()); builder.append_variant(VariantDecimal8::try_new(1234, 3).unwrap().into()); builder.append_variant(VariantDecimal8::try_new(1234, 0).unwrap().into()); builder.append_null();
4066 builder.append_variant(
4067 VariantDecimal16::try_new((VariantDecimal8::MAX_UNSCALED_VALUE as i128) + 1, 3)
4068 .unwrap()
4069 .into(),
4070 ); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4072
4073 let field = Field::new("result", DataType::Decimal64(18, 2), true);
4074 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4075 let result = variant_get(&variant_array, options).unwrap();
4076 let result = result.as_any().downcast_ref::<Decimal64Array>().unwrap();
4077
4078 assert_eq!(result.precision(), 18);
4079 assert_eq!(result.scale(), 2);
4080 assert_eq!(result.value(0), 1234);
4081 assert_eq!(result.value(1), 123);
4082 assert_eq!(result.value(2), 123400);
4083 assert!(result.is_null(3));
4084 assert_eq!(
4085 result.value(4),
4086 VariantDecimal8::MAX_UNSCALED_VALUE / 10 + 1
4087 ); }
4089
4090 #[test]
4091 fn get_decimal64_scale_down_rounding() {
4092 let mut builder = crate::VariantArrayBuilder::new(7);
4093 builder.append_variant(VariantDecimal8::try_new(1235, 0).unwrap().into());
4094 builder.append_variant(VariantDecimal8::try_new(1245, 0).unwrap().into());
4095 builder.append_variant(VariantDecimal8::try_new(-1235, 0).unwrap().into());
4096 builder.append_variant(VariantDecimal8::try_new(-1245, 0).unwrap().into());
4097 builder.append_variant(VariantDecimal8::try_new(1235, 2).unwrap().into()); builder.append_variant(VariantDecimal8::try_new(1235, 3).unwrap().into()); builder.append_variant(VariantDecimal8::try_new(5235, 3).unwrap().into()); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4101
4102 let field = Field::new("result", DataType::Decimal64(18, -1), true);
4103 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4104 let result = variant_get(&variant_array, options).unwrap();
4105 let result = result.as_any().downcast_ref::<Decimal64Array>().unwrap();
4106
4107 assert_eq!(result.precision(), 18);
4108 assert_eq!(result.scale(), -1);
4109 assert_eq!(result.value(0), 124);
4110 assert_eq!(result.value(1), 125);
4111 assert_eq!(result.value(2), -124);
4112 assert_eq!(result.value(3), -125);
4113 assert_eq!(result.value(4), 1);
4114 assert!(result.is_valid(5));
4115 assert_eq!(result.value(5), 0);
4116 assert_eq!(result.value(6), 1);
4117 }
4118
4119 #[test]
4120 fn get_decimal64_large_scale_reduction() {
4121 let mut builder = crate::VariantArrayBuilder::new(2);
4122 builder.append_variant(
4123 VariantDecimal8::try_new(-VariantDecimal8::MAX_UNSCALED_VALUE, 0)
4124 .unwrap()
4125 .into(),
4126 );
4127 builder.append_variant(
4128 VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0)
4129 .unwrap()
4130 .into(),
4131 );
4132 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4133
4134 let field = Field::new("result", DataType::Decimal64(18, -18), true);
4135 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4136 let result = variant_get(&variant_array, options).unwrap();
4137 let result = result.as_any().downcast_ref::<Decimal64Array>().unwrap();
4138
4139 assert_eq!(result.precision(), 18);
4140 assert_eq!(result.scale(), -18);
4141 assert_eq!(result.value(0), -1);
4142 assert_eq!(result.value(1), 1);
4143
4144 let field = Field::new("result", DataType::Decimal64(18, -19), true);
4145 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4146 let result = variant_get(&variant_array, options).unwrap();
4147 let result = result.as_any().downcast_ref::<Decimal64Array>().unwrap();
4148
4149 assert_eq!(result.precision(), 18);
4150 assert_eq!(result.scale(), -19);
4151 assert!(result.is_valid(0));
4152 assert_eq!(result.value(0), 0);
4153 assert!(result.is_valid(1));
4154 assert_eq!(result.value(1), 0);
4155 }
4156
4157 #[test]
4158 fn get_decimal64_precision_overflow_safe() {
4159 let mut builder = crate::VariantArrayBuilder::new(2);
4161 builder.append_variant(
4162 VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0)
4163 .unwrap()
4164 .into(),
4165 );
4166 builder.append_variant(
4167 VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 18)
4168 .unwrap()
4169 .into(),
4170 ); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4172
4173 let field = Field::new("result", DataType::Decimal64(2, 2), true);
4174 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4175 let result = variant_get(&variant_array, options).unwrap();
4176 let result = result.as_any().downcast_ref::<Decimal64Array>().unwrap();
4177
4178 assert!(result.is_null(0));
4179 assert!(result.is_null(1));
4180 }
4181
4182 #[test]
4183 fn get_decimal64_precision_overflow_unsafe_errors() {
4184 let mut builder = crate::VariantArrayBuilder::new(1);
4185 builder.append_variant(
4186 VariantDecimal8::try_new(VariantDecimal8::MAX_UNSCALED_VALUE, 0)
4187 .unwrap()
4188 .into(),
4189 );
4190 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4191
4192 let field = Field::new("result", DataType::Decimal64(18, 2), true);
4193 let cast_options = CastOptions {
4194 safe: false,
4195 ..Default::default()
4196 };
4197 let options = GetOptions::new()
4198 .with_as_type(Some(FieldRef::from(field)))
4199 .with_cast_options(cast_options);
4200 let err = variant_get(&variant_array, options).unwrap_err();
4201
4202 assert!(
4203 err.to_string().contains(
4204 "Failed to cast to Decimal64(precision=18, scale=2) from variant Decimal8"
4205 )
4206 );
4207 }
4208
4209 #[test]
4210 fn get_decimal128_rescaled_to_scale2() {
4211 let mut builder = crate::VariantArrayBuilder::new(4);
4212 builder.append_variant(VariantDecimal16::try_new(1234, 2).unwrap().into());
4213 builder.append_variant(VariantDecimal16::try_new(1234, 3).unwrap().into());
4214 builder.append_variant(VariantDecimal16::try_new(1234, 0).unwrap().into());
4215 builder.append_null();
4216 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4217
4218 let field = Field::new("result", DataType::Decimal128(38, 2), true);
4219 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4220 let result = variant_get(&variant_array, options).unwrap();
4221 let result = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
4222
4223 assert_eq!(result.precision(), 38);
4224 assert_eq!(result.scale(), 2);
4225 assert_eq!(result.value(0), 1234);
4226 assert_eq!(result.value(1), 123);
4227 assert_eq!(result.value(2), 123400);
4228 assert!(result.is_null(3));
4229 }
4230
4231 #[test]
4232 fn get_decimal128_scale_down_rounding() {
4233 let mut builder = crate::VariantArrayBuilder::new(7);
4234 builder.append_variant(VariantDecimal16::try_new(1235, 0).unwrap().into());
4235 builder.append_variant(VariantDecimal16::try_new(1245, 0).unwrap().into());
4236 builder.append_variant(VariantDecimal16::try_new(-1235, 0).unwrap().into());
4237 builder.append_variant(VariantDecimal16::try_new(-1245, 0).unwrap().into());
4238 builder.append_variant(VariantDecimal16::try_new(1235, 2).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(1235, 3).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(5235, 3).unwrap().into()); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4242
4243 let field = Field::new("result", DataType::Decimal128(38, -1), true);
4244 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4245 let result = variant_get(&variant_array, options).unwrap();
4246 let result = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
4247
4248 assert_eq!(result.precision(), 38);
4249 assert_eq!(result.scale(), -1);
4250 assert_eq!(result.value(0), 124);
4251 assert_eq!(result.value(1), 125);
4252 assert_eq!(result.value(2), -124);
4253 assert_eq!(result.value(3), -125);
4254 assert_eq!(result.value(4), 1);
4255 assert!(result.is_valid(5));
4256 assert_eq!(result.value(5), 0);
4257 assert_eq!(result.value(6), 1);
4258 }
4259
4260 #[test]
4261 fn get_decimal128_precision_overflow_safe() {
4262 let mut builder = crate::VariantArrayBuilder::new(2);
4264 builder.append_variant(
4265 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0)
4266 .unwrap()
4267 .into(),
4268 );
4269 builder.append_variant(
4270 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 38)
4271 .unwrap()
4272 .into(),
4273 ); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4275
4276 let field = Field::new("result", DataType::Decimal128(2, 2), true);
4277 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4278 let result = variant_get(&variant_array, options).unwrap();
4279 let result = result.as_any().downcast_ref::<Decimal128Array>().unwrap();
4280
4281 assert!(result.is_null(0));
4282 assert!(result.is_null(1)); }
4284
4285 #[test]
4286 fn get_decimal128_precision_overflow_unsafe_errors() {
4287 let mut builder = crate::VariantArrayBuilder::new(1);
4288 builder.append_variant(
4289 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0)
4290 .unwrap()
4291 .into(),
4292 );
4293 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4294
4295 let field = Field::new("result", DataType::Decimal128(38, 2), true);
4296 let cast_options = CastOptions {
4297 safe: false,
4298 ..Default::default()
4299 };
4300 let options = GetOptions::new()
4301 .with_as_type(Some(FieldRef::from(field)))
4302 .with_cast_options(cast_options);
4303 let err = variant_get(&variant_array, options).unwrap_err();
4304
4305 assert!(err.to_string().contains(
4306 "Failed to cast to Decimal128(precision=38, scale=2) from variant Decimal16"
4307 ));
4308 }
4309
4310 #[test]
4311 fn get_decimal256_rescaled_to_scale2() {
4312 let mut builder = crate::VariantArrayBuilder::new(4);
4314 builder.append_variant(VariantDecimal16::try_new(1234, 2).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(1234, 3).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(1234, 0).unwrap().into()); builder.append_null();
4318 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4319
4320 let field = Field::new("result", DataType::Decimal256(76, 2), true);
4321 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4322 let result = variant_get(&variant_array, options).unwrap();
4323 let result = result.as_any().downcast_ref::<Decimal256Array>().unwrap();
4324
4325 assert_eq!(result.precision(), 76);
4326 assert_eq!(result.scale(), 2);
4327 assert_eq!(result.value(0), i256::from_i128(1234));
4328 assert_eq!(result.value(1), i256::from_i128(123));
4329 assert_eq!(result.value(2), i256::from_i128(123400));
4330 assert!(result.is_null(3));
4331 }
4332
4333 #[test]
4334 fn get_decimal256_scale_down_rounding() {
4335 let mut builder = crate::VariantArrayBuilder::new(7);
4336 builder.append_variant(VariantDecimal16::try_new(1235, 0).unwrap().into());
4337 builder.append_variant(VariantDecimal16::try_new(1245, 0).unwrap().into());
4338 builder.append_variant(VariantDecimal16::try_new(-1235, 0).unwrap().into());
4339 builder.append_variant(VariantDecimal16::try_new(-1245, 0).unwrap().into());
4340 builder.append_variant(VariantDecimal16::try_new(1235, 2).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(1235, 3).unwrap().into()); builder.append_variant(VariantDecimal16::try_new(5235, 3).unwrap().into()); let variant_array: ArrayRef = ArrayRef::from(builder.build());
4344
4345 let field = Field::new("result", DataType::Decimal256(76, -1), true);
4346 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4347 let result = variant_get(&variant_array, options).unwrap();
4348 let result = result.as_any().downcast_ref::<Decimal256Array>().unwrap();
4349
4350 assert_eq!(result.precision(), 76);
4351 assert_eq!(result.scale(), -1);
4352 assert_eq!(result.value(0), i256::from_i128(124));
4353 assert_eq!(result.value(1), i256::from_i128(125));
4354 assert_eq!(result.value(2), i256::from_i128(-124));
4355 assert_eq!(result.value(3), i256::from_i128(-125));
4356 assert_eq!(result.value(4), i256::from_i128(1));
4357 assert!(result.is_valid(5));
4358 assert_eq!(result.value(5), i256::from_i128(0));
4359 assert_eq!(result.value(6), i256::from_i128(1));
4360 }
4361
4362 #[test]
4363 fn get_decimal256_precision_overflow_safe() {
4364 let mut builder = crate::VariantArrayBuilder::new(2);
4366 builder.append_variant(
4367 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1)
4368 .unwrap()
4369 .into(),
4370 );
4371 builder.append_variant(
4372 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0)
4373 .unwrap()
4374 .into(),
4375 );
4376 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4377
4378 let field = Field::new("result", DataType::Decimal256(76, 39), true);
4379 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4380 let result = variant_get(&variant_array, options).unwrap();
4381 let result = result.as_any().downcast_ref::<Decimal256Array>().unwrap();
4382
4383 let base = i256::from_i128(10);
4386 let factor = base.checked_pow(38).unwrap();
4387 let expected = i256::from_i128(VariantDecimal16::MAX_UNSCALED_VALUE)
4388 .checked_mul(factor)
4389 .unwrap();
4390 assert_eq!(result.value(0), expected);
4391 assert!(result.is_null(1));
4392 }
4393
4394 #[test]
4395 fn get_decimal256_precision_overflow_unsafe_errors() {
4396 let mut builder = crate::VariantArrayBuilder::new(2);
4398 builder.append_variant(
4399 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 1)
4400 .unwrap()
4401 .into(),
4402 );
4403 builder.append_variant(
4404 VariantDecimal16::try_new(VariantDecimal16::MAX_UNSCALED_VALUE, 0)
4405 .unwrap()
4406 .into(),
4407 );
4408 let variant_array: ArrayRef = ArrayRef::from(builder.build());
4409
4410 let field = Field::new("result", DataType::Decimal256(76, 39), true);
4411 let cast_options = CastOptions {
4412 safe: false,
4413 ..Default::default()
4414 };
4415 let options = GetOptions::new()
4416 .with_as_type(Some(FieldRef::from(field)))
4417 .with_cast_options(cast_options);
4418 let err = variant_get(&variant_array, options).unwrap_err();
4419
4420 assert!(err.to_string().contains(
4421 "Failed to cast to Decimal256(precision=76, scale=39) from variant Decimal16"
4422 ));
4423 }
4424
4425 #[test]
4426 fn get_non_supported_temporal_types_error() {
4427 let values = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse)];
4428 let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(values));
4429
4430 let test_cases = vec![
4431 FieldRef::from(Field::new(
4432 "result",
4433 DataType::Duration(TimeUnit::Microsecond),
4434 true,
4435 )),
4436 FieldRef::from(Field::new(
4437 "result",
4438 DataType::Interval(IntervalUnit::YearMonth),
4439 true,
4440 )),
4441 ];
4442
4443 for field in test_cases {
4444 let options = GetOptions::new().with_as_type(Some(field));
4445 let err = variant_get(&variant_array, options).unwrap_err();
4446 assert!(
4447 err.to_string()
4448 .contains("Casting Variant to duration/interval types is not supported")
4449 );
4450 }
4451 }
4452
4453 #[test]
4454 fn get_variant_as_dictionary() {
4455 let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![
4456 Some(Variant::from("apple")),
4457 Some(Variant::from("banana")),
4458 None,
4459 Some(Variant::from("apple")),
4460 ]));
4461 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
4462 let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
4463 "dict",
4464 data_type.clone(),
4465 true,
4466 ))));
4467
4468 let result = variant_get(&variant_array, options).unwrap();
4469 assert_eq!(result.data_type(), &data_type);
4470
4471 let decoded = cast(result.as_ref(), &DataType::Utf8).unwrap();
4472 let expected = StringArray::from(vec![Some("apple"), Some("banana"), None, Some("apple")]);
4473 assert_eq!(decoded.as_ref(), &expected);
4474 }
4475
4476 #[test]
4477 fn get_variant_as_numeric_dictionary() {
4478 let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![
4479 Some(Variant::from(42)),
4480 Some(Variant::from(7)),
4481 None,
4482 Some(Variant::from(42)),
4483 ]));
4484 let data_type = DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Int32));
4485 let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
4486 "dict",
4487 data_type.clone(),
4488 true,
4489 ))));
4490
4491 let result = variant_get(&variant_array, options).unwrap();
4492 assert_eq!(result.data_type(), &data_type);
4493
4494 let decoded = cast(result.as_ref(), &DataType::Int32).unwrap();
4495 let expected = Int32Array::from(vec![Some(42), Some(7), None, Some(42)]);
4496 assert_eq!(decoded.as_ref(), &expected);
4497 }
4498
4499 #[test]
4500 fn get_variant_as_run_end_encoded() {
4501 let variant_array: ArrayRef = ArrayRef::from(VariantArray::from_iter(vec![
4502 Some(Variant::from("apple")),
4503 Some(Variant::from("apple")),
4504 None,
4505 Some(Variant::from("banana")),
4506 Some(Variant::from("banana")),
4507 ]));
4508 let run_ends = Arc::new(Field::new(
4509 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4510 DataType::Int32,
4511 false,
4512 ));
4513 let values = Arc::new(Field::new(
4514 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4515 DataType::Utf8,
4516 true,
4517 ));
4518 let data_type = DataType::RunEndEncoded(run_ends, values);
4519 let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
4520 "ree",
4521 data_type.clone(),
4522 true,
4523 ))));
4524
4525 let result = variant_get(&variant_array, options).unwrap();
4526 assert_eq!(result.data_type(), &data_type);
4527
4528 let decoded = cast(result.as_ref(), &DataType::Utf8).unwrap();
4529 let expected = StringArray::from(vec![
4530 Some("apple"),
4531 Some("apple"),
4532 None,
4533 Some("banana"),
4534 Some("banana"),
4535 ]);
4536 assert_eq!(decoded.as_ref(), &expected);
4537 }
4538
4539 fn map_data_type(value_type: DataType) -> DataType {
4542 DataType::Map(
4543 Arc::new(Field::new(
4544 "entries",
4545 DataType::Struct(Fields::from(vec![
4546 Field::new("keys", DataType::Utf8, false),
4547 Field::new("values", value_type, true),
4548 ])),
4549 false,
4550 )),
4551 false,
4552 )
4553 }
4554
4555 fn map_get_options(data_type: &DataType) -> GetOptions<'static> {
4556 GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
4557 "map",
4558 data_type.clone(),
4559 true,
4560 ))))
4561 }
4562
4563 #[test]
4564 fn get_variant_as_map() {
4565 let input: ArrayRef = Arc::new(StringArray::from(vec![
4566 Some(r#"{"a": 1, "b": 2}"#),
4567 Some(r#"{"c": 3}"#),
4568 None,
4569 Some("{}"),
4570 Some(r#"{"d": null}"#),
4571 ]));
4572 let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
4573
4574 let data_type = map_data_type(DataType::Int64);
4575 let result = variant_get(&variant_array, map_get_options(&data_type)).unwrap();
4576 assert_eq!(result.data_type(), &data_type);
4577
4578 let mut expected = MapBuilder::new(None, StringBuilder::new(), Int64Builder::new());
4579 expected.keys().append_value("a");
4580 expected.values().append_value(1);
4581 expected.keys().append_value("b");
4582 expected.values().append_value(2);
4583 expected.append(true).unwrap();
4584 expected.keys().append_value("c");
4585 expected.values().append_value(3);
4586 expected.append(true).unwrap();
4587 expected.append(false).unwrap(); expected.append(true).unwrap(); expected.keys().append_value("d");
4590 expected.values().append_null(); expected.append(true).unwrap();
4592 let expected = expected.finish();
4593 assert_eq!(result.as_ref(), &expected);
4594 }
4595
4596 #[test]
4597 fn get_variant_as_map_of_lists() {
4598 let input: ArrayRef = Arc::new(StringArray::from(vec![
4599 Some(r#"{"a": [1, 2], "b": []}"#),
4600 Some(r#"{"c": [3]}"#),
4601 ]));
4602 let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
4603
4604 let data_type = map_data_type(DataType::List(Arc::new(Field::new(
4605 "item",
4606 DataType::Int64,
4607 true,
4608 ))));
4609 let result = variant_get(&variant_array, map_get_options(&data_type)).unwrap();
4610 assert_eq!(result.data_type(), &data_type);
4611
4612 let mut expected = MapBuilder::new(
4613 None,
4614 StringBuilder::new(),
4615 ListBuilder::new(Int64Builder::new()),
4616 );
4617 expected.keys().append_value("a");
4618 expected.values().append_value([Some(1), Some(2)]);
4619 expected.keys().append_value("b");
4620 expected.values().append_value([]);
4621 expected.append(true).unwrap();
4622 expected.keys().append_value("c");
4623 expected.values().append_value([Some(3)]);
4624 expected.append(true).unwrap();
4625 let expected = expected.finish();
4626 assert_eq!(result.as_ref(), &expected);
4627 }
4628
4629 #[test]
4630 fn get_variant_as_map_non_object_rows() {
4631 let input: ArrayRef = Arc::new(StringArray::from(vec![
4632 Some(r#"{"a": 1}"#),
4633 Some("42"), ]));
4635 let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
4636 let data_type = map_data_type(DataType::Int64);
4637
4638 let result = variant_get(&variant_array, map_get_options(&data_type)).unwrap();
4640 let mut expected = MapBuilder::new(None, StringBuilder::new(), Int64Builder::new());
4641 expected.keys().append_value("a");
4642 expected.values().append_value(1);
4643 expected.append(true).unwrap();
4644 expected.append(false).unwrap();
4645 let expected = expected.finish();
4646 assert_eq!(result.as_ref(), &expected);
4647
4648 let options = map_get_options(&data_type).with_cast_options(CastOptions {
4650 safe: false,
4651 format_options: FormatOptions::default(),
4652 });
4653 let err = variant_get(&variant_array, options).unwrap_err();
4654 assert!(
4655 err.to_string().contains("Failed to extract object"),
4656 "unexpected error: {err}"
4657 );
4658 }
4659
4660 #[test]
4661 fn get_variant_as_map_invalid_entries() {
4662 let input: ArrayRef = Arc::new(StringArray::from(vec![Some(r#"{"a": 1}"#)]));
4663 let variant_array = ArrayRef::from(json_to_variant(&input).unwrap());
4664
4665 let data_type = DataType::Map(
4667 Arc::new(Field::new("entries", DataType::Int32, false)),
4668 false,
4669 );
4670 let err = variant_get(&variant_array, map_get_options(&data_type)).unwrap_err();
4671 assert!(
4672 err.to_string().contains("Map entries must be Struct"),
4673 "unexpected error: {err}"
4674 );
4675
4676 let data_type = DataType::Map(
4678 Arc::new(Field::new(
4679 "entries",
4680 DataType::Struct(Fields::from(vec![Field::new(
4681 "keys",
4682 DataType::Utf8,
4683 false,
4684 )])),
4685 false,
4686 )),
4687 false,
4688 );
4689 let err = variant_get(&variant_array, map_get_options(&data_type)).unwrap_err();
4690 assert!(
4691 err.to_string()
4692 .contains("Map entries must have exactly two fields"),
4693 "unexpected error: {err}"
4694 );
4695 }
4696
4697 fn invalid_time_variant_array() -> ArrayRef {
4698 let mut builder = VariantArrayBuilder::new(3);
4699 builder.append_variant(Variant::Int64(86401000000));
4701 builder.append_variant(Variant::Int64(86401000000));
4702 builder.append_variant(Variant::Int64(86401000000));
4703 Arc::new(builder.build().into_inner())
4704 }
4705
4706 #[test]
4707 fn test_variant_get_error_when_cast_failure_and_safe_false() {
4708 let variant_array = invalid_time_variant_array();
4709
4710 let field = Field::new("result", DataType::Time64(TimeUnit::Microsecond), true);
4711 let cast_options = CastOptions {
4712 safe: false, ..Default::default()
4714 };
4715 let options = GetOptions::new()
4716 .with_as_type(Some(FieldRef::from(field)))
4717 .with_cast_options(cast_options);
4718 let err = variant_get(&variant_array, options).unwrap_err();
4719 assert!(
4720 err.to_string().contains(
4721 "Cast error: Failed to extract primitive of type Time64(µs) from variant Int64(86401000000) at path VariantPath([])"
4722 ),
4723 "actual: {err}",
4724 );
4725 }
4726
4727 #[test]
4728 fn test_variant_get_return_null_when_cast_failure_and_safe_true() {
4729 let variant_array = invalid_time_variant_array();
4730
4731 let field = Field::new("result", DataType::Time64(TimeUnit::Microsecond), true);
4732 let cast_options = CastOptions {
4733 safe: true, ..Default::default()
4735 };
4736 let options = GetOptions::new()
4737 .with_as_type(Some(FieldRef::from(field)))
4738 .with_cast_options(cast_options);
4739 let result = variant_get(&variant_array, options).unwrap();
4740 assert_eq!(3, result.len());
4741
4742 for i in 0..3 {
4743 assert!(result.is_null(i));
4744 }
4745 }
4746
4747 #[test]
4748 fn test_perfect_shredding_returns_same_arc_ptr() {
4749 let variant_array = perfectly_shredded_int32_variant_array();
4750
4751 let variant_array_ref = VariantArray::try_new(&variant_array).unwrap();
4752 let typed_value_arc = variant_array_ref.typed_value_column().unwrap().clone();
4753
4754 let field = Field::new("result", DataType::Int32, true);
4755 let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
4756 let result = variant_get(&variant_array, options).unwrap();
4757
4758 assert!(Arc::ptr_eq(&typed_value_arc, &result));
4759 }
4760
4761 #[test]
4762 fn test_perfect_shredding_three_typed_value_columns() {
4763 let all_nulls_values: Arc<Int32Array> = Arc::new(Int32Array::from(vec![
4765 Option::<i32>::None,
4766 Option::<i32>::None,
4767 Option::<i32>::None,
4768 ]));
4769 let all_nulls_erased: ArrayRef = all_nulls_values.clone();
4770 let all_nulls_field =
4771 ShreddedVariantFieldArray::perfectly_shredded(all_nulls_erased.clone());
4772 let all_nulls_type = all_nulls_field.data_type().clone();
4773 let all_nulls_struct: ArrayRef = ArrayRef::from(all_nulls_field);
4774
4775 let some_nulls_values: Arc<Int32Array> =
4777 Arc::new(Int32Array::from(vec![Some(10), None, Some(30)]));
4778 let some_nulls_erased: ArrayRef = some_nulls_values.clone();
4779 let some_nulls_field =
4780 ShreddedVariantFieldArray::perfectly_shredded(some_nulls_erased.clone());
4781 let some_nulls_type = some_nulls_field.data_type().clone();
4782 let some_nulls_struct: ArrayRef = ArrayRef::from(some_nulls_field);
4783
4784 let inner_values: Arc<Int32Array> =
4786 Arc::new(Int32Array::from(vec![Some(111), None, Some(333)]));
4787 let inner_erased: ArrayRef = inner_values.clone();
4788 let inner_field = ShreddedVariantFieldArray::perfectly_shredded(inner_erased.clone());
4789 let inner_field_type = inner_field.data_type().clone();
4790 let inner_struct_array: ArrayRef = ArrayRef::from(inner_field);
4791
4792 let nested_struct = Arc::new(
4793 StructArray::try_new(
4794 Fields::from(vec![Field::new("inner", inner_field_type, true)]),
4795 vec![inner_struct_array],
4796 None,
4797 )
4798 .unwrap(),
4799 );
4800 let nested_struct_erased: ArrayRef = nested_struct.clone();
4801 let struct_field =
4802 ShreddedVariantFieldArray::perfectly_shredded(nested_struct_erased.clone());
4803 let struct_field_type = struct_field.data_type().clone();
4804 let struct_field_struct: ArrayRef = ArrayRef::from(struct_field);
4805
4806 let typed_value_struct = StructArray::try_new(
4808 Fields::from(vec![
4809 Field::new("all_nulls", all_nulls_type, true),
4810 Field::new("some_nulls", some_nulls_type, true),
4811 Field::new("struct_field", struct_field_type, true),
4812 ]),
4813 vec![all_nulls_struct, some_nulls_struct, struct_field_struct],
4814 None,
4815 )
4816 .unwrap();
4817
4818 let metadata = BinaryViewArray::from_iter_values(std::iter::repeat_n(
4819 EMPTY_VARIANT_METADATA_BYTES,
4820 all_nulls_values.len(),
4821 ));
4822 let variant_array: ArrayRef = VariantArray::perfectly_shredded(
4823 Arc::new(metadata),
4824 Arc::new(typed_value_struct),
4825 None,
4826 )
4827 .into();
4828
4829 let all_nulls_field_ref = FieldRef::from(Field::new("result", DataType::Int32, true));
4831 let all_nulls_result = variant_get(
4832 &variant_array,
4833 GetOptions::new_with_path(VariantPath::try_from("all_nulls").unwrap())
4834 .with_as_type(Some(all_nulls_field_ref)),
4835 )
4836 .unwrap();
4837 assert!(Arc::ptr_eq(&all_nulls_result, &all_nulls_erased));
4838
4839 let some_nulls_field_ref = FieldRef::from(Field::new("result", DataType::Int32, true));
4841 let some_nulls_result = variant_get(
4842 &variant_array,
4843 GetOptions::new_with_path(VariantPath::try_from("some_nulls").unwrap())
4844 .with_as_type(Some(some_nulls_field_ref)),
4845 )
4846 .unwrap();
4847 assert!(Arc::ptr_eq(&some_nulls_result, &some_nulls_erased));
4848
4849 let struct_child_fields = Fields::from(vec![Field::new("inner", DataType::Int32, true)]);
4851 let struct_field_ref = FieldRef::from(Field::new(
4852 "result",
4853 DataType::Struct(struct_child_fields.clone()),
4854 true,
4855 ));
4856 let struct_result = variant_get(
4857 &variant_array,
4858 GetOptions::new_with_path(VariantPath::try_from("struct_field").unwrap())
4859 .with_as_type(Some(struct_field_ref)),
4860 )
4861 .unwrap();
4862 let struct_array = struct_result
4863 .as_any()
4864 .downcast_ref::<StructArray>()
4865 .unwrap();
4866 assert_eq!(struct_array.len(), 3);
4867 assert_eq!(struct_array.null_count(), 0);
4868
4869 let inner_values_result = struct_array
4870 .column(0)
4871 .as_any()
4872 .downcast_ref::<Int32Array>()
4873 .unwrap();
4874 assert_eq!(inner_values_result.len(), 3);
4875 assert_eq!(inner_values_result.value(0), 111);
4876 assert!(inner_values_result.is_null(1));
4877 assert_eq!(inner_values_result.value(2), 333);
4878 }
4879
4880 #[test]
4881 fn test_variant_get_list_like_safe_cast() {
4882 let string_array: ArrayRef = Arc::new(StringArray::from(vec![
4883 r#"{"outer":{"list":[1, "two", 3]}}"#,
4884 r#"{"outer":{"list":"not a list"}}"#,
4885 ]));
4886 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
4887
4888 let element_array: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(3)]));
4889 let field = Arc::new(Field::new("item", Int64, true));
4890
4891 let expectations = vec![
4892 (
4893 DataType::List(field.clone()),
4894 Arc::new(ListArray::new(
4895 field.clone(),
4896 OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 3])),
4897 element_array.clone(),
4898 Some(NullBuffer::from(vec![true, false])),
4899 )) as ArrayRef,
4900 ),
4901 (
4902 DataType::LargeList(field.clone()),
4903 Arc::new(LargeListArray::new(
4904 field.clone(),
4905 OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 3])),
4906 element_array.clone(),
4907 Some(NullBuffer::from(vec![true, false])),
4908 )) as ArrayRef,
4909 ),
4910 (
4911 DataType::ListView(field.clone()),
4912 Arc::new(ListViewArray::new(
4913 field.clone(),
4914 ScalarBuffer::from(vec![0, 3]),
4915 ScalarBuffer::from(vec![3, 0]),
4916 element_array.clone(),
4917 Some(NullBuffer::from(vec![true, false])),
4918 )) as ArrayRef,
4919 ),
4920 (
4921 DataType::LargeListView(field.clone()),
4922 Arc::new(LargeListViewArray::new(
4923 field.clone(),
4924 ScalarBuffer::from(vec![0, 3]),
4925 ScalarBuffer::from(vec![3, 0]),
4926 element_array,
4927 Some(NullBuffer::from(vec![true, false])),
4928 )) as ArrayRef,
4929 ),
4930 (
4931 DataType::FixedSizeList(field.clone(), 3),
4932 Arc::new(FixedSizeListArray::new(
4933 field,
4934 3,
4935 Arc::new(Int64Array::from(vec![
4936 Some(1),
4937 None,
4938 Some(3),
4939 None,
4940 None,
4941 None,
4942 ])),
4943 Some(NullBuffer::from(vec![true, false])),
4944 )) as ArrayRef,
4945 ),
4946 ];
4947
4948 for (request_type, expected) in expectations {
4949 let options =
4950 GetOptions::new_with_path(VariantPath::try_from("outer").unwrap().join("list"))
4951 .with_as_type(Some(FieldRef::from(Field::new(
4952 "result",
4953 request_type.clone(),
4954 true,
4955 ))));
4956
4957 let result = variant_get(&variant_array, options).unwrap();
4958 assert_eq!(result.data_type(), expected.data_type());
4959 assert_eq!(&result, &expected);
4960 }
4961
4962 for (idx, expected) in [
4963 (0, vec![Some(1), None]),
4964 (1, vec![None, None]),
4965 (2, vec![Some(3), None]),
4966 ] {
4967 let index_options = GetOptions::new_with_path(
4968 VariantPath::try_from("outer")
4969 .unwrap()
4970 .join("list")
4971 .join(idx),
4972 )
4973 .with_as_type(Some(FieldRef::from(Field::new(
4974 "result",
4975 DataType::Int64,
4976 true,
4977 ))));
4978 let index_result = variant_get(&variant_array, index_options).unwrap();
4979 let index_expected: ArrayRef = Arc::new(Int64Array::from(expected));
4980 assert_eq!(&index_result, &index_expected);
4981 }
4982 }
4983
4984 #[test]
4985 fn test_variant_get_nested_list() {
4986 use arrow::datatypes::Int64Type;
4987
4988 let string_array: ArrayRef = Arc::new(StringArray::from(vec![
4989 "[[1, 2], [3]]",
4990 r#"[[4], "not a list", [5, 6]]"#,
4991 ]));
4992 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
4993
4994 let inner_field = Arc::new(Field::new("item", Int64, true));
4995 let outer_field = Arc::new(Field::new(
4996 "item",
4997 DataType::List(inner_field.clone()),
4998 true,
4999 ));
5000 let request_type = DataType::List(outer_field.clone());
5001
5002 let options = GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
5003 "result",
5004 request_type,
5005 true,
5006 ))));
5007 let result = variant_get(&variant_array, options).unwrap();
5008 let outer = result.as_list::<i32>();
5009
5010 let row0 = outer.value(0);
5012 let row0 = row0.as_list::<i32>();
5013 assert_eq!(row0.len(), 2);
5014 let elem0 = row0.value(0);
5015 assert_eq!(elem0.as_primitive::<Int64Type>().values(), &[1, 2]);
5016 let elem1 = row0.value(1);
5017 assert_eq!(elem1.as_primitive::<Int64Type>().values(), &[3]);
5018
5019 let row1 = outer.value(1);
5021 let row1 = row1.as_list::<i32>();
5022 assert_eq!(row1.len(), 3);
5023 let elem0 = row1.value(0);
5024 assert_eq!(elem0.as_primitive::<Int64Type>().values(), &[4]);
5025 assert!(row1.is_null(1));
5026 let elem2 = row1.value(2);
5027 assert_eq!(elem2.as_primitive::<Int64Type>().values(), &[5, 6]);
5028 }
5029
5030 #[test]
5031 fn test_variant_get_list_like_unsafe_cast_errors_on_element_mismatch() {
5032 let string_array: ArrayRef =
5033 Arc::new(StringArray::from(vec![r#"[1, "two", 3]"#, "[4, 5]"]));
5034 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
5035 let cast_options = CastOptions {
5036 safe: false,
5037 ..Default::default()
5038 };
5039
5040 let item_field = Arc::new(Field::new("item", DataType::Int64, true));
5041 let request_types = vec![
5042 DataType::List(item_field.clone()),
5043 DataType::LargeList(item_field.clone()),
5044 DataType::ListView(item_field.clone()),
5045 DataType::LargeListView(item_field),
5046 ];
5047
5048 for request_type in request_types {
5049 let options = GetOptions::new()
5050 .with_as_type(Some(FieldRef::from(Field::new(
5051 "result",
5052 request_type.clone(),
5053 true,
5054 ))))
5055 .with_cast_options(cast_options.clone());
5056
5057 let err = variant_get(&variant_array, options).unwrap_err();
5058 assert!(
5059 err.to_string()
5060 .contains("Failed to extract primitive of type Int64")
5061 );
5062 }
5063 }
5064
5065 #[test]
5066 fn test_variant_get_list_like_unsafe_cast_preserves_null_elements() {
5067 let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, null, 3]"]));
5068 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
5069 let cast_options = CastOptions {
5070 safe: false,
5071 ..Default::default()
5072 };
5073 let options = GetOptions::new()
5074 .with_as_type(Some(FieldRef::from(Field::new(
5075 "result",
5076 DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
5077 true,
5078 ))))
5079 .with_cast_options(cast_options);
5080
5081 let result = variant_get(&variant_array, options).unwrap();
5082 let list_array = result.as_any().downcast_ref::<ListArray>().unwrap();
5083 let values = list_array
5084 .values()
5085 .as_any()
5086 .downcast_ref::<Int64Array>()
5087 .unwrap();
5088
5089 assert_eq!(values.len(), 3);
5090 assert_eq!(values.value(0), 1);
5091 assert!(values.is_null(1));
5092 assert_eq!(values.value(2), 3);
5093 }
5094
5095 #[test]
5096 fn test_variant_get_list_like_unsafe_cast_errors_on_non_list() {
5097 let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2]", "\"not a list\""]));
5098 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
5099 let cast_options = CastOptions {
5100 safe: false,
5101 ..Default::default()
5102 };
5103 let item_field = Arc::new(Field::new("item", Int64, true));
5104 let data_types = vec![
5105 DataType::List(item_field.clone()),
5106 DataType::LargeList(item_field.clone()),
5107 DataType::ListView(item_field.clone()),
5108 DataType::LargeListView(item_field.clone()),
5109 DataType::FixedSizeList(item_field, 2),
5110 ];
5111
5112 for data_type in data_types {
5113 let options = GetOptions::new()
5114 .with_as_type(Some(FieldRef::from(Field::new("result", data_type, true))))
5115 .with_cast_options(cast_options.clone());
5116
5117 let err = variant_get(&variant_array, options).unwrap_err();
5118 assert!(
5119 err.to_string()
5120 .contains("Failed to extract list from variant"),
5121 );
5122 }
5123 }
5124
5125 #[test]
5126 fn test_variant_get_fixed_size_list_wrong_size() {
5127 let string_array: ArrayRef = Arc::new(StringArray::from(vec!["[1, 2, 3]"]));
5128 let variant_array = ArrayRef::from(json_to_variant(&string_array).unwrap());
5129 let item_field = Arc::new(Field::new("item", Int64, true));
5130
5131 let options = GetOptions::new()
5133 .with_as_type(Some(FieldRef::from(Field::new(
5134 "result",
5135 DataType::FixedSizeList(item_field.clone(), 2),
5136 true,
5137 ))))
5138 .with_cast_options(CastOptions {
5139 safe: true,
5140 ..Default::default()
5141 });
5142 let result = variant_get(&variant_array, options).unwrap();
5143 let fixed_size_list = result
5144 .as_any()
5145 .downcast_ref::<FixedSizeListArray>()
5146 .expect("Expected FixedSizeListArray");
5147 assert_eq!(fixed_size_list.len(), 1);
5148 assert!(fixed_size_list.is_null(0));
5149
5150 let options = GetOptions::new()
5152 .with_as_type(Some(FieldRef::from(Field::new(
5153 "result",
5154 DataType::FixedSizeList(item_field.clone(), 2),
5155 true,
5156 ))))
5157 .with_cast_options(CastOptions {
5158 safe: false,
5159 ..Default::default()
5160 });
5161 let err = variant_get(&variant_array, options).unwrap_err();
5162 assert!(
5163 err.to_string()
5164 .contains("Expected fixed size list of size 2, got size 3"),
5165 "got: {err}",
5166 );
5167 }
5168
5169 macro_rules! perfectly_shredded_preserves_top_level_nulls_test {
5170 ($name:ident, $result_type:expr, $typed_value:expr, $expected_array:expr) => {
5171 perfectly_shredded_preserves_top_level_nulls_test!(
5172 $name,
5173 $result_type,
5174 $typed_value,
5175 Some(NullBuffer::from(vec![true, false, true])),
5176 $expected_array
5177 );
5178 };
5179 ($name:ident, $result_type:expr, $typed_value:expr, $parent_nulls:expr, $expected_array:expr) => {
5180 #[test]
5181 fn $name() {
5182 let metadata = Arc::new(BinaryViewArray::from_iter_values(std::iter::repeat_n(
5183 EMPTY_VARIANT_METADATA_BYTES,
5184 3,
5185 )));
5186 let typed_value: ArrayRef = Arc::new($typed_value);
5187 let variant_array: ArrayRef =
5188 VariantArray::perfectly_shredded(metadata, typed_value, $parent_nulls).into();
5189
5190 let result = variant_get(
5191 &variant_array,
5192 GetOptions::new().with_as_type(Some(FieldRef::from(Field::new(
5193 "result",
5194 $result_type,
5195 true,
5196 )))),
5197 )
5198 .unwrap();
5199
5200 let expected_array: ArrayRef = Arc::new($expected_array);
5201 assert_eq!(&result, &expected_array);
5202 }
5203 };
5204 }
5205
5206 perfectly_shredded_preserves_top_level_nulls_test!(
5207 test_variant_get_perfectly_shredded_integer_preserves_top_level_nulls,
5208 DataType::Int32,
5209 Int32Array::from(vec![Some(0_i32), Some(1_i32), Some(2_i32)]),
5210 Int32Array::from(vec![Some(0_i32), None, Some(2_i32)])
5211 );
5212
5213 perfectly_shredded_preserves_top_level_nulls_test!(
5214 test_variant_get_perfectly_shredded_integer_unions_child_and_top_level_nulls,
5215 DataType::Int32,
5216 Int32Array::from(vec![None, Some(1_i32), Some(2_i32)]),
5217 Some(NullBuffer::from(vec![true, false, true])),
5218 Int32Array::from(vec![None, None, Some(2_i32)])
5219 );
5220
5221 perfectly_shredded_preserves_top_level_nulls_test!(
5222 test_variant_get_perfectly_shredded_null_preserves_top_level_nulls,
5223 DataType::Null,
5224 NullArray::new(3),
5225 NullArray::new(3)
5226 );
5227
5228 perfectly_shredded_preserves_top_level_nulls_test!(
5229 test_variant_get_perfectly_shredded_binary_view_preserves_top_level_nulls,
5230 DataType::BinaryView,
5231 BinaryViewArray::from(vec![
5232 Some(b"Apache" as &[u8]),
5233 Some(b"masked-null" as &[u8]),
5234 Some(b"Parquet-variant" as &[u8]),
5235 ]),
5236 BinaryViewArray::from(vec![
5237 Some(b"Apache" as &[u8]),
5238 None,
5239 Some(b"Parquet-variant" as &[u8]),
5240 ])
5241 );
5242
5243 perfectly_shredded_preserves_top_level_nulls_test!(
5244 test_variant_get_perfectly_shredded_binary_preserves_top_level_nulls,
5245 DataType::Binary,
5246 BinaryArray::from(vec![
5247 Some(b"Apache" as &[u8]),
5248 Some(b"masked-null" as &[u8]),
5249 Some(b"Parquet-variant" as &[u8]),
5250 ]),
5251 BinaryArray::from(vec![
5252 Some(b"Apache" as &[u8]),
5253 None,
5254 Some(b"Parquet-variant" as &[u8]),
5255 ])
5256 );
5257
5258 perfectly_shredded_preserves_top_level_nulls_test!(
5259 test_variant_get_perfectly_shredded_decimal4_preserves_top_level_nulls,
5260 DataType::Decimal32(5, 2),
5261 Decimal32Array::from(vec![Some(12345), Some(23400), Some(-12342)])
5262 .with_precision_and_scale(5, 2)
5263 .unwrap(),
5264 Decimal32Array::from(vec![Some(12345), None, Some(-12342)])
5265 .with_precision_and_scale(5, 2)
5266 .unwrap()
5267 );
5268
5269 perfectly_shredded_preserves_top_level_nulls_test!(
5270 test_variant_get_perfectly_shredded_decimal8_preserves_top_level_nulls,
5271 DataType::Decimal64(10, 1),
5272 Decimal64Array::from(vec![Some(1234567809), Some(1456787000), Some(-1234561203)])
5273 .with_precision_and_scale(10, 1)
5274 .unwrap(),
5275 Decimal64Array::from(vec![Some(1234567809), None, Some(-1234561203)])
5276 .with_precision_and_scale(10, 1)
5277 .unwrap()
5278 );
5279
5280 perfectly_shredded_preserves_top_level_nulls_test!(
5281 test_variant_get_perfectly_shredded_decimal16_preserves_top_level_nulls,
5282 DataType::Decimal128(20, 3),
5283 Decimal128Array::from(vec![
5284 Some(i128::from_str("12345678901234567899").unwrap()),
5285 Some(i128::from_str("23445677483748324300").unwrap()),
5286 Some(i128::from_str("-12345678901234567899").unwrap()),
5287 ])
5288 .with_precision_and_scale(20, 3)
5289 .unwrap(),
5290 Decimal128Array::from(vec![
5291 Some(i128::from_str("12345678901234567899").unwrap()),
5292 None,
5293 Some(i128::from_str("-12345678901234567899").unwrap()),
5294 ])
5295 .with_precision_and_scale(20, 3)
5296 .unwrap()
5297 );
5298
5299 fn union_get_options(fields: &UnionFields, mode: UnionMode) -> GetOptions<'static> {
5300 let field = Field::new("union", DataType::Union(fields.clone(), mode), true);
5301 GetOptions::new().with_as_type(Some(FieldRef::from(field)))
5302 }
5303
5304 fn int_str_bool_union_fields() -> UnionFields {
5305 UnionFields::try_new(
5306 vec![0, 1, 2],
5307 vec![
5308 Field::new("int", DataType::Int64, true),
5309 Field::new("str", DataType::Utf8, true),
5310 Field::new("bool", DataType::Boolean, true),
5311 ],
5312 )
5313 .unwrap()
5314 }
5315
5316 fn mixed_variant_array() -> ArrayRef {
5318 let mut builder = VariantArrayBuilder::new(7);
5319 builder.append_variant(Variant::Int8(1));
5320 builder.append_variant(Variant::from("hello"));
5321 builder.append_variant(Variant::from(true));
5322 builder.append_null();
5323 builder.append_variant(Variant::Null);
5324 builder.append_variant(Variant::Double(2.5));
5325 builder.append_variant(Variant::Int64(5_000_000_000));
5326 ArrayRef::from(builder.build())
5327 }
5328
5329 #[test]
5330 fn get_variant_as_dense_union() {
5331 let fields = int_str_bool_union_fields();
5332 let array = mixed_variant_array();
5333 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5334
5335 let expected: ArrayRef = Arc::new(
5337 UnionArray::try_new(
5338 fields,
5339 ScalarBuffer::from(vec![0i8, 1, 2, 0, 0, 0, 0]),
5340 Some(ScalarBuffer::from(vec![0i32, 0, 0, 1, 2, 3, 4])),
5341 vec![
5342 Arc::new(Int64Array::from(vec![
5343 Some(1),
5344 None,
5345 None,
5346 None,
5347 Some(5_000_000_000),
5348 ])),
5349 Arc::new(StringArray::from(vec!["hello"])),
5350 Arc::new(BooleanArray::from(vec![true])),
5351 ],
5352 )
5353 .unwrap(),
5354 );
5355 assert_eq!(&result, &expected);
5356 }
5357
5358 #[test]
5359 fn get_variant_as_sparse_union() {
5360 let fields = int_str_bool_union_fields();
5361 let array = mixed_variant_array();
5362 let result = variant_get(&array, union_get_options(&fields, UnionMode::Sparse)).unwrap();
5363
5364 let expected: ArrayRef = Arc::new(
5365 UnionArray::try_new(
5366 fields,
5367 ScalarBuffer::from(vec![0i8, 1, 2, 0, 0, 0, 0]),
5368 None,
5369 vec![
5370 Arc::new(Int64Array::from(vec![
5371 Some(1),
5372 None,
5373 None,
5374 None,
5375 None,
5376 None,
5377 Some(5_000_000_000),
5378 ])),
5379 Arc::new(StringArray::from(vec![
5380 None,
5381 Some("hello"),
5382 None,
5383 None,
5384 None,
5385 None,
5386 None,
5387 ])),
5388 Arc::new(BooleanArray::from(vec![
5389 None,
5390 None,
5391 Some(true),
5392 None,
5393 None,
5394 None,
5395 None,
5396 ])),
5397 ],
5398 )
5399 .unwrap(),
5400 );
5401 assert_eq!(&result, &expected);
5402 }
5403
5404 #[test]
5405 fn get_variant_as_union_prefers_most_exact_field() {
5406 let fields = UnionFields::try_new(
5408 vec![0, 1],
5409 vec![
5410 Field::new("big", DataType::Int64, true),
5411 Field::new("small", DataType::Int32, true),
5412 ],
5413 )
5414 .unwrap();
5415 let mut builder = VariantArrayBuilder::new(3);
5416 builder.append_variant(Variant::Int8(1));
5417 builder.append_variant(Variant::Int32(2));
5418 builder.append_variant(Variant::Int64(3));
5419 let array = ArrayRef::from(builder.build());
5420
5421 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5422
5423 let expected: ArrayRef = Arc::new(
5424 UnionArray::try_new(
5425 fields,
5426 ScalarBuffer::from(vec![1i8, 1, 0]),
5427 Some(ScalarBuffer::from(vec![0i32, 1, 0])),
5428 vec![
5429 Arc::new(Int64Array::from(vec![3])),
5430 Arc::new(Int32Array::from(vec![1, 2])),
5431 ],
5432 )
5433 .unwrap(),
5434 );
5435 assert_eq!(&result, &expected);
5436 }
5437
5438 #[test]
5439 fn get_variant_as_union_with_encoded_children() {
5440 let encoded_types = [
5441 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
5442 DataType::RunEndEncoded(
5443 Arc::new(Field::new(
5444 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
5445 DataType::Int32,
5446 false,
5447 )),
5448 Arc::new(Field::new(
5449 Field::REE_VALUES_FIELD_DEFAULT_NAME,
5450 DataType::Utf8,
5451 true,
5452 )),
5453 ),
5454 ];
5455
5456 for data_type in encoded_types {
5457 let fields = UnionFields::try_new(
5458 vec![0],
5459 vec![Field::new("encoded", data_type.clone(), true)],
5460 )
5461 .unwrap();
5462 let mut builder = VariantArrayBuilder::new(2);
5463 builder.append_variant(Variant::from("apple"));
5464 builder.append_variant(Variant::from("banana"));
5465 let array = ArrayRef::from(builder.build());
5466 let options =
5467 union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions {
5468 safe: false,
5469 ..Default::default()
5470 });
5471
5472 let result = variant_get(&array, options).unwrap();
5473 let union = result.as_any().downcast_ref::<UnionArray>().unwrap();
5474 assert_eq!(union.type_ids(), &[0i8, 0]);
5475 assert_eq!(union.child(0).data_type(), &data_type);
5476
5477 let decoded = cast(union.child(0).as_ref(), &DataType::Utf8).unwrap();
5478 let expected = StringArray::from(vec!["apple", "banana"]);
5479 assert_eq!(decoded.as_ref(), &expected);
5480 }
5481 }
5482
5483 #[test]
5484 fn get_variant_as_union_with_fixed_size_list_child() {
5485 let item = Arc::new(Field::new("item", DataType::Int64, true));
5486 let fields = UnionFields::try_new(
5487 vec![0],
5488 vec![Field::new("fixed", DataType::FixedSizeList(item, 2), true)],
5489 )
5490 .unwrap();
5491 let json = StringArray::from(vec!["[1, 2]"]);
5492 let array = ArrayRef::from(json_to_variant(&(Arc::new(json) as ArrayRef)).unwrap());
5493
5494 for safe in [true, false] {
5495 let options =
5496 union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions {
5497 safe,
5498 ..Default::default()
5499 });
5500 let result = variant_get(&array, options).unwrap();
5501 let union = result.as_any().downcast_ref::<UnionArray>().unwrap();
5502 assert_eq!(union.type_ids(), &[0i8]);
5503 let list = union
5504 .child(0)
5505 .as_any()
5506 .downcast_ref::<FixedSizeListArray>()
5507 .unwrap();
5508 assert_eq!(
5509 list.value(0)
5510 .as_primitive::<arrow::datatypes::Int64Type>()
5511 .values(),
5512 &[1, 2]
5513 );
5514 }
5515 }
5516
5517 #[test]
5518 fn get_variant_as_union_skips_decimal_that_cannot_fit() {
5519 let fields = UnionFields::try_new(
5520 vec![0, 1],
5521 vec![
5522 Field::new("too_narrow", DataType::Decimal32(3, 2), true),
5523 Field::new("fits", DataType::Decimal32(5, 2), true),
5524 ],
5525 )
5526 .unwrap();
5527 let mut builder = VariantArrayBuilder::new(1);
5528 builder.append_variant(VariantDecimal4::try_new(12_345, 2).unwrap().into());
5529 let array = ArrayRef::from(builder.build());
5530
5531 for safe in [true, false] {
5532 let options =
5533 union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions {
5534 safe,
5535 ..Default::default()
5536 });
5537 let result = variant_get(&array, options).unwrap();
5538 let union = result.as_any().downcast_ref::<UnionArray>().unwrap();
5539 assert_eq!(union.type_ids(), &[1i8]);
5540 let decimal = union
5541 .child(1)
5542 .as_any()
5543 .downcast_ref::<Decimal32Array>()
5544 .unwrap();
5545 assert_eq!(decimal.value(0), 12_345);
5546 }
5547 }
5548
5549 #[test]
5550 fn get_variant_as_union_with_null_field() {
5551 let fields = UnionFields::try_new(
5553 vec![0, 1],
5554 vec![
5555 Field::new("int", DataType::Int64, true),
5556 Field::new("null", DataType::Null, true),
5557 ],
5558 )
5559 .unwrap();
5560 let mut builder = VariantArrayBuilder::new(4);
5561 builder.append_variant(Variant::Int8(1));
5562 builder.append_null();
5563 builder.append_variant(Variant::Null);
5564 builder.append_variant(Variant::from("no matching field"));
5565 let array = ArrayRef::from(builder.build());
5566
5567 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5568
5569 let expected: ArrayRef = Arc::new(
5570 UnionArray::try_new(
5571 fields,
5572 ScalarBuffer::from(vec![0i8, 1, 1, 1]),
5573 Some(ScalarBuffer::from(vec![0i32, 0, 1, 2])),
5574 vec![
5575 Arc::new(Int64Array::from(vec![1])),
5576 Arc::new(NullArray::new(3)),
5577 ],
5578 )
5579 .unwrap(),
5580 );
5581 assert_eq!(&result, &expected);
5582 }
5583
5584 #[test]
5585 fn get_variant_as_union_of_nested_types() {
5586 let fields = UnionFields::try_new(
5587 vec![0, 1, 2],
5588 vec![
5589 Field::new(
5590 "struct",
5591 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int64, true)])),
5592 true,
5593 ),
5594 Field::new(
5595 "list",
5596 DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
5597 true,
5598 ),
5599 Field::new("str", DataType::Utf8, true),
5600 ],
5601 )
5602 .unwrap();
5603 let json = StringArray::from(vec![r#"{"a": 1}"#, "[1, 2, 3]", "\"s\""]);
5604 let array = ArrayRef::from(json_to_variant(&(Arc::new(json) as ArrayRef)).unwrap());
5605
5606 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5607
5608 let mut list_builder = ListBuilder::new(Int64Builder::new());
5609 list_builder.append_value([Some(1), Some(2), Some(3)]);
5610 let expected: ArrayRef = Arc::new(
5611 UnionArray::try_new(
5612 fields,
5613 ScalarBuffer::from(vec![0i8, 1, 2]),
5614 Some(ScalarBuffer::from(vec![0i32, 0, 0])),
5615 vec![
5616 Arc::new(StructArray::from(vec![(
5617 Arc::new(Field::new("a", DataType::Int64, true)),
5618 Arc::new(Int64Array::from(vec![1])) as ArrayRef,
5619 )])),
5620 Arc::new(list_builder.finish()),
5621 Arc::new(StringArray::from(vec!["s"])),
5622 ],
5623 )
5624 .unwrap(),
5625 );
5626 assert_eq!(&result, &expected);
5627 }
5628
5629 #[test]
5630 fn get_variant_as_union_with_map_field() {
5631 let fields = UnionFields::try_new(
5633 vec![0, 1],
5634 vec![
5635 Field::new("map", map_data_type(DataType::Int64), true),
5636 Field::new("str", DataType::Utf8, true),
5637 ],
5638 )
5639 .unwrap();
5640 let json = StringArray::from(vec![r#"{"a": 1, "b": 2}"#, "\"hi\""]);
5641 let array = ArrayRef::from(json_to_variant(&(Arc::new(json) as ArrayRef)).unwrap());
5642
5643 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5644
5645 let mut map_builder = MapBuilder::new(None, StringBuilder::new(), Int64Builder::new());
5646 map_builder.keys().append_value("a");
5647 map_builder.values().append_value(1);
5648 map_builder.keys().append_value("b");
5649 map_builder.values().append_value(2);
5650 map_builder.append(true).unwrap();
5651 let expected: ArrayRef = Arc::new(
5652 UnionArray::try_new(
5653 fields,
5654 ScalarBuffer::from(vec![0i8, 1]),
5655 Some(ScalarBuffer::from(vec![0i32, 0])),
5656 vec![
5657 Arc::new(map_builder.finish()),
5658 Arc::new(StringArray::from(vec!["hi"])),
5659 ],
5660 )
5661 .unwrap(),
5662 );
5663 assert_eq!(&result, &expected);
5664 }
5665
5666 #[test]
5667 fn get_variant_as_union_prefers_struct_over_map() {
5668 let fields = UnionFields::try_new(
5671 vec![0, 1],
5672 vec![
5673 Field::new("map", map_data_type(DataType::Int64), true),
5674 Field::new(
5675 "struct",
5676 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int64, true)])),
5677 true,
5678 ),
5679 ],
5680 )
5681 .unwrap();
5682 let json = StringArray::from(vec![r#"{"a": 1}"#]);
5683 let array = ArrayRef::from(json_to_variant(&(Arc::new(json) as ArrayRef)).unwrap());
5684
5685 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5686 let union = result.as_any().downcast_ref::<UnionArray>().unwrap();
5687 assert_eq!(union.type_ids(), &[1i8]);
5689 }
5690
5691 #[test]
5692 fn get_variant_as_union_no_matching_field() {
5693 let fields =
5695 UnionFields::try_new(vec![0], vec![Field::new("str", DataType::Utf8, false)]).unwrap();
5696 let mut builder = VariantArrayBuilder::new(2);
5697 builder.append_variant(Variant::from("kept"));
5698 builder.append_variant(Variant::Int8(1));
5699 let array = ArrayRef::from(builder.build());
5700
5701 let result = variant_get(&array, union_get_options(&fields, UnionMode::Dense)).unwrap();
5703 let expected: ArrayRef = Arc::new(
5704 UnionArray::try_new(
5705 fields.clone(),
5706 ScalarBuffer::from(vec![0i8, 0]),
5707 Some(ScalarBuffer::from(vec![0i32, 1])),
5708 vec![Arc::new(StringArray::from(vec![Some("kept"), None]))],
5709 )
5710 .unwrap(),
5711 );
5712 assert_eq!(&result, &expected);
5713
5714 let options = union_get_options(&fields, UnionMode::Dense).with_cast_options(CastOptions {
5716 safe: false,
5717 ..Default::default()
5718 });
5719 let err = variant_get(&array, options).unwrap_err();
5720 assert!(
5721 err.to_string().contains("no field can represent it"),
5722 "unexpected error: {err}"
5723 );
5724 }
5725
5726 #[test]
5727 fn get_variant_as_union_empty_fields_errors() {
5728 let mut builder = VariantArrayBuilder::new(1);
5729 builder.append_variant(Variant::Int8(1));
5730 let array = ArrayRef::from(builder.build());
5731
5732 let err = variant_get(
5733 &array,
5734 union_get_options(&UnionFields::empty(), UnionMode::Dense),
5735 )
5736 .unwrap_err();
5737 assert!(
5738 err.to_string().contains("at least one union field"),
5739 "unexpected error: {err}"
5740 );
5741 }
5742}