always include offset in groupby_dynamic (#3779)

pola-rs · Jun 23, 2022 · bdb1dd0 · bdb1dd0
1 parent cd710c4
commit bdb1dd0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/polars/polars-time/src/windows/test.rs b/polars/polars-time/src/windows/test.rs
@@ -550,3 +550,30 @@ fn test_groupby_windows_duplicates_2931() {
     );
     assert_eq!(groups, [[0, 1], [1, 2], [3, 2]]);
 }
+
+#[test]
+fn test_groupby_windows_offsets_3776() {
+    let dates = &[
+        NaiveDate::from_ymd(2020, 12, 1),
+        NaiveDate::from_ymd(2021, 2, 1),
+        NaiveDate::from_ymd(2021, 5, 1),
+    ];
+    let ts = dates
+        .iter()
+        .map(|d| d.and_hms(0, 0, 0).timestamp_nanos())
+        .collect::<Vec<_>>();
+
+    let window = Window::new(
+        Duration::parse("2d"),
+        Duration::parse("2d"),
+        Duration::parse("-2d"),
+    );
+    let (groups, _, _) = groupby_windows(
+        window,
+        &ts,
+        false,
+        ClosedWindow::Right,
+        TimeUnit::Nanoseconds,
+    );
+    assert_eq!(groups, [[0, 1], [1, 1], [2, 1]]);
+}
diff --git a/polars/polars-time/src/windows/window.rs b/polars/polars-time/src/windows/window.rs
@@ -71,10 +71,9 @@ impl Window {
         let start = if !self.every.months_only()
             && self.every.duration_ns() > NANOSECONDS * SECONDS_IN_DAY
         {
-            t
+            self.offset.add_ns(t)
         } else {
-            // original code translates offset here
-            // we don't. Seems unintuitive to me.
+            // offset is translated in the truncate
             self.truncate_ns(t)
         };
 
@@ -87,7 +86,7 @@ impl Window {
         let start = if !self.every.months_only()
             && self.every.duration_us() > MICROSECONDS * SECONDS_IN_DAY
         {
-            t
+            self.offset.add_us(t)
         } else {
             self.truncate_us(t)
         };
@@ -100,7 +99,7 @@ impl Window {
         let start = if !self.every.months_only()
             && self.every.duration_ms() > MILLISECONDS * SECONDS_IN_DAY
         {
-            t
+            self.offset.add_ms(t)
         } else {
             self.truncate_ms(t)
         };