fix explode for datetime

pola-rs · Jan 15, 2022 · 42d09c1 · 42d09c1
1 parent 2d48db4
commit 42d09c1
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 8 deletions.
diff --git a/polars/polars-core/src/series/mod.rs b/polars/polars-core/src/series/mod.rs
@@ -746,20 +746,35 @@ impl Series {
             Ok(s)
         }
     }
+
     #[cfg(feature = "dtype-time")]
     pub(crate) fn into_time(self) -> Series {
-        self.i64()
-            .expect("impl error")
-            .clone()
-            .into_time()
-            .into_series()
+        match self.dtype() {
+            DataType::Int64 => self.i64().unwrap().clone().into_time().into_series(),
+            DataType::Time => self
+                .time()
+                .unwrap()
+                .as_ref()
+                .clone()
+                .into_time()
+                .into_series(),
+            dt => panic!("date not implemented for {:?}", dt),
+        }
     }
 
     pub(crate) fn into_date(self) -> Series {
         match self.dtype() {
             #[cfg(feature = "dtype-date")]
             DataType::Int32 => self.i32().unwrap().clone().into_date().into_series(),
-            _ => unimplemented!(),
+            #[cfg(feature = "dtype-date")]
+            DataType::Date => self
+                .date()
+                .unwrap()
+                .as_ref()
+                .clone()
+                .into_date()
+                .into_series(),
+            dt => panic!("date not implemented for {:?}", dt),
         }
     }
     pub(crate) fn into_datetime(self, timeunit: TimeUnit, tz: Option<TimeZone>) -> Series {
@@ -771,7 +786,15 @@ impl Series {
                 .clone()
                 .into_datetime(timeunit, tz)
                 .into_series(),
-            _ => unimplemented!(),
+            #[cfg(feature = "dtype-datetime")]
+            DataType::Datetime(_, _) => self
+                .datetime()
+                .unwrap()
+                .as_ref()
+                .clone()
+                .into_datetime(timeunit, tz)
+                .into_series(),
+            dt => panic!("into_datetime not implemented for {:?}", dt),
         }
     }
 
@@ -784,7 +807,15 @@ impl Series {
                 .clone()
                 .into_duration(timeunit)
                 .into_series(),
-            _ => unimplemented!(),
+            #[cfg(feature = "dtype-duration")]
+            DataType::Duration(_) => self
+                .duration()
+                .unwrap()
+                .as_ref()
+                .clone()
+                .into_duration(timeunit)
+                .into_series(),
+            dt => panic!("into_duration not implemented for {:?}", dt),
         }
     }
 

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -378,3 +378,32 @@ def test_non_exact_strptime() -> None:
 
     with pytest.raises(Exception):
         a.str.strptime(pl.Date, fmt, strict=True, exact=True)
+
+
+def test_explode_date() -> None:
+    datetimes = [
+        datetime(2021, 12, 1, 0, 0),
+        datetime(2021, 12, 1, 0, 0),
+        datetime(2021, 12, 1, 0, 0),
+        datetime(2021, 12, 1, 0, 0),
+    ]
+    dates = [
+        date(2021, 12, 1),
+        date(2021, 12, 1),
+        date(2021, 12, 1),
+        date(2021, 12, 1),
+    ]
+    for d in [dates, datetimes]:
+        df = pl.DataFrame(
+            {
+                "a": d,
+                "b": ["a", "b", "a", "b"],
+                "c": [1.0, 2.0, 1.1, 2.2],
+            }
+        )
+        out = (
+            df.groupby("b")
+            .agg([pl.col("a"), pl.col("c").pct_change()])
+            .explode(["a", "c"])
+        )
+        assert out.shape == (4, 3)