fix duration computation (#3790)

pola-rs · Jun 24, 2022 · 128550b · 128550b
1 parent cb81c86
commit 128550b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 6 deletions.
diff --git a/polars/polars-lazy/src/dsl/functions.rs b/polars/polars-lazy/src/dsl/functions.rs
@@ -475,22 +475,22 @@ pub fn duration(args: DurationArgs) -> Expr {
             nanoseconds = nanoseconds.expand_at_index(0, max_len);
         }
         if condition(&days) {
-            nanoseconds = (&nanoseconds + &days) * NANOSECONDS * SECONDS_IN_DAY;
+            nanoseconds = nanoseconds + days * NANOSECONDS * SECONDS_IN_DAY;
         }
         if condition(&seconds) {
-            nanoseconds = (&nanoseconds + &seconds) * NANOSECONDS;
+            nanoseconds = nanoseconds + &seconds * NANOSECONDS;
         }
         if condition(&milliseconds) {
-            nanoseconds = (&nanoseconds + &milliseconds) * 1_000_000;
+            nanoseconds = nanoseconds + milliseconds * 1_000_000;
         }
         if condition(&minutes) {
-            nanoseconds = (&nanoseconds + &minutes) * NANOSECONDS * 60;
+            nanoseconds = nanoseconds + minutes * NANOSECONDS * 60;
         }
         if condition(&hours) {
-            nanoseconds = (&nanoseconds + &hours) * NANOSECONDS * 60 * 60;
+            nanoseconds = nanoseconds + hours * NANOSECONDS * 60 * 60;
         }
         if condition(&weeks) {
-            nanoseconds = (&nanoseconds + &weeks) * NANOSECONDS * SECONDS_IN_DAY * 7;
+            nanoseconds = nanoseconds + weeks * NANOSECONDS * SECONDS_IN_DAY * 7;
         }
 
         nanoseconds.cast(&DataType::Duration(TimeUnit::Nanoseconds))

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -1053,3 +1053,31 @@ def test_groupby_rolling_by_ordering() -> None:
         ],
         "sum val": [2, 2, 2, 2, 1, 1, 1],
     }
+
+
+def test_add_duration_3786() -> None:
+    df = pl.DataFrame(
+        {
+            "datetime": [datetime(2022, 1, 1), datetime(2022, 1, 2)],
+            "add": [1, 2],
+        }
+    )
+    assert df.slice(0, 1).with_columns(
+        [
+            (pl.col("datetime") + pl.duration(weeks="add")).alias("add_weeks"),
+            (pl.col("datetime") + pl.duration(days="add")).alias("add_days"),
+            (pl.col("datetime") + pl.duration(seconds="add")).alias("add_seconds"),
+            (pl.col("datetime") + pl.duration(milliseconds="add")).alias(
+                "add_milliseconds"
+            ),
+            (pl.col("datetime") + pl.duration(hours="add")).alias("add_hours"),
+        ]
+    ).to_dict(False) == {
+        "datetime": [datetime(2022, 1, 1, 0, 0)],
+        "add": [1],
+        "add_weeks": [datetime(2022, 1, 8, 0, 0)],
+        "add_days": [datetime(2022, 1, 2, 0, 0)],
+        "add_seconds": [datetime(2022, 1, 1, 0, 0, 1)],
+        "add_milliseconds": [datetime(2022, 1, 1, 0, 0, 0, 1000)],
+        "add_hours": [datetime(2022, 1, 1, 1, 0)],
+    }