From 965433052afe75e54869cefd6d212aa8bd73e45f Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com> Date: Wed, 22 Mar 2023 10:01:53 +0000 Subject: [PATCH] Return ScalarBuffer from PrimitiveArray::values (#3879) (#3896) * Return ScalarBuffer from PrimitiveArray::values (#3879) * Fix docs * Review feedback --- arrow-array/src/array/primitive_array.rs | 4 +-- arrow-buffer/src/buffer/scalar.rs | 33 ++++++++++++++++++++++++ arrow/src/lib.rs | 6 ++--- 3 files changed, 38 insertions(+), 5 deletions(-) diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index 78859bd5956f..241e2a051197 100644 --- a/arrow-array/src/array/primitive_array.rs +++ b/arrow-array/src/array/primitive_array.rs @@ -275,9 +275,9 @@ impl PrimitiveArray { self.data.is_empty() } - /// Returns a slice of the values of this array + /// Returns the values of this array #[inline] - pub fn values(&self) -> &[T::Native] { + pub fn values(&self) -> &ScalarBuffer { &self.raw_values } diff --git a/arrow-buffer/src/buffer/scalar.rs b/arrow-buffer/src/buffer/scalar.rs index 9b3a47785098..04c6d9dcc7ac 100644 --- a/arrow-buffer/src/buffer/scalar.rs +++ b/arrow-buffer/src/buffer/scalar.rs @@ -114,6 +114,39 @@ impl From> for ScalarBuffer { } } +impl<'a, T: ArrowNativeType> IntoIterator for &'a ScalarBuffer { + type Item = &'a T; + type IntoIter = std::slice::Iter<'a, T>; + + fn into_iter(self) -> Self::IntoIter { + self.as_ref().iter() + } +} + +impl + ?Sized> PartialEq for ScalarBuffer { + fn eq(&self, other: &S) -> bool { + self.as_ref().eq(other.as_ref()) + } +} + +impl PartialEq> for [T; N] { + fn eq(&self, other: &ScalarBuffer) -> bool { + self.as_ref().eq(other.as_ref()) + } +} + +impl PartialEq> for [T] { + fn eq(&self, other: &ScalarBuffer) -> bool { + self.as_ref().eq(other.as_ref()) + } +} + +impl PartialEq> for Vec { + fn eq(&self, other: &ScalarBuffer) -> bool { + self.as_slice().eq(other.as_ref()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/arrow/src/lib.rs b/arrow/src/lib.rs index 3d1bced298c9..4b1251ebcd2b 100644 --- a/arrow/src/lib.rs +++ b/arrow/src/lib.rs @@ -69,7 +69,7 @@ //! //! let collected: Vec<_> = array.iter().collect(); //! assert_eq!(collected, vec![Some(1), None, Some(3)]); -//! assert_eq!(array.values(), [1, 0, 3]) +//! assert_eq!(array.values(), &[1, 0, 3]) //! ``` //! //! It is also possible to write generic code. For example, the following is generic over @@ -168,7 +168,7 @@ //! //! let array = parse_strings(["1", "2", "3"], DataType::Int32); //! let integers = array.as_any().downcast_ref::().unwrap(); -//! assert_eq!(integers.values(), [1, 2, 3]) +//! assert_eq!(integers.values(), &[1, 2, 3]) //! ``` //! //! # Compute Kernels @@ -192,7 +192,7 @@ //! //! let array = parse_strings(["1", "2", "3"], &DataType::UInt32).unwrap(); //! let integers = array.as_any().downcast_ref::().unwrap(); -//! assert_eq!(integers.values(), [1, 2, 3]) +//! assert_eq!(integers.values(), &[1, 2, 3]) //! ``` //! //! This module also implements many common vertical operations: