ENH: Enable .mode to sort with NA values (pandas-dev#60702)

Author: rhshadrach

pandas/core/algorithms.py

@@ -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}",

pandas/tests/frame/test_reductions.py

@@ -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):

pandas/tests/reductions/test_reductions.py

@@ -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):