diff --git a/rust/arrow/src/array/array_primitive.rs b/rust/arrow/src/array/array_primitive.rs index 3dcb187495f9a..984454ce779ed 100644 --- a/rust/arrow/src/array/array_primitive.rs +++ b/rust/arrow/src/array/array_primitive.rs @@ -94,6 +94,21 @@ impl PrimitiveArray { let offset = i + self.offset(); unsafe { *self.raw_values.as_ptr().add(offset) } } + + /// Creates a PrimitiveArray based on an iterator of values without nulls + pub fn from_iter_values>(iter: I) -> Self { + let val_buf: Buffer = iter.into_iter().collect(); + let data = ArrayData::new( + T::DATA_TYPE, + val_buf.len() / mem::size_of::<::Native>(), + None, + None, + 0, + vec![val_buf], + vec![], + ); + PrimitiveArray::from(Arc::new(data)) + } } impl Array for PrimitiveArray { @@ -820,6 +835,18 @@ mod tests { } } + #[test] + fn test_primitive_from_iter_values() { + // Test building a primitive array with from_iter_values + + let arr: PrimitiveArray = PrimitiveArray::from_iter_values(0..10); + assert_eq!(10, arr.len()); + assert_eq!(0, arr.null_count()); + for i in 0..10i32 { + assert_eq!(i, arr.value(i as usize)); + } + } + #[test] #[should_panic(expected = "PrimitiveArray data should contain a single buffer only \ (values buffer)")] diff --git a/rust/arrow/src/array/array_string.rs b/rust/arrow/src/array/array_string.rs index f1523e35c3bb6..4b98fafdeafeb 100644 --- a/rust/arrow/src/array/array_string.rs +++ b/rust/arrow/src/array/array_string.rs @@ -146,6 +146,36 @@ impl GenericStringArray { pub(crate) fn from_opt_vec(v: Vec>) -> Self { v.into_iter().collect() } + + /// Creates a `GenericStringArray` based on an iterator of values without nulls + pub fn from_iter_values>(iter: I) -> Self + where + Ptr: AsRef, + { + let iter = iter.into_iter(); + let (_, data_len) = iter.size_hint(); + let data_len = data_len.expect("Iterator must be sized"); // panic if no upper bound. + + let mut offsets = + MutableBuffer::new((data_len + 1) * std::mem::size_of::()); + let mut values = MutableBuffer::new(0); + + let mut length_so_far = OffsetSize::zero(); + offsets.push(length_so_far); + + for i in iter { + let s = i.as_ref(); + length_so_far += OffsetSize::from_usize(s.len()).unwrap(); + offsets.push(length_so_far); + values.extend_from_slice(s.as_bytes()); + } + let array_data = ArrayData::builder(OffsetSize::DATA_TYPE) + .len(data_len) + .add_buffer(offsets.into()) + .add_buffer(values.into()) + .build(); + Self::from(array_data) + } } impl<'a, Ptr, OffsetSize: StringOffsetSizeTrait> FromIterator> @@ -411,6 +441,7 @@ mod tests { ); } + #[test] fn test_string_array_from_iter() { let data = vec![Some("hello"), None, Some("arrow")]; // from Vec> @@ -424,4 +455,13 @@ mod tests { assert_eq!(array1, array2); assert_eq!(array2, array3); } + + #[test] + fn test_string_array_from_iter_values() { + let data = vec!["hello", "hello2"]; + let array1 = StringArray::from_iter_values(data.iter()); + + assert_eq!(array1.value(0), "hello"); + assert_eq!(array1.value(1), "hello2"); + } }