Mark MutableBuffer::typed_data_mut unsafe

arrow/src/buffer/mutable.rs

@@ -273,15 +273,20 @@ impl MutableBuffer {
     }
 
     /// View this buffer asa slice of a specific type.
+    ///
     /// # Safety
-    /// This function must only be used when this buffer was extended with items of type `T`.
-    /// Failure to do so results in undefined behavior.
-    pub fn typed_data_mut<T: ArrowNativeType>(&mut self) -> &mut [T] {
-        unsafe {
-            let (prefix, offsets, suffix) = self.as_slice_mut().align_to_mut::<T>();
-            assert!(prefix.is_empty() && suffix.is_empty());
-            offsets
-        }
+    ///
+    /// This function must only be used with buffers which are treated
+    /// as type `T` (e.g.  extended with items of type `T`).
+    ///
+    /// # Panics
+    ///
+    /// This function panics if the underlying buffer is not aligned
+    /// correctly for type `T`.
+    pub unsafe fn typed_data_mut<T: ArrowNativeType>(&mut self) -> &mut [T] {
+        let (prefix, offsets, suffix) = self.as_slice_mut().align_to_mut::<T>();
+        assert!(prefix.is_empty() && suffix.is_empty());
+        offsets
     }
 
     /// Extends this buffer from a slice of items that can be represented in bytes, increasing its capacity if needed.

arrow/src/buffer/ops.rs

@@ -168,7 +168,10 @@ where
         MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false);
 
     let left_chunks = left.bit_chunks(offset_in_bits, len_in_bits);
-    let result_chunks = result.typed_data_mut::<u64>().iter_mut();
+
+    // Safety: buffer is always treated as type `u64` in the code
+    // below.
+    let result_chunks = unsafe { result.typed_data_mut::<u64>().iter_mut() };
 
     result_chunks
         .zip(left_chunks.iter())

arrow/src/compute/kernels/sort.rs

@@ -471,7 +471,8 @@ fn sort_boolean(
     let mut result = MutableBuffer::new(result_capacity);
     // sets len to capacity so we can access the whole buffer as a typed slice
     result.resize(result_capacity, 0);
-    let result_slice: &mut [u32] = result.typed_data_mut();
+    // Safety: the buffer is always treated as `u32` in the code below
+    let result_slice: &mut [u32] = unsafe { result.typed_data_mut() };
 
     if options.nulls_first {
         let size = nulls_len.min(len);
@@ -559,7 +560,8 @@ where
     let mut result = MutableBuffer::new(result_capacity);
     // sets len to capacity so we can access the whole buffer as a typed slice
     result.resize(result_capacity, 0);
-    let result_slice: &mut [u32] = result.typed_data_mut();
+    // Safety: the buffer is always treated as `u32` in the code below
+    let result_slice: &mut [u32] = unsafe { result.typed_data_mut() };
 
     if options.nulls_first {
         let size = nulls_len.min(len);

arrow/src/compute/kernels/take.rs

@@ -632,7 +632,8 @@ where
     let bytes_offset = (data_len + 1) * std::mem::size_of::<OffsetSize>();
     let mut offsets_buffer = MutableBuffer::from_len_zeroed(bytes_offset);
 
-    let offsets = offsets_buffer.typed_data_mut();
+    // Safety: the buffer is always treated as as a type of `OffsetSize` in the code below
+    let offsets = unsafe { offsets_buffer.typed_data_mut() };
     let mut values = MutableBuffer::new(0);
     let mut length_so_far = OffsetSize::zero();
     offsets[0] = length_so_far;

parquet/src/arrow/array_reader.rs

@@ -1154,7 +1154,8 @@ impl ArrayReader for StructArrayReader {
         let mut def_level_data_buffer = MutableBuffer::new(buffer_size);
         def_level_data_buffer.resize(buffer_size, 0);
 
-        let def_level_data = def_level_data_buffer.typed_data_mut();
+        // Safety: the buffer is always treated as `u16` in the code below
+        let def_level_data = unsafe { def_level_data_buffer.typed_data_mut() };
 
         def_level_data
             .iter_mut()