You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2022/11/03 08:28:54 UTC

[arrow-rs] branch master updated: Add macro downcast_temporal_array (#3007)

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 24afac490 Add macro downcast_temporal_array (#3007)
24afac490 is described below

commit 24afac49063d4ee92c0a2be84fd7c305d6088e14
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Thu Nov 3 01:28:48 2022 -0700

    Add macro downcast_temporal_array (#3007)
    
    * Add macro downcast_temporal_array
    
    * Fix TimeUnit import
    
    * Use downcast_temporal in downcast_primitive
    
    * Fix typo
---
 arrow-array/src/cast.rs | 137 ++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 110 insertions(+), 27 deletions(-)

diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs
index 4436dc77c..7bc627137 100644
--- a/arrow-array/src/cast.rs
+++ b/arrow-array/src/cast.rs
@@ -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)*)
             }
@@ -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)*)
             }
@@ -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)*)
             }
@@ -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,)*
+                }
+            }
         }
     };
 }