Skip to content

REF: standardize __array_ufunc__ patterns #45113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Dec 30, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion pandas/core/arrays/categorical.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,10 @@
notna,
)

from pandas.core import ops
from pandas.core import (
arraylike,
ops,
)
from pandas.core.accessor import (
PandasDelegate,
delegate_names,
Expand Down Expand Up @@ -1516,6 +1519,14 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
if result is not NotImplemented:
return result

if method == "reduce":
# e.g. TestCategoricalAnalytics::test_min_max_ordered
result = arraylike.dispatch_reduction_ufunc(
self, ufunc, method, *inputs, **kwargs
)
if result is not NotImplemented:
return result

# for all other cases, raise for now (similarly as what happens in
# Series.__array_prepare__)
raise TypeError(
Expand Down
4 changes: 0 additions & 4 deletions pandas/core/arrays/numpy_.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
from __future__ import annotations

import numbers

import numpy as np

from pandas._libs import lib
Expand Down Expand Up @@ -130,8 +128,6 @@ def dtype(self) -> PandasDtype:
def __array__(self, dtype: NpDtype | None = None) -> np.ndarray:
return np.asarray(self._ndarray, dtype=dtype)

_HANDLED_TYPES = (np.ndarray, numbers.Number)

def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
# Lightly modified version of
# https://numpy.org/doc/stable/reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html
Expand Down
28 changes: 26 additions & 2 deletions pandas/core/arrays/sparse/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@
notna,
)

from pandas.core import arraylike
import pandas.core.algorithms as algos
from pandas.core.arraylike import OpsMixin
from pandas.core.arrays import ExtensionArray
Expand Down Expand Up @@ -1415,7 +1416,9 @@ def any(self, axis=0, *args, **kwargs):

return values.any().item()

def sum(self, axis: int = 0, min_count: int = 0, *args, **kwargs) -> Scalar:
def sum(
self, axis: int = 0, min_count: int = 0, skipna: bool = True, *args, **kwargs
) -> Scalar:
"""
Sum of non-NA/null values

Expand All @@ -1437,6 +1440,11 @@ def sum(self, axis: int = 0, min_count: int = 0, *args, **kwargs) -> Scalar:
nv.validate_sum(args, kwargs)
valid_vals = self._valid_sp_values
sp_sum = valid_vals.sum()
has_na = self.sp_index.ngaps > 0 and not self._null_fill_value

if has_na and not skipna:
return na_value_for_dtype(self.dtype.subtype, compat=False)

if self._null_fill_value:
if check_below_min_count(valid_vals.shape, None, min_count):
return na_value_for_dtype(self.dtype.subtype, compat=False)
Expand Down Expand Up @@ -1589,6 +1597,21 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
if result is not NotImplemented:
return result

if "out" in kwargs:
# e.g. tests.arrays.sparse.test_arithmetics.test_ndarray_inplace
res = arraylike.dispatch_ufunc_with_out(
self, ufunc, method, *inputs, **kwargs
)
return res

if method == "reduce":
result = arraylike.dispatch_reduction_ufunc(
self, ufunc, method, *inputs, **kwargs
)
if result is not NotImplemented:
# e.g. tests.series.test_ufunc.TestNumpyReductions
return result

if len(inputs) == 1:
# No alignment necessary.
sp_values = getattr(ufunc, method)(self.sp_values, **kwargs)
Expand All @@ -1611,7 +1634,8 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
sp_values, self.sp_index, SparseDtype(sp_values.dtype, fill_value)
)

result = getattr(ufunc, method)(*(np.asarray(x) for x in inputs), **kwargs)
new_inputs = tuple(np.asarray(x) for x in inputs)
result = getattr(ufunc, method)(*new_inputs, **kwargs)
if out:
if len(out) == 1:
out = out[0]
Expand Down
6 changes: 6 additions & 0 deletions pandas/core/indexes/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -878,6 +878,12 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str_t, *inputs, **kwargs):
if result is not NotImplemented:
return result

if "out" in kwargs:
# e.g. test_dti_isub_tdi
return arraylike.dispatch_ufunc_with_out(
self, ufunc, method, *inputs, **kwargs
)

if method == "reduce":
result = arraylike.dispatch_reduction_ufunc(
self, ufunc, method, *inputs, **kwargs
Expand Down
4 changes: 1 addition & 3 deletions pandas/tests/arithmetic/test_datetime64.py
Original file line number Diff line number Diff line change
Expand Up @@ -2107,7 +2107,7 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
np.subtract(out, tdi, out=out)
tm.assert_datetime_array_equal(out, expected._data)

msg = "cannot subtract .* from a TimedeltaArray"
msg = "cannot subtract a datelike from a TimedeltaArray"
with pytest.raises(TypeError, match=msg):
tdi -= dti

Expand All @@ -2116,11 +2116,9 @@ def test_dti_isub_tdi(self, tz_naive_fixture):
result -= tdi.values
tm.assert_index_equal(result, expected)

msg = "cannot subtract DatetimeArray from ndarray"
with pytest.raises(TypeError, match=msg):
tdi.values -= dti

msg = "cannot subtract a datelike from a TimedeltaArray"
with pytest.raises(TypeError, match=msg):
tdi._values -= dti

Expand Down
22 changes: 17 additions & 5 deletions pandas/tests/arrays/categorical/test_analytics.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,20 +29,32 @@ def test_min_max_not_ordered_raises(self, aggregation):
with pytest.raises(TypeError, match=msg):
agg_func()

def test_min_max_ordered(self):
ufunc = np.minimum if aggregation == "min" else np.maximum
with pytest.raises(TypeError, match=msg):
ufunc.reduce(cat)

def test_min_max_ordered(self, index_or_series_or_array):
cat = Categorical(["a", "b", "c", "d"], ordered=True)
_min = cat.min()
_max = cat.max()
obj = index_or_series_or_array(cat)
_min = obj.min()
_max = obj.max()
assert _min == "a"
assert _max == "d"

assert np.minimum.reduce(obj) == "a"
assert np.maximum.reduce(obj) == "d"
# TODO: raises if we pass axis=0 (on Index and Categorical, not Series)

cat = Categorical(
["a", "b", "c", "d"], categories=["d", "c", "b", "a"], ordered=True
)
_min = cat.min()
_max = cat.max()
obj = index_or_series_or_array(cat)
_min = obj.min()
_max = obj.max()
assert _min == "d"
assert _max == "a"
assert np.minimum.reduce(obj) == "d"
assert np.maximum.reduce(obj) == "a"

@pytest.mark.parametrize(
"categories,expected",
Expand Down
8 changes: 8 additions & 0 deletions pandas/tests/extension/decimal/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
is_list_like,
is_scalar,
)
from pandas.core import arraylike
from pandas.core.arraylike import OpsMixin
from pandas.core.arrays import (
ExtensionArray,
Expand Down Expand Up @@ -121,6 +122,13 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
inputs = tuple(x._data if isinstance(x, DecimalArray) else x for x in inputs)
result = getattr(ufunc, method)(*inputs, **kwargs)

if method == "reduce":
result = arraylike.dispatch_reduction_ufunc(
self, ufunc, method, *inputs, **kwargs
)
if result is not NotImplemented:
return result

def reconstruct(x):
if isinstance(x, (decimal.Decimal, numbers.Number)):
return x
Expand Down