Validate arguments to ArrayData::new: null bit buffer and buffers

apache · Oct 14, 2021 · b270574 · b270574
1 parent f14d213
commit b270574
Show file tree

Hide file tree

Showing 9 changed files with 587 additions and 57 deletions.
diff --git a/arrow-flight/src/arrow.flight.protocol.rs b/arrow-flight/src/arrow.flight.protocol.rs
@@ -242,7 +242,7 @@ pub mod flight_service_client {
             interceptor: F,
         ) -> FlightServiceClient<InterceptedService<T, F>>
         where
-            F: tonic::service::Interceptor,
+            F: FnMut(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status>,
             T: tonic::codegen::Service<
                 http::Request<tonic::body::BoxBody>,
                 Response = http::Response<
@@ -666,7 +666,7 @@ pub mod flight_service_server {
             interceptor: F,
         ) -> InterceptedService<Self, F>
         where
-            F: tonic::service::Interceptor,
+            F: FnMut(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status>,
         {
             InterceptedService::new(Self::new(inner), interceptor)
         }

diff --git a/arrow/src/array/array_binary.rs b/arrow/src/array/array_binary.rs
@@ -891,10 +891,18 @@ mod tests {
             assert!(binary_array.is_valid(i));
             assert!(!binary_array.is_null(i));
         }
+    }
+
+    #[test]
+    fn test_binary_array_with_offsets() {
+        let values: [u8; 12] = [
+            b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
+        ];
+        let offsets: [i32; 4] = [0, 5, 5, 12];
 
         // Test binary array with offset
         let array_data = ArrayData::builder(DataType::Binary)
-            .len(4)
+            .len(2)
             .offset(1)
             .add_buffer(Buffer::from_slice_ref(&offsets))
             .add_buffer(Buffer::from_slice_ref(&values))
@@ -947,10 +955,18 @@ mod tests {
             assert!(binary_array.is_valid(i));
             assert!(!binary_array.is_null(i));
         }
+    }
+
+    #[test]
+    fn test_large_binary_array_with_offsets() {
+        let values: [u8; 12] = [
+            b'h', b'e', b'l', b'l', b'o', b'p', b'a', b'r', b'q', b'u', b'e', b't',
+        ];
+        let offsets: [i64; 4] = [0, 5, 5, 12];
 
         // Test binary array with offset
         let array_data = ArrayData::builder(DataType::LargeBinary)
-            .len(4)
+            .len(2)
             .offset(1)
             .add_buffer(Buffer::from_slice_ref(&offsets))
             .add_buffer(Buffer::from_slice_ref(&values))
@@ -1196,15 +1212,13 @@ mod tests {
 
     #[test]
     #[should_panic(
-        expected = "FixedSizeBinaryArray can only be created from list array of u8 values \
-                    (i.e. FixedSizeList<PrimitiveArray<u8>>)."
+        expected = "FixedSizeBinaryArray can only be created from FixedSizeList<u8> arrays"
     )]
     fn test_fixed_size_binary_array_from_incorrect_list_array() {
         let values: [u32; 12] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
         let values_data = ArrayData::builder(DataType::UInt32)
             .len(12)
             .add_buffer(Buffer::from_slice_ref(&values))
-            .add_child_data(ArrayData::builder(DataType::Boolean).build().unwrap())
             .build()
             .unwrap();
 

diff --git a/arrow/src/array/array_boolean.rs b/arrow/src/array/array_boolean.rs
@@ -329,13 +329,17 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "BooleanArray data should contain a single buffer only \
-                               (values buffer)")]
+    #[should_panic(
+        expected = "Invalid argument error: Expected 1 buffers in array of type Boolean, got 2"
+    )]
     fn test_boolean_array_invalid_buffer_len() {
         let data = ArrayData::builder(DataType::Boolean)
             .len(5)
+            .add_buffer(Buffer::from([100_u8]))
+            .add_buffer(Buffer::from([200_u8]))
             .build()
             .unwrap();
+
         BooleanArray::from(data);
     }
 }
diff --git a/arrow/src/array/array_list.rs b/arrow/src/array/array_list.rs
@@ -1034,9 +1034,7 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(
-        expected = "ListArray data should contain a single buffer only (value offsets)"
-    )]
+    #[should_panic(expected = "Expected 1 buffers in array of type List")]
     fn test_list_array_invalid_buffer_len() {
         let value_data = ArrayData::builder(DataType::Int32)
             .len(8)

diff --git a/arrow/src/array/array_primitive.rs b/arrow/src/array/array_primitive.rs
@@ -894,7 +894,7 @@ mod tests {
     #[test]
     fn test_primitive_array_builder() {
         // Test building a primitive array with ArrayData builder and offset
-        let buf = Buffer::from_slice_ref(&[0, 1, 2, 3, 4]);
+        let buf = Buffer::from_slice_ref(&[0i32, 1, 2, 3, 4, 5, 6]);
         let buf2 = buf.clone();
         let data = ArrayData::builder(DataType::Int32)
             .len(5)
@@ -947,10 +947,16 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "PrimitiveArray data should contain a single buffer only \
-                               (values buffer)")]
+    #[should_panic(expected = "Expected 1 buffers in array of type Int32, got 2")]
     fn test_primitive_array_invalid_buffer_len() {
-        let data = ArrayData::builder(DataType::Int32).len(5).build().unwrap();
+        let buffer = Buffer::from_slice_ref(&[0i32, 1, 2, 3, 4]);
+        let data = ArrayData::builder(DataType::Int32)
+            .add_buffer(buffer.clone())
+            .add_buffer(buffer.clone())
+            .len(5)
+            .build()
+            .unwrap();
+
         Int32Array::from(data);
     }