Skip to content

Commit

Permalink
ENH: Enable .mode to sort with NA values (#60702)
Browse files Browse the repository at this point in the history
  • Loading branch information
rhshadrach authored Jan 13, 2025
1 parent 221ad46 commit 1708e90
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 26 deletions.
2 changes: 1 addition & 1 deletion pandas/core/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -1012,7 +1012,7 @@ def mode(
return npresult, res_mask # type: ignore[return-value]

try:
npresult = np.sort(npresult)
npresult = safe_sort(npresult)
except TypeError as err:
warnings.warn(
f"Unable to sort modes: {err}",
Expand Down
17 changes: 2 additions & 15 deletions pandas/tests/frame/test_reductions.py
Original file line number Diff line number Diff line change
Expand Up @@ -672,23 +672,10 @@ def test_mode_dropna(self, dropna, expected):
expected = DataFrame(expected)
tm.assert_frame_equal(result, expected)

def test_mode_sortwarning(self, using_infer_string):
# Check for the warning that is raised when the mode
# results cannot be sorted

def test_mode_sort_with_na(self, using_infer_string):
df = DataFrame({"A": [np.nan, np.nan, "a", "a"]})
expected = DataFrame({"A": ["a", np.nan]})

# TODO(infer_string) avoid this UserWarning for python storage
warning = (
None
if using_infer_string and df.A.dtype.storage == "pyarrow"
else UserWarning
)
with tm.assert_produces_warning(warning, match="Unable to sort modes"):
result = df.mode(dropna=False)
result = result.sort_values(by="A").reset_index(drop=True)

result = df.mode(dropna=False)
tm.assert_frame_equal(result, expected)

def test_mode_empty_df(self):
Expand Down
13 changes: 3 additions & 10 deletions pandas/tests/reductions/test_reductions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1607,17 +1607,10 @@ def test_mode_intoverflow(self, dropna, expected1, expected2):
expected2 = Series(expected2, dtype=np.uint64)
tm.assert_series_equal(result, expected2)

def test_mode_sortwarning(self):
# Check for the warning that is raised when the mode
# results cannot be sorted

expected = Series(["foo", np.nan], dtype=object)
def test_mode_sort_with_na(self):
s = Series([1, "foo", "foo", np.nan, np.nan])

with tm.assert_produces_warning(UserWarning, match="Unable to sort modes"):
result = s.mode(dropna=False)
result = result.sort_values().reset_index(drop=True)

expected = Series(["foo", np.nan], dtype=object)
result = s.mode(dropna=False)
tm.assert_series_equal(result, expected)

def test_mode_boolean_with_na(self):
Expand Down

0 comments on commit 1708e90

Please sign in to comment.