diff --git a/py-polars/polars/internals/expr/expr.py b/py-polars/polars/internals/expr/expr.py index 76ccdb0069f3..d8cd94118989 100644 --- a/py-polars/polars/internals/expr/expr.py +++ b/py-polars/polars/internals/expr/expr.py @@ -290,6 +290,12 @@ def __array_ufunc__( ) -> Expr: """Numpy universal functions.""" + num_expr = sum(isinstance(inp, Expr) for inp in inputs) + if num_expr > 1: + raise ValueError( + f"Numpy ufunc can only be used with one expression, {num_expr} given. Use `pl.reduce` to call numpy functions over multiple expressions." + ) + def function(s: pli.Series) -> pli.Series: # pragma: no cover args = [inp if not isinstance(inp, Expr) else s for inp in inputs] return ufunc(*args, **kwargs) diff --git a/py-polars/tests/unit/test_lazy.py b/py-polars/tests/unit/test_lazy.py index 281d19d5a8b2..92e6d5b9f6a6 100644 --- a/py-polars/tests/unit/test_lazy.py +++ b/py-polars/tests/unit/test_lazy.py @@ -973,6 +973,20 @@ def test_ufunc_expr_not_first() -> None: assert_frame_equal(out, expected) +def test_ufunc_multiple_expr() -> None: + df = pl.DataFrame( + [ + pl.Series("a", [1, 2, 3], dtype=pl.Float64), + pl.Series("b", [4, 5, 6], dtype=pl.Float64), + ] + ) + + with pytest.raises( + ValueError, match="Numpy ufunc can only be used with one expression, 2 given" + ): + df.select(np.arctan2(pl.col("a"), pl.col("b"))) # type: ignore[call-overload] + + def test_clip() -> None: df = pl.DataFrame({"a": [1, 2, 3, 4, 5]}) assert df.select(pl.col("a").clip(2, 4))["a"].to_list() == [2, 2, 3, 4, 4]