fix: Do not panic if time is invalid (#12466)

pola-rs · Nov 15, 2023 · 38b5322 · 38b5322
1 parent 5017d14
commit 38b5322
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 12 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/range/date_range.rs b/crates/polars-plan/src/dsl/function_expr/range/date_range.rs
@@ -48,8 +48,12 @@ fn date_range(s: &[Series], interval: Duration, closed: ClosedWindow) -> PolarsR
     ensure_range_bounds_contain_exactly_one_value(start, end)?;
 
     let dtype = DataType::Date;
-    let start = temporal_series_to_i64_scalar(start) * MILLISECONDS_IN_DAY;
-    let end = temporal_series_to_i64_scalar(end) * MILLISECONDS_IN_DAY;
+    let start = temporal_series_to_i64_scalar(start)
+        .ok_or_else(|| polars_err!(ComputeError: "start is an out-of-range time."))?
+        * MILLISECONDS_IN_DAY;
+    let end = temporal_series_to_i64_scalar(end)
+        .ok_or_else(|| polars_err!(ComputeError: "end is an out-of-range time."))?
+        * MILLISECONDS_IN_DAY;
 
     let result = datetime_range_impl(
         "date",

diff --git a/crates/polars-plan/src/dsl/function_expr/range/datetime_range.rs b/crates/polars-plan/src/dsl/function_expr/range/datetime_range.rs
@@ -67,8 +67,10 @@ pub(super) fn datetime_range(
         _ => {},
     };
 
-    let start = temporal_series_to_i64_scalar(&start);
-    let end = temporal_series_to_i64_scalar(&end);
+    let start = temporal_series_to_i64_scalar(&start)
+        .ok_or_else(|| polars_err!(ComputeError: "start is an out-of-range time."))?;
+    let end = temporal_series_to_i64_scalar(&end)
+        .ok_or_else(|| polars_err!(ComputeError: "end is an out-of-range time."))?;
 
     let result = match dtype {
         DataType::Datetime(tu, ref tz) => {

diff --git a/crates/polars-plan/src/dsl/function_expr/range/time_range.rs b/crates/polars-plan/src/dsl/function_expr/range/time_range.rs
@@ -17,8 +17,10 @@ pub(super) fn time_range(
     ensure_range_bounds_contain_exactly_one_value(start, end)?;
 
     let dtype = DataType::Time;
-    let start = temporal_series_to_i64_scalar(&start.cast(&dtype)?);
-    let end = temporal_series_to_i64_scalar(&end.cast(&dtype)?);
+    let start = temporal_series_to_i64_scalar(&start.cast(&dtype)?)
+        .ok_or_else(|| polars_err!(ComputeError: "start is an out-of-range time."))?;
+    let end = temporal_series_to_i64_scalar(&end.cast(&dtype)?)
+        .ok_or_else(|| polars_err!(ComputeError: "end is an out-of-range time."))?;
 
     let out = time_range_impl("time", start, end, interval, closed)?;
     Ok(out.cast(&dtype).unwrap().into_series())

diff --git a/crates/polars-plan/src/dsl/function_expr/range/utils.rs b/crates/polars-plan/src/dsl/function_expr/range/utils.rs
@@ -1,12 +1,8 @@
 use polars_core::prelude::{polars_bail, polars_ensure, PolarsResult};
 use polars_core::series::Series;
 
-pub(super) fn temporal_series_to_i64_scalar(s: &Series) -> i64 {
-    s.to_physical_repr()
-        .get(0)
-        .unwrap()
-        .extract::<i64>()
-        .unwrap()
+pub(super) fn temporal_series_to_i64_scalar(s: &Series) -> Option<i64> {
+    s.to_physical_repr().get(0).unwrap().extract::<i64>()
 }
 
 pub(super) fn ensure_range_bounds_contain_exactly_one_value(

diff --git a/py-polars/tests/unit/functions/range/test_date_range.py b/py-polars/tests/unit/functions/range/test_date_range.py
@@ -29,6 +29,11 @@ def test_date_range_invalid_time_unit() -> None:
         )
 
 
+def test_date_range_invalid_time() -> None:
+    with pytest.raises(pl.ComputeError, match="end is an out-of-range time"):
+        pl.date_range(pl.date(2024, 1, 1), pl.date(2024, 2, 30), eager=True)
+
+
 def test_date_range_lazy_with_literals() -> None:
     df = pl.DataFrame({"misc": ["x"]}).with_columns(
         pl.date_ranges(