Skip to content

Commit 8974a95

Browse files
author
Khor Chean Wei
authored
BUG: array ufuncs with more than 2 inputs fails with 2 or more pd.Series and scalar arguments (#48280)
1 parent 04db721 commit 8974a95

File tree

4 files changed

+52
-3
lines changed

4 files changed

+52
-3
lines changed

doc/source/whatsnew/v1.6.0.rst

+1-1
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ Timezones
190190

191191
Numeric
192192
^^^^^^^
193-
-
193+
- Bug in :meth:`DataFrame.add` cannot apply ufunc when inputs contain mixed DataFrame type and Series type (:issue:`39853`)
194194
-
195195

196196
Conversion

pandas/core/arraylike.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,10 @@ def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any)
250250
--------
251251
numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__
252252
"""
253+
from pandas.core.frame import (
254+
DataFrame,
255+
Series,
256+
)
253257
from pandas.core.generic import NDFrame
254258
from pandas.core.internals import BlockManager
255259

@@ -295,8 +299,8 @@ def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any)
295299
# At the moment, there aren't any ufuncs with more than two inputs
296300
# so this ends up just being x1.index | x2.index, but we write
297301
# it to handle *args.
298-
299-
if len(set(types)) > 1:
302+
set_types = set(types)
303+
if len(set_types) > 1 and {DataFrame, Series}.issubset(set_types):
300304
# We currently don't handle ufunc(DataFrame, Series)
301305
# well. Previously this raised an internal ValueError. We might
302306
# support it someday, so raise a NotImplementedError.

pandas/tests/frame/test_ufunc.py

+22
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from functools import partial
2+
import re
23

34
import numpy as np
45
import pytest
@@ -301,3 +302,24 @@ def my_ufunc(x, y, z):
301302
result = my_ufunc(df1.values, df2, df3)
302303
expected = expected.set_axis(["b", "c"], axis=1)
303304
tm.assert_frame_equal(result, expected)
305+
306+
307+
def test_array_ufuncs_for_many_arguments():
308+
# GH39853
309+
def add3(x, y, z):
310+
return x + y + z
311+
312+
ufunc = np.frompyfunc(add3, 3, 1)
313+
df = pd.DataFrame([[1, 2], [3, 4]])
314+
315+
result = ufunc(df, df, 1)
316+
expected = pd.DataFrame([[3, 5], [7, 9]], dtype=object)
317+
tm.assert_frame_equal(result, expected)
318+
319+
ser = pd.Series([1, 2])
320+
msg = (
321+
"Cannot apply ufunc <ufunc 'add3 (vectorized)'> "
322+
"to mixed DataFrame and Series inputs."
323+
)
324+
with pytest.raises(NotImplementedError, match=re.escape(msg)):
325+
ufunc(df, df, ser)

pandas/tests/series/test_ufunc.py

+23
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from collections import deque
2+
import re
23
import string
34

45
import numpy as np
@@ -453,3 +454,25 @@ def test_np_matmul():
453454
expected_result,
454455
result,
455456
)
457+
458+
459+
def test_array_ufuncs_for_many_arguments():
460+
# GH39853
461+
def add3(x, y, z):
462+
return x + y + z
463+
464+
ufunc = np.frompyfunc(add3, 3, 1)
465+
ser = pd.Series([1, 2])
466+
467+
result = ufunc(ser, ser, 1)
468+
expected = pd.Series([3, 5], dtype=object)
469+
tm.assert_series_equal(result, expected)
470+
471+
df = pd.DataFrame([[1, 2]])
472+
473+
msg = (
474+
"Cannot apply ufunc <ufunc 'add3 (vectorized)'> "
475+
"to mixed DataFrame and Series inputs."
476+
)
477+
with pytest.raises(NotImplementedError, match=re.escape(msg)):
478+
ufunc(ser, ser, df)

0 commit comments

Comments
 (0)