diff --git a/arrow/src/array/cast.rs b/arrow/src/array/cast.rs index dd6c71357101..fd57133bc843 100644 --- a/arrow/src/array/cast.rs +++ b/arrow/src/array/cast.rs @@ -15,12 +15,35 @@ // specific language governing permissions and limitations // under the License. -//! Defines helper functions for force Array type downcast +//! Defines helper functions for force [`Array`] downcasts use crate::array::*; use crate::datatypes::*; -/// Force downcast ArrayRef to PrimitiveArray +/// Force downcast of an [`Array`], such as an [`ArrayRef`], to +/// [`PrimitiveArray`], panic'ing on failure. +/// +/// # Example +/// +/// ``` +/// # use arrow::array::*; +/// # use arrow::datatypes::*; +/// # use std::sync::Arc; +/// let arr: ArrayRef = Arc::new(Int32Array::from(vec![Some(1)])); +/// +/// // Downcast an `ArrayRef` to Int32Array / PrimiveArray: +/// let primitive_array: &Int32Array = as_primitive_array(&arr); +/// +/// // Equivalently: +/// let primitive_array = as_primitive_array::(&arr); +/// +/// // This is the equivalent of: +/// let primitive_array = arr +/// .as_any() +/// .downcast_ref::() +/// .unwrap(); +/// ``` + pub fn as_primitive_array(arr: &dyn Array) -> &PrimitiveArray where T: ArrowPrimitiveType, @@ -30,7 +53,19 @@ where .expect("Unable to downcast to primitive array") } -/// Force downcast ArrayRef to DictionaryArray +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`DictionaryArray`], panic'ing on failure. +/// +/// # Example +/// +/// ``` +/// # use arrow::array::*; +/// # use arrow::datatypes::*; +/// # use std::sync::Arc; +/// let arr: DictionaryArray = vec![Some("foo")].into_iter().collect(); +/// let arr: ArrayRef = std::sync::Arc::new(arr); +/// let dict_array: &DictionaryArray = as_dictionary_array::(&arr); +/// ``` pub fn as_dictionary_array(arr: &dyn Array) -> &DictionaryArray where T: ArrowDictionaryKeyType, @@ -40,7 +75,8 @@ where .expect("Unable to downcast to dictionary array") } -#[doc = "Force downcast ArrayRef to GenericListArray"] +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`GenericListArray`], panic'ing on failure. pub fn as_generic_list_array( arr: &dyn Array, ) -> &GenericListArray { @@ -49,19 +85,22 @@ pub fn as_generic_list_array( .expect("Unable to downcast to list array") } -#[doc = "Force downcast ArrayRef to ListArray"] +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`ListArray`], panic'ing on failure. #[inline] pub fn as_list_array(arr: &dyn Array) -> &ListArray { as_generic_list_array::(arr) } -#[doc = "Force downcast ArrayRef to LargeListArray"] +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`LargeListArray`], panic'ing on failure. #[inline] pub fn as_large_list_array(arr: &dyn Array) -> &LargeListArray { as_generic_list_array::(arr) } -#[doc = "Force downcast ArrayRef to GenericBinaryArray"] +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`GenericBinaryArray`], panic'ing on failure. #[inline] pub fn as_generic_binary_array( arr: &dyn Array, @@ -71,9 +110,43 @@ pub fn as_generic_binary_array( .expect("Unable to downcast to binary array") } +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`StringArray`], panic'ing on failure. +/// +/// # Example +/// +/// ``` +/// # use arrow::array::*; +/// # use std::sync::Arc; +/// let arr: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo")])); +/// let string_array = as_string_array(&arr); +/// ``` +pub fn as_string_array(arr: &dyn Array) -> &StringArray { + arr.as_any() + .downcast_ref::() + .expect("Unable to downcast to StringArray") +} + +/// Force downcast of an [`Array`], such as an [`ArrayRef`] to +/// [`BooleanArray`], panic'ing on failure. +/// +/// # Example +/// +/// ``` +/// # use arrow::array::*; +/// # use std::sync::Arc; +/// let arr: ArrayRef = Arc::new(BooleanArray::from_iter(vec![Some(true)])); +/// let boolean_array = as_boolean_array(&arr); +/// ``` +pub fn as_boolean_array(arr: &dyn Array) -> &BooleanArray { + arr.as_any() + .downcast_ref::() + .expect("Unable to downcast to BooleanArray") +} + macro_rules! array_downcast_fn { ($name: ident, $arrty: ty, $arrty_str:expr) => { - #[doc = "Force downcast ArrayRef to "] + #[doc = "Force downcast of an [`Array`], such as an [`ArrayRef`] to "] #[doc = $arrty_str] pub fn $name(arr: &dyn Array) -> &$arrty { arr.as_any().downcast_ref::<$arrty>().expect(concat!( @@ -85,13 +158,15 @@ macro_rules! array_downcast_fn { // use recursive macro to generate dynamic doc string for a given array type ($name: ident, $arrty: ty) => { - array_downcast_fn!($name, $arrty, stringify!($arrty)); + array_downcast_fn!( + $name, + $arrty, + concat!("[`", stringify!($arrty), "`], panic'ing on failure.") + ); }; } -array_downcast_fn!(as_string_array, StringArray); array_downcast_fn!(as_largestring_array, LargeStringArray); -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); diff --git a/arrow/src/array/mod.rs b/arrow/src/array/mod.rs index e533f9e75d58..d805710ccf5d 100644 --- a/arrow/src/array/mod.rs +++ b/arrow/src/array/mod.rs @@ -24,9 +24,11 @@ //! Arrays are often passed around as a dynamically typed [`&dyn Array`] or [`ArrayRef`]. //! For example, [`RecordBatch`](`crate::record_batch::RecordBatch`) stores columns as [`ArrayRef`]. //! -//! Whilst these arrays can be passed directly to the [`compute`](crate::compute), -//! [`csv`](crate::csv), [`json`](crate::json), etc... APIs, it is often the case that you wish -//! to interact with the data directly. This requires downcasting to the concrete type of the array: +//! Whilst these arrays can be passed directly to the +//! [`compute`](crate::compute), [`csv`](crate::csv), +//! [`json`](crate::json), etc... APIs, it is often the case that you +//! wish to interact with the data directly. This requires downcasting +//! to the concrete type of the array: //! //! ``` //! # use arrow::array::{Array, Float32Array, Int32Array}; @@ -42,6 +44,19 @@ //! } //! ``` //! +//! Additionally, there are convenient functions to do this casting +//! such as [`as_primitive_array`] and [`as_string_array`]: +//! +//! ``` +//! # use arrow::array::*; +//! # use arrow::datatypes::*; +//! # +//! fn as_f32_slice(array: &dyn Array) -> &[f32] { +//! // use as_primtive_array +//! as_primitive_array::(array).values() +//! } +//! ``` + //! # Building an Array //! //! Most [`Array`] implementations can be constructed directly from iterators or [`Vec`]