diff --git a/rust/arrow/src/array/builder.rs b/rust/arrow/src/array/builder.rs index b0b97c2210740..b647ae1eac813 100644 --- a/rust/arrow/src/array/builder.rs +++ b/rust/arrow/src/array/builder.rs @@ -705,7 +705,7 @@ mod tests { assert_eq!(0, b.len()); assert_eq!(16, b.capacity()); let a = b.finish(); - assert_eq!(0, a.len()); + assert_eq!(64, a.len()); } #[test] @@ -713,7 +713,7 @@ mod tests { let mut b = Int32BufferBuilder::new(0); b.append(123).unwrap(); let a = b.finish(); - assert_eq!(4, a.len()); + assert_eq!(64, a.len()); } #[test] @@ -724,7 +724,7 @@ mod tests { } assert_eq!(16, b.capacity()); let a = b.finish(); - assert_eq!(20, a.len()); + assert_eq!(64, a.len()); // 16 * size_of(i32) } #[test] @@ -736,7 +736,7 @@ mod tests { } assert_eq!(32, b.capacity()); let a = b.finish(); - assert_eq!(80, a.len()); + assert_eq!(128, a.len()); // 32 * size_of(i32) } #[test] @@ -747,7 +747,7 @@ mod tests { b.append(i).unwrap(); } let mut a = b.finish(); - assert_eq!(40, a.len()); + assert_eq!(64, a.len()); assert_eq!(0, b.len()); assert_eq!(0, b.capacity()); @@ -757,7 +757,7 @@ mod tests { } assert_eq!(32, b.capacity()); a = b.finish(); - assert_eq!(80, a.len()); + assert_eq!(128, a.len()); } #[test] @@ -783,12 +783,12 @@ mod tests { b.append_slice("Hello, ".as_bytes()).unwrap(); b.append_slice("World!".as_bytes()).unwrap(); let buffer = b.finish(); - assert_eq!(13, buffer.len()); + assert_eq!(64, buffer.len()); let mut b = Int32BufferBuilder::new(0); b.append_slice(&[32, 54]).unwrap(); let buffer = b.finish(); - assert_eq!(8, buffer.len()); + assert_eq!(64, buffer.len()); } #[test] @@ -801,14 +801,14 @@ mod tests { assert_eq!(4, b.len()); assert_eq!(512, b.capacity()); let buffer = b.finish(); - assert_eq!(1, buffer.len()); + assert_eq!(64, buffer.len()); let mut b = BooleanBufferBuilder::new(4); b.append_slice(&[false, true, false, true]).unwrap(); assert_eq!(4, b.len()); assert_eq!(512, b.capacity()); let buffer = b.finish(); - assert_eq!(1, buffer.len()); + assert_eq!(64, buffer.len()); } #[test] @@ -819,7 +819,7 @@ mod tests { assert_eq!(4, b.len()); assert_eq!(16, b.capacity()); let buffer = b.finish(); - assert_eq!(16, buffer.len()); + assert_eq!(64, buffer.len()); } #[test] diff --git a/rust/arrow/src/array/data.rs b/rust/arrow/src/array/data.rs index ac6ad4fa9ce56..48055b78ea781 100644 --- a/rust/arrow/src/array/data.rs +++ b/rust/arrow/src/array/data.rs @@ -275,7 +275,7 @@ mod tests { assert_eq!(10, arr_data.null_count()); assert_eq!(5, arr_data.offset()); assert_eq!(1, arr_data.buffers().len()); - assert_eq!(&[0, 1, 2, 3], arr_data.buffers()[0].data()); + assert_eq!([0, 1, 2, 3], arr_data.buffers()[0].data()[..4]); assert_eq!(1, arr_data.child_data().len()); assert_eq!(child_arr_data, arr_data.child_data()[0]); } diff --git a/rust/arrow/src/buffer.rs b/rust/arrow/src/buffer.rs index eea000fe5c28a..f4660d293f73d 100644 --- a/rust/arrow/src/buffer.rs +++ b/rust/arrow/src/buffer.rs @@ -178,7 +178,7 @@ impl> From for Buffer { unsafe { memory::memcpy(buffer, slice.as_ptr(), len); } - Buffer::from_raw_parts(buffer, len) + Buffer::from_raw_parts(buffer, capacity) } } @@ -448,7 +448,7 @@ impl MutableBuffer { pub fn freeze(self) -> Buffer { let buffer_data = BufferData { ptr: self.data, - len: self.len, + len: self.capacity, }; ::std::mem::forget(self); Buffer { @@ -534,26 +534,26 @@ mod tests { assert!(buf.raw_data().is_null()); let buf = Buffer::from(&[0, 1, 2, 3, 4]); - assert_eq!(5, buf.len()); + assert_eq!(64, buf.len()); assert!(!buf.raw_data().is_null()); - assert_eq!(&[0, 1, 2, 3, 4], buf.data()); + assert_eq!([0, 1, 2, 3, 4], buf.data()[..5]); } #[test] fn test_from_vec() { let buf = Buffer::from(&[0, 1, 2, 3, 4]); - assert_eq!(5, buf.len()); + assert_eq!(64, buf.len()); assert!(!buf.raw_data().is_null()); - assert_eq!(&[0, 1, 2, 3, 4], buf.data()); + assert_eq!([0, 1, 2, 3, 4], buf.data()[..5]); } #[test] fn test_copy() { let buf = Buffer::from(&[0, 1, 2, 3, 4]); let buf2 = buf.clone(); - assert_eq!(5, buf2.len()); + assert_eq!(64, buf2.len()); assert!(!buf2.raw_data().is_null()); - assert_eq!(&[0, 1, 2, 3, 4], buf2.data()); + assert_eq!([0, 1, 2, 3, 4], buf2.data()[..5]); } #[test] @@ -561,20 +561,17 @@ mod tests { let buf = Buffer::from(&[2, 4, 6, 8, 10]); let buf2 = buf.slice(2); - assert_eq!(&[6, 8, 10], buf2.data()); - assert_eq!(3, buf2.len()); + assert_eq!([6, 8, 10], buf2.data()[0..3]); + assert_eq!(62, buf2.len()); assert_eq!(unsafe { buf.raw_data().offset(2) }, buf2.raw_data()); let buf3 = buf2.slice(1); - assert_eq!(&[8, 10], buf3.data()); - assert_eq!(2, buf3.len()); + assert_eq!([8, 10], buf3.data()[..2]); + assert_eq!(61, buf3.len()); assert_eq!(unsafe { buf.raw_data().offset(3) }, buf3.raw_data()); let buf4 = buf.slice(5); - let empty_slice: [u8; 0] = []; - assert_eq!(empty_slice, buf4.data()); - assert_eq!(0, buf4.len()); - assert!(buf4.is_empty()); + assert_eq!(59, buf4.len()); } #[test] @@ -583,7 +580,7 @@ mod tests { )] fn test_slice_offset_out_of_bound() { let buf = Buffer::from(&[2, 4, 6, 8, 10]); - buf.slice(6); + buf.slice(65); } #[test] @@ -627,14 +624,15 @@ mod tests { #[test] fn test_bitwise_not() { let buf = Buffer::from([0b01101010]); - assert_eq!(Buffer::from([0b10010101]), !&buf); + assert_eq!(Buffer::from([0b10010101]).data()[..1], (!&buf).data()[..1]); } #[test] #[should_panic(expected = "Buffers must be the same size to apply Bitwise OR.")] fn test_buffer_bitand_different_sizes() { - let buf1 = Buffer::from([1_u8, 1_u8]); - let buf2 = Buffer::from([0b01001110]); + let size = memory::ALIGNMENT * 2; + let buf1 = Buffer::from_raw_parts(memory::allocate_aligned(size), size); + let buf2 = Buffer::from([0x00]); let _buf3 = (&buf1 | &buf2).unwrap(); } @@ -722,18 +720,22 @@ mod tests { buf.write("aaaa bbbb cccc dddd".as_bytes()) .expect("write should be OK"); assert_eq!(19, buf.len()); + assert_eq!(64, buf.capacity()); assert_eq!("aaaa bbbb cccc dddd".as_bytes(), buf.data()); let immutable_buf = buf.freeze(); - assert_eq!(19, immutable_buf.len()); - assert_eq!("aaaa bbbb cccc dddd".as_bytes(), immutable_buf.data()); + assert_eq!(64, immutable_buf.len()); + assert_eq!( + "aaaa bbbb cccc dddd".as_bytes(), + &immutable_buf.data()[..19] + ); } #[test] fn test_access_concurrently() { let buffer = Buffer::from(vec![1, 2, 3, 4, 5]); let buffer2 = buffer.clone(); - assert_eq!(&[1, 2, 3, 4, 5], buffer.data()); + assert_eq!([1, 2, 3, 4, 5], buffer.data()[..5]); let buffer_copy = thread::spawn(move || { // access buffer in another thread. diff --git a/rust/arrow/src/memory.rs b/rust/arrow/src/memory.rs index d73ddd11676c8..32c9f726b5fd0 100644 --- a/rust/arrow/src/memory.rs +++ b/rust/arrow/src/memory.rs @@ -26,7 +26,7 @@ pub const ALIGNMENT: usize = 64; pub fn allocate_aligned(size: usize) -> *mut u8 { unsafe { let layout = Layout::from_size_align_unchecked(size, ALIGNMENT); - ::std::alloc::alloc(layout) + ::std::alloc::alloc_zeroed(layout) } } diff --git a/rust/arrow/src/tensor.rs b/rust/arrow/src/tensor.rs index 5382ca19c8807..ef69e35165142 100644 --- a/rust/arrow/src/tensor.rs +++ b/rust/arrow/src/tensor.rs @@ -86,11 +86,6 @@ impl<'a, T: ArrowPrimitiveType> Tensor<'a, T> { ) -> Self { match &shape { None => { - assert_eq!( - buffer.len(), - mem::size_of::(), - "underlying buffer should only contain a single tensor element" - ); assert_eq!(None, strides); assert_eq!(None, names); } @@ -192,7 +187,10 @@ impl<'a, T: ArrowPrimitiveType> Tensor<'a, T> { /// The total number of elements in the `Tensor` pub fn size(&self) -> usize { - (self.buffer.len() / mem::size_of::()) + match self.shape { + None => 0, + Some(ref s) => s.iter().fold(1, |a, b| a * b), + } } /// Indicates if the data is laid out contiguously in memory @@ -260,7 +258,7 @@ mod tests { fn test_zero_dim() { let buf = Buffer::from(&[1]); let tensor = UInt8Tensor::new(buf, None, None, None); - assert_eq!(1, tensor.size()); + assert_eq!(0, tensor.size()); assert_eq!(None, tensor.shape()); assert_eq!(None, tensor.names()); assert_eq!(0, tensor.ndim()); @@ -270,7 +268,7 @@ mod tests { let buf = Buffer::from(&[1, 2, 2, 2]); let tensor = Int32Tensor::new(buf, None, None, None); - assert_eq!(1, tensor.size()); + assert_eq!(0, tensor.size()); assert_eq!(None, tensor.shape()); assert_eq!(None, tensor.names()); assert_eq!(0, tensor.ndim());