Skip to content

REF: share NumericArray._arith_method with BooleanArray #45849

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 1 commit into from
Feb 9, 2022
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
37 changes: 0 additions & 37 deletions pandas/core/arrays/boolean.py
Original file line number Diff line number Diff line change
Expand Up @@ -382,42 +382,5 @@ def _logical_method(self, other, op):
# expected "ndarray"
return BooleanArray(result, mask) # type: ignore[arg-type]

def _arith_method(self, other, op):
mask = None
op_name = op.__name__

if isinstance(other, BaseMaskedArray):
other, mask = other._data, other._mask

elif is_list_like(other):
other = np.asarray(other)
if other.ndim > 1:
raise NotImplementedError("can only perform ops with 1-d structures")
if len(self) != len(other):
raise ValueError("Lengths must match")

mask = self._propagate_mask(mask, other)

if other is libmissing.NA:
# if other is NA, the result will be all NA and we can't run the
# actual op, so we need to choose the resulting dtype manually
if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
dtype = "int8"
elif op_name in {"truediv", "rtruediv"}:
dtype = "float64"
else:
dtype = "bool"
result = np.zeros(len(self._data), dtype=dtype)
else:
if op_name in {"pow", "rpow"} and isinstance(other, np.bool_):
# Avoid DeprecationWarning: In future, it will be an error
# for 'np.bool_' scalars to be interpreted as an index
other = bool(other)

with np.errstate(all="ignore"):
result = op(self._data, other)

return self._maybe_mask_result(result, mask, other, op_name)

def __abs__(self):
return self.copy()
80 changes: 80 additions & 0 deletions pandas/core/arrays/masked.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@
from pandas.core.array_algos.quantile import quantile_with_mask
from pandas.core.arraylike import OpsMixin
from pandas.core.arrays import ExtensionArray
from pandas.core.construction import ensure_wrapped_if_datetimelike
from pandas.core.indexers import check_array_indexer
from pandas.core.ops import invalid_comparison

Expand Down Expand Up @@ -593,6 +594,85 @@ def _propagate_mask(
mask = self._mask | mask
return mask

def _arith_method(self, other, op):
op_name = op.__name__
omask = None

if isinstance(other, BaseMaskedArray):
other, omask = other._data, other._mask

elif is_list_like(other):
if not isinstance(other, ExtensionArray):
other = np.asarray(other)
if other.ndim > 1:
raise NotImplementedError("can only perform ops with 1-d structures")

# We wrap the non-masked arithmetic logic used for numpy dtypes
# in Series/Index arithmetic ops.
other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
pd_op = ops.get_array_op(op)
other = ensure_wrapped_if_datetimelike(other)

if op_name in {"pow", "rpow"} and isinstance(other, np.bool_):
# Avoid DeprecationWarning: In future, it will be an error
# for 'np.bool_' scalars to be interpreted as an index
# e.g. test_array_scalar_like_equivalence
other = bool(other)

mask = self._propagate_mask(omask, other)

if other is libmissing.NA:
result = np.ones_like(self._data)
if self.dtype.kind == "b":
if op_name in {"floordiv", "rfloordiv", "mod", "rmod", "pow", "rpow"}:
dtype = "int8"
elif op_name in {"truediv", "rtruediv"}:
dtype = "float64"
else:
dtype = "bool"
result = result.astype(dtype)
elif "truediv" in op_name and self.dtype.kind != "f":
# The actual data here doesn't matter since the mask
# will be all-True, but since this is division, we want
# to end up with floating dtype.
result = result.astype(np.float64)
else:
# Make sure we do this before the "pow" mask checks
# to get an expected exception message on shape mismatch.
if self.dtype.kind in ["i", "u"] and op_name in ["floordiv", "mod"]:
# TODO(GH#30188) ATM we don't match the behavior of non-masked
# types with respect to floordiv-by-zero
pd_op = op

elif self.dtype.kind == "b" and (
"div" in op_name or "pow" in op_name or "mod" in op_name
):
# TODO(GH#41165): should these be disallowed?
pd_op = op

with np.errstate(all="ignore"):
result = pd_op(self._data, other)

if op_name == "pow":
# 1 ** x is 1.
mask = np.where((self._data == 1) & ~self._mask, False, mask)
# x ** 0 is 1.
if omask is not None:
mask = np.where((other == 0) & ~omask, False, mask)
elif other is not libmissing.NA:
mask = np.where(other == 0, False, mask)

elif op_name == "rpow":
# 1 ** x is 1.
if omask is not None:
mask = np.where((other == 1) & ~omask, False, mask)
elif other is not libmissing.NA:
mask = np.where(other == 1, False, mask)
# x ** 0 is 1.
mask = np.where((self._data == 0) & ~self._mask, False, mask)

return self._maybe_mask_result(result, mask, other, op_name)

def _cmp_method(self, other, op) -> BooleanArray:
from pandas.core.arrays import BooleanArray

Expand Down
63 changes: 0 additions & 63 deletions pandas/core/arrays/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,15 @@
is_bool_dtype,
is_float_dtype,
is_integer_dtype,
is_list_like,
is_object_dtype,
is_string_dtype,
pandas_dtype,
)

