Return ScalarBuffer from PrimitiveArray::values (apache#3879)

tustvold · Mar 21, 2023 · fdcd643 · fdcd643
1 parent 3a6f8bd
commit fdcd643
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 2 deletions.
diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs
@@ -877,7 +877,7 @@ mod tests {
         let keys = array.keys();
         assert_eq!(&DataType::Int8, keys.data_type());
         assert_eq!(0, keys.null_count());
-        assert_eq!(&[0, 1, 2, 0], keys.values());
+        assert_eq!(keys.values(), &[0, 1, 2, 0]);
     }
 
     #[test]

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
@@ -277,7 +277,7 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
 
     /// Returns a slice of the values of this array
     #[inline]
-    pub fn values(&self) -> &[T::Native] {
+    pub fn values(&self) -> &ScalarBuffer<T::Native> {
         &self.raw_values
     }
 

diff --git a/arrow-buffer/src/buffer/scalar.rs b/arrow-buffer/src/buffer/scalar.rs
@@ -114,6 +114,39 @@ impl<T: ArrowNativeType> From<Vec<T>> for ScalarBuffer<T> {
     }
 }
 
+impl<'a, T: ArrowNativeType> IntoIterator for &'a ScalarBuffer<T> {
+    type Item = &'a T;
+    type IntoIter = std::slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.as_ref().iter()
+    }
+}
+
+impl<T: ArrowNativeType, S: AsRef<[T]> + ?Sized> PartialEq<S> for ScalarBuffer<T> {
+    fn eq(&self, other: &S) -> bool {
+        self.as_ref().eq(other.as_ref())
+    }
+}
+
+impl<T: ArrowNativeType, const N: usize> PartialEq<ScalarBuffer<T>> for [T; N] {
+    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
+        self.as_ref().eq(other.as_ref())
+    }
+}
+
+impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for [T] {
+    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
+        self.as_ref().eq(other.as_ref())
+    }
+}
+
+impl<T: ArrowNativeType> PartialEq<ScalarBuffer<T>> for Vec<T> {
+    fn eq(&self, other: &ScalarBuffer<T>) -> bool {
+        self.as_slice().eq(other.as_ref())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;