Skip to content

Commit

Permalink
fix: ensure binview doesn't OOB (#13876)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored Jan 20, 2024
1 parent 9849361 commit 79ae3b3
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 5 deletions.
2 changes: 1 addition & 1 deletion crates/polars-arrow/src/array/binview/mutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ impl<T: ViewType + ?Sized> MutableBinaryViewArray<T> {
let (data_ptr, data_len) = *buffers.get_unchecked_release(buffer_idx as usize);
let data = std::slice::from_raw_parts(data_ptr, data_len);
let offset = offset as usize;
let bytes = data.get_unchecked(offset..offset + len as usize);
let bytes = data.get_unchecked_release(offset..offset + len as usize);
let t = T::from_bytes_unchecked(bytes);
self.push_value_ignore_validity(t)
}
Expand Down
13 changes: 9 additions & 4 deletions crates/polars-arrow/src/array/growable/binview.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ pub struct GrowableBinaryViewArray<'a, T: ViewType + ?Sized> {
validity: Option<MutableBitmap>,
views: Vec<u128>,
buffers: Vec<Buffer<u8>>,
buffers_offsets: Vec<u32>,
buffers_idx_offsets: Vec<u32>,
total_bytes_len: usize,
total_buffer_len: usize,
}
Expand Down Expand Up @@ -63,7 +63,7 @@ impl<'a, T: ViewType + ?Sized> GrowableBinaryViewArray<'a, T> {
validity: prepare_validity(use_validity, capacity),
views: Vec::with_capacity(capacity),
buffers,
buffers_offsets: cum_offset,
buffers_idx_offsets: cum_offset,
total_bytes_len: 0,
total_buffer_len,
}
Expand Down Expand Up @@ -101,9 +101,14 @@ impl<'a, T: ViewType + ?Sized> GrowableBinaryViewArray<'a, T> {
self.total_bytes_len += len;

if len > 12 {
let buffer_offset = *self.buffers_offsets.get_unchecked(index);
// Take the buffer index of the View.
let current_buffer_idx = (view >> 64) as u32;
// And add the offset of the buffers.
let buffer_idx =
*self.buffers_idx_offsets.get_unchecked(index) + current_buffer_idx;
// Mask out the old buffer-idx and OR in the new.
let mask = (u32::MAX as u128) << 64;
(view & !mask) | ((buffer_offset as u128) << 64)
(view & !mask) | ((buffer_idx as u128) << 64)
} else {
view
}
Expand Down

0 comments on commit 79ae3b3

Please sign in to comment.