feat: add NaiveDate and TimeDelta from chrono (#230)

mbrobbel · Sep 2, 2024 · 5972c2b · 5972c2b
1 parent ceeb3ad
commit 5972c2b
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 3 deletions.
diff --git a/examples/parquet.rs b/examples/parquet.rs
@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use chrono::{DateTime, NaiveTime, Utc};
+use chrono::{DateTime, NaiveDate, NaiveTime, TimeDelta, Utc};
 
 #[rustversion::attr(nightly, allow(non_local_definitions))]
 fn main() {
@@ -32,6 +32,8 @@ fn main() {
         j: DateTime<Utc>,
         k: NaiveTime,
         l: Option<HashMap<String, Vec<u8>>>,
+        m: NaiveDate,
+        n: TimeDelta,
     }
     let input = [
         Foo {
@@ -50,6 +52,8 @@ fn main() {
                 "a".to_string(),
                 vec![1, 2, 3, 4, 42],
             )])),
+            m: NaiveDate::MAX,
+            n: TimeDelta::max_value(),
         },
         Foo {
             a: 42,
@@ -64,6 +68,8 @@ fn main() {
             j: Utc::now(),
             k: Utc::now().time(),
             l: None,
+            m: NaiveDate::MIN,
+            n: TimeDelta::min_value(),
         },
     ];
 

diff --git a/src/logical/chrono.rs b/src/logical/chrono.rs
@@ -1,4 +1,4 @@
-use chrono::{DateTime, NaiveDateTime, NaiveTime, Timelike, Utc};
+use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, TimeDelta, Timelike, Utc};
 
 use crate::{
     array::{ArrayType, UnionType},
@@ -60,6 +60,32 @@ impl LogicalArrayType<NaiveDateTime> for NaiveDateTime {
 pub type NaiveDateTimeArray<const NULLABLE: bool = false, Buffer = crate::buffer::VecBuffer> =
     LogicalArray<NaiveDateTime, NULLABLE, Buffer, crate::offset::NA, crate::array::union::NA>;
 
+impl ArrayType<NaiveDate> for NaiveDate {
+    type Array<Buffer: BufferType, OffsetItem: OffsetElement, UnionLayout: UnionType> =
+        LogicalArray<Self, false, Buffer, OffsetItem, UnionLayout>;
+}
+
+impl ArrayType<NaiveDate> for Option<NaiveDate> {
+    type Array<Buffer: BufferType, OffsetItem: OffsetElement, UnionLayout: UnionType> =
+        LogicalArray<NaiveDate, true, Buffer, OffsetItem, UnionLayout>;
+}
+
+impl LogicalArrayType<NaiveDate> for NaiveDate {
+    type ArrayType = i32;
+
+    fn from_array_type(item: Self::ArrayType) -> Self {
+        NaiveDate::from_num_days_from_ce_opt(item).expect("out of range")
+    }
+
+    fn into_array_type(self) -> Self::ArrayType {
+        self.num_days_from_ce()
+    }
+}
+
+/// An array for [`NaiveDate`] items.
+pub type NaiveDateArray<const NULLABLE: bool = false, Buffer = crate::buffer::VecBuffer> =
+    LogicalArray<NaiveDate, NULLABLE, Buffer, crate::offset::NA, crate::array::union::NA>;
+
 impl ArrayType<NaiveTime> for NaiveTime {
     type Array<Buffer: BufferType, OffsetItem: OffsetElement, UnionLayout: UnionType> =
         LogicalArray<Self, false, Buffer, OffsetItem, UnionLayout>;
@@ -94,13 +120,53 @@ impl LogicalArrayType<NaiveTime> for NaiveTime {
 pub type NaiveTimeArray<const NULLABLE: bool = false, Buffer = crate::buffer::VecBuffer> =
     LogicalArray<NaiveTime, NULLABLE, Buffer, crate::offset::NA, crate::array::union::NA>;
 
+impl ArrayType<TimeDelta> for TimeDelta {
+    type Array<Buffer: BufferType, OffsetItem: OffsetElement, UnionLayout: UnionType> =
+        LogicalArray<Self, false, Buffer, OffsetItem, UnionLayout>;
+}
+
+impl ArrayType<TimeDelta> for Option<TimeDelta> {
+    type Array<Buffer: BufferType, OffsetItem: OffsetElement, UnionLayout: UnionType> =
+        LogicalArray<TimeDelta, true, Buffer, OffsetItem, UnionLayout>;
+}
+
+impl LogicalArrayType<TimeDelta> for TimeDelta {
+    type ArrayType = i64;
+
+    fn from_array_type(item: Self::ArrayType) -> Self {
+        Self::nanoseconds(item)
+    }
+
+    fn into_array_type(self) -> Self::ArrayType {
+        self.num_nanoseconds().expect("out of range")
+    }
+}
+
+/// An array for [`TimeDelta`] items.
+pub type TimeDeltaArray<const NULLABLE: bool = false, Buffer = crate::buffer::VecBuffer> =
+    LogicalArray<TimeDelta, NULLABLE, Buffer, crate::offset::NA, crate::array::union::NA>;
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::Length;
 
     #[test]
-    fn round_trip() {
+    fn round_trip_naivedate() {
+        for value in [
+            NaiveDate::from_yo_opt(2024, 7)
+                .expect("out of range")
+                .num_days_from_ce(),
+            NaiveDate::from_yo_opt(2020, 6)
+                .expect("out of range")
+                .num_days_from_ce(),
+        ] {
+            assert_eq!(NaiveDate::from_array_type(value).into_array_type(), value);
+        }
+    }
+
+    #[test]
+    fn round_trip_naivetime() {
         for value in [
             0,
             1234,
@@ -111,6 +177,18 @@ mod tests {
         }
     }
 
+    #[test]
+    fn round_trip_timedelta() {
+        for value in [
+            0,
+            1234,
+            1234 * NANO_SECONDS,
+            86_398 * NANO_SECONDS + 1_999_999_999,
+        ] {
+            assert_eq!(TimeDelta::nanoseconds(value).into_array_type(), value);
+        }
+    }
+
     #[test]
     fn from_iter() {
         let array = [DateTime::<Utc>::UNIX_EPOCH, DateTime::<Utc>::UNIX_EPOCH]