fix: Propagate error instead of panicking when calling product on a…

…n invalid type (pola-rs#15093)
mcrumiller · Mar 16, 2024 · be92470 · be92470
1 parent 5d449cc
commit be92470
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 8 deletions.
diff --git a/crates/polars-core/src/series/mod.rs b/crates/polars-core/src/series/mod.rs
@@ -599,7 +599,7 @@ impl Series {
     ///
     /// If the [`DataType`] is one of `{Int8, UInt8, Int16, UInt16}` the `Series` is
     /// first cast to `Int64` to prevent overflow issues.
-    pub fn product(&self) -> Series {
+    pub fn product(&self) -> PolarsResult<Series> {
         #[cfg(feature = "product")]
         {
             use DataType::*;
@@ -609,11 +609,13 @@ impl Series {
                     let s = self.cast(&Int64).unwrap();
                     s.product()
                 },
-                Int64 => self.i64().unwrap().prod_as_series(),
-                UInt64 => self.u64().unwrap().prod_as_series(),
-                Float32 => self.f32().unwrap().prod_as_series(),
-                Float64 => self.f64().unwrap().prod_as_series(),
-                dt => panic!("product not supported for dtype: {dt:?}"),
+                Int64 => Ok(self.i64().unwrap().prod_as_series()),
+                UInt64 => Ok(self.u64().unwrap().prod_as_series()),
+                Float32 => Ok(self.f32().unwrap().prod_as_series()),
+                Float64 => Ok(self.f64().unwrap().prod_as_series()),
+                dt => {
+                    polars_bail!(InvalidOperation: "`product` operation not supported for dtype `{dt}`")
+                },
             }
         }
         #[cfg(not(feature = "product"))]

diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs
@@ -737,7 +737,7 @@ impl Expr {
         };
 
         self.function_with_options(
-            move |s: Series| Ok(Some(s.product())),
+            move |s: Series| Some(s.product()).transpose(),
             GetOutput::map_dtype(|dt| {
                 use DataType::*;
                 match dt {

diff --git a/py-polars/src/series/aggregation.rs b/py-polars/src/series/aggregation.rs
@@ -101,7 +101,14 @@ impl PySeries {
     }
 
     fn product(&self, py: Python) -> PyResult<PyObject> {
-        Ok(Wrap(self.series.product().get(0).map_err(PyPolarsErr::from)?).into_py(py))
+        Ok(Wrap(
+            self.series
+                .product()
+                .map_err(PyPolarsErr::from)?
+                .get(0)
+                .map_err(PyPolarsErr::from)?,
+        )
+        .into_py(py))
     }
 
     fn quantile(

diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -706,3 +706,11 @@ def test_raise_not_found_in_simplify_14974() -> None:
     df = pl.DataFrame()
     with pytest.raises(pl.ColumnNotFoundError):
         df.select(1 / (1 + pl.col("a")))
+
+
+def test_invalid_product_type() -> None:
+    with pytest.raises(
+        pl.InvalidOperationError,
+        match="`product` operation not supported for dtype",
+    ):
+        pl.Series([[1, 2, 3]]).product()