diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index d0c8ffa3978..09ef8b69dce 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -241,6 +241,18 @@ impl Buffer { "the offset of the new Buffer cannot exceed the existing length: slice offset={offset} length={length} selflen={}", self.length ); + // Safety: + // offset + length <= self.length + unsafe { self.slice_with_length_unchecked(offset, length) } + } + + /// Returns a new [Buffer] that is a slice of this buffer starting at `offset`, + /// with `length` bytes. + /// Doing so allows the same memory region to be shared between buffers. + /// + /// # Safety + /// `(offset + length)` MUST be `<=` [`Self::len`]. + pub unsafe fn slice_with_length_unchecked(&self, offset: usize, length: usize) -> Self { // Safety: // offset + length <= self.length let ptr = unsafe { self.ptr.add(offset) }; diff --git a/arrow-buffer/src/buffer/scalar.rs b/arrow-buffer/src/buffer/scalar.rs index ab6c87168e5..fed3fbf2226 100644 --- a/arrow-buffer/src/buffer/scalar.rs +++ b/arrow-buffer/src/buffer/scalar.rs @@ -64,7 +64,7 @@ impl ScalarBuffer { /// /// * `offset` or `len` would result in overflow /// * `buffer` is not aligned to a multiple of `std::mem::align_of::` - /// * `bytes` is not large enough for the requested slice + /// * `buffer` is not large enough for the requested slice pub fn new(buffer: Buffer, offset: usize, len: usize) -> Self { let size = std::mem::size_of::(); let byte_offset = offset.checked_mul(size).expect("offset overflow"); @@ -72,6 +72,33 @@ impl ScalarBuffer { buffer.slice_with_length(byte_offset, byte_len).into() } + /// Create a new [`ScalarBuffer`] from a [`Buffer`], and an `offset` + /// and `length` in units of `T` + /// + /// # Safety + /// + /// This method will be safe UNLESS any of the following: + /// + /// * `offset` or `len` would result in overflow + /// * `buffer` is not aligned to a multiple of `std::mem::align_of::` + /// * `buffer` is not large enough for the requested slice + pub unsafe fn new_unchecked(buffer: Buffer, offset: usize, len: usize) -> Self { + let size = std::mem::size_of::(); + let byte_offset = offset * size; + let byte_len = len * size; + unsafe { + buffer + .slice_with_length_unchecked(byte_offset, byte_len) + .into() + } + } + + /// The length of the scalar buffer, in units of `T`. + #[inline] + pub fn len_in_elements(&self) -> usize { + self.buffer.len() / std::mem::size_of::() + } + /// Free up unused memory. pub fn shrink_to_fit(&mut self) { self.buffer.shrink_to_fit(); @@ -82,6 +109,14 @@ impl ScalarBuffer { Self::new(self.buffer.clone(), offset, len) } + /// Returns a zero-copy slice of this buffer with length `len` and starting at `offset` + /// + /// # Safety + /// `(offset + len)` <= [`Self::len_in_elements`] + pub unsafe fn slice_unchecked(&self, offset: usize, len: usize) -> Self { + unsafe { Self::new_unchecked(self.buffer.clone(), offset, len) } + } + /// Returns the inner [`Buffer`] pub fn inner(&self) -> &Buffer { &self.buffer