from pandas.core import ops
from pandas.core.arrays.base import ExtensionArray
from pandas.core.arrays.masked import (
BaseMaskedArray,
BaseMaskedDtype,
)
from pandas.core.construction import ensure_wrapped_if_datetimelike

if TYPE_CHECKING:
import pyarrow
Expand Down Expand Up @@ -214,65 +210,6 @@ def _from_sequence_of_strings(
scalars = to_numeric(strings, errors="raise")
return cls._from_sequence(scalars, dtype=dtype, copy=copy)

def _arith_method(self, other, op):
op_name = op.__name__
omask = None

if isinstance(other, BaseMaskedArray):
other, omask = other._data, other._mask

elif is_list_like(other):
if not isinstance(other, ExtensionArray):
other = np.asarray(other)
if other.ndim > 1:
raise NotImplementedError("can only perform ops with 1-d structures")

# We wrap the non-masked arithmetic logic used for numpy dtypes
# in Series/Index arithmetic ops.
other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
pd_op = ops.get_array_op(op)
other = ensure_wrapped_if_datetimelike(other)

mask = self._propagate_mask(omask, other)

if other is libmissing.NA:
result = np.ones_like(self._data)
if "truediv" in op_name and self.dtype.kind != "f":
# The actual data here doesn't matter since the mask
# will be all-True, but since this is division, we want
# to end up with floating dtype.
result = result.astype(np.float64)
else:
# Make sure we do this before the "pow" mask checks
# to get an expected exception message on shape mismatch.
if self.dtype.kind in ["i", "u"] and op_name in ["floordiv", "mod"]:
# ATM we don't match the behavior of non-masked types with
# respect to floordiv-by-zero
pd_op = op

with np.errstate(all="ignore"):
result = pd_op(self._data, other)

if op_name == "pow":
# 1 ** x is 1.
mask = np.where((self._data == 1) & ~self._mask, False, mask)
# x ** 0 is 1.
if omask is not None:
mask = np.where((other == 0) & ~omask, False, mask)
elif other is not libmissing.NA:
mask = np.where(other == 0, False, mask)

elif op_name == "rpow":
# 1 ** x is 1.
if omask is not None:
mask = np.where((other == 1) & ~omask, False, mask)
elif other is not libmissing.NA:
mask = np.where(other == 1, False, mask)
# x ** 0 is 1.
mask = np.where((self._data == 0) & ~self._mask, False, mask)

return self._maybe_mask_result(result, mask, other, op_name)

_HANDLED_TYPES = (np.ndarray, numbers.Number)

def __neg__(self):
Expand Down
4 changes: 1 addition & 3 deletions pandas/tests/arrays/boolean/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,7 @@ def test_div(left_array, right_array):
[
"floordiv",
"mod",
pytest.param(
"pow", marks=pytest.mark.xfail(reason="TODO follow int8 behaviour? GH34686")
),
"pow",
],
)
def test_op_int8(left_array, right_array, opname):
Expand Down
17 changes: 12 additions & 5 deletions pandas/tests/arrays/masked/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,22 +153,29 @@ def test_error_len_mismatch(data, all_arithmetic_operators):

other = [scalar] * (len(data) - 1)

err = ValueError
msg = "|".join(
[
r"operands could not be broadcast together with shapes \(3,\) \(4,\)",
r"operands could not be broadcast together with shapes \(4,\) \(3,\)",
]
)

if data.dtype.kind == "b":
msg = "Lengths must match"
if data.dtype.kind == "b" and all_arithmetic_operators.strip("_") in [
"sub",
"rsub",
]:
err = TypeError
msg = (
r"numpy boolean subtract, the `\-` operator, is not supported, use "
r"the bitwise_xor, the `\^` operator, or the logical_xor function instead"
)

for other in [other, np.array(other)]:
with pytest.raises(ValueError, match=msg):
with pytest.raises(err, match=msg):
op(data, other)

s = pd.Series(data)
with pytest.raises(ValueError, match=msg):
with pytest.raises(err, match=msg):
op(s, other)


Expand Down