ENH: Implement unary operators for FloatingArray class (#39916)

Author: zitorelova

doc/source/whatsnew/v1.3.0.rst

@@ -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`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/arrays/integer.py

@@ -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

pandas/core/arrays/numeric.py

@@ -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())

pandas/tests/arrays/floating/test_arithmetic.py

@@ -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)

pandas/tests/arrays/integer/test_arithmetic.py

@@ -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)

pandas/tests/arrays/masked/test_arithmetic.py

@@ -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

pandas/tests/series/test_unary.py

@@ -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)