Skip to content

Commit

Permalink
Add unsafe/unchecked slice functions
Browse files Browse the repository at this point in the history
  • Loading branch information
emilk committed Dec 18, 2024
1 parent fc814bc commit a28e444
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 1 deletion.
11 changes: 11 additions & 0 deletions arrow-buffer/src/buffer/immutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,17 @@ 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::length`].
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) };
Expand Down
31 changes: 30 additions & 1 deletion arrow-buffer/src/buffer/scalar.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,14 +64,35 @@ impl<T: ArrowNativeType> ScalarBuffer<T> {
///
/// * `offset` or `len` would result in overflow
/// * `buffer` is not aligned to a multiple of `std::mem::align_of::<T>`
/// * `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::<T>();
let byte_offset = offset.checked_mul(size).expect("offset overflow");
let byte_len = len.checked_mul(size).expect("length overflow");
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::<T>`
/// * `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::<T>();
let byte_offset = offset * size;
let byte_len = len * size;
unsafe {
buffer
.slice_with_length_unchecked(byte_offset, byte_len)
.into()
}
}

/// Free up unused memory.
pub fn shrink_to_fit(&mut self) {
self.buffer.shrink_to_fit();
Expand All @@ -82,6 +103,14 @@ impl<T: ArrowNativeType> ScalarBuffer<T> {
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::length`]
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
Expand Down

0 comments on commit a28e444

Please sign in to comment.