diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs index e7fd7262b29..acab58649b0 100644 --- a/src/array/fixed_size_list/mutable.rs +++ b/src/array/fixed_size_list/mutable.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use crate::{ array::{Array, MutableArray, TryExtend, TryPush}, bitmap::MutableBitmap, - datatypes::DataType, + datatypes::{DataType, Field}, error::{ArrowError, Result}, }; @@ -41,6 +41,21 @@ impl MutableFixedSizeListArray { } } + /// Creates a new [`MutableFixedSizeListArray`] from a [`MutableArray`] and size. + pub fn new_with_field(values: M, name: &str, nullable: bool, size: usize) -> Self { + let data_type = DataType::FixedSizeList( + Box::new(Field::new(name, values.data_type().clone(), nullable)), + size, + ); + assert_eq!(values.len(), 0); + Self { + size, + data_type, + values, + validity: None, + } + } + /// The inner values pub fn values(&self) -> &M { &self.values @@ -61,7 +76,9 @@ impl MutableFixedSizeListArray { } #[inline] - fn try_push_valid(&mut self) -> Result<()> { + /// Needs to be called when a valid value was extended to this array. + /// This is a relatively low level function, prefer `try_push` when you can. + pub fn try_push_valid(&mut self) -> Result<()> { if self.values.len() % self.size != 0 { return Err(ArrowError::Overflow); }; diff --git a/tests/it/array/fixed_size_list/mutable.rs b/tests/it/array/fixed_size_list/mutable.rs index 2fd437144cc..39826645ba7 100644 --- a/tests/it/array/fixed_size_list/mutable.rs +++ b/tests/it/array/fixed_size_list/mutable.rs @@ -1,4 +1,5 @@ use arrow2::array::*; +use arrow2::datatypes::{DataType, Field}; #[test] fn primitive() { @@ -24,3 +25,41 @@ fn primitive() { let expected = Int32Array::from(vec![None, None, None]); assert_eq!(a, &expected) } + +#[test] +fn new_with_field() { + let data = vec![ + Some(vec![Some(1i32), Some(2), Some(3)]), + Some(vec![None, None, None]), + Some(vec![Some(4), None, Some(6)]), + ]; + + let mut list = MutableFixedSizeListArray::new_with_field( + MutablePrimitiveArray::::new(), + "custom_items", + false, + 3, + ); + list.try_extend(data).unwrap(); + let list: FixedSizeListArray = list.into(); + + assert_eq!( + list.data_type(), + &DataType::FixedSizeList( + Box::new(Field::new("custom_items", DataType::Int32, false)), + 3 + ) + ); + + let a = list.value(0); + let a = a.as_any().downcast_ref::().unwrap(); + + let expected = Int32Array::from(vec![Some(1i32), Some(2), Some(3)]); + assert_eq!(a, &expected); + + let a = list.value(1); + let a = a.as_any().downcast_ref::().unwrap(); + + let expected = Int32Array::from(vec![None, None, None]); + assert_eq!(a, &expected) +}