Skip to content

Commit

Permalink
ENH: Implement unary operators for FloatingArray class (#39916)
Browse files Browse the repository at this point in the history
  • Loading branch information
zitorelova authored Feb 28, 2021
1 parent 879cd22 commit 4c5e6fa
Show file tree
Hide file tree
Showing 7 changed files with 78 additions and 73 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v1.3.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ Other enhancements
- :meth:`Series.loc.__getitem__` and :meth:`Series.loc.__setitem__` with :class:`MultiIndex` now raising helpful error message when indexer has too many dimensions (:issue:`35349`)
- :meth:`pandas.read_stata` and :class:`StataReader` support reading data from compressed files.
- Add support for parsing ``ISO 8601``-like timestamps with negative signs to :meth:`pandas.Timedelta` (:issue:`37172`)
- Add support for unary operators in :class:`FloatingArray` (:issue:`38749`)

.. ---------------------------------------------------------------------------
Expand Down
9 changes: 0 additions & 9 deletions pandas/core/arrays/integer.py
Original file line number Diff line number Diff line change
Expand Up @@ -315,15 +315,6 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False):
)
super().__init__(values, mask, copy=copy)

def __neg__(self):
return type(self)(-self._data, self._mask.copy())

def __pos__(self):
return self

def __abs__(self):
return type(self)(np.abs(self._data), self._mask.copy())

@classmethod
def _from_sequence(
cls, scalars, *, dtype: Optional[Dtype] = None, copy: bool = False
Expand Down
9 changes: 9 additions & 0 deletions pandas/core/arrays/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,3 +199,12 @@ def reconstruct(x):
return tuple(reconstruct(x) for x in result)
else:
return reconstruct(result)

def __neg__(self):
return type(self)(-self._data, self._mask.copy())

def __pos__(self):
return self

def __abs__(self):
return type(self)(abs(self._data), self._mask.copy())
21 changes: 21 additions & 0 deletions pandas/tests/arrays/floating/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,3 +180,24 @@ def test_cross_type_arithmetic():
result = df.A + df.B
expected = pd.Series([2, np.nan, np.nan], dtype="Float64")
tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
"source, neg_target, abs_target",
[
([1.1, 2.2, 3.3], [-1.1, -2.2, -3.3], [1.1, 2.2, 3.3]),
([1.1, 2.2, None], [-1.1, -2.2, None], [1.1, 2.2, None]),
([-1.1, 0.0, 1.1], [1.1, 0.0, -1.1], [1.1, 0.0, 1.1]),
],
)
def test_unary_float_operators(float_ea_dtype, source, neg_target, abs_target):
# GH38794
dtype = float_ea_dtype
arr = pd.array(source, dtype=dtype)
neg_result, pos_result, abs_result = -arr, +arr, abs(arr)
neg_target = pd.array(neg_target, dtype=dtype)
abs_target = pd.array(abs_target, dtype=dtype)

tm.assert_extension_array_equal(neg_result, neg_target)
tm.assert_extension_array_equal(pos_result, arr)
tm.assert_extension_array_equal(abs_result, abs_target)
40 changes: 13 additions & 27 deletions pandas/tests/arrays/integer/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,36 +284,22 @@ def test_reduce_to_float(op):


@pytest.mark.parametrize(
"source, target",
"source, neg_target, abs_target",
[
([1, 2, 3], [-1, -2, -3]),
([1, 2, None], [-1, -2, None]),
([-1, 0, 1], [1, 0, -1]),
([1, 2, 3], [-1, -2, -3], [1, 2, 3]),
([1, 2, None], [-1, -2, None], [1, 2, None]),
([-1, 0, 1], [1, 0, -1], [1, 0, 1]),
],
)
def test_unary_minus_nullable_int(any_signed_nullable_int_dtype, source, target):
def test_unary_int_operators(
any_signed_nullable_int_dtype, source, neg_target, abs_target
):
dtype = any_signed_nullable_int_dtype
arr = pd.array(source, dtype=dtype)
result = -arr
expected = pd.array(target, dtype=dtype)
tm.assert_extension_array_equal(result, expected)


@pytest.mark.parametrize("source", [[1, 2, 3], [1, 2, None], [-1, 0, 1]])
def test_unary_plus_nullable_int(any_signed_nullable_int_dtype, source):
dtype = any_signed_nullable_int_dtype
expected = pd.array(source, dtype=dtype)
result = +expected
tm.assert_extension_array_equal(result, expected)
neg_result, pos_result, abs_result = -arr, +arr, abs(arr)
neg_target = pd.array(neg_target, dtype=dtype)
abs_target = pd.array(abs_target, dtype=dtype)


