From e260bf9970269a689afa9d1027ee36a9bb8ce383 Mon Sep 17 00:00:00 2001 From: Brock Date: Mon, 21 Feb 2022 12:34:13 -0800 Subject: [PATCH] ENH: IntegerArray bitwise operations --- doc/source/whatsnew/v1.5.0.rst | 2 +- pandas/core/arrays/boolean.py | 8 +++--- pandas/core/arrays/masked.py | 2 ++ .../tests/arrays/floating/test_arithmetic.py | 12 +++++++++ .../tests/arrays/integer/test_arithmetic.py | 27 +++++++++++++++++++ 5 files changed, 46 insertions(+), 5 deletions(-) diff --git a/doc/source/whatsnew/v1.5.0.rst b/doc/source/whatsnew/v1.5.0.rst index e3f772ac026ab..5f22592b412fa 100644 --- a/doc/source/whatsnew/v1.5.0.rst +++ b/doc/source/whatsnew/v1.5.0.rst @@ -39,7 +39,7 @@ Other enhancements - :meth:`.GroupBy.min` and :meth:`.GroupBy.max` now supports `Numba `_ execution with the ``engine`` keyword (:issue:`45428`) - Implemented a ``bool``-dtype :class:`Index`, passing a bool-dtype array-like to ``pd.Index`` will now retain ``bool`` dtype instead of casting to ``object`` (:issue:`45061`) - Implemented a complex-dtype :class:`Index`, passing a complex-dtype array-like to ``pd.Index`` will now retain complex dtype instead of casting to ``object`` (:issue:`45845`) - +- :class:`Series` and :class:`DataFrame` with ``IntegerDtype`` now supports bitwise operations (:issue:`34463`) - .. --------------------------------------------------------------------------- diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py index 10cb7683e57d4..10469f2aef9ea 100644 --- a/pandas/core/arrays/boolean.py +++ b/pandas/core/arrays/boolean.py @@ -375,9 +375,9 @@ def _logical_method(self, other, op): result, mask = ops.kleene_or(self._data, other, self._mask, mask) elif op.__name__ in {"and_", "rand_"}: result, mask = ops.kleene_and(self._data, other, self._mask, mask) - elif op.__name__ in {"xor", "rxor"}: + else: + # i.e. xor, rxor result, mask = ops.kleene_xor(self._data, other, self._mask, mask) - # error: Argument 2 to "BooleanArray" has incompatible type "Optional[Any]"; - # expected "ndarray" - return BooleanArray(result, mask) # type: ignore[arg-type] + # i.e. BooleanArray + return self._maybe_mask_result(result, mask) diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 21f44dbc6a1cd..8ba213f0d9a41 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -675,6 +675,8 @@ def _arith_method(self, other, op): return self._maybe_mask_result(result, mask) + _logical_method = _arith_method + def _cmp_method(self, other, op) -> BooleanArray: from pandas.core.arrays import BooleanArray diff --git a/pandas/tests/arrays/floating/test_arithmetic.py b/pandas/tests/arrays/floating/test_arithmetic.py index d0f25d2cd5026..ec7419d6346a8 100644 --- a/pandas/tests/arrays/floating/test_arithmetic.py +++ b/pandas/tests/arrays/floating/test_arithmetic.py @@ -218,3 +218,15 @@ def test_unary_float_operators(float_ea_dtype, source, neg_target, abs_target): tm.assert_extension_array_equal(pos_result, arr) assert not tm.shares_memory(pos_result, arr) tm.assert_extension_array_equal(abs_result, abs_target) + + +def test_bitwise(dtype): + left = pd.array([1, None, 3, 4], dtype=dtype) + right = pd.array([None, 3, 5, 4], dtype=dtype) + + with pytest.raises(TypeError, match="unsupported operand type"): + left | right + with pytest.raises(TypeError, match="unsupported operand type"): + left & right + with pytest.raises(TypeError, match="unsupported operand type"): + left ^ right diff --git a/pandas/tests/arrays/integer/test_arithmetic.py b/pandas/tests/arrays/integer/test_arithmetic.py index d14a6d089ecd6..b7ab00468048d 100644 --- a/pandas/tests/arrays/integer/test_arithmetic.py +++ b/pandas/tests/arrays/integer/test_arithmetic.py @@ -314,3 +314,30 @@ def test_unary_int_operators(any_signed_int_ea_dtype, source, neg_target, abs_ta tm.assert_extension_array_equal(pos_result, arr) assert not tm.shares_memory(pos_result, arr) tm.assert_extension_array_equal(abs_result, abs_target) + + +def test_bitwise(dtype): + left = pd.array([1, None, 3, 4], dtype=dtype) + right = pd.array([None, 3, 5, 4], dtype=dtype) + + result = left | right + expected = pd.array([None, None, 3 | 5, 4 | 4], dtype=dtype) + tm.assert_extension_array_equal(result, expected) + + result = left & right + expected = pd.array([None, None, 3 & 5, 4 & 4], dtype=dtype) + tm.assert_extension_array_equal(result, expected) + + result = left ^ right + expected = pd.array([None, None, 3 ^ 5, 4 ^ 4], dtype=dtype) + tm.assert_extension_array_equal(result, expected) + + # TODO: desired behavior when operating with boolean? defer? + + floats = right.astype("Float64") + with pytest.raises(TypeError, match="unsupported operand type"): + left | floats + with pytest.raises(TypeError, match="unsupported operand type"): + left & floats + with pytest.raises(TypeError, match="unsupported operand type"): + left ^ floats