-
-
Notifications
You must be signed in to change notification settings - Fork 18.6k
Fix #60766:.map,.apply would convert element type for extension array #61396
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
f8153f3
bf6aaef
e8edcea
d845306
ef9812e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
import pandas as pd | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can't we have this in an existing tests file? |
||
|
||
|
||
def test_basemaskedarray_map(): | ||
for dtype, data, expected_data in [ | ||
("Int32", [1, 2, None, 4], [2, 3, pd.NA, 5]), | ||
]: | ||
Comment on lines
+5
to
+7
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why to loop over a single value? |
||
s = pd.Series(data, dtype=dtype) | ||
|
||
def transform(x): | ||
if x is None: | ||
return x | ||
return x + 1 | ||
|
||
result = s.map(transform) | ||
expected = pd.Series(expected_data, dtype=result.dtype) | ||
|
||
assert result.tolist() == expected.tolist() | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Any reason why not to compare the Series directly instead of converting them to lists? You can check other tests, there is a function assert_series_equal in case you're not aware. |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -171,6 +171,12 @@ class TestMaskedArrays(base.ExtensionTests): | |
@pytest.mark.parametrize("na_action", [None, "ignore"]) | ||
def test_map(self, data_missing, na_action): | ||
result = data_missing.map(lambda x: x, na_action=na_action) | ||
if data_missing.dtype.kind != "b": | ||
for i in range(len(result)): | ||
if result[i] is pd.NA: | ||
result[i] = "nan" | ||
result = result.astype("float64") | ||
|
||
if data_missing.dtype == Float32Dtype(): | ||
# map roundtrips through objects, which converts to float64 | ||
expected = data_missing.to_numpy(dtype="float64", na_value=np.nan) | ||
|
@@ -181,10 +187,15 @@ def test_map(self, data_missing, na_action): | |
def test_map_na_action_ignore(self, data_missing_for_sorting): | ||
zero = data_missing_for_sorting[2] | ||
result = data_missing_for_sorting.map(lambda x: zero, na_action="ignore") | ||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Better to avoid this unrelated changes |
||
if data_missing_for_sorting.dtype.kind == "b": | ||
expected = np.array([False, pd.NA, False], dtype=object) | ||
else: | ||
expected = np.array([zero, np.nan, zero]) | ||
for i in range(len(result)): | ||
if result[i] is pd.NA: | ||
result[i] = "nan" | ||
result = result.astype("float64") | ||
Comment on lines
+195
to
+198
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Are you changing the result to match the expected value? Why not change the expected value if what you are proposing here is that? |
||
tm.assert_numpy_array_equal(result, expected) | ||
|
||
def _get_expected_exception(self, op_name, obj, other): | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you remove this new line please?