pivot index maintain logical type (#3367)

pola-rs · May 11, 2022 · 272368b · 272368b
1 parent c075b5a
commit 272368b
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 6 deletions.
diff --git a/polars/polars-core/src/frame/groupby/pivot.rs b/polars/polars-core/src/frame/groupby/pivot.rs
@@ -140,10 +140,15 @@ impl DataFrame {
                 .collect::<Vec<_>>();
 
             let row_index = match count {
-                0 => Some(vec![Series::new(
-                    &index[0],
-                    row_to_idx.into_iter().map(|(k, _)| k).collect::<Vec<_>>(),
-                )]),
+                0 => {
+                    let s = Series::new(
+                        &index[0],
+                        row_to_idx.into_iter().map(|(k, _)| k).collect::<Vec<_>>(),
+                    );
+                    // restore logical type
+                    let s = s.cast(index_s.dtype()).unwrap();
+                    Some(vec![s])
+                }
                 _ => None,
             };
 
@@ -186,7 +191,7 @@ impl DataFrame {
                         .iter()
                         .enumerate()
                         .map(|(i, name)| {
-                            Series::new(
+                            let s = Series::new(
                                 name,
                                 row_to_idx
                                     .iter()
@@ -195,7 +200,9 @@ impl DataFrame {
                                         unsafe { k.get_unchecked(i).clone() }
                                     })
                                     .collect::<Vec<_>>(),
-                            )
+                            );
+                            // restore logical type
+                            s.cast(index_s[i].dtype()).unwrap()
                         })
                         .collect::<Vec<_>>(),
                 ),

diff --git a/polars/tests/it/core/pivot.rs b/polars/tests/it/core/pivot.rs
@@ -1,3 +1,4 @@
+use polars::export::chrono::{NaiveDate, NaiveDateTime};
 use polars::prelude::*;
 
 #[test]
@@ -167,3 +168,24 @@ fn test_pivot_2() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(feature = "dtype-datetime")]
+fn test_pivot_datetime() -> Result<()> {
+    let dt = NaiveDate::from_ymd(2021, 1, 1).and_hms(12, 15, 0);
+    let df = df![
+        "dt" => [dt, dt, dt, dt],
+        "key" => ["x", "x", "y", "y"],
+        "val" => [100, 50, 500, -80]
+    ]?;
+
+    let out = df.pivot(["val"], ["dt"], ["key"], PivotAgg::Sum, false)?;
+    let expected = df![
+        "dt" => [dt],
+        "x" => [150],
+        "y" => [420]
+    ]?;
+    assert!(out.frame_equal(&expected));
+
+    Ok(())
+}