add Variant::{as_timestamp_micros/as_timestamp_nanos} and simplify the impl of for timestamp types

Author: klion26

parquet-variant-compute/src/type_conversion.rs

@@ -17,10 +17,7 @@
 
 //! Module for transforming a typed arrow `Array` to `VariantArray`.
 
-use arrow::datatypes::{
-    self, ArrowPrimitiveType, ArrowTimestampType, Date32Type, TimestampMicrosecondType,
-    TimestampNanosecondType,
-};
+use arrow::datatypes::{self, ArrowPrimitiveType, ArrowTimestampType, Date32Type};
 use parquet_variant::Variant;
 
 /// Options for controlling the behavior of `cast_to_variant_with_options`.
@@ -41,6 +38,13 @@ pub(crate) trait PrimitiveFromVariant: ArrowPrimitiveType {
     fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
 }
 
+/// Extension trait for Arrow timestamp types that can extract their native value from a Variant
+/// We can't use [`PrimitiveFromVariant`] directly because we might need to use methods that
+/// are only available on [`ArrowTimestampType`] (such as with_timezone_opt)
+pub(crate) trait TimestampFromVariant: ArrowTimestampType {
+    fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
+}
+
 /// Macro to generate PrimitiveFromVariant implementations for Arrow primitive types
 macro_rules! impl_primitive_from_variant {
     ($arrow_type:ty, $variant_method:ident $(, $cast_fn:expr)?) => {
@@ -52,6 +56,18 @@ macro_rules! impl_primitive_from_variant {
             }
         }
     };
+    ($arrow_type:ty $(, $variant_method:ident => $cast_fn:expr )+ ) => {
+        impl TimestampFromVariant for $arrow_type {
+            fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native> {
+                $(
+                    if let Some(value) = variant.$variant_method() {
+                        return Some($cast_fn(value));
+                    }
+                )+
+                None
+            }
+        }
+    };
 }
 
 impl_primitive_from_variant!(datatypes::Int32Type, as_int32);
@@ -70,41 +86,13 @@ impl_primitive_from_variant!(
     as_naive_date,
     Date32Type::from_naive_date
 );
-
-pub(crate) trait TimestampFromVariant: ArrowTimestampType {
-    fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native>;
-}
-
-macro_rules! impl_timestamp_from_variant {
-    ($timestamp_type:ty,
-    $( $variant_pattern:pat => $conversion:expr ),+ $(,)?
-    ) => {
-        impl TimestampFromVariant for $timestamp_type {
-            fn from_variant(variant: &Variant<'_, '_>) -> Option<Self::Native> {
-                match variant {
-                    $(
-                        $variant_pattern => $conversion,
-                    )+
-                    _ => None,
-                }
-            }
-        }
-    };
-}
-
-impl_timestamp_from_variant!(
-    TimestampMicrosecondType,
-    Variant::TimestampMicros(t) => Some(t.timestamp_micros()),
-    Variant::TimestampNtzMicros(t) => Some(t.and_utc().timestamp_micros()),
-);
-
-impl_timestamp_from_variant!(
-    TimestampNanosecondType,
-    Variant::TimestampMicros(t) => Some(t.timestamp_micros()).map(|t| t * 1000),
-    Variant::TimestampNtzMicros(t) => Some(t.and_utc().timestamp_micros()).map(|t| t * 1000),
-    Variant::TimestampNanos(t) => t.timestamp_nanos_opt(),
-    Variant::TimestampNtzNanos(t) => t.and_utc().timestamp_nanos_opt(),
-);
+impl_primitive_from_variant!(
+    datatypes::TimestampMicrosecondType,
+    as_timestamp_micros => |t| t);
+impl_primitive_from_variant!(
+    datatypes::TimestampNanosecondType,
+    as_timestamp_micros => |t| 1000 * t,
+    as_timestamp_nanos => |t| t);
 
 /// Convert the value at a specific index in the given array into a `Variant`.
 macro_rules! non_generic_conversion_single_value {

parquet-variant/src/variant.rs

@@ -561,6 +561,72 @@ impl<'m, 'v> Variant<'m, 'v> {
         }
     }
 
