Skip to content

Commit

Permalink
check null buffer length before new_unchecked (#1714)
Browse files Browse the repository at this point in the history
Signed-off-by: remzi <13716567376yh@gmail.com>
  • Loading branch information
HaoYang670 authored May 21, 2022
1 parent a30e787 commit 8855d22
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion arrow/src/array/data.rs
Original file line number Diff line number Diff line change
Expand Up @@ -327,6 +327,19 @@ impl ArrayData {
buffers: Vec<Buffer>,
child_data: Vec<ArrayData>,
) -> Result<Self> {
// we must check the length of `null_bit_buffer` first
// because we use this buffer to calculate `null_count`
// in `Self::new_unchecked`.
if let Some(null_bit_buffer) = null_bit_buffer.as_ref() {
let needed_len = bit_util::ceil(len + offset, 8);
if null_bit_buffer.len() < needed_len {
return Err(ArrowError::InvalidArgumentError(format!(
"null_bit_buffer size too small. got {} needed {}",
null_bit_buffer.len(),
needed_len
)));
}
}
// Safety justification: `validate_full` is called below
let new_self = unsafe {
Self::new_unchecked(
Expand Down Expand Up @@ -1733,7 +1746,7 @@ mod tests {

#[test]
#[should_panic(expected = "null_bit_buffer size too small. got 1 needed 2")]
fn test_bitmap_too_small() {
fn test_bitmap_too_small_with_null_count() {
let buffer = make_i32_buffer(9);
let null_bit_buffer = Buffer::from(vec![0b11111111]);

Expand All @@ -1749,6 +1762,24 @@ mod tests {
.unwrap();
}

#[test]
#[should_panic(expected = "null_bit_buffer size too small. got 1 needed 2")]
fn test_bitmap_too_small_without_null_count() {
let buffer = make_i32_buffer(9);
let null_bit_buffer = Buffer::from(vec![0b11111111]);

ArrayData::try_new(
DataType::Int32,
9,
None,
Some(null_bit_buffer),
0,
vec![buffer],
vec![],
)
.unwrap();
}

#[test]
#[should_panic(expected = "null_count 3 for an array exceeds length of 2 elements")]
fn test_bad_null_count() {
Expand Down

0 comments on commit 8855d22

Please sign in to comment.