fix[rust]: serde timeunit of datetime type (#4598)

pola-rs · Aug 28, 2022 · bd91e1d · bd91e1d
1 parent ae218d4
commit bd91e1d
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/polars/polars-core/src/serde/mod.rs b/polars/polars-core/src/serde/mod.rs
@@ -38,7 +38,7 @@ enum DeDataType<'a> {
     Float64,
     Utf8,
     Date,
-    Datetime,
+    Datetime(TimeUnit, Option<TimeZone>),
     #[serde(with = "TimeUnitDef")]
     Time64(ArrowTimeUnit),
     List,
@@ -55,7 +55,7 @@ impl From<&DataType> for DeDataType<'_> {
             DataType::Int64 => DeDataType::Int64,
             DataType::UInt64 => DeDataType::UInt64,
             DataType::Date => DeDataType::Date,
-            DataType::Datetime(_, _) => DeDataType::Datetime,
+            DataType::Datetime(tu, tz) => DeDataType::Datetime(*tu, tz.clone()),
             DataType::Float32 => DeDataType::Float32,
             DataType::Float64 => DeDataType::Float64,
             DataType::Utf8 => DeDataType::Utf8,

diff --git a/polars/polars-core/src/serde/series.rs b/polars/polars-core/src/serde/series.rs
@@ -42,10 +42,7 @@ impl Serialize for Series {
                 }
                 #[cfg(feature = "dtype-datetime")]
                 DataType::Datetime(_, _) => {
-                    let s = self
-                        .cast(&DataType::Datetime(TimeUnit::Microseconds, None))
-                        .unwrap();
-                    let ca = s.datetime().unwrap();
+                    let ca = self.datetime().unwrap();
                     ca.serialize(serializer)
                 }
                 #[cfg(feature = "dtype-categorical")]
@@ -151,10 +148,10 @@ impl<'de> Deserialize<'de> for Series {
                         Ok(Series::new(&name, values).cast(&DataType::Date).unwrap())
                     }
                     #[cfg(feature = "dtype-datetime")]
-                    DeDataType::Datetime => {
+                    DeDataType::Datetime(tu, tz) => {
                         let values: Vec<Option<i64>> = map.next_value()?;
                         Ok(Series::new(&name, values)
-                            .cast(&DataType::Datetime(TimeUnit::Microseconds, None))
+                            .cast(&DataType::Datetime(tu, tz))
                             .unwrap())
                     }
                     DeDataType::Boolean => {

diff --git a/py-polars/tests/test_serde.py b/py-polars/tests/test_serde.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pickle
+from datetime import datetime, timedelta
 
 import polars as pl
 
@@ -21,3 +22,13 @@ def serde_lazy_frame_lp() -> None:
         .to_series()
         .series_equal(pl.Series("a", [1, 2, 3]))
     )
+
+
+def test_serde_time_unit() -> None:
+    assert pickle.loads(
+        pickle.dumps(
+            pl.Series(
+                [datetime(2022, 1, 1) + timedelta(days=1) for _ in range(3)]
+            ).cast(pl.Datetime("ns"))
+        )
+    ).dtype == pl.Datetime("ns")