From f100521c4a619a7d931e790d96f9d40996d1d4aa Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Thu, 29 Sep 2022 22:13:17 +0100 Subject: [PATCH] Simplify downcast_primitive_array --- arrow-array/src/cast.rs | 442 ++++++++++------------------------------ 1 file changed, 105 insertions(+), 337 deletions(-) diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs index 653836b8d4e..6eb5407966f 100644 --- a/arrow-array/src/cast.rs +++ b/arrow-array/src/cast.rs @@ -20,6 +20,18 @@ use crate::array::*; use crate::types::*; +/// Repeats the provided pattern based on the number of comma separated identifiers +#[doc(hidden)] +#[macro_export] +macro_rules! repeat_pat { + ($e:pat, $v_:ident) => { + $e + }; + ($e:pat, $v_:ident $(, $tail:ident)+) => { + ($e, $crate::repeat_pat!($e $(, $tail)+)) + } +} + /// Downcast an [`Array`] to a [`PrimitiveArray`] based on its [`DataType`] /// accepts a number of subsequent patterns to match the data type /// @@ -50,427 +62,183 @@ macro_rules! downcast_primitive_array { ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => { downcast_primitive_array!($values => {$e} $($p => $fallback)*) }; - - ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => { - match $values.data_type() { - arrow_schema::DataType::Int8 => { - let $values = $crate::cast::as_primitive_array::< + (($($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 $(,)*)*) => { + match ($($values.data_type()),+) { + $crate::repeat_pat!(arrow_schema::DataType::Int8, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Int8Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Int16 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Int16, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Int16Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Int32 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Int32, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Int32Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Int64 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Int64, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Int64Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::UInt8 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::UInt8, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::UInt8Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::UInt16 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::UInt16, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::UInt16Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::UInt32 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::UInt32, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::UInt32Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::UInt64 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::UInt64, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::UInt64Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Float16 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Float16, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Float16Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Float32 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Float32, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Float32Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Float64 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Float64, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Float64Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Date32 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Date32, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Date32Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Date64 => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Date64, $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Date64Type, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Time32SecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Millisecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Millisecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Time32MillisecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Microsecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Microsecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Time64MicrosecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Nanosecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Nanosecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::Time64NanosecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Second, _) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Second, _), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::TimestampSecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::TimestampMillisecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::TimestampMicrosecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::TimestampNanosecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::IntervalYearMonthType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::DayTime) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::DayTime), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::IntervalDayTimeType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::IntervalMonthDayNanoType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Second) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Second), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::DurationSecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Millisecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Millisecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::DurationMillisecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Microsecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Microsecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::DurationMicrosecondType, - >($values); + >($values);)+ $e } - arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond) => { - let $values = $crate::cast::as_primitive_array::< + $crate::repeat_pat!(arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond), $($values),+) => { + $(let $values = $crate::cast::as_primitive_array::< $crate::types::DurationNanosecondType, - >($values); - $e - } - $($p => $fallback,)* - } - }; - - (($values1:ident, $values2:ident) => $e:block $($p:pat => $fallback:expr $(,)*)*) => { - match ($values1.data_type(), $values2.data_type()) { - (arrow_schema::DataType::Int8, arrow_schema::DataType::Int8) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Int8Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Int8Type, - >($values2); - $e - } - (arrow_schema::DataType::Int16, arrow_schema::DataType::Int16) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Int16Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Int16Type, - >($values2); - $e - } - (arrow_schema::DataType::Int32, arrow_schema::DataType::Int32) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Int32Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Int32Type, - >($values2); - $e - } - (arrow_schema::DataType::Int64, arrow_schema::DataType::Int64) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Int64Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Int64Type, - >($values2); - $e - } - (arrow_schema::DataType::UInt8, arrow_schema::DataType::UInt8) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::UInt8Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::UInt8Type, - >($values2); - $e - } - (arrow_schema::DataType::UInt16, arrow_schema::DataType::UInt16) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::UInt16Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::UInt16Type, - >($values2); - $e - } - (arrow_schema::DataType::UInt32, arrow_schema::DataType::UInt32) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::UInt32Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::UInt32Type, - >($values2); - $e - } - (arrow_schema::DataType::UInt64, arrow_schema::DataType::UInt64) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::UInt64Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::UInt64Type, - >($values2); - $e - } - (arrow_schema::DataType::Float32, arrow_schema::DataType::Float32) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Float32Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Float32Type, - >($values2); - $e - } - (arrow_schema::DataType::Float64, arrow_schema::DataType::Float64) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Float64Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Float64Type, - >($values2); - $e - } - (arrow_schema::DataType::Date32, arrow_schema::DataType::Date32) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Date32Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Date32Type, - >($values2); - $e - } - (arrow_schema::DataType::Date64, arrow_schema::DataType::Date64) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Date64Type, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Date64Type, - >($values2); + >($values);)+ $e } - (arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second), arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Second)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Time32SecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Time32SecondType, - >($values2); - $e - } - (arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Millisecond), arrow_schema::DataType::Time32(arrow_schema::TimeUnit::Millisecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Time32MillisecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Time32MillisecondType, - >($values2); - $e - } - (arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Microsecond), arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Microsecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Time64MicrosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Time64MicrosecondType, - >($values2); - $e - } - (arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Nanosecond), arrow_schema::DataType::Time64(arrow_schema::TimeUnit::Nanosecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::Time64NanosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::Time64NanosecondType, - >($values2); - $e - } - (arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Second, _), arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Second, _)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::TimestampSecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::TimestampSecondType, - >($values2); - $e - } - (arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _), arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::TimestampMillisecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::TimestampMillisecondType, - >($values2); - $e - } - (arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _), arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, _)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::TimestampMicrosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::TimestampMicrosecondType, - >($values2); - $e - } - (arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _), arrow_schema::DataType::Timestamp(arrow_schema::TimeUnit::Nanosecond, _)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::TimestampNanosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::TimestampNanosecondType, - >($values2); - $e - } - (arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth), arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::IntervalYearMonthType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::IntervalYearMonthType, - >($values2); - $e - } - (arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::DayTime), arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::DayTime)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::IntervalDayTimeType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::IntervalDayTimeType, - >($values2); - $e - } - (arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano), arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::IntervalMonthDayNanoType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::IntervalMonthDayNanoType, - >($values2); - $e - } - (arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Second), arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Second)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::DurationSecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::DurationSecondType, - >($values2); - $e - } - (arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Millisecond), arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Millisecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::DurationMillisecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::DurationMillisecondType, - >($values2); - $e - } - (arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Microsecond), arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Microsecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::DurationMicrosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::DurationMicrosecondType, - >($values2); - $e - } - (arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond), arrow_schema::DataType::Duration(arrow_schema::TimeUnit::Nanosecond)) => { - let $values1 = $crate::cast::as_primitive_array::< - $crate::types::DurationNanosecondType, - >($values1); - let $values2 = $crate::cast::as_primitive_array::< - $crate::types::DurationNanosecondType, - >($values2); - $e - } - $($p => $fallback,)* + $(($($p),+) => $fallback,)* } }; }