Return ComputeError intead of panic

pola-rs · Feb 1, 2025 · d0803dc · d0803dc
1 parent 3812022
commit d0803dc
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 24 deletions.
diff --git a/crates/polars-time/src/chunkedarray/date.rs b/crates/polars-time/src/chunkedarray/date.rs
@@ -84,21 +84,32 @@ pub trait DateMethods: AsDate {
         day: &Int8Chunked,
         name: PlSmallStr,
     ) -> PolarsResult<DateChunked> {
+        let mut error_values: Option<(i32, i8, i8)> = None;
         let mut ca: Int32Chunked = year
             .into_iter()
             .zip(month)
             .zip(day)
             .map(|((y, m), d)| {
                 if let (Some(y), Some(m), Some(d)) = (y, m, d) {
-                    let Some(ns) = NaiveDate::from_ymd_opt(y, m as u32, d as u32) else {
-                        panic!("Invalid date components ({}, {}, {}) supplied", y, m, d)
-                    };
-                    Some(ns.num_days_from_ce() - EPOCH_DAYS_FROM_CE)
+                    NaiveDate::from_ymd_opt(y, m as u32, d as u32).map_or_else(
+                        // If None is returned, then we have an invalid date.
+                        // We save the faulty values and move on.
+                        || {
+                            error_values = Some((y, m, d));
+                            None
+                        },
+                        // We have a valid date.
+                        |ns| Some(ns.num_days_from_ce() - EPOCH_DAYS_FROM_CE),
+                    )
                 } else {
                     None
                 }
             })
             .collect_trusted();
+        if let Some(values) = error_values {
+            // An invalid y/m/d was detected.
+            polars_bail!(ComputeError: format!("Invalid date components ({}, {}, {}) supplied", values.0, values.1, values.2))
+        };
         ca.rename(name);
         Ok(ca.into_date())
     }

diff --git a/crates/polars-time/src/chunkedarray/datetime.rs b/crates/polars-time/src/chunkedarray/datetime.rs
@@ -167,6 +167,8 @@ pub trait DatetimeMethods: AsDatetime {
         time_zone: Option<&str>,
         name: PlSmallStr,
     ) -> PolarsResult<DatetimeChunked> {
+        let mut error_date_values: Option<(i32, i8, i8)> = None;
+        let mut error_time_values: Option<(i8, i8, i8, i32)> = None;
         let ca: Int64Chunked = year
             .into_iter()
             .zip(month)
@@ -179,28 +181,54 @@ pub trait DatetimeMethods: AsDatetime {
                 if let (Some(y), Some(m), Some(d), Some(h), Some(mnt), Some(s), Some(ns)) =
                     (y, m, d, h, mnt, s, ns)
                 {
-                    let Some(t) = NaiveDate::from_ymd_opt(y, m as u32, d as u32) else {
-                        panic!("Invalid date components ({}, {}, {}) supplied", y, m, d)
-                    };
-                    let Some(ndt) = t.and_hms_nano_opt(h as u32, mnt as u32, s as u32, ns as u32)
-                    else {
-                        panic!(
-                            "Invalid time components ({}, {}, {}, {}) supplied",
-                            h, mnt, s, ns
-                        )
-                    };
-                    Some(match time_unit {
-                        TimeUnit::Milliseconds => ndt.and_utc().timestamp_millis(),
-                        TimeUnit::Microseconds => ndt.and_utc().timestamp_micros(),
-                        TimeUnit::Nanoseconds => ndt.and_utc().timestamp_nanos_opt().unwrap(),
-                    })
+                    NaiveDate::from_ymd_opt(y, m as u32, d as u32).map_or_else(
+                        // If None is returned, then we have an invalid date.
+                        // We save the faulty values and move on.
+                        || {
+                            error_date_values = Some((y, m, d));
+                            None
+                        },
+                        // We have a valid date.
+                        |date| {
+                            date.and_hms_nano_opt(h as u32, mnt as u32, s as u32, ns as u32)
+                                .map_or_else(
+                                    // If None is returned, then we have invalid time components for the
+                                    // specified date. We save the faulty values and move on.
+                                    || {
+                                        error_time_values = Some((h, mnt, s, ns));
+                                        None
+                                    },
+                                    // We have a valid date.
+                                    |ndt| {
+                                        Some(match time_unit {
+                                            TimeUnit::Milliseconds => {
+                                                ndt.and_utc().timestamp_millis()
+                                            },
+                                            TimeUnit::Microseconds => {
+                                                ndt.and_utc().timestamp_micros()
+                                            },
+                                            TimeUnit::Nanoseconds => {
+                                                ndt.and_utc().timestamp_nanos_opt().unwrap()
+                                            },
+                                        })
+                                    },
+                                )
+                        },
+                    )
                 } else {
                     None
                 }
             })
             .collect_trusted();
 
-        println!("here");
+        if let Some(values) = error_date_values {
+            // An invalid date was detected.
+            polars_bail!(ComputeError: format!("Invalid date components ({}, {}, {}) supplied", values.0, values.1, values.2))
+        };
+        if let Some(values) = error_time_values {
+            // An invalid time was detected.
+            polars_bail!(ComputeError: format!("Invalid time components ({}, {}, {}, {}) supplied", values.0, values.1, values.2, values.3))
+        };
         let mut ca = match time_zone {
             #[cfg(feature = "timezones")]
             Some(_) => {

diff --git a/py-polars/tests/unit/functions/as_datatype/test_datetime.py b/py-polars/tests/unit/functions/as_datatype/test_datetime.py
@@ -7,7 +7,7 @@
 import pytest
 
 import polars as pl
-from polars.exceptions import ComputeError, PanicException
+from polars.exceptions import ComputeError
 from polars.testing import assert_series_equal
 
 if TYPE_CHECKING:
@@ -43,7 +43,7 @@ def test_date_datetime() -> None:
 def test_date_invalid_component(components: list[int]) -> None:
     y, m, d = components
     msg = rf"Invalid date components \({y}, {m}, {d}\) supplied"
-    with pytest.raises(PanicException, match=msg):
+    with pytest.raises(ComputeError, match=msg):
         pl.select(pl.date(*components))
 
 
@@ -58,7 +58,7 @@ def test_date_invalid_component(components: list[int]) -> None:
 def test_datetime_invalid_date_component(components: list[int]) -> None:
     y, m, d = components[0:3]
     msg = rf"Invalid date components \({y}, {m}, {d}\) supplied"
-    with pytest.raises(PanicException, match=msg):
+    with pytest.raises(ComputeError, match=msg):
         pl.select(pl.datetime(*components))
 
 
@@ -75,7 +75,7 @@ def test_datetime_invalid_time_component(components: list[int]) -> None:
     h, mnt, s, us = components[3:]
     ns = us * 1_000
     msg = rf"Invalid time components \({h}, {mnt}, {s}, {ns}\) supplied"
-    with pytest.raises(PanicException, match=msg):
+    with pytest.raises(ComputeError, match=msg):
         pl.select(pl.datetime(*components))