From 40ce3d24d7011903eaec22c14243a93dc66082fc Mon Sep 17 00:00:00 2001 From: Weijie Guo Date: Wed, 15 Nov 2023 14:59:20 +0800 Subject: [PATCH 1/2] fix: Do not panic if time is invalid --- .../src/dsl/function_expr/range/date_range.rs | 8 ++++++-- .../src/dsl/function_expr/range/datetime_range.rs | 6 ++++-- .../src/dsl/function_expr/range/time_range.rs | 6 ++++-- .../polars-plan/src/dsl/function_expr/range/utils.rs | 10 +++------- .../tests/unit/functions/range/test_date_range.py | 5 +++++ 5 files changed, 22 insertions(+), 13 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 index 84406821f493..15338e03a9c6 100644 --- 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 index 874656cd55fc..d0d64d7304b5 100644 --- 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 index 3dd3e44fb904..d543191e3885 100644 --- 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 index 1ff8da6b2ba3..49d942550182 100644 --- 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::prelude::{polars_bail, polars_ensure, polars_err, 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::() - .unwrap() +pub(super) fn temporal_series_to_i64_scalar(s: &Series) -> Option { + s.to_physical_repr().get(0).unwrap().extract::() } 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 index 16a746aac4fe..b5fc889ab008 100644 --- 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( From 59943f0e89334bf0f44de39b0328457e730b2846 Mon Sep 17 00:00:00 2001 From: Weijie Guo Date: Wed, 15 Nov 2023 15:04:30 +0800 Subject: [PATCH 2/2] clippy --- crates/polars-plan/src/dsl/function_expr/range/utils.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/polars-plan/src/dsl/function_expr/range/utils.rs b/crates/polars-plan/src/dsl/function_expr/range/utils.rs index 49d942550182..c395f2b3c768 100644 --- a/crates/polars-plan/src/dsl/function_expr/range/utils.rs +++ b/crates/polars-plan/src/dsl/function_expr/range/utils.rs @@ -1,4 +1,4 @@ -use polars_core::prelude::{polars_bail, polars_ensure, polars_err, PolarsResult}; +use polars_core::prelude::{polars_bail, polars_ensure, PolarsResult}; use polars_core::series::Series; pub(super) fn temporal_series_to_i64_scalar(s: &Series) -> Option {