Skip to content

Commit ea264c4

Browse files
Remove ceremony from iterator of variants into VariantArray
1 parent 751b082 commit ea264c4

File tree

1 file changed

+60
-0
lines changed

1 file changed

+60
-0
lines changed

parquet-variant-compute/src/variant_array.rs

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
//! [`VariantArray`] implementation
1919
20+
use crate::VariantArrayBuilder;
2021
use crate::type_conversion::{generic_conversion_single_value, primitive_conversion_single_value};
2122
use arrow::array::{Array, ArrayRef, AsArray, BinaryViewArray, StructArray};
2223
use arrow::buffer::NullBuffer;
@@ -439,6 +440,20 @@ impl From<VariantArray> for ArrayRef {
439440
}
440441
}
441442

443+
impl<'m, 'v> FromIterator<Option<Variant<'m, 'v>>> for VariantArray {
444+
fn from_iter<T: IntoIterator<Item = Option<Variant<'m, 'v>>>>(iter: T) -> Self {
445+
let mut b = VariantArrayBuilder::new(0);
446+
b.extend(iter);
447+
b.build()
448+
}
449+
}
450+
451+
impl<'m, 'v> FromIterator<Variant<'m, 'v>> for VariantArray {
452+
fn from_iter<T: IntoIterator<Item = Variant<'m, 'v>>>(iter: T) -> Self {
453+
Self::from_iter(iter.into_iter().map(Some))
454+
}
455+
}
456+
442457
/// An iterator over [`VariantArray`]
443458
///
444459
/// This iterator returns `Option<Option<Variant<'a, 'a>>>` where:
@@ -1141,6 +1156,7 @@ mod test {
11411156
use super::*;
11421157
use arrow::array::{BinaryViewArray, Int32Array};
11431158
use arrow_schema::{Field, Fields};
1159+
use parquet_variant::ShortString;
11441160

11451161
#[test]
11461162
fn invalid_not_a_struct_array() {
@@ -1405,4 +1421,48 @@ mod test {
14051421
assert!(i.next().is_none());
14061422
assert!(i.next_back().is_none());
14071423
}
1424+
1425+
#[test]
1426+
fn test_from_variant_opts_into_variant_array() {
1427+
let v = vec![None, Some(Variant::Null), Some(Variant::BooleanFalse), None];
1428+
1429+
let variant_array = VariantArray::from_iter(v);
1430+
1431+
assert_eq!(variant_array.len(), 4);
1432+
1433+
assert!(variant_array.is_null(0));
1434+
1435+
assert!(!variant_array.is_null(1));
1436+
assert_eq!(variant_array.value(1), Variant::Null);
1437+
1438+
assert!(!variant_array.is_null(2));
1439+
assert_eq!(variant_array.value(2), Variant::BooleanFalse);
1440+
1441+
assert!(variant_array.is_null(3));
1442+
}
1443+
1444+
#[test]
1445+
fn test_from_variants_into_variant_array() {
1446+
let v = vec![
1447+
Variant::Null,
1448+
Variant::BooleanFalse,
1449+
Variant::ShortString(ShortString::try_new("norm").unwrap()),
1450+
];
1451+
1452+
let variant_array = VariantArray::from_iter(v);
1453+
1454+
assert_eq!(variant_array.len(), 3);
1455+
1456+
assert!(!variant_array.is_null(0));
1457+
assert_eq!(variant_array.value(0), Variant::Null);
1458+
1459+
assert!(!variant_array.is_null(1));
1460+
assert_eq!(variant_array.value(1), Variant::BooleanFalse);
1461+
1462+
assert!(!variant_array.is_null(3));
1463+
assert_eq!(
1464+
variant_array.value(2),
1465+
Variant::ShortString(ShortString::try_new("norm").unwrap())
1466+
);
1467+
}
14081468
}

0 commit comments

Comments
 (0)