Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Check overflow in MutableArrayData extend offsets (#3123) #3157

Merged
merged 2 commits into from Nov 22, 2022
Merged

Check overflow in MutableArrayData extend offsets (#3123) #3157

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f808ca9
Check overflow in MutableArrayData extend offsets (#3123)
tustvold Nov 22, 2022
83e9bb8
Update arrow-data/src/transform/utils.rs
tustvold Nov 22, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 4 additions & 2 deletions arrow-data/src/transform/list.rs
View file
Open in desktop
Expand Up @@ -21,9 +21,11 @@ use super::{
};
use crate::ArrayData;
use arrow_buffer::ArrowNativeType;
use num::Integer;
use num::{CheckedAdd, Integer};

pub(super) fn build_extend<T: ArrowNativeType + Integer>(array: &ArrayData) -> Extend {
pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd>(
array: &ArrayData,
) -> Extend {
let offsets = array.buffer::<T>(0);
if array.null_count() == 0 {
// fast case where we can copy regions without nullability checks
Expand Down
19 changes: 16 additions & 3 deletions arrow-data/src/transform/utils.rs
View file
Open in desktop
Expand Up @@ -16,7 +16,7 @@
// under the License.

use arrow_buffer::{bit_util, ArrowNativeType, MutableBuffer};
use num::Integer;
use num::{CheckedAdd, Integer};

/// extends the `buffer` to be able to hold `len` bits, setting all bits of the new size to zero.
#[inline]
Expand All @@ -27,7 +27,7 @@ pub(super) fn resize_for_bits(buffer: &mut MutableBuffer, len: usize) {
}
}

pub(super) fn extend_offsets<T: ArrowNativeType + Integer>(
pub(super) fn extend_offsets<T: ArrowNativeType + Integer + CheckedAdd>(
buffer: &mut MutableBuffer,
mut last_offset: T,
offsets: &[T],
Expand All @@ -36,7 +36,7 @@ pub(super) fn extend_offsets<T: ArrowNativeType + Integer>(
offsets.windows(2).for_each(|offsets| {
// compute the new offset
let length = offsets[1] - offsets[0];
last_offset = last_offset + length;
last_offset = last_offset.checked_add(&length).expect("offset overflow");
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Extend is infallible, so we can't return an Error. We also fairly consistently panic on offset overflow in other kernels, and I wish to eventually remove MutableArrayData, and so I think this is fine for now

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think panic'ing rather than segfaulting is a much better practice

tustvold marked this conversation as resolved.
Show resolved Hide resolved
buffer.push(last_offset);
});
}
Expand All @@ -55,3 +55,16 @@ pub(super) unsafe fn get_last_offset<T: ArrowNativeType>(
debug_assert!(prefix.is_empty() && suffix.is_empty());
*offsets.get_unchecked(offsets.len() - 1)
}

#[cfg(test)]
mod tests {
use crate::transform::utils::extend_offsets;
use arrow_buffer::MutableBuffer;

#[test]
#[should_panic(expected = "offset overflow")]
fn test_overflow() {
let mut buffer = MutableBuffer::new(10);
extend_offsets(&mut buffer, i32::MAX - 4, &[0, 5]);
}
}
6 changes: 4 additions & 2 deletions arrow-data/src/transform/variable_size.rs
View file
Open in desktop
Expand Up @@ -18,7 +18,7 @@
use crate::ArrayData;
use arrow_buffer::{ArrowNativeType, MutableBuffer};
use num::traits::AsPrimitive;
use num::Integer;
use num::{CheckedAdd, Integer};

use super::{
Extend, _MutableArrayData,
Expand All @@ -39,7 +39,9 @@ fn extend_offset_values<T: ArrowNativeType + AsPrimitive<usize>>(
buffer.extend_from_slice(new_values);
}

pub(super) fn build_extend<T: ArrowNativeType + Integer + AsPrimitive<usize>>(
pub(super) fn build_extend<
T: ArrowNativeType + Integer + CheckedAdd + AsPrimitive<usize>,
>(
array: &ArrayData,
) -> Extend {
let offsets = array.buffer::<T>(0);
Expand Down