Remove ceremony from iterator of variants into VariantArray

Author: friendlymatthew

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<VariantArray> for ArrayRef {
     }
 }
 
+impl<'m, 'v> FromIterator<Option<Variant<'m, 'v>>> for VariantArray {
+    fn from_iter<T: IntoIterator<Item = Option<Variant<'m, 'v>>>>(iter: T) -> Self {
+        let mut b = VariantArrayBuilder::new(0);
+        b.extend(iter);
+        b.build()
+    }
+}
+
+impl<'m, 'v> FromIterator<Variant<'m, 'v>> for VariantArray {
+    fn from_iter<T: IntoIterator<Item = Variant<'m, 'v>>>(iter: T) -> Self {
+        let mut b = VariantArrayBuilder::new(0);
+        b.extend(iter.into_iter().map(Some));
+        b.build()
+    }
+}
+
 /// An iterator over [`VariantArray`]
 ///
 /// This iterator returns `Option<Option<Variant<'a, 'a>>>` 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.into_iter());
+
+        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())
+        );
+    }
 }