+    /// Converts this variant to a `i64` representing microseconds since the Unix epoch if possible.
+    /// This is useful when convert the variant to arrow types.
+    ///
+    /// Returns Some(i64) for [`Variant::TimestampMicros`] and [`Variant::TimestampNtzMicros`],
+    /// None for the other variant types.
+    ///
+    /// ```
+    /// use parquet_variant::Variant;
+    /// use chrono::NaiveDate;
+    ///
+    /// // you can extract an i64 from Variant::TimestampMicros
+    /// let datetime = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_milli_opt(12, 34, 56, 789).unwrap().and_utc();
+    /// let v1 = Variant::from(datetime);
+    /// assert_eq!(v1.as_timestamp_micros(), Some(1759494896789000));
+    ///
+    /// // or Variant::TimestampNtzMicros
+    /// let datetime_ntz = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_milli_opt(12, 34, 56, 789).unwrap();
+    /// let v2 = Variant::from(datetime_ntz);
+    /// assert_eq!(v1.as_timestamp_micros(), Some(1759494896789000));
+    ///
+    /// // but not from other variants
+    /// let datetime_nanos = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_nano_opt(12, 34, 56, 789123456).unwrap().and_utc();
+    /// let v3 = Variant::from(datetime_nanos);
+    /// assert_eq!(v3.as_timestamp_micros(), None);
+    /// ```
+    pub fn as_timestamp_micros(&self) -> Option<i64> {
+        match *self {
+            Variant::TimestampMicros(d) => Some(d.timestamp_micros()),
+            Variant::TimestampNtzMicros(d) => Some(d.and_utc().timestamp_micros()),
+            _ => None,
+        }
+    }
+
+    /// Converts this variant to a `i64` representing nanoseconds since the Unix epoch if possible.
+    /// This is useful when convert the variant to arrow types.
+    ///
+    /// Returns Some(i64) for [`Variant::TimestampNanos`] and [`Variant::TimestampNtzNanos`],
+    /// None for the other variant types.
+    ///
+    /// ```
+    /// use parquet_variant::Variant;
+    /// use chrono::NaiveDate;
+    ///
+    /// // you can extract an i64 from Variant::TimestampNanos
+    /// let datetime = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_nano_opt(12, 34, 56, 789123456).unwrap().and_utc();
+    /// let v1 = Variant::from(datetime);
+    /// assert_eq!(v1.as_timestamp_nanos(), Some(1759494896789123456));
+    ///
+    /// // or Variant::TimestampNtzNanos
+    /// let datetime_ntz = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_nano_opt(12, 34, 56, 789123456).unwrap();
+    /// let v2 = Variant::from(datetime_ntz);
+    /// assert_eq!(v1.as_timestamp_nanos(), Some(1759494896789123456));
+    ///
+    /// // but not from other variants
+    /// let datetime_micros = NaiveDate::from_ymd_opt(2025, 10, 03).unwrap().and_hms_micro_opt(12, 34, 56, 789).unwrap().and_utc();
+    /// let v3 = Variant::from(datetime_micros);
+    /// assert_eq!(v3.as_timestamp_nanos(), None);
+    /// ```
+    pub fn as_timestamp_nanos(&self) -> Option<i64> {
+        match *self {
+            Variant::TimestampNanos(d) => d.timestamp_nanos_opt(),
+            Variant::TimestampNtzNanos(d) => d.and_utc().timestamp_nanos_opt(),
+            _ => None,
+        }
+    }
+
     /// Converts this variant to a `NaiveDateTime` if possible.
     ///
     /// Returns `Some(NaiveDateTime)` for timestamp variants,