@pytest.mark.parametrize(
"source, target",
[([1, 2, 3], [1, 2, 3]), ([1, -2, None], [1, 2, None]), ([-1, 0, 1], [1, 0, 1])],
)
def test_abs_nullable_int(any_signed_nullable_int_dtype, source, target):
dtype = any_signed_nullable_int_dtype
s = pd.array(source, dtype=dtype)
result = abs(s)
expected = pd.array(target, dtype=dtype)
tm.assert_extension_array_equal(result, expected)
tm.assert_extension_array_equal(neg_result, neg_target)
tm.assert_extension_array_equal(pos_result, arr)
tm.assert_extension_array_equal(abs_result, abs_target)
12 changes: 7 additions & 5 deletions pandas/tests/arrays/masked/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,12 +165,14 @@ def test_error_len_mismatch(data, all_arithmetic_operators):


@pytest.mark.parametrize("op", ["__neg__", "__abs__", "__invert__"])
@pytest.mark.parametrize(
"values, dtype", [([1, 2, 3], "Int64"), ([True, False, True], "boolean")]
)
def test_unary_op_does_not_propagate_mask(op, values, dtype):
def test_unary_op_does_not_propagate_mask(data, op, request):
# https://github.com/pandas-dev/pandas/issues/39943
s = pd.Series(values, dtype=dtype)
data, _ = data
if data.dtype in ["Float32", "Float64"] and op == "__invert__":
request.node.add_marker(
pytest.mark.xfail(reason="invert is not implemented for float ea dtypes")
)
s = pd.Series(data)
result = getattr(s, op)()
expected = result.copy(deep=True)
s[0] = None
Expand Down
59 changes: 27 additions & 32 deletions pandas/tests/series/test_unary.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,40 +18,35 @@ def test_invert(self):
tm.assert_series_equal(-(ser < 0), ~(ser < 0))

@pytest.mark.parametrize(
"source, target",
"source, neg_target, abs_target",
[
([1, 2, 3], [-1, -2, -3]),
([1, 2, None], [-1, -2, None]),
([-1, 0, 1], [1, 0, -1]),
([1, 2, 3], [-1, -2, -3], [1, 2, 3]),
([1, 2, None], [-1, -2, None], [1, 2, None]),
],
)
def test_unary_minus_nullable_int(
self, any_signed_nullable_int_dtype, source, target
def test_all_numeric_unary_operators(
self, any_nullable_numeric_dtype, source, neg_target, abs_target
):
dtype = any_signed_nullable_int_dtype
# GH38794
dtype = any_nullable_numeric_dtype
ser = Series(source, dtype=dtype)
result = -ser
expected = Series(target, dtype=dtype)
tm.assert_series_equal(result, expected)

@pytest.mark.parametrize("source", [[1, 2, 3], [1, 2, None], [-1, 0, 1]])
def test_unary_plus_nullable_int(self, any_signed_nullable_int_dtype, source):
dtype = any_signed_nullable_int_dtype
expected = Series(source, dtype=dtype)
result = +expected
tm.assert_series_equal(result, expected)

@pytest.mark.parametrize(
"source, target",
[
([1, 2, 3], [1, 2, 3]),
([1, -2, None], [1, 2, None]),
([-1, 0, 1], [1, 0, 1]),
],
)
def test_abs_nullable_int(self, any_signed_nullable_int_dtype, source, target):
dtype = any_signed_nullable_int_dtype
ser = Series(source, dtype=dtype)
result = abs(ser)
expected = Series(target, dtype=dtype)
tm.assert_series_equal(result, expected)
neg_result, pos_result, abs_result = -ser, +ser, abs(ser)
if dtype.startswith("U"):
neg_target = -Series(source, dtype=dtype)
else:
neg_target = Series(neg_target, dtype=dtype)

abs_target = Series(abs_target, dtype=dtype)

tm.assert_series_equal(neg_result, neg_target)
tm.assert_series_equal(pos_result, ser)
tm.assert_series_equal(abs_result, abs_target)

@pytest.mark.parametrize("op", ["__neg__", "__abs__"])
def test_unary_float_op_mask(self, float_ea_dtype, op):
dtype = float_ea_dtype
ser = Series([1.1, 2.2, 3.3], dtype=dtype)
result = getattr(ser, op)()
target = result.copy(deep=True)
ser[0] = None
tm.assert_series_equal(result, target)

0 comments on commit 4c5e6fa

Please sign in to comment.