python: improve error (#4223)

py-polars/polars/internals/frame.py

@@ -1993,7 +1993,10 @@ def __getitem__(
                 )
 
         # if no data has been returned, the operation is not supported
-        raise NotImplementedError
+        raise ValueError(
+            f"Cannot __getitem__ on DataFrame with item: '{item}'"
+            f" of type: '{type(item)}'."
+        )
 
     def __setitem__(
         self, key: str | list | tuple[Any, str | int], value: Any
@@ -2048,7 +2051,11 @@ def __setitem__(
             elif isinstance(col_selection, str):
                 self.replace(col_selection, s)
         else:
-            raise NotImplementedError
+            raise ValueError(
+                f"Cannot __setitem__ on DataFrame with key: '{key}' "
+                f"of type: '{type(key)}' and value: '{value}' "
+                f"of type: '{type(value)}'."
+            )
 
     def __len__(self) -> int:
         return self.height

py-polars/polars/internals/series.py

@@ -579,7 +579,10 @@ def __getitem__(self, item: int | Series | range | slice | np.ndarray) -> Any:
         if isinstance(item, slice):
             return PolarsSlice(self).apply(item)
 
-        raise NotImplementedError
+        raise ValueError(
+            f"Cannot __getitem__ on Series of dtype: '{self.dtype}' "
+            f"with argument: '{item}' of type: '{type(item)}'."
+        )
 
     def __setitem__(
         self, key: int | Series | np.ndarray | list | tuple, value: Any

py-polars/src/lazy/apply.rs

@@ -49,7 +49,7 @@ pub(crate) fn call_lambda_with_series(
     s: Series,
     lambda: &PyObject,
     polars_module: &PyObject,
-) -> PyObject {
+) -> PyResult<PyObject> {
     let pypolars = polars_module.cast_as::<PyModule>(py).unwrap();
 
     // create a PySeries struct/object for Python
@@ -61,10 +61,7 @@ pub(crate) fn call_lambda_with_series(
         .call1((pyseries,))
         .unwrap();
     // call the lambda and get a python side Series wrapper
-    match lambda.call1(py, (python_series_wrapper,)) {
-        Ok(pyobj) => pyobj,
-        Err(e) => panic!("python apply failed: {}", e.value(py)),
-    }
+    lambda.call1(py, (python_series_wrapper,))
 }
 
 /// A python lambda taking two Series
@@ -159,7 +156,8 @@ pub fn map_single(
         let py = gil.python();
 
         // this is a python Series
-        let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars);
+        let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars)
+            .map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;
 
         Ok(out.to_series(py, &pypolars, s.name()))
     };

py-polars/src/lazy/dsl.rs

@@ -912,7 +912,8 @@ impl PyExpr {
             let gil = Python::acquire_gil();
             let py = gil.python();
 
-            let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars);
+            let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars)
+                .expect("python function failed");
             match out.getattr(py, "_s") {
                 Ok(pyseries) => {
                     let pyseries = pyseries.extract::<PySeries>(py).unwrap();

py-polars/tests/test_df.py

@@ -612,7 +612,7 @@ def test_set() -> None:
             df[(1, 2, 3)] = 1  # type: ignore[index]
 
         # we cannot index with any type, such as bool
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             df[True] = 1  # type: ignore[index]
 
 
@@ -1660,7 +1660,7 @@ def test_get_item() -> None:
 
     # note that we cannot use floats (even if they could be casted to integer without
     # loss)
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         _ = df[np.array([1.0])]
 
     # using boolean masks with numpy is deprecated

py-polars/tests/test_errors.py

@@ -117,3 +117,26 @@ def test_not_found_on_rename() -> None:
                 "does_not_exist": "exists",
             }
         )
+
+
+@typing.no_type_check
+def test_getitem_errs() -> None:
+    df = pl.DataFrame({"a": [1, 2, 3]})
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot __getitem__ on DataFrame with item: "
+        r"'{'some'}' of type: '<class 'set'>'.",
+    ):
+        df[{"some"}]
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot __getitem__ on Series of dtype: "
+        r"'<class 'polars.datatypes.Int64'>' with argument: "
+        r"'{'strange'}' of type: '<class 'set'>'.",
+    ):
+        df["a"][{"strange"}]
+
+    with pytest.raises(ValueError, match="Cannot __setitem__ on DataFrame with key:.*"):
+        df[{"some"}] = "foo"