ARROW-10989: [Rust] Iterate primitive buffers by slice

rust/arrow/src/array/array_dictionary.rs

@@ -386,7 +386,7 @@ mod tests {
  let keys = array.keys_array();
  assert_eq!(&DataType::Int8, keys.data_type());
  assert_eq!(0, keys.null_count());
- assert_eq!(&[0, 1, 2, 0], keys.value_slice(0, keys.len()));
+ assert_eq!(&[0, 1, 2, 0], keys.values());
  }
 
  #[test]

rust/arrow/src/array/array_primitive.rs

@@ -43,17 +43,21 @@ const NANOSECONDS: i64 = 1_000_000_000;
 
 /// Array whose elements are of primitive types.
 pub struct PrimitiveArray<T: ArrowPrimitiveType> {
+ /// Underlying ArrayData
+ /// # Safety
+ /// must have exactly one buffer, aligned to type T
  data: ArrayDataRef,
  /// Pointer to the value array. The lifetime of this must be <= to the value buffer
  /// stored in `data`, so it's safe to store.
- /// Also note that boolean arrays are bit-packed, so although the underlying pointer
- /// is of type bool it should be cast back to u8 before being used.
- /// i.e. `self.raw_values.get() as *const u8`
+ /// # Safety
+ /// raw_values must have a value equivalent to data.buffers()[0].raw_data()
+ /// raw_values must have alignment for type T::NativeType
  raw_values: RawPtrBox<T::Native>,
 }
 
 impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
  /// Returns the length of this array.
+ #[inline]
  pub fn len(&self) -> usize {
  self.data.len()
  }
@@ -63,35 +67,43 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
  self.data.is_empty()
  }
 
- /// Returns a raw pointer to the values of this array.
- pub fn raw_values(&self) -> *const T::Native {
- unsafe { self.raw_values.get().add(self.data.offset()) }
- }
-
  /// Returns a slice for the given offset and length
  ///
  /// Note this doesn't do any bound checking, for performance reason.
