From 9b400145620dbdb00c5ab5d0752e22de377ab791 Mon Sep 17 00:00:00 2001
From: kikkon <nian920@outlook.com>
Date: Sun, 3 Mar 2024 16:29:23 +0800
Subject: [PATCH] Enhancement: Replace unsafe code

---
 arrow-buffer/src/buffer/immutable.rs | 23 +++--------------------
 1 file changed, 3 insertions(+), 20 deletions(-)

diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs
index ca3ee6260ba7..10b62a2ce473 100644
--- a/arrow-buffer/src/buffer/immutable.rs
+++ b/arrow-buffer/src/buffer/immutable.rs
@@ -423,25 +423,8 @@ impl Buffer {
 
 impl<T: ArrowNativeType> FromIterator<T> for Buffer {
     fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
-        let mut iterator = iter.into_iter();
-        let size = std::mem::size_of::<T>();
-
-        // first iteration, which will likely reserve sufficient space for the buffer.
-        let mut buffer = match iterator.next() {
-            None => MutableBuffer::new(0),
-            Some(element) => {
-                let (lower, _) = iterator.size_hint();
-                let mut buffer = MutableBuffer::new(lower.saturating_add(1).saturating_mul(size));
-                unsafe {
-                    std::ptr::write(buffer.as_mut_ptr() as *mut T, element);
-                    buffer.set_len(size);
-                }
-                buffer
-            }
-        };
-
-        buffer.extend_from_iter(iterator);
-        buffer.into()
+        let vec = Vec::from_iter(iter);
+        Buffer::from_vec(vec)
     }
 }
 
@@ -822,7 +805,7 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "failed to round to next highest power of 2")]
+    #[should_panic(expected = "capacity overflow")]
     fn test_from_iter_overflow() {
         let iter_len = usize::MAX / std::mem::size_of::<u64>() + 1;
         let _ = Buffer::from_iter(std::iter::repeat(0_u64).take(iter_len));