Skip to content

Commit

Permalink
Add append_nulls and append_trusted_len_iter to PrimitiveBuilder (
Browse files Browse the repository at this point in the history
#728)

* stub out impl

* mark unsafe

* add tests
  • Loading branch information
bjchambers authored Sep 9, 2021
1 parent c4e20da commit 7675164
Showing 1 changed file with 77 additions and 0 deletions.
77 changes: 77 additions & 0 deletions arrow/src/array/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,23 @@ impl<T: ArrowNativeType> BufferBuilder<T> {
self.len += slice.len();
}

/// # Safety
/// This requires the iterator be a trusted length. This could instead require
/// the iterator implement `TrustedLen` once that is stabilized.
#[inline]
pub unsafe fn append_trusted_len_iter(&mut self, iter: impl IntoIterator<Item = T>) {
let iter = iter.into_iter();
let len = iter
.size_hint()
.1
.expect("append_trusted_len_iter expects upper bound");
self.reserve(len);
for v in iter {
self.buffer.push(v)
}
self.len += len;
}

/// Resets this builder and returns an immutable [`Buffer`](crate::buffer::Buffer).
///
/// # Example:
Expand Down Expand Up @@ -695,6 +712,14 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
Ok(())
}

#[inline]
pub fn append_nulls(&mut self, n: usize) -> Result<()> {
self.materialize_bitmap_builder();
self.bitmap_builder.as_mut().unwrap().append_n(n, false);
self.values_builder.advance(n);
Ok(())
}

/// Appends an `Option<T>` into the builder
#[inline]
pub fn append_option(&mut self, v: Option<T::Native>) -> Result<()> {
Expand Down Expand Up @@ -737,6 +762,29 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
Ok(())
}

/// Appends values from a trusted length iterator.
///
/// # Safety
/// This requires the iterator be a trusted length. This could instead require
/// the iterator implement `TrustedLen` once that is stabilized.
#[inline]
pub unsafe fn append_trusted_len_iter(
&mut self,
iter: impl IntoIterator<Item = T::Native>,
) -> Result<()> {
let iter = iter.into_iter();
let len = iter
.size_hint()
.1
.expect("append_trusted_len_iter requires an upper bound");

if let Some(b) = self.bitmap_builder.as_mut() {
b.append_n(len, true);
}
self.values_builder.append_trusted_len_iter(iter);
Ok(())
}

/// Builds the `PrimitiveArray` and reset this builder.
pub fn finish(&mut self) -> PrimitiveArray<T> {
let len = self.len();
Expand Down Expand Up @@ -2739,6 +2787,35 @@ mod tests {
}
}

#[test]
fn test_primitive_array_builder_i32_append_iter() {
let mut builder = Int32Array::builder(5);
unsafe { builder.append_trusted_len_iter(0..5) }.unwrap();
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(0, arr.null_count());
for i in 0..5 {
assert!(!arr.is_null(i));
assert!(arr.is_valid(i));
assert_eq!(i as i32, arr.value(i));
}
}

#[test]
fn test_primitive_array_builder_i32_append_nulls() {
let mut builder = Int32Array::builder(5);
builder.append_nulls(5).unwrap();
let arr = builder.finish();
assert_eq!(5, arr.len());
assert_eq!(0, arr.offset());
assert_eq!(5, arr.null_count());
for i in 0..5 {
assert!(arr.is_null(i));
assert!(!arr.is_valid(i));
}
}

#[test]
fn test_primitive_array_builder_date32() {
let mut builder = Date32Array::builder(5);
Expand Down

0 comments on commit 7675164

Please sign in to comment.