diff --git a/arrow-array/src/array/list_view_array.rs b/arrow-array/src/array/list_view_array.rs index cd1c0b99002e..35e0c4ed9402 100644 --- a/arrow-array/src/array/list_view_array.rs +++ b/arrow-array/src/array/list_view_array.rs @@ -23,8 +23,12 @@ use std::ops::Add; use std::sync::Arc; use crate::array::{make_array, print_long_array}; +use crate::builder::{GenericListViewBuilder, PrimitiveBuilder}; use crate::iterator::GenericListViewArrayIter; -use crate::{Array, ArrayAccessor, ArrayRef, FixedSizeListArray, OffsetSizeTrait, new_empty_array}; +use crate::{ + Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, FixedSizeListArray, OffsetSizeTrait, + new_empty_array, +}; /// A [`GenericListViewArray`] of variable size lists, storing offsets as `i32`. pub type ListViewArray = GenericListViewArray; @@ -357,6 +361,46 @@ impl GenericListViewArray { value_sizes: self.value_sizes.slice(offset, length), } } + + /// Creates a [`GenericListViewArray`] from an iterator of primitive values + /// # Example + /// ``` + /// # use arrow_array::ListViewArray; + /// # use arrow_array::types::Int32Type; + /// + /// let data = vec![ + /// Some(vec![Some(0), Some(1), Some(2)]), + /// None, + /// Some(vec![Some(3), None, Some(5)]), + /// Some(vec![Some(6), Some(7)]), + /// ]; + /// let list_array = ListViewArray::from_iter_primitive::(data); + /// println!("{:?}", list_array); + /// ``` + pub fn from_iter_primitive(iter: I) -> Self + where + T: ArrowPrimitiveType, + P: IntoIterator::Native>>, + I: IntoIterator>, + { + let iter = iter.into_iter(); + let size_hint = iter.size_hint().0; + let mut builder = + GenericListViewBuilder::with_capacity(PrimitiveBuilder::::new(), size_hint); + + for i in iter { + match i { + Some(p) => { + for t in p { + builder.values().append_option(t); + } + builder.append(true); + } + None => builder.append(false), + } + } + builder.finish() + } } impl ArrayAccessor for &GenericListViewArray { @@ -559,7 +603,7 @@ impl GenericListViewArray { #[cfg(test)] mod tests { - use arrow_buffer::{BooleanBuffer, Buffer, ScalarBuffer, bit_util}; + use arrow_buffer::{BooleanBuffer, Buffer, NullBufferBuilder, ScalarBuffer, bit_util}; use arrow_schema::Field; use crate::builder::{FixedSizeListBuilder, Int32Builder}; @@ -1120,4 +1164,29 @@ mod tests { .collect(); assert_eq!(values, vec![Some(vec![]), Some(vec![]), Some(vec![])]); } + + #[test] + fn test_from_iter_primitive() { + let data = vec![ + Some(vec![Some(0), Some(1), Some(2)]), + None, + Some(vec![Some(3), Some(4), Some(5)]), + Some(vec![Some(6), Some(7)]), + ]; + let list_array = ListViewArray::from_iter_primitive::(data); + + // [[0, 1, 2], NULL, [3, 4, 5], [6, 7]] + let values = Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7]); + let offsets = ScalarBuffer::from(vec![0, 3, 3, 6]); + let sizes = ScalarBuffer::from(vec![3, 0, 3, 2]); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + + let mut nulls = NullBufferBuilder::new(4); + nulls.append(true); + nulls.append(false); + nulls.append_n_non_nulls(2); + let another = ListViewArray::new(field, offsets, sizes, Arc::new(values), nulls.finish()); + + assert_eq!(list_array, another) + } }