Implement Array for ArrayRef, Improve as_* kernels to take `&dyn Arra…

…y` (#1129) * Implement Array for ArrayRef * Improve as_* kernels to take &dyn Array * remove uneeded pyarrow binding
apache · Jan 5, 2022 · ae2e39c · ae2e39c
1 parent 3c12ebe
commit ae2e39c
Show file tree

Hide file tree

Showing 4 changed files with 124 additions and 18 deletions.
diff --git a/arrow/src/array/array.rs b/arrow/src/array/array.rs
@@ -227,6 +227,69 @@ pub trait Array: fmt::Debug + Send + Sync + JsonEqual {
 /// A reference-counted reference to a generic `Array`.
 pub type ArrayRef = Arc<dyn Array>;
 
+/// Ergonomics: Allow use of an ArrayRef as an `&dyn Array`
+impl Array for ArrayRef {
+    fn as_any(&self) -> &dyn Any {
+        self.as_ref().as_any()
+    }
+
+    fn data(&self) -> &ArrayData {
+        self.as_ref().data()
+    }
+
+    fn data_ref(&self) -> &ArrayData {
+        self.as_ref().data_ref()
+    }
+
+    fn data_type(&self) -> &DataType {
+        self.as_ref().data_type()
+    }
+
+    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
+        self.as_ref().slice(offset, length)
+    }
+
+    fn len(&self) -> usize {
+        self.as_ref().len()
+    }
+
+    fn is_empty(&self) -> bool {
+        self.as_ref().is_empty()
+    }
+
+    fn offset(&self) -> usize {
+        self.as_ref().offset()
+    }
+
+    fn is_null(&self, index: usize) -> bool {
+        self.as_ref().is_null(index)
+    }
+
+    fn is_valid(&self, index: usize) -> bool {
+        self.as_ref().is_valid(index)
+    }
+
+    fn null_count(&self) -> usize {
+        self.as_ref().null_count()
+    }
+
+    fn get_buffer_memory_size(&self) -> usize {
+        self.as_ref().get_buffer_memory_size()
+    }
+
+    fn get_array_memory_size(&self) -> usize {
+        self.as_ref().get_array_memory_size()
+    }
+
+    fn to_raw(
+        &self,
+    ) -> Result<(*const ffi::FFI_ArrowArray, *const ffi::FFI_ArrowSchema)> {
+        let data = self.data().clone();
+        let array = ffi::ArrowArray::try_from(data)?;
+        Ok(ffi::ArrowArray::into_raw(array))
+    }
+}
+
 /// Constructs an array using the input `data`.
 /// Returns a reference-counted `Array` instance.
 pub fn make_array(data: ArrayData) -> ArrayRef {
@@ -843,4 +906,22 @@ mod tests {
             expected_size
         );
     }
+
+    /// Test function that takes an &dyn Array
+    fn compute_my_thing(arr: &dyn Array) -> bool {
+        !arr.is_empty()
+    }
+
+    #[test]
+    fn test_array_ref_as_array() {
+        let arr: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+
+        // works well!
+        assert!(compute_my_thing(&arr));
+
+        // Should also work when wrapped as an ArrayRef
+        let arr: ArrayRef = Arc::new(arr);
+        assert!(compute_my_thing(&arr));
+        assert!(compute_my_thing(arr.as_ref()));
+    }
 }
diff --git a/arrow/src/array/cast.rs b/arrow/src/array/cast.rs
@@ -21,7 +21,7 @@ use crate::array::*;
 use crate::datatypes::*;
 
 /// Force downcast ArrayRef to PrimitiveArray<T>
-pub fn as_primitive_array<T>(arr: &ArrayRef) -> &PrimitiveArray<T>
+pub fn as_primitive_array<T>(arr: &dyn Array) -> &PrimitiveArray<T>
 where
     T: ArrowPrimitiveType,
 {
@@ -31,7 +31,7 @@ where
 }
 
 /// Force downcast ArrayRef to DictionaryArray<T>
-pub fn as_dictionary_array<T>(arr: &ArrayRef) -> &DictionaryArray<T>
+pub fn as_dictionary_array<T>(arr: &dyn Array) -> &DictionaryArray<T>
 where
     T: ArrowDictionaryKeyType,
 {
@@ -41,28 +41,30 @@ where
 }
 
 #[doc = "Force downcast ArrayRef to GenericListArray"]
