[arrow-rs] branch master updated: check null buffer length before new_unchecked (#1714)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 8855d2201 check null buffer length before new_unchecked (#1714)
8855d2201 is described below

commit 8855d220104003fe440f068c7bff30a04d1ef54b
Author: Remzi Yang <59...@users.noreply.github.com>
AuthorDate: Sun May 22 06:13:32 2022 +0800

    check null buffer length before new_unchecked (#1714)
    
    Signed-off-by: remzi <13...@gmail.com>
---
 arrow/src/array/data.rs | 33 ++++++++++++++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/arrow/src/array/data.rs b/arrow/src/array/data.rs
index 15cd9cc5f..22536ca58 100644
--- a/arrow/src/array/data.rs
+++ b/arrow/src/array/data.rs
@@ -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(
@@ -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]);
 
@@ -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() {