diff --git a/doc/source/whatsnew/v1.6.0.rst b/doc/source/whatsnew/v1.6.0.rst index bdaea89776b7c..ce60844f68b51 100644 --- a/doc/source/whatsnew/v1.6.0.rst +++ b/doc/source/whatsnew/v1.6.0.rst @@ -189,7 +189,7 @@ Timezones Numeric ^^^^^^^ -- +- Bug in :meth:`DataFrame.add` cannot apply ufunc when inputs contain mixed DataFrame type and Series type (:issue:`39853`) - Conversion diff --git a/pandas/core/arraylike.py b/pandas/core/arraylike.py index 4e8e4ea7e8d87..ca9fee3cf8896 100644 --- a/pandas/core/arraylike.py +++ b/pandas/core/arraylike.py @@ -250,6 +250,10 @@ def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any) -------- numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__ """ + from pandas.core.frame import ( + DataFrame, + Series, + ) from pandas.core.generic import NDFrame from pandas.core.internals import BlockManager @@ -295,8 +299,8 @@ def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any) # At the moment, there aren't any ufuncs with more than two inputs # so this ends up just being x1.index | x2.index, but we write # it to handle *args. - - if len(set(types)) > 1: + set_types = set(types) + if len(set_types) > 1 and {DataFrame, Series}.issubset(set_types): # We currently don't handle ufunc(DataFrame, Series) # well. Previously this raised an internal ValueError. We might # support it someday, so raise a NotImplementedError. diff --git a/pandas/tests/frame/test_ufunc.py b/pandas/tests/frame/test_ufunc.py index b884224315a61..8c79611003013 100644 --- a/pandas/tests/frame/test_ufunc.py +++ b/pandas/tests/frame/test_ufunc.py @@ -1,4 +1,5 @@ from functools import partial +import re import numpy as np import pytest @@ -301,3 +302,24 @@ def my_ufunc(x, y, z): result = my_ufunc(df1.values, df2, df3) expected = expected.set_axis(["b", "c"], axis=1) tm.assert_frame_equal(result, expected) + + +def test_array_ufuncs_for_many_arguments(): + # GH39853 + def add3(x, y, z): + return x + y + z + + ufunc = np.frompyfunc(add3, 3, 1) + df = pd.DataFrame([[1, 2], [3, 4]]) + + result = ufunc(df, df, 1) + expected = pd.DataFrame([[3, 5], [7, 9]], dtype=object) + tm.assert_frame_equal(result, expected) + + ser = pd.Series([1, 2]) + msg = ( + "Cannot apply ufunc " + "to mixed DataFrame and Series inputs." + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + ufunc(df, df, ser) diff --git a/pandas/tests/series/test_ufunc.py b/pandas/tests/series/test_ufunc.py index 8e1384db024b0..924980b62a51b 100644 --- a/pandas/tests/series/test_ufunc.py +++ b/pandas/tests/series/test_ufunc.py @@ -1,4 +1,5 @@ from collections import deque +import re import string import numpy as np @@ -453,3 +454,25 @@ def test_np_matmul(): expected_result, result, ) + + +def test_array_ufuncs_for_many_arguments(): + # GH39853 + def add3(x, y, z): + return x + y + z + + ufunc = np.frompyfunc(add3, 3, 1) + ser = pd.Series([1, 2]) + + result = ufunc(ser, ser, 1) + expected = pd.Series([3, 5], dtype=object) + tm.assert_series_equal(result, expected) + + df = pd.DataFrame([[1, 2]]) + + msg = ( + "Cannot apply ufunc " + "to mixed DataFrame and Series inputs." + ) + with pytest.raises(NotImplementedError, match=re.escape(msg)): + ufunc(ser, ser, df)