Skip to content

Commit

Permalink
Improve docstrings + examples for as_primitive_array cast functions (…
Browse files Browse the repository at this point in the history 
…#2114)

* Improve doc strings for as_primtive_array

* Add more docs for cast

* improvements

* improvements

* Update arrow/src/array/mod.rs

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>

* fixups

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>
  • Loading branch information
@alamb @viirya
alamb and viirya committed Jul 21, 2022
1 parent db91031 commit f62f26a
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 14 deletions.
97 changes: 86 additions & 11 deletions arrow/src/array/cast.rs
View file
Expand Up @@ -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<T>
/// Force downcast of an [`Array`], such as an [`ArrayRef`], to
/// [`PrimitiveArray<T>`], 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<Int32>:
/// let primitive_array: &Int32Array = as_primitive_array(&arr);
///
/// // Equivalently:
/// let primitive_array = as_primitive_array::<Int32Type>(&arr);
///
/// // This is the equivalent of:
/// let primitive_array = arr
/// .as_any()
/// .downcast_ref::<Int32Array>()
/// .unwrap();
/// ```

pub fn as_primitive_array<T>(arr: &dyn Array) -> &PrimitiveArray<T>
where
T: ArrowPrimitiveType,
Expand All @@ -30,7 +53,19 @@ where
.expect("Unable to downcast to primitive array")
}

/// Force downcast ArrayRef to DictionaryArray<T>
/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
/// [`DictionaryArray<T>`], panic'ing on failure.
///
/// # Example
///
/// ```
/// # use arrow::array::*;
/// # use arrow::datatypes::*;
/// # use std::sync::Arc;
/// let arr: DictionaryArray<Int32Type> = vec![Some("foo")].into_iter().collect();
/// let arr: ArrayRef = std::sync::Arc::new(arr);
/// let dict_array: &DictionaryArray<Int32Type> = as_dictionary_array::<Int32Type>(&arr);
/// ```
pub fn as_dictionary_array<T>(arr: &dyn Array) -> &DictionaryArray<T>
where
T: ArrowDictionaryKeyType,
Expand All @@ -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<T>`], panic'ing on failure.
pub fn as_generic_list_array<S: OffsetSizeTrait>(
arr: &dyn Array,
) -> &GenericListArray<S> {
Expand All @@ -49,19 +85,22 @@ pub fn as_generic_list_array<S: OffsetSizeTrait>(
.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::<i32>(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::<i64>(arr)
}

#[doc = "Force downcast ArrayRef to GenericBinaryArray"]
/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
/// [`GenericBinaryArray<S>`], panic'ing on failure.
#[inline]
pub fn as_generic_binary_array<S: OffsetSizeTrait>(
arr: &dyn Array,
Expand All @@ -71,9 +110,43 @@ pub fn as_generic_binary_array<S: OffsetSizeTrait>(
.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::<StringArray>()
.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::<BooleanArray>()
.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!(
Expand All @@ -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);
Expand Down
21 changes: 18 additions & 3 deletions arrow/src/array/mod.rs
View file
Expand Up @@ -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};
Expand All @@ -42,6 +44,19 @@
//! }
//! ```
//!
//! Additionally, there are convenient functions to do this casting
//! such as [`as_primitive_array<T>`] 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::<Float32Type>(array).values()
//! }
//! ```

//! # Building an Array
//!
//! Most [`Array`] implementations can be constructed directly from iterators or [`Vec`]
Expand Down

0 comments on commit f62f26a

Please sign in to comment.