fix date_range w/ nominal (i.e. monthly) interval

pola-rs · Jan 11, 2022 · efda57b · efda57b
1 parent ec1873c
commit efda57b
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 8 deletions.
diff --git a/polars/polars-time/src/calendar.rs b/polars/polars-time/src/calendar.rs
@@ -48,34 +48,34 @@ pub fn date_range(
 
     let mut t = start;
     let f = match tu {
-        TimeUnit::Nanoseconds => <Duration>::duration_ns,
-        TimeUnit::Milliseconds => <Duration>::duration_ms,
+        TimeUnit::Nanoseconds => <Duration>::add_ns,
+        TimeUnit::Milliseconds => <Duration>::add_ms,
     };
     match closed {
         ClosedWindow::Both => {
             while t <= stop {
                 ts.push(t);
-                t += f(&every)
+                t = f(&every, t)
             }
         }
         ClosedWindow::Left => {
             while t < stop {
                 ts.push(t);
-                t += f(&every)
+                t = f(&every, t)
             }
         }
         ClosedWindow::Right => {
-            t += f(&every);
+            t = f(&every, t);
             while t <= stop {
                 ts.push(t);
-                t += f(&every)
+                t = f(&every, t)
             }
         }
         ClosedWindow::None => {
-            t += f(&every);
+            t = f(&every, t);
             while t < stop {
                 ts.push(t);
-                t += f(&every)
+                t = f(&every, t)
             }
         }
     }

diff --git a/polars/polars-time/src/test.rs b/polars/polars-time/src/test.rs
@@ -5,6 +5,30 @@ use crate::window::Window;
 use chrono::prelude::*;
 use polars_arrow::export::arrow::temporal_conversions::timestamp_ns_to_datetime;
 
+#[test]
+fn test_date_range() {
+    // Test month as interval in date range
+    let start = NaiveDate::from_ymd(2022, 1, 1).and_hms(0, 0, 0);
+    let end = NaiveDate::from_ymd(2022, 4, 1).and_hms(0, 0, 0);
+    let dates = date_range(
+        start.timestamp_nanos(),
+        end.timestamp_nanos(),
+        Duration::parse("1mo"),
+        ClosedWindow::Both,
+        TimeUnit::Nanoseconds,
+    );
+    let expected = [
+        NaiveDate::from_ymd(2022, 1, 1),
+        NaiveDate::from_ymd(2022, 2, 1),
+        NaiveDate::from_ymd(2022, 3, 1),
+        NaiveDate::from_ymd(2022, 4, 1),
+    ]
+    .iter()
+    .map(|d| d.and_hms(0, 0, 0).timestamp_nanos())
+    .collect::<Vec<_>>();
+    assert_eq!(dates, expected);
+}
+
 fn print_ns(ts: &[i64]) {
     for ts in ts {
         println!("{}", timestamp_ns_to_datetime(*ts));