diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index 5f7f826819d8..6ce84f232aae 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -17,6 +17,7 @@ //! [`VariantArray`] implementation +use crate::VariantArrayBuilder; use crate::type_conversion::{generic_conversion_single_value, primitive_conversion_single_value}; use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray, StructArray}; use arrow::buffer::NullBuffer; @@ -439,6 +440,22 @@ impl From for ArrayRef { } } +impl<'m, 'v> FromIterator>> for VariantArray { + fn from_iter>>>(iter: T) -> Self { + let iter = iter.into_iter(); + + let mut b = VariantArrayBuilder::new(iter.size_hint().0); + b.extend(iter); + b.build() + } +} + +impl<'m, 'v> FromIterator> for VariantArray { + fn from_iter>>(iter: T) -> Self { + Self::from_iter(iter.into_iter().map(Some)) + } +} + /// An iterator over [`VariantArray`] /// /// This iterator returns `Option>>` where: @@ -1141,6 +1158,7 @@ mod test { use super::*; use arrow::array::{BinaryViewArray, Int32Array}; use arrow_schema::{Field, Fields}; + use parquet_variant::ShortString; #[test] fn invalid_not_a_struct_array() { @@ -1405,4 +1423,48 @@ mod test { assert!(i.next().is_none()); assert!(i.next_back().is_none()); } + + #[test] + fn test_from_variant_opts_into_variant_array() { + let v = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse), None]; + + let variant_array = VariantArray::from_iter(v); + + assert_eq!(variant_array.len(), 4); + + assert!(variant_array.is_null(0)); + + assert!(!variant_array.is_null(1)); + assert_eq!(variant_array.value(1), Variant::Null); + + assert!(!variant_array.is_null(2)); + assert_eq!(variant_array.value(2), Variant::BooleanFalse); + + assert!(variant_array.is_null(3)); + } + + #[test] + fn test_from_variants_into_variant_array() { + let v = vec![ + Variant::Null, + Variant::BooleanFalse, + Variant::ShortString(ShortString::try_new("norm").unwrap()), + ]; + + let variant_array = VariantArray::from_iter(v); + + assert_eq!(variant_array.len(), 3); + + assert!(!variant_array.is_null(0)); + assert_eq!(variant_array.value(0), Variant::Null); + + assert!(!variant_array.is_null(1)); + assert_eq!(variant_array.value(1), Variant::BooleanFalse); + + assert!(!variant_array.is_null(3)); + assert_eq!( + variant_array.value(2), + Variant::ShortString(ShortString::try_new("norm").unwrap()) + ); + } }