Skip to content

Commit

Permalink
ENH: IntegerArray bitwise operations (pandas-dev#46104)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Feb 27, 2022
1 parent de98025 commit ac2f746
Show file tree
Hide file tree
Showing 5 changed files with 46 additions and 5 deletions.
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.5.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ Other enhancements
- :meth:`.GroupBy.min` and :meth:`.GroupBy.max` now supports `Numba <https://numba.pydata.org/>`_ 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`)
-

.. ---------------------------------------------------------------------------
Expand Down
8 changes: 4 additions & 4 deletions pandas/core/arrays/boolean.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
2 changes: 2 additions & 0 deletions pandas/core/arrays/masked.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
12 changes: 12 additions & 0 deletions pandas/tests/arrays/floating/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
27 changes: 27 additions & 0 deletions pandas/tests/arrays/integer/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -323,3 +323,30 @@ def test_values_multiplying_large_series_by_NA():
expected = pd.Series([pd.NA] * 10001)

tm.assert_series_equal(result, expected)


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

0 comments on commit ac2f746

Please sign in to comment.