Modify nargsort

Author: makbigc

pandas/core/sorting.py

@@ -9,8 +9,10 @@
 from pandas.core.dtypes.cast import infer_dtype_from_array
 from pandas.core.dtypes.common import (
     ensure_int64, ensure_platform_int, is_categorical_dtype,
-    is_extension_array_dtype, is_list_like)
+    is_extension_array_dtype, is_list_like, is_sparse)
 from pandas.core.dtypes.missing import isna
+from pandas.core.dtypes.generic import ABCIndexClass
+
 
 import pandas.core.algorithms as algorithms
 
@@ -239,12 +241,12 @@ def nargsort(items, kind='quicksort', ascending=True, na_position='last'):
     GH #6399, #5231
     """
 
+    mask = isna(items)
     # specially handle Categorical
     if is_categorical_dtype(items):
         if na_position not in {'first', 'last'}:
             raise ValueError('invalid na_position: {!r}'.format(na_position))
 
-        mask = isna(items)
         cnt_null = mask.sum()
         sorted_idx = items.argsort(ascending=ascending, kind=kind)
         if ascending and na_position == 'last':
@@ -255,15 +257,19 @@ def nargsort(items, kind='quicksort', ascending=True, na_position='last'):
             sorted_idx = np.roll(sorted_idx, cnt_null)
         return sorted_idx
 
-    with warnings.catch_warnings():
-        # https://github.com/pandas-dev/pandas/issues/25439
-        # can be removed once ExtensionArrays are properly handled by nargsort
-        warnings.filterwarnings(
-            "ignore", category=FutureWarning,
-            message="Converting timezone-aware DatetimeArray to")
+    if (not isinstance(items, ABCIndexClass)
+            and is_extension_array_dtype(items)):
+
+        if is_sparse(items):
+            # The conversion to np.ndarray is the fact that
+            # SparseArray.isna() is also a SparseArray
+            mask = np.array(isna(items))
+
+        items = items._values_for_argsort()
+    else:
         items = np.asanyarray(items)
+
     idx = np.arange(len(items))
-    mask = isna(items)
     non_nans = items[~mask]
     non_nan_idx = idx[~mask]
     nan_idx = np.nonzero(mask)[0]