102
102
Categorical ,
103
103
DataFrame ,
104
104
Index ,
105
- MultiIndex ,
106
105
Series ,
107
106
)
108
107
from pandas .core .arrays import (
@@ -1768,7 +1767,7 @@ def safe_sort(
1768
1767
na_sentinel : int | None = - 1 ,
1769
1768
assume_unique : bool = False ,
1770
1769
verify : bool = True ,
1771
- ) -> np . ndarray | MultiIndex | tuple [np . ndarray | MultiIndex , np .ndarray ]:
1770
+ ) -> AnyArrayLike | tuple [AnyArrayLike , np .ndarray ]:
1772
1771
"""
1773
1772
Sort ``values`` and reorder corresponding ``codes``.
1774
1773
@@ -1797,7 +1796,7 @@ def safe_sort(
1797
1796
1798
1797
Returns
1799
1798
-------
1800
- ordered : ndarray or MultiIndex
1799
+ ordered : AnyArrayLike
1801
1800
Sorted ``values``
1802
1801
new_codes : ndarray
1803
1802
Reordered ``codes``; returned when ``codes`` is not None.
@@ -1816,7 +1815,7 @@ def safe_sort(
1816
1815
"Only list-like objects are allowed to be passed to safe_sort as values"
1817
1816
)
1818
1817
1819
- if not isinstance (values , ( np . ndarray , ABCExtensionArray , ABCMultiIndex ) ):
1818
+ if not is_array_like (values ):
1820
1819
# don't convert to string types
1821
1820
dtype , _ = infer_dtype_from_array (values )
1822
1821
# error: Argument "dtype" to "asarray" has incompatible type "Union[dtype[Any],
@@ -1826,7 +1825,7 @@ def safe_sort(
1826
1825
values = np .asarray (values , dtype = dtype ) # type: ignore[arg-type]
1827
1826
1828
1827
sorter = None
1829
- ordered : np . ndarray | MultiIndex
1828
+ ordered : AnyArrayLike
1830
1829
1831
1830
if (
1832
1831
not is_extension_array_dtype (values )
@@ -1894,15 +1893,19 @@ def safe_sort(
1894
1893
return ordered , ensure_platform_int (new_codes )
1895
1894
1896
1895
1897
- def _sort_mixed (values ) -> np . ndarray :
1896
+ def _sort_mixed (values ) -> AnyArrayLike :
1898
1897
"""order ints before strings before nulls in 1d arrays"""
1899
1898
str_pos = np .array ([isinstance (x , str ) for x in values ], dtype = bool )
1900
1899
null_pos = np .array ([isna (x ) for x in values ], dtype = bool )
1901
- nums = np .sort (values [~ str_pos & ~ null_pos ])
1902
- strs = np .sort (values [str_pos ])
1903
- return np .concatenate (
1904
- [nums , np .asarray (strs , dtype = object ), np .array (values [null_pos ])]
1905
- )
1900
+ num_pos = ~ str_pos & ~ null_pos
1901
+ str_argsort = np .argsort (values [str_pos ])
1902
+ num_argsort = np .argsort (values [num_pos ])
1903
+ # convert boolean arrays to positional indices, then order by underlying values
1904
+ str_locs = str_pos .nonzero ()[0 ].take (str_argsort )
1905
+ num_locs = num_pos .nonzero ()[0 ].take (num_argsort )
1906
+ null_locs = null_pos .nonzero ()[0 ]
1907
+ locs = np .concatenate ([num_locs , str_locs , null_locs ])
1908
+ return values .take (locs )
1906
1909
1907
1910
1908
1911
def _sort_tuples (values : np .ndarray ) -> np .ndarray :
0 commit comments