fix(python): fix ufunc in agg (change __ufunc_array__ so it uses is_elementwise=True parameter)

py-polars/polars/exceptions.py

@@ -121,6 +121,10 @@ class ArrowError(Exception):
     """Deprecated: will be removed."""
 
 
+class CustomUFuncWarning(PolarsWarning):  # type: ignore[misc]
+    """Warning issued when a custom ufunc is handled differently than numpy ufunc would."""  # noqa: W505
+
+
 __all__ = [
     "ArrowError",
     "ColumnNotFoundError",

py-polars/polars/expr/expr.py

@@ -31,7 +31,7 @@
 )
 from polars.dependencies import _check_for_numpy
 from polars.dependencies import numpy as np
-from polars.exceptions import PolarsInefficientMapWarning
+from polars.exceptions import CustomUFuncWarning, PolarsInefficientMapWarning
 from polars.expr.array import ExprArrayNameSpace
 from polars.expr.binary import ExprBinaryNameSpace
 from polars.expr.categorical import ExprCatNameSpace
@@ -58,6 +58,7 @@
 from polars.utils.meta import threadpool_size
 from polars.utils.unstable import issue_unstable_warning, unstable
 from polars.utils.various import (
+    find_stacklevel,
     no_default,
     sphinx_accessor,
     warn_null_comparison,
@@ -286,6 +287,7 @@ def __array_ufunc__(
         self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any
     ) -> Self:
         """Numpy universal functions."""
+        is_custom_ufunc = ufunc.__class__ != np.ufunc
         num_expr = sum(isinstance(inp, Expr) for inp in inputs)
         if num_expr > 1:
             if num_expr < len(inputs):
@@ -302,7 +304,20 @@ def function(s: Series) -> Series:  # pragma: no cover
             args = [inp if not isinstance(inp, Expr) else s for inp in inputs]
             return ufunc(*args, **kwargs)
 
-        return self.map_batches(function)
+        if is_custom_ufunc is True:
+            msg = (
+                "Native numpy ufuncs are dispatched using `map_batches(ufunc, is_elementwise=True)` which "
+                "is safe for native Numpy and Scipy ufuncs but custom ufuncs in a group_by "
+                "context won't be properly grouped. Custom ufuncs are dispatched with is_elementwise=False. "
+                f"If {ufunc.__name__} needs elementwise then please use map_batches directly."
+            )
+            warnings.warn(
+                msg,
+                CustomUFuncWarning,
+                stacklevel=find_stacklevel(),
+            )
+            return self.map_batches(function, is_elementwise=False)
+        return self.map_batches(function, is_elementwise=True)
 
     @classmethod
     def from_json(cls, value: str) -> Self:

py-polars/tests/unit/operations/map/test_map_batches.py

@@ -77,3 +77,13 @@ def test_map_deprecated() -> None:
         pl.col("a").map(lambda x: x)
     with pytest.deprecated_call():
         pl.LazyFrame({"a": [1, 2]}).map(lambda x: x)
+
+
+def test_ufunc_recognition() -> None:
+    df = pl.DataFrame({"a": [1, 1, 2, 2], "b": [1.1, 2.2, 3.3, 4.4]})
+    assert_frame_equal(df.select(np.exp(pl.col("b"))), df.select(pl.col("b").exp()))
+
+
+def test_grouped_ufunc() -> None:
+    df = pl.DataFrame({"id": ["a", "a", "b", "b"], "values": [0.1, 0.1, -0.1, -0.1]})
+    df.group_by("id").agg(pl.col("values").log1p().sum().pipe(np.expm1))