diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index fc001afe5a07..0bbe3781585e 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -106,11 +106,18 @@ impl Buffer { /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` /// /// self.ptr and self.data can be different after slicing or advancing the buffer. + /// + /// to check if the buffer is sliced you can call [`Self::is_sliced`] pub fn ptr_offset(&self) -> usize { // Safety: `ptr` is always in bounds of `data`. unsafe { self.ptr.offset_from(self.data.ptr().as_ptr()) as usize } } + /// Returns whether the buffer is sliced (does not point to entire original data) + pub fn is_sliced(&self) -> bool { + self.length != self.data.len() + } + /// Returns the pointer to the start of the buffer without the offset. pub fn data_ptr(&self) -> NonNull { self.data.ptr() @@ -362,6 +369,11 @@ impl Buffer { /// Returns `Err` if this is shared or its allocation is from an external source or /// it is not allocated with alignment [`ALIGNMENT`] /// + /// If the buffer is sliced, the returned [`MutableBuffer`] will be a view of the original buffer + /// (include values outside the sliced data). + /// + /// you can check if the buffer is sliced by calling [`Self::is_sliced`] + /// /// # Example: Creating a [`MutableBuffer`] from a [`Buffer`] /// ``` /// # use arrow_buffer::buffer::{Buffer, MutableBuffer}; @@ -382,11 +394,13 @@ impl Buffer { /// [`ALIGNMENT`]: crate::alloc::ALIGNMENT pub fn into_mutable(self) -> Result { let ptr = self.ptr; + let data_ptr = self.data_ptr(); let length = self.length; + Arc::try_unwrap(self.data) .and_then(|bytes| { - // The pointer of underlying buffer should not be offset. - assert_eq!(ptr, bytes.ptr().as_ptr()); + // The pointer of underlying buffer should be the same + assert_eq!(data_ptr, bytes.ptr()); MutableBuffer::from_bytes(bytes).map_err(Arc::new) }) .map_err(|bytes| Buffer { @@ -1080,4 +1094,56 @@ mod tests { drop(capture); assert_eq!(buffer2.strong_count(), 1); } + + #[test] + fn test_is_sliced() { + let buffer = Buffer::from(&[1, 2, 3, 4]); + assert!(!buffer.is_sliced()); + assert!(!buffer.clone().is_sliced()); + { + let mut advanced = buffer.clone(); + advanced.advance(0); + assert!(!advanced.is_sliced()); + } + { + let mut advanced = buffer.clone(); + advanced.advance(1); + assert!(advanced.is_sliced()); + } + + assert!(!buffer.slice(0).is_sliced()); + assert!(buffer.slice(1).is_sliced()); + + assert!(buffer.slice_with_length(1, 3).is_sliced()); + assert!(!buffer.slice_with_length(0, 4).is_sliced()); + assert!(buffer.slice_with_length(0, 3).is_sliced()); + assert!(buffer.slice_with_length(0, 0).is_sliced()); + } + + #[test] + fn into_mutable_should_return_the_entire_data_regardless_of_slicing() { + let original_buffer_data = [1_u8, 2, 3, 4, 5, 6, 7, 8]; + for (slice_from, slice_length) in [ + (0, 0), + (0, original_buffer_data.len()), + (original_buffer_data.len(), 0), + (2, 4), + (2, original_buffer_data.len() - 2), + ] { + let buffer = Buffer::from(original_buffer_data); + let original_buffer_len = buffer.len(); + let original_data_ptr = buffer.data_ptr(); + let sliced = buffer.slice_with_length(slice_from, slice_length); + drop(buffer); // Keep only 1 owner + + let mutable = sliced.into_mutable().expect("should convert to mutable"); + assert_eq!(mutable.len(), original_buffer_len); + let new_buffer = Buffer::from(mutable); + assert_eq!(new_buffer.data_ptr(), original_data_ptr); + assert_eq!(new_buffer.len(), original_buffer_len); + assert!(!new_buffer.is_sliced()); + + assert_eq!(new_buffer.as_slice(), &original_buffer_data); + } + } }