- pub fn value_slice(&self, offset: usize, len: usize) -> &[T::Native] {
- let raw =
- unsafe { std::slice::from_raw_parts(self.raw_values().add(offset), len) };
- &raw[..]
+ /// # Safety
+ /// caller must ensure that the passed in offset + len are less than the array len()
+ #[deprecated(note = "Please use values() instead")]
+ pub unsafe fn value_slice(&self, offset: usize, len: usize) -> &[T::Native] {
+ std::slice::from_raw_parts(
+ self.raw_values.get().add(self.data.offset()).add(offset),
+ len,
+ )
+ }
+
+ /// Returns a slice of the values of this array
+ #[inline]
+ pub fn values(&self) -> &[T::Native] {
+ // Soundness
+ // raw_values alignment & location is ensured by fn from(ArrayDataRef)
+ // buffer bounds/offset is ensured by the ArrayData instance.
+ unsafe {
+ std::slice::from_raw_parts(
+ self.raw_values.get().add(self.data.offset()),
+ self.len(),
+ )
+ }
  }
 
  // Returns a new primitive array builder
  pub fn builder(capacity: usize) -> PrimitiveBuilder<T> {
  PrimitiveBuilder::<T>::new(capacity)
  }
 
- /// Returns a `Buffer` holding all the values of this array.
- ///
- /// Note this doesn't take the offset of this array into account.
- pub fn values(&self) -> Buffer {
- self.data.buffers()[0].clone()
- }
-
  /// Returns the primitive value at index `i`.
  ///
  /// Note this doesn't do any bound checking, for performance reason.
+ /// # Safety
+ /// caller must ensure that the passed in offset is less than the array len()
  pub fn value(&self, i: usize) -> T::Native {
  let offset = i + self.offset();
  unsafe { *self.raw_values.get().add(offset) }
@@ -455,8 +467,8 @@ mod tests {
  fn test_primitive_array_from_vec() {
  let buf = Buffer::from(&[0, 1, 2, 3, 4].to_byte_slice());
  let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
- let slice = unsafe { std::slice::from_raw_parts(arr.raw_values(), 5) };
- assert_eq!(buf, arr.values());
+ let slice = arr.values();
+ assert_eq!(buf, arr.data.buffers()[0]);
  assert_eq!(&[0, 1, 2, 3, 4], slice);
  assert_eq!(5, arr.len());
  assert_eq!(0, arr.offset());
@@ -738,12 +750,6 @@ mod tests {
  assert_eq!(false, bool_arr.value(4));
  }
 
- #[test]
- fn test_value_slice_no_bounds_check() {
- let arr = Int32Array::from(vec![2, 3, 4]);
- let _slice = arr.value_slice(0, 4);
- }
-
  #[test]
  fn test_int32_fmt_debug() {
  let arr = Int32Array::from(vec![0, 1, 2, 3, 4]);
@@ -823,7 +829,7 @@ mod tests {
  .add_buffer(buf)
  .build();
  let arr = Int32Array::from(data);
- assert_eq!(buf2, arr.values());
+ assert_eq!(buf2, arr.data.buffers()[0]);
  assert_eq!(5, arr.len());
  assert_eq!(0, arr.null_count());
  for i in 0..3 {

rust/arrow/src/array/builder.rs

@@ -3344,7 +3344,7 @@ mod tests {
  // Values are polymorphic and so require a downcast.
  let av = array.values();
  let ava: &UInt32Array = av.as_any().downcast_ref::<UInt32Array>().unwrap();
- let avs: &[u32] = ava.value_slice(0, array.values().len());
+ let avs: &[u32] = ava.values();
 
  assert_eq!(array.is_null(0), false);
  assert_eq!(array.is_null(1), true);

rust/arrow/src/array/mod.rs

@@ -75,7 +75,7 @@
 //! assert_eq!(2, array.value(2), "Get the value with index 2");
 //!
 //! assert_eq!(
-//! array.value_slice(3, 2),
+//! &array.values()[3..5],
 //! &[3, 4],
 //! "Get slice of len 2 starting at idx 3"
 //! )

rust/arrow/src/compute/kernels/aggregate.rs

@@ -114,7 +114,7 @@ where
  }
 
  let data = array.data();
- let m = array.value_slice(0, data.len());
+ let m = array.values();
  let mut n;
 
  if null_count == 0 {
@@ -150,7 +150,7 @@ where
  return None;
  }
 
- let data: &[T::Native] = array.value_slice(0, array.len());
+ let data: &[T::Native] = array.values();
 
  match array.data().null_buffer() {
  None => {
@@ -478,7 +478,7 @@ mod simd {
  return None;
  }
 
- let data: &[T::Native] = array.value_slice(0, array.len());
+ let data: &[T::Native] = array.values();
 
  let mut chunk_acc = A::init_accumulator_chunk();
  let mut rem_acc = A::init_accumulator_scalar();

rust/arrow/src/compute/kernels/arithmetic.rs

@@ -51,8 +51,10 @@ where
  T::Native: Neg<Output = T::Native>,
  F: Fn(T::Native) -> T::Native,
 {
- let values = (0..array.len())
- .map(|i| op(array.value(i)))
+ let values = array
+ .values()
+ .iter()
+ .map(|v| op(*v))
  .collect::<Vec<T::Native>>();
 
  let data = ArrayData::new(
@@ -84,7 +86,7 @@ where
  let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
 
  let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut array_chunks = array.value_slice(0, array.len()).chunks_exact(lanes);
+ let mut array_chunks = array.values().chunks_exact(lanes);
 
  result_chunks
  .borrow_mut()
@@ -141,8 +143,11 @@ where
  let null_bit_buffer =
  combine_option_bitmap(left.data_ref(), right.data_ref(), left.len())?;
 
- let values = (0..left.len())
- .map(|i| op(left.value(i), right.value(i)))
+ let values = left
+ .values()
+ .iter()
+ .zip(right.values().iter())
+ .map(|(l, r)| op(*l, *r))
  .collect::<Vec<T::Native>>();
 
  let data = ArrayData::new(
@@ -248,8 +253,8 @@ where
  let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
 
  let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut left_chunks = left.value_slice(0, left.len()).chunks_exact(lanes);
- let mut right_chunks = right.value_slice(0, left.len()).chunks_exact(lanes);
+ let mut left_chunks = left.values().chunks_exact(lanes);
+ let mut right_chunks = right.values().chunks_exact(lanes);
 
  result_chunks
  .borrow_mut()
@@ -385,8 +390,8 @@ where
 
  // process data in chunks of 64 elements since we also get 64 bits of validity information at a time
  let mut result_chunks = result.typed_data_mut().chunks_exact_mut(64);
- let mut left_chunks = left.value_slice(0, left.len()).chunks_exact(64);
- let mut right_chunks = right.value_slice(0, left.len()).chunks_exact(64);
+ let mut left_chunks = left.values().chunks_exact(64);
+ let mut right_chunks = right.values().chunks_exact(64);
 
  valid_chunks
  .iter()
@@ -430,8 +435,8 @@ where
  }
  None => {
  let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes);
- let mut left_chunks = left.value_slice(0, left.len()).chunks_exact(lanes);
- let mut right_chunks = right.value_slice(0, left.len()).chunks_exact(lanes);
+ let mut left_chunks = left.values().chunks_exact(lanes);
+ let mut right_chunks = right.values().chunks_exact(lanes);
 
  result_chunks
  .borrow_mut()

rust/arrow/src/compute/kernels/comparison.rs

@@ -398,17 +398,19 @@ where
  let rem = len % lanes;
 
  for i in (0..len - rem).step_by(lanes) {
- let simd_left = T::load(left.value_slice(i, lanes));
- let simd_right = T::load(right.value_slice(i, lanes));
+ let simd_left = T::load(unsafe { left.value_slice(i, lanes) });
+ let simd_right = T::load(unsafe { right.value_slice(i, lanes) });
  let simd_result = op(simd_left, simd_right);
  T::bitmask(&simd_result, |b| {
  result.extend_from_slice(b);
  });
  }
 
  if rem > 0 {
- let simd_left = T::load(left.value_slice(len - rem, lanes));
- let simd_right = T::load(right.value_slice(len - rem, lanes));
+ //Soundness
+ //	This is not sound because it can read past the end of PrimitiveArray buffer (lanes is always greater than rem), see ARROW-10990
+ let simd_left = T::load(unsafe { left.value_slice(len - rem, lanes) });
+ let simd_right = T::load(unsafe { right.value_slice(len - rem, lanes) });
  let simd_result = op(simd_left, simd_right);
  let rem_buffer_size = (rem as f32 / 8f32).ceil() as usize;
  T::bitmask(&simd_result, |b| {
@@ -451,15 +453,17 @@ where
  let rem = len % lanes;
 
  for i in (0..len - rem).step_by(lanes) {
- let simd_left = T::load(left.value_slice(i, lanes));
+ let simd_left = T::load(unsafe { left.value_slice(i, lanes) });
  let simd_result = op(simd_left, simd_right);
  T::bitmask(&simd_result, |b| {
  result.extend_from_slice(b);
  });
  }
 
  if rem > 0 {
- let simd_left = T::load(left.value_slice(len - rem, lanes));
+ //Soundness
+ //	This is not sound because it can read past the end of PrimitiveArray buffer (lanes is always greater than rem), see ARROW-10990
+ let simd_left = T::load(unsafe { left.value_slice(len - rem, lanes) });
  let simd_result = op(simd_left, simd_right);
  let rem_buffer_size = (rem as f32 / 8f32).ceil() as usize;
  T::bitmask(&simd_result, |b| {