Skip to content

Commit

Permalink
Add macro downcast_temporal_array (#3007)
Browse files Browse the repository at this point in the history 
* Add macro downcast_temporal_array

* Fix TimeUnit import

* Use downcast_temporal in downcast_primitive

* Fix typo
  • Loading branch information
@viirya
viirya committed Nov 3, 2022
1 parent a2c6647 commit 24afac4
Showing 1 changed file with 110 additions and 27 deletions.
137 changes: 110 additions & 27 deletions arrow-array/src/cast.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,48 +98,31 @@ macro_rules! downcast_integer {
/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
///
/// ```
/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType};
/// # use arrow_array::{downcast_temporal, ArrowPrimitiveType};
/// # use arrow_schema::DataType;
///
/// macro_rules! primitive_size_helper {
/// macro_rules! temporal_size_helper {
/// ($t:ty, $o:ty) => {
/// std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
/// };
/// }
///
/// fn primitive_size(t: &DataType) -> u8 {
/// downcast_primitive! {
/// t => (primitive_size_helper, u8),
/// fn temporal_size(t: &DataType) -> u8 {
/// downcast_temporal! {
/// t => (temporal_size_helper, u8),
/// _ => u8::MAX
/// }
/// }
///
/// assert_eq!(primitive_size(&DataType::Int32), 4);
/// assert_eq!(primitive_size(&DataType::Int64), 8);
/// assert_eq!(primitive_size(&DataType::Float16), 2);
/// assert_eq!(temporal_size(&DataType::Date32), 4);
/// assert_eq!(temporal_size(&DataType::Date64), 8);
/// ```
///
/// [`DataType`]: arrow_schema::DataType
#[macro_export]
macro_rules! downcast_primitive {
macro_rules! downcast_temporal {
($($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::Date32, $($data_type),+) => {
$m!($crate::types::Date32Type $(, $args)*)
}
$crate::repeat_pat!(arrow_schema::DataType::Date64, $($data_type),+) => {
$m!($crate::types::Date64Type $(, $args)*)
}
match ($($data_type),+) {
$crate::repeat_pat!(arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second), $($data_type),+) => {
$m!($crate::types::Time32SecondType $(, $args)*)
}
Expand All @@ -152,6 +135,12 @@ macro_rules! downcast_primitive {
$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)*)
}
Expand All @@ -164,6 +153,95 @@ macro_rules! downcast_primitive {
$crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _), $($data_type),+) => {
$m!($crate::types::TimestampNanosecondType $(, $args)*)
}
$(($($p),+) => $fallback,)*
}
};
}

/// Downcast an [`Array`] to a temporal [`PrimitiveArray`] based on its [`DataType`]
/// accepts a number of subsequent patterns to match the data type
///
/// ```
/// # use arrow_array::{Array, downcast_temporal_array, cast::as_string_array};
/// # use arrow_schema::DataType;
///
/// fn print_temporal(array: &dyn Array) {
/// downcast_temporal_array!(
/// array => {
/// for v in array {
/// println!("{:?}", v);
/// }
/// }
/// DataType::Utf8 => {
/// for v in as_string_array(array) {
/// println!("{:?}", v);
/// }
/// }
/// t => println!("Unsupported datatype {}", t)
/// )
/// }
/// ```
///
/// [`DataType`]: arrow_schema::DataType
#[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,)*
}
};
}

/// Given one or more expressions evaluating to primitive [`DataType`] invokes the provided macro
/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
///
/// ```
/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType};
/// # use arrow_schema::DataType;
///
/// macro_rules! primitive_size_helper {
/// ($t:ty, $o:ty) => {
/// std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
/// };
/// }
///
/// fn primitive_size(t: &DataType) -> u8 {
/// downcast_primitive! {
/// t => (primitive_size_helper, u8),
/// _ => u8::MAX
/// }
/// }
///
/// assert_eq!(primitive_size(&DataType::Int32), 4);
/// assert_eq!(primitive_size(&DataType::Int64), 8);
/// assert_eq!(primitive_size(&DataType::Float16), 2);
/// ```
///
/// [`DataType`]: arrow_schema::DataType
#[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::Interval(arrow_schema::IntervalUnit::YearMonth), $($data_type),+) => {
$m!($crate::types::IntervalYearMonthType $(, $args)*)
}
Expand All @@ -185,7 +263,12 @@ macro_rules! downcast_primitive {
$crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond), $($data_type),+) => {
$m!($crate::types::DurationNanosecondType $(, $args)*)
}
$($($p),+ => $fallback,)*
_ => {
$crate::downcast_temporal! {
$($data_type),+ => ($m $(, $args)*),
$($($p),+ => $fallback,)*
}
}
}
};
}
Expand Down

0 comments on commit 24afac4

Please sign in to comment.