fix: Raise proper error message when too small interval is passed to …

…datetime_range (#19955)
pola-rs · Nov 24, 2024 · 5516536 · 5516536
1 parent 8d9ed64
commit 5516536
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 8 deletions.
diff --git a/crates/polars-time/src/date_range.rs b/crates/polars-time/src/date_range.rs
@@ -100,8 +100,8 @@ pub(crate) fn datetime_range_i64(
     end: i64,
     interval: Duration,
     closed: ClosedWindow,
-    tu: TimeUnit,
-    tz: Option<&Tz>,
+    time_unit: TimeUnit,
+    time_zone: Option<&Tz>,
 ) -> PolarsResult<Vec<i64>> {
     if start > end {
         return Ok(Vec::new());
@@ -111,12 +111,12 @@ pub(crate) fn datetime_range_i64(
         ComputeError: "`interval` must be positive"
     );
 
-    let duration = match tu {
+    let duration = match time_unit {
         TimeUnit::Nanoseconds => interval.duration_ns(),
         TimeUnit::Microseconds => interval.duration_us(),
         TimeUnit::Milliseconds => interval.duration_ms(),
     };
-    let time_zone_opt_string: Option<String> = match tz {
+    let time_zone_opt_string: Option<String> = match time_zone {
         #[cfg(feature = "timezones")]
         Some(tz) => Some(tz.to_string()),
         _ => None,
@@ -126,6 +126,12 @@ pub(crate) fn datetime_range_i64(
         let step: usize = duration.try_into().map_err(
             |_err| polars_err!(ComputeError: "Could not convert {:?} to usize", duration),
         )?;
+        polars_ensure!(
+            step != 0,
+            InvalidOperation: "interval {} is too small for time unit {} and got rounded down to zero",
+            interval,
+            time_unit,
+        );
         return match closed {
             ClosedWindow::Both => Ok((start..=end).step_by(step).collect::<Vec<i64>>()),
             ClosedWindow::None => Ok((start + duration..end).step_by(step).collect::<Vec<i64>>()),
@@ -135,7 +141,7 @@ pub(crate) fn datetime_range_i64(
     }
 
     let size = ((end - start) / duration + 1) as usize;
-    let offset_fn = match tu {
+    let offset_fn = match time_unit {
         TimeUnit::Nanoseconds => Duration::add_ns,
         TimeUnit::Microseconds => Duration::add_us,
         TimeUnit::Milliseconds => Duration::add_ms,
@@ -145,20 +151,20 @@ pub(crate) fn datetime_range_i64(
         ClosedWindow::Both | ClosedWindow::Left => 0,
         ClosedWindow::Right | ClosedWindow::None => 1,
     };
-    let mut t = offset_fn(&(interval * i), start, tz)?;
+    let mut t = offset_fn(&(interval * i), start, time_zone)?;
     i += 1;
     match closed {
         ClosedWindow::Both | ClosedWindow::Right => {
             while t <= end {
                 ts.push(t);
-                t = offset_fn(&(interval * i), start, tz)?;
+                t = offset_fn(&(interval * i), start, time_zone)?;
                 i += 1;
             }
         },
         ClosedWindow::Left | ClosedWindow::None => {
             while t < end {
                 ts.push(t);
-                t = offset_fn(&(interval * i), start, tz)?;
+                t = offset_fn(&(interval * i), start, time_zone)?;
                 i += 1;
             }
         },

diff --git a/py-polars/tests/unit/functions/range/test_datetime_range.py b/py-polars/tests/unit/functions/range/test_datetime_range.py
@@ -620,3 +620,17 @@ def test_datetime_range_fast_slow_paths(
         eager=True,
     )
     assert_series_equal(result_slow, result_fast)
+
+
+def test_dt_range_with_nanosecond_interval_19931() -> None:
+    with pytest.raises(
+        InvalidOperationError, match="interval 1ns is too small for time unit ms"
+    ):
+        pl.datetime_range(
+            pl.date(2022, 1, 1),
+            pl.date(2022, 1, 1),
+            time_zone="Asia/Kathmandu",
+            interval="1ns",
+            time_unit="ms",
+            eager=True,
+        )