Implement FromIterator for Builders

arrow/src/array/array_boolean.rs

@@ -222,38 +222,7 @@ impl<'a> BooleanArray {
 
 impl<Ptr: Borrow<Option<bool>>> FromIterator<Ptr> for BooleanArray {
     fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
-        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 num_bytes = bit_util::ceil(data_len, 8);
-        let mut null_buf = MutableBuffer::from_len_zeroed(num_bytes);
-        let mut val_buf = MutableBuffer::from_len_zeroed(num_bytes);
-
-        let data = val_buf.as_slice_mut();
-
-        let null_slice = null_buf.as_slice_mut();
-        iter.enumerate().for_each(|(i, item)| {
-            if let Some(a) = item.borrow() {
-                bit_util::set_bit(null_slice, i);
-                if *a {
-                    bit_util::set_bit(data, i);
-                }
-            }
-        });
-
-        let data = unsafe {
-            ArrayData::new_unchecked(
-                DataType::Boolean,
-                data_len,
-                None,
-                Some(null_buf.into()),
-                0,
-                vec![val_buf.into()],
-                vec![],
-            )
-        };
-        BooleanArray::from(data)
+        iter.into_iter().collect::<BooleanBuilder>().finish()
     }
 }
 

arrow/src/array/builder/boolean_builder.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 use std::any::Any;
+use std::borrow::Borrow;
 use std::sync::Arc;
 
 use crate::array::ArrayBuilder;
@@ -57,6 +58,24 @@ use super::BooleanBufferBuilder;
 ///     assert!(arr.is_valid(3));
 ///     assert!(!arr.is_null(3));
 /// ```
+///
+/// Using `from_iter`
+/// ```
+///     use arrow::array::{Array, BooleanBuilder};
+///     let v = vec![Some(false), Some(true), Some(false), Some(true)];
+///     let arr = v.into_iter().collect::<BooleanBuilder>().finish();
+///     assert_eq!(4, arr.len());
+///     assert_eq!(0, arr.offset());
+///     assert_eq!(0, arr.null_count());
+///     assert!(arr.is_valid(0));
+///     assert_eq!(false, arr.value(0));
+///     assert!(arr.is_valid(1));
+///     assert_eq!(true, arr.value(1));
+///     assert!(arr.is_valid(2));
+///     assert_eq!(false, arr.value(2));
+///     assert!(arr.is_valid(3));
+///     assert_eq!(true, arr.value(3));
+/// ```
 #[derive(Debug)]
 pub struct BooleanBuilder {
     values_builder: BooleanBufferBuilder,
@@ -140,6 +159,21 @@ impl BooleanBuilder {
 }
 
 impl ArrayBuilder for BooleanBuilder {
+    /// Returns the number of array slots in the builder
+    fn len(&self) -> usize {
+        self.values_builder.len()
+    }
+
+    /// Returns whether the number of array slots is zero
+    fn is_empty(&self) -> bool {
+        self.values_builder.is_empty()
+    }
+
+    /// Builds the array and reset this builder.
+    fn finish(&mut self) -> ArrayRef {
+        Arc::new(self.finish())
+    }
+
     /// Returns the builder as a non-mutable `Any` reference.
     fn as_any(&self) -> &dyn Any {
         self
@@ -154,26 +188,38 @@ impl ArrayBuilder for BooleanBuilder {
     fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
         self
     }
+}
 
-    /// Returns the number of array slots in the builder
-    fn len(&self) -> usize {
-        self.values_builder.len()
-    }
+impl<Ptr: Borrow<Option<bool>>> FromIterator<Ptr> for BooleanBuilder {
+    fn from_iter<T: IntoIterator<Item = Ptr>>(iter: T) -> Self {
+        let iter = iter.into_iter();
+        let (lower, upper) = iter.size_hint();
+        let size_hint = upper.unwrap_or(lower);
 
-    /// Returns whether the number of array slots is zero
-    fn is_empty(&self) -> bool {
-        self.values_builder.is_empty()
-    }
+        let mut bitmap_builder = BooleanBufferBuilder::new(size_hint);
+        let mut values_builder = BooleanBufferBuilder::new(size_hint);
 
-    /// Builds the array and reset this builder.
-    fn finish(&mut self) -> ArrayRef {
-        Arc::new(self.finish())
+        iter.for_each(|item| {
+            if let Some(a) = item.borrow() {
+                bitmap_builder.append(true);
+                values_builder.append(*a);
+            } else {
+                bitmap_builder.append(false);
+                values_builder.append(false);
+            }
+        });
+
+        BooleanBuilder {
+            values_builder,
+            bitmap_builder,
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::array::Array;
 
     #[test]
     fn test_boolean_array_builder_append_slice() {
@@ -200,4 +246,18 @@ mod tests {
 
         assert_eq!(arr1, arr2);
     }
+
+    #[test]
+    fn test_boolean_array_builder_from_iter() {
+        let v = vec![Some(false), Some(true), Some(false), Some(true)];
+        let arr = v.into_iter().collect::<BooleanBuilder>().finish();
+        assert_eq!(4, arr.len());
+        assert_eq!(0, arr.offset());
+        assert_eq!(0, arr.null_count());
+        for i in 0..3 {
+            assert!(!arr.is_null(i));
+            assert!(arr.is_valid(i));
+            assert_eq!(i == 1 || i == 3, arr.value(i), "failed at {}", i)
+        }
+    }
 }