fix(rust, python): correct output dtype for cummin/cumsum/cummax (#6062)

pola-rs · Jan 5, 2023 · a609f25 · a609f25
1 parent 5bf8e31
commit a609f25
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 16 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/mod.rs
@@ -880,14 +880,18 @@ impl Expr {
             move |s: Series| Ok(s.cumsum(reverse)),
             GetOutput::map_dtype(|dt| {
                 use DataType::*;
-                match dt {
-                    Boolean => UInt32,
-                    Int32 => Int32,
-                    UInt32 => UInt32,
-                    UInt64 => UInt64,
-                    Float32 => Float32,
-                    Float64 => Float64,
-                    _ => Int64,
+                if dt.is_logical() {
+                    dt.clone()
+                } else {
+                    match dt {
+                        Boolean => UInt32,
+                        Int32 => Int32,
+                        UInt32 => UInt32,
+                        UInt64 => UInt64,
+                        Float32 => Float32,
+                        Float64 => Float64,
+                        _ => Int64,
+                    }
                 }
             }),
         )
@@ -928,14 +932,7 @@ impl Expr {
     pub fn cummax(self, reverse: bool) -> Self {
         self.apply(
             move |s: Series| Ok(s.cummax(reverse)),
-            GetOutput::map_dtype(|dt| {
-                use DataType::*;
-                match dt {
-                    Float32 => Float32,
-                    Float64 => Float64,
-                    _ => Int64,
-                }
-            }),
+            GetOutput::same_type(),
         )
         .with_fmt("cummax")
     }

diff --git a/py-polars/tests/unit/test_datelike.py b/py-polars/tests/unit/test_datelike.py
@@ -2337,3 +2337,58 @@ def test_tz_aware_day_weekday() -> None:
         "tyo_weekday": [1, 4, 7],
         "ny_weekday": [7, 3, 6],
     }
+
+
+def test_datetime_cum_agg_schema() -> None:
+    df = pl.DataFrame(
+        {
+            "timestamp": [
+                datetime(2023, 1, 1),
+                datetime(2023, 1, 2),
+                datetime(2023, 1, 3),
+            ]
+        }
+    )
+    # Exactly the same as above but with lazy() and collect() later
+    assert (
+        df.lazy()
+        .with_columns(
+            [
+                (pl.col("timestamp").cummin()).alias("cummin"),
+                (pl.col("timestamp").cummax()).alias("cummax"),
+            ]
+        )
+        .with_columns(
+            [
+                (pl.col("cummin") + pl.duration(hours=24)).alias("cummin+24"),
+                (pl.col("cummax") + pl.duration(hours=24)).alias("cummax+24"),
+            ]
+        )
+        .collect()
+    ).to_dict(False) == {
+        "timestamp": [
+            datetime(2023, 1, 1, 0, 0),
+            datetime(2023, 1, 2, 0, 0),
+            datetime(2023, 1, 3, 0, 0),
+        ],
+        "cummin": [
+            datetime(2023, 1, 1, 0, 0),
+            datetime(2023, 1, 1, 0, 0),
+            datetime(2023, 1, 1, 0, 0),
+        ],
+        "cummax": [
+            datetime(2023, 1, 1, 0, 0),
+            datetime(2023, 1, 2, 0, 0),
+            datetime(2023, 1, 3, 0, 0),
+        ],
+        "cummin+24": [
+            datetime(2023, 1, 2, 0, 0),
+            datetime(2023, 1, 2, 0, 0),
+            datetime(2023, 1, 2, 0, 0),
+        ],
+        "cummax+24": [
+            datetime(2023, 1, 2, 0, 0),
+            datetime(2023, 1, 3, 0, 0),
+            datetime(2023, 1, 4, 0, 0),
+        ],
+    }