-pub fn as_generic_list_array<S: OffsetSizeTrait>(arr: &ArrayRef) -> &GenericListArray<S> {
+pub fn as_generic_list_array<S: OffsetSizeTrait>(
+    arr: &dyn Array,
+) -> &GenericListArray<S> {
     arr.as_any()
         .downcast_ref::<GenericListArray<S>>()
         .expect("Unable to downcast to list array")
 }
 
 #[doc = "Force downcast ArrayRef to ListArray"]
 #[inline]
-pub fn as_list_array(arr: &ArrayRef) -> &ListArray {
+pub fn as_list_array(arr: &dyn Array) -> &ListArray {
     as_generic_list_array::<i32>(arr)
 }
 
 #[doc = "Force downcast ArrayRef to LargeListArray"]
 #[inline]
-pub fn as_large_list_array(arr: &ArrayRef) -> &LargeListArray {
+pub fn as_large_list_array(arr: &dyn Array) -> &LargeListArray {
     as_generic_list_array::<i64>(arr)
 }
 
 #[doc = "Force downcast ArrayRef to GenericBinaryArray"]
 #[inline]
 pub fn as_generic_binary_array<S: BinaryOffsetSizeTrait>(
-    arr: &ArrayRef,
+    arr: &dyn Array,
 ) -> &GenericBinaryArray<S> {
     arr.as_any()
         .downcast_ref::<GenericBinaryArray<S>>()
@@ -73,7 +75,7 @@ macro_rules! array_downcast_fn {
     ($name: ident, $arrty: ty, $arrty_str:expr) => {
         #[doc = "Force downcast ArrayRef to "]
         #[doc = $arrty_str]
-        pub fn $name(arr: &ArrayRef) -> &$arrty {
+        pub fn $name(arr: &dyn Array) -> &$arrty {
             arr.as_any().downcast_ref::<$arrty>().expect(concat!(
                 "Unable to downcast to typed array through ",
                 stringify!($name)
@@ -93,3 +95,30 @@ array_downcast_fn!(as_boolean_array, BooleanArray);
 array_downcast_fn!(as_null_array, NullArray);
 array_downcast_fn!(as_struct_array, StructArray);
 array_downcast_fn!(as_union_array, UnionArray);
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::*;
+
+    #[test]
+    fn test_as_primitive_array_ref() {
+        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
+
+        // should also work when wrapped in an Arc
+        let array: ArrayRef = Arc::new(array);
+        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
+    }
+
+    #[test]
+    fn test_as_string_array_ref() {
+        let array: StringArray = vec!["foo", "bar"].into_iter().map(Some).collect();
+        assert!(!as_string_array(&array).is_empty());
+
+        // should also work when wrapped in an Arc
+        let array: ArrayRef = Arc::new(array);
+        assert!(!as_string_array(&array).is_empty())
+    }
+}
diff --git a/arrow/src/array/equal_json.rs b/arrow/src/array/equal_json.rs
@@ -441,6 +441,12 @@ impl PartialEq<Value> for NullArray {
     }
 }
 
+impl JsonEqual for ArrayRef {
+    fn equals_json(&self, json: &[&Value]) -> bool {
+        self.as_ref().equals_json(json)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/arrow/src/pyarrow.rs b/arrow/src/pyarrow.rs
@@ -26,7 +26,7 @@ use pyo3::import_exception;
 use pyo3::prelude::*;
 use pyo3::types::PyList;
 
-use crate::array::{make_array, Array, ArrayData, ArrayRef};
+use crate::array::{Array, ArrayData, ArrayRef};
 use crate::datatypes::{DataType, Field, Schema};
 use crate::error::ArrowError;
 use crate::ffi;
@@ -148,16 +148,6 @@ impl PyArrowConvert for ArrayData {
     }
 }
 
-impl PyArrowConvert for ArrayRef {
-    fn from_pyarrow(value: &PyAny) -> PyResult<Self> {
-        Ok(make_array(ArrayData::from_pyarrow(value)?))
-    }
-
-    fn to_pyarrow(&self, py: Python) -> PyResult<PyObject> {
-        self.data().to_pyarrow(py)
-    }
-}
-
 impl<T> PyArrowConvert for T
 where
     T: Array + From<ArrayData>,