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