diff --git a/arrow/src/array/array.rs b/arrow/src/array/array.rs index ce3751def8a4..795439ec80a5 100644 --- 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; +/// 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 index e4284ef2218b..c96e49e8ee00 100644 --- 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 -pub fn as_primitive_array(arr: &ArrayRef) -> &PrimitiveArray +pub fn as_primitive_array(arr: &dyn Array) -> &PrimitiveArray where T: ArrowPrimitiveType, { @@ -31,7 +31,7 @@ where } /// Force downcast ArrayRef to DictionaryArray -pub fn as_dictionary_array(arr: &ArrayRef) -> &DictionaryArray +pub fn as_dictionary_array(arr: &dyn Array) -> &DictionaryArray where T: ArrowDictionaryKeyType, { @@ -41,7 +41,9 @@ where } #[doc = "Force downcast ArrayRef to GenericListArray"] -pub fn as_generic_list_array(arr: &ArrayRef) -> &GenericListArray { +pub fn as_generic_list_array( + arr: &dyn Array, +) -> &GenericListArray { arr.as_any() .downcast_ref::>() .expect("Unable to downcast to list array") @@ -49,20 +51,20 @@ pub fn as_generic_list_array(arr: &ArrayRef) -> &GenericList #[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::(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::(arr) } #[doc = "Force downcast ArrayRef to GenericBinaryArray"] #[inline] pub fn as_generic_binary_array( - arr: &ArrayRef, + arr: &dyn Array, ) -> &GenericBinaryArray { arr.as_any() .downcast_ref::>() @@ -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::(&array).is_empty()); + + // should also work when wrapped in an Arc + let array: ArrayRef = Arc::new(array); + assert!(!as_primitive_array::(&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 index adc33a7a1cd3..0350fc962b4e 100644 --- a/arrow/src/array/equal_json.rs +++ b/arrow/src/array/equal_json.rs @@ -441,6 +441,12 @@ impl PartialEq 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 index 12d4c0d1f92d..62e6316b621c 100644 --- 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 { - Ok(make_array(ArrayData::from_pyarrow(value)?)) - } - - fn to_pyarrow(&self, py: Python) -> PyResult { - self.data().to_pyarrow(py) - } -} - impl PyArrowConvert for T where T: Array + From,