Optimize PrimitiveBuilder::append_trusted_len_iter

Author: alamb

arrow-array/src/builder/primitive_builder.rs

@@ -284,14 +284,10 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
     /// the iterator implement `TrustedLen` once that is stabilized.
     #[inline]
     pub unsafe fn append_trusted_len_iter(&mut self, iter: impl IntoIterator<Item = T::Native>) {
-        let iter = iter.into_iter();
-        let len = iter
-            .size_hint()
-            .1
-            .expect("append_trusted_len_iter requires an upper bound");
-
-        self.null_buffer_builder.append_n_non_nulls(len);
+        let starting_len = self.len();
         self.values_builder.append_trusted_len_iter(iter);
+        self.null_buffer_builder
+            .append_n_non_nulls(self.len() - starting_len);
     }
 
     /// Builds the [`PrimitiveArray`] and reset this builder.

arrow-buffer/src/buffer/mutable.rs

@@ -543,8 +543,42 @@ impl MutableBuffer {
         iterator.for_each(|item| self.push(item));
     }
 
+    /// Extends a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length.
+    ///
+    /// See [`MutableBuffer::from_trusted_len_iter`] for more details.
+    ///
+    /// # Safety
+    /// This method assumes that the iterator's size is correct and is undefined
+    /// behavior to use it on an iterator that reports an incorrect length.
+    #[inline]
+    pub unsafe fn extend_from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
+        &mut self,
+        iterator: I,
+    ) {
+        let item_size = std::mem::size_of::<T>();
+        let (lower, _) = iterator.size_hint();
+        let additional = lower * item_size;
+        self.reserve(additional);
+
+        // this is necessary because of https://github.com/rust-lang/rust/issues/32155
+        let mut dst = unsafe { self.data.as_ptr().add(self.len) };
+        for item in iterator {
+            // note how there is no reserve here (compared with `extend_from_iter`)
+            let src = item.to_byte_slice().as_ptr();
+            std::ptr::copy_nonoverlapping(src, dst, item_size);
+            dst = dst.add(item_size);
+        }
+        self.len += additional;
+        assert_eq!(
+            dst.offset_from(self.data.as_ptr()) as usize,
+            self.len,
+            "Trusted iterator length was not accurately reported"
+        );
+    }
+
     /// Creates a [`MutableBuffer`] from an [`Iterator`] with a trusted (upper) length.
-    /// Prefer this to `collect` whenever possible, as it is faster ~60% faster.
+    /// Prefer this to `collect` whenever possible, as it is ~60% faster.
+    ///
     /// # Example
     /// ```
     /// # use arrow_buffer::buffer::MutableBuffer;
@@ -554,16 +588,18 @@ impl MutableBuffer {
     /// assert_eq!(buffer.len(), 4) // u32 has 4 bytes
     /// ```
     /// # Safety
-    /// This method assumes that the iterator's size is correct and is undefined behavior
-    /// to use it on an iterator that reports an incorrect length.
-    // This implementation is required for two reasons:
-    // 1. there is no trait `TrustedLen` in stable rust and therefore
-    //    we can't specialize `extend` for `TrustedLen` like `Vec` does.
-    // 2. `from_trusted_len_iter` is faster.
+    /// This method assumes that the iterator's size is correct and is undefined
+    /// behavior to use it on an iterator that reports an incorrect length.
+    ///
+    /// This implementation is required for two reasons:
+    /// 1. there is no trait `TrustedLen` in stable Rust and therefore
+    ///    we can't specialize `extend` for `TrustedLen` like `Vec` does.
+    /// 2. `from_trusted_len_iter` is faster.
     #[inline]
     pub unsafe fn from_trusted_len_iter<T: ArrowNativeType, I: Iterator<Item = T>>(
         iterator: I,
     ) -> Self {
+        // TODO: reduce duplication with `extend_from_trusted_len_iter`
         let item_size = std::mem::size_of::<T>();
         let (_, upper) = iterator.size_hint();
         let upper = upper.expect("from_trusted_len_iter requires an upper limit");

arrow-buffer/src/builder/mod.rs

@@ -330,13 +330,13 @@ impl<T: ArrowNativeType> BufferBuilder<T> {
     /// 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);
-        self.extend(iter);
+        let starting_buffer_len = self.len;
+        debug_assert_eq!(
+            starting_buffer_len,
+            self.buffer.len() / std::mem::size_of::<T>()
+        );
+        self.buffer.extend_from_trusted_len_iter(iter.into_iter());
+        self.len = self.buffer.len() / std::mem::size_of::<T>();
     }
 
     /// Resets this builder and returns an immutable [Buffer].