use crate::array::*;
use crate::types::*;
use arrow_data::ArrayData;
#[doc(hidden)]
#[macro_export]
macro_rules! repeat_pat {
($e:pat, $v_:expr) => {
$e
};
($e:pat, $v_:expr $(, $tail:expr)+) => {
($e, $crate::repeat_pat!($e $(, $tail)+))
}
}
#[macro_export]
macro_rules! downcast_integer {
($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($($p:pat),+ => $fallback:expr $(,)*)*) => {
match ($($data_type),+) {
$crate::repeat_pat!(arrow_schema::DataType::Int8, $($data_type),+) => {
$m!($crate::types::Int8Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Int16, $($data_type),+) => {
$m!($crate::types::Int16Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Int32, $($data_type),+) => {
$m!($crate::types::Int32Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Int64, $($data_type),+) => {
$m!($crate::types::Int64Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::UInt8, $($data_type),+) => {
$m!($crate::types::UInt8Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::UInt16, $($data_type),+) => {
$m!($crate::types::UInt16Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::UInt32, $($data_type),+) => {
$m!($crate::types::UInt32Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::UInt64, $($data_type),+) => {
$m!($crate::types::UInt64Type $(, $args)*)
}
$(($($p),+) => $fallback,)*
}
};
}
#[macro_export]
macro_rules! downcast_temporal {
($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($($p:pat),+ => $fallback:expr $(,)*)*) => {
match ($($data_type),+) {
$crate::repeat_pat!(arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second), $($data_type),+) => {
$m!($crate::types::Time32SecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Millisecond), $($data_type),+) => {
$m!($crate::types::Time32MillisecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Microsecond), $($data_type),+) => {
$m!($crate::types::Time64MicrosecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Nanosecond), $($data_type),+) => {
$m!($crate::types::Time64NanosecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Date32, $($data_type),+) => {
$m!($crate::types::Date32Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Date64, $($data_type),+) => {
$m!($crate::types::Date64Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Second, _), $($data_type),+) => {
$m!($crate::types::TimestampSecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _), $($data_type),+) => {
$m!($crate::types::TimestampMillisecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _), $($data_type),+) => {
$m!($crate::types::TimestampMicrosecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _), $($data_type),+) => {
$m!($crate::types::TimestampNanosecondType $(, $args)*)
}
$(($($p),+) => $fallback,)*
}
};
}
#[macro_export]
macro_rules! downcast_temporal_array {
($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
$crate::downcast_temporal_array!($values => {$e} $($p => $fallback)*)
};
(($($values:ident),+) => $e:block $($($p:pat),+ => $fallback:expr $(,)*)*) => {
$crate::downcast_temporal_array!($($values),+ => $e $($($p),+ => $fallback)*)
};
(($($values:ident),+) => $e:block $(($($p:pat),+) => $fallback:expr $(,)*)*) => {
$crate::downcast_temporal_array!($($values),+ => $e $($($p),+ => $fallback)*)
};
($($values:ident),+ => $e:block $($($p:pat),+ => $fallback:expr $(,)*)*) => {
$crate::downcast_temporal!{
$($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
$($($p),+ => $fallback,)*
}
};
}
#[macro_export]
macro_rules! downcast_primitive {
($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($($p:pat),+ => $fallback:expr $(,)*)*) => {
$crate::downcast_integer! {
$($data_type),+ => ($m $(, $args)*),
$crate::repeat_pat!(arrow_schema::DataType::Float16, $($data_type),+) => {
$m!($crate::types::Float16Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Float32, $($data_type),+) => {
$m!($crate::types::Float32Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Float64, $($data_type),+) => {
$m!($crate::types::Float64Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Decimal128(_, _), $($data_type),+) => {
$m!($crate::types::Decimal128Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Decimal256(_, _), $($data_type),+) => {
$m!($crate::types::Decimal256Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth), $($data_type),+) => {
$m!($crate::types::IntervalYearMonthType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::DayTime), $($data_type),+) => {
$m!($crate::types::IntervalDayTimeType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano), $($data_type),+) => {
$m!($crate::types::IntervalMonthDayNanoType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Second), $($data_type),+) => {
$m!($crate::types::DurationSecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Millisecond), $($data_type),+) => {
$m!($crate::types::DurationMillisecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Microsecond), $($data_type),+) => {
$m!($crate::types::DurationMicrosecondType $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond), $($data_type),+) => {
$m!($crate::types::DurationNanosecondType $(, $args)*)
}
_ => {
$crate::downcast_temporal! {
$($data_type),+ => ($m $(, $args)*),
$($($p),+ => $fallback,)*
}
}
}
};
}
#[macro_export]
#[doc(hidden)]
macro_rules! downcast_primitive_array_helper {
($t:ty, $($values:ident),+, $e:block) => {{
$(let $values = $crate::cast::as_primitive_array::<$t>($values);)+
$e
}};
}
#[macro_export]
macro_rules! downcast_primitive_array {
($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
$crate::downcast_primitive_array!($values => {$e} $($p => $fallback)*)
};
(($($values:ident),+) => $e:block $($($p:pat),+ => $fallback:expr $(,)*)*) => {
$crate::downcast_primitive_array!($($values),+ => $e $($($p),+ => $fallback)*)
};
(($($values:ident),+) => $e:block $(($($p:pat),+) => $fallback:expr $(,)*)*) => {
$crate::downcast_primitive_array!($($values),+ => $e $($($p),+ => $fallback)*)
};
($($values:ident),+ => $e:block $($($p:pat),+ => $fallback:expr $(,)*)*) => {
$crate::downcast_primitive!{
$($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
$($($p),+ => $fallback,)*
}
};
}
pub fn as_primitive_array<T>(arr: &dyn Array) -> &PrimitiveArray<T>
where
T: ArrowPrimitiveType,
{
arr.as_any()
.downcast_ref::<PrimitiveArray<T>>()
.expect("Unable to downcast to primitive array")
}
#[macro_export]
#[doc(hidden)]
macro_rules! downcast_dictionary_array_helper {
($t:ty, $($values:ident),+, $e:block) => {{
$(let $values = $crate::cast::as_dictionary_array::<$t>($values);)+
$e
}};
}
#[macro_export]
macro_rules! downcast_dictionary_array {
($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
downcast_dictionary_array!($values => {$e} $($p => $fallback)*)
};
($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
match $values.data_type() {
arrow_schema::DataType::Dictionary(k, _) => {
$crate::downcast_integer! {
k.as_ref() => ($crate::downcast_dictionary_array_helper, $values, $e),
k => unreachable!("unsupported dictionary key type: {}", k)
}
}
$($p => $fallback,)*
}
}
}
pub fn as_dictionary_array<T>(arr: &dyn Array) -> &DictionaryArray<T>
where
T: ArrowDictionaryKeyType,
{
arr.as_any()
.downcast_ref::<DictionaryArray<T>>()
.expect("Unable to downcast to dictionary array")
}
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")
}
#[inline]
pub fn as_list_array(arr: &dyn Array) -> &ListArray {
as_generic_list_array::<i32>(arr)
}
#[inline]
pub fn as_large_list_array(arr: &dyn Array) -> &LargeListArray {
as_generic_list_array::<i64>(arr)
}
#[inline]
pub fn as_generic_binary_array<S: OffsetSizeTrait>(
arr: &dyn Array,
) -> &GenericBinaryArray<S> {
arr.as_any()
.downcast_ref::<GenericBinaryArray<S>>()
.expect("Unable to downcast to binary array")
}
pub fn as_string_array(arr: &dyn Array) -> &StringArray {
arr.as_any()
.downcast_ref::<StringArray>()
.expect("Unable to downcast to StringArray")
}
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 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!(
"Unable to downcast to typed array through ",
stringify!($name)
))
}
};
($name: ident, $arrty: ty) => {
array_downcast_fn!(
$name,
$arrty,
concat!("[`", stringify!($arrty), "`], panic'ing on failure.")
);
};
}
array_downcast_fn!(as_largestring_array, LargeStringArray);
array_downcast_fn!(as_null_array, NullArray);
array_downcast_fn!(as_struct_array, StructArray);
array_downcast_fn!(as_union_array, UnionArray);
array_downcast_fn!(as_map_array, MapArray);
array_downcast_fn!(as_decimal_array, Decimal128Array);
pub fn downcast_array<T>(array: &dyn Array) -> T
where
T: From<ArrayData>,
{
T::from(array.data().clone())
}
#[cfg(test)]
mod tests {
use arrow_buffer::i256;
use std::sync::Arc;
use super::*;
#[test]
fn test_as_decimal_array_ref() {
let array: Decimal128Array = vec![Some(123), None, Some(1111)]
.into_iter()
.collect::<Decimal128Array>()
.with_precision_and_scale(10, 2)
.unwrap();
assert!(!as_decimal_array(&array).is_empty());
let result_decimal = as_decimal_array(&array);
assert_eq!(result_decimal, &array);
}
#[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());
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());
let array: ArrayRef = Arc::new(array);
assert!(!as_string_array(&array).is_empty())
}
#[test]
fn test_decimal128array() {
let a = Decimal128Array::from_iter_values([1, 2, 4, 5]);
assert!(!as_primitive_array::<Decimal128Type>(&a).is_empty());
}
#[test]
fn test_decimal256array() {
let a = Decimal256Array::from_iter_values(
[1, 2, 4, 5].into_iter().map(i256::from_i128),
);
assert!(!as_primitive_array::<Decimal256Type>(&a).is_empty());
}
}