|
21 | 21 | BytesIO,
|
22 | 22 | StringIO,
|
23 | 23 | )
|
| 24 | +import operator |
24 | 25 | import pickle
|
25 | 26 | import re
|
26 | 27 |
|
@@ -1216,7 +1217,7 @@ def test_add_series_with_extension_array(self, data, request):
|
1216 | 1217 |
|
1217 | 1218 |
|
1218 | 1219 | class TestBaseComparisonOps(base.BaseComparisonOpsTests):
|
1219 |
| - def test_compare_array(self, data, comparison_op, na_value, request): |
| 1220 | + def test_compare_array(self, data, comparison_op, na_value): |
1220 | 1221 | ser = pd.Series(data)
|
1221 | 1222 | # pd.Series([ser.iloc[0]] * len(ser)) may not return ArrowExtensionArray
|
1222 | 1223 | # since ser.iloc[0] is a python scalar
|
@@ -1255,6 +1256,20 @@ def test_invalid_other_comp(self, data, comparison_op):
|
1255 | 1256 | ):
|
1256 | 1257 | comparison_op(data, object())
|
1257 | 1258 |
|
| 1259 | + @pytest.mark.parametrize("masked_dtype", ["boolean", "Int64", "Float64"]) |
| 1260 | + def test_comp_masked_numpy(self, masked_dtype, comparison_op): |
| 1261 | + # GH 52625 |
| 1262 | + data = [1, 0, None] |
| 1263 | + ser_masked = pd.Series(data, dtype=masked_dtype) |
| 1264 | + ser_pa = pd.Series(data, dtype=f"{masked_dtype.lower()}[pyarrow]") |
| 1265 | + result = comparison_op(ser_pa, ser_masked) |
| 1266 | + if comparison_op in [operator.lt, operator.gt, operator.ne]: |
| 1267 | + exp = [False, False, None] |
| 1268 | + else: |
| 1269 | + exp = [True, True, None] |
| 1270 | + expected = pd.Series(exp, dtype=ArrowDtype(pa.bool_())) |
| 1271 | + tm.assert_series_equal(result, expected) |
| 1272 | + |
1258 | 1273 |
|
1259 | 1274 | class TestLogicalOps:
|
1260 | 1275 | """Various Series and DataFrame logical ops methods."""
|
@@ -1399,6 +1414,23 @@ def test_kleene_xor_scalar(self, other, expected):
|
1399 | 1414 | a, pd.Series([True, False, None], dtype="boolean[pyarrow]")
|
1400 | 1415 | )
|
1401 | 1416 |
|
| 1417 | + @pytest.mark.parametrize( |
| 1418 | + "op, exp", |
| 1419 | + [ |
| 1420 | + ["__and__", True], |
| 1421 | + ["__or__", True], |
| 1422 | + ["__xor__", False], |
| 1423 | + ], |
| 1424 | + ) |
| 1425 | + def test_logical_masked_numpy(self, op, exp): |
| 1426 | + # GH 52625 |
| 1427 | + data = [True, False, None] |
| 1428 | + ser_masked = pd.Series(data, dtype="boolean") |
| 1429 | + ser_pa = pd.Series(data, dtype="boolean[pyarrow]") |
| 1430 | + result = getattr(ser_pa, op)(ser_masked) |
| 1431 | + expected = pd.Series([exp, False, None], dtype=ArrowDtype(pa.bool_())) |
| 1432 | + tm.assert_series_equal(result, expected) |
| 1433 | + |
1402 | 1434 |
|
1403 | 1435 | def test_arrowdtype_construct_from_string_type_with_unsupported_parameters():
|
1404 | 1436 | with pytest.raises(NotImplementedError, match="Passing pyarrow type"):
|
|
0 commit comments