Add mask as return value of _values_for_argsort

Author: makbigc

pandas/core/arrays/base.py

@@ -361,7 +361,7 @@ def isna(self) -> ArrayLike:
         """
         raise AbstractMethodError(self)
 
-    def _values_for_argsort(self) -> np.ndarray:
+    def _values_for_argsort(self) -> Tuple[np.ndarray, np.ndarray]:
         """
         Return values for sorting.
 
@@ -376,7 +376,7 @@ def _values_for_argsort(self) -> np.ndarray:
         ExtensionArray.argsort
         """
         # Note: this is used in `ExtensionArray.argsort`.
-        return np.array(self)
+        return np.array(self), self.isna()
 
     def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         """
@@ -406,8 +406,29 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         # 1. _values_for_argsort : construct the values passed to np.argsort
         # 2. argsort : total control over sorting.
         ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
-        values = self._values_for_argsort()
-        result = np.argsort(values, kind=kind, **kwargs)
+        values, mask = self._values_for_argsort()
+
+        def permutation(mask):
+            # Return a permutation which maps the indices of the
+            # subarray without nan to the indices of the original array.
+            permu = np.arange(len(mask))
+            nan_loc = np.arange(len(mask))[mask]
+            offset = 0
+            for x in nan_loc:
+                permu[x - offset:] += 1
+                offset += 1
+            return permu
+
+        if mask.any():
+            notmask = ~mask
+            notnull = np.argsort(values[notmask], kind=kind, **kwargs)
+            permu = permutation(mask)
+            notnull = permu[notnull]
+            allnan = np.arange(len(self))[mask]
+            result = np.append(notnull, allnan)
+        else:
+            result = np.argsort(values, kind=kind, **kwargs)
+
         if not ascending:
             result = result[::-1]
         return result

pandas/core/arrays/categorical.py

@@ -1521,7 +1521,7 @@ def check_for_ordered(self, op):
                             "Categorical to an ordered one\n".format(op=op))
 
     def _values_for_argsort(self):
-        return self._codes.copy()
+        return self._codes.copy(), self.isna()
 
     def argsort(self, *args, **kwargs):
         # TODO(PY2): use correct signature

pandas/core/arrays/datetimelike.py

@@ -621,7 +621,7 @@ def _from_factorized(cls, values, original):
         return cls(values, dtype=original.dtype)
 
     def _values_for_argsort(self):
-        return self._data
+        return self._data, self._isnan
 
     # ------------------------------------------------------------------
     # Additional array methods

pandas/core/arrays/integer.py

@@ -512,7 +512,8 @@ def value_counts(self, dropna=True):
 
         return Series(array, index=index)
 
-    def _values_for_argsort(self) -> np.ndarray:
+    #def _values_for_argsort(self) -> np.ndarray:
+    def _values_for_argsort(self):
         """Return values for sorting.
 
         Returns
@@ -526,8 +527,9 @@ def _values_for_argsort(self) -> np.ndarray:
         ExtensionArray.argsort
         """
         data = self._data.copy()
+        mask = self._mask
         data[self._mask] = data.min() - 1
-        return data
+        return data, mask
 
     @classmethod
     def _create_comparison_method(cls, op):

pandas/core/arrays/numpy_.py

@@ -281,7 +281,7 @@ def copy(self, deep=False):
         return type(self)(self._ndarray.copy())
 
     def _values_for_argsort(self):
-        return self._ndarray
+        return self._ndarray, self.isna()
 
     def _values_for_factorize(self):
         return self._ndarray, -1

pandas/core/arrays/period.py

@@ -674,7 +674,7 @@ def _check_timedeltalike_freq_compat(self, other):
         _raise_on_incompatible(self, other)
 
     def _values_for_argsort(self):
-        return self._data
+        return self._data, self._isnan
 
 
 PeriodArray._add_comparison_ops()

pandas/tests/extension/base/methods.py

@@ -44,6 +44,11 @@ def test_argsort_missing(self, data_missing_for_sorting):
         expected = pd.Series(np.array([1, -1, 0], dtype=np.int64))
         self.assert_series_equal(result, expected)
 
+    def test_argsort_nan_loc(self, data_multiple_nan):
+        result = data_multiple_nan.argsort()
+        expected = np.array([3, 9, 7, 1, 0, 6, 2, 4, 5, 8])
+        tm.assert_numpy_array_equal(result, expected)
+
     @pytest.mark.parametrize('ascending', [True, False])
     def test_sort_values(self, data_for_sorting, ascending):
         ser = pd.Series(data_for_sorting)

pandas/tests/extension/decimal/test_decimal.py

@@ -46,6 +46,12 @@ def data_missing_for_sorting():
                          decimal.Decimal('0')])
 
 
+@pytest.fixture
+def data_multiple_nan():
+    return DecimalArray([decimal.Decimal(x) for x in
+                        [5, 4, np.nan, 1, np.nan, np.nan, 6, 3, np.nan, 2]])
+
+
 @pytest.fixture
 def na_cmp():
     return lambda x, y: x.is_nan() and y.is_nan()

pandas/tests/extension/json/array.py

@@ -171,7 +171,7 @@ def _concat_same_type(cls, to_concat):
         return cls(data)
 
     def _values_for_factorize(self):
-        frozen = self._values_for_argsort()
+        frozen, _ = self._values_for_argsort()
         if len(frozen) == 0:
             # _factorize_array expects 1-d array, this is a len-0 2-d array.
             frozen = frozen.ravel()
@@ -182,7 +182,7 @@ def _values_for_argsort(self):
         # If all the elemnts of self are the same size P, NumPy will
         # cast them to an (N, P) array, instead of an (N,) array of tuples.
         frozen = [()] + [tuple(x.items()) for x in self]
-        return np.array(frozen, dtype=object)[1:]
+        return np.array(frozen, dtype=object)[1:], self.isna()
 
 
 def make_data():

pandas/tests/extension/json/test_json.py

@@ -50,6 +50,12 @@ def data_missing_for_sorting():
     return JSONArray([{'b': 1}, {}, {'a': 4}])
 
 
+@pytest.fixture
+def data_multiple_nan():
+    return JSONArray([{'e': 5}, {'d': 4}, {}, {'a': 1}, {},
+                      {}, {'f': 6}, {'c': 3}, {}, {'b': 2}])
+
+
 @pytest.fixture
 def na_value(dtype):
     return dtype.na_value

pandas/tests/extension/test_categorical.py

@@ -68,6 +68,12 @@ def data_missing_for_sorting():
                        ordered=True)
 
 
+@pytest.fixture
+def data_multiple_nan():
+    return Categorical(['E', 'D', None, 'A', None,
+                        None, 'F', 'C', None, 'B'], ordered=True)
+
+
 @pytest.fixture
 def na_value():
     return np.nan

pandas/tests/extension/test_datetime.py

@@ -45,6 +45,15 @@ def data_missing_for_sorting(dtype):
                          dtype=dtype)
 
 
+@pytest.fixture
+def data_multiple_nan(dtype):
+    idx = pd.date_range(start='2000-01-01', end='2000-01-06')
+    return DatetimeArray(np.array([idx[4], idx[3], 'NaT', idx[0], 'NaT',
+                                   'NaT', idx[5], idx[2], 'NaT', idx[1]],
+                                  dtype='datetime64[ns]'),
+                         dtype=dtype)
+
+
 @pytest.fixture
 def data_for_grouping(dtype):
     """

pandas/tests/extension/test_integer.py

@@ -57,6 +57,12 @@ def data_for_sorting(dtype):
     return integer_array([1, 2, 0], dtype=dtype)
 
 
+@pytest.fixture
+def data_multiple_nan(dtype):
+    return integer_array([5, 4, np.nan, 1, np.nan,
+                          np.nan, 6, 3, np.nan, 2], dtype=dtype)
+
+
 @pytest.fixture
 def data_missing_for_sorting(dtype):
     return integer_array([1, np.nan, 0], dtype=dtype)

pandas/tests/extension/test_interval.py

@@ -57,6 +57,13 @@ def data_missing_for_sorting():
     return IntervalArray.from_tuples([(1, 2), None, (0, 1)])
 
 
+@pytest.fixture
+def data_multiple_nan():
+    return IntervalArray.from_tuples([(5, 6), (4, 5), None, (1, 2),
+                                      None, None, (6, 7), (3, 4),
+                                      None, (2, 3)])
+
+
 @pytest.fixture
 def na_value():
     return np.nan

pandas/tests/extension/test_numpy.py

@@ -101,6 +101,23 @@ def data_missing_for_sorting(allow_in_pandas, dtype):
     )
 
 
+@pytest.fixture
+def data_multiple_nan(allow_in_pandas, dtype):
+    """Length-10 array with a known sort order.
+
+    This should be three items [B, NA, A] with
+    A < B and NA missing.
+    """
+    if dtype.numpy_dtype == 'object':
+        return PandasArray(
+            np.array([(5,), (4,), np.nan, (1,), np.nan,
+                      np.nan, (6,), (3,), np.nan, (2,)])
+        )
+    return PandasArray(
+        np.array([5, 4, np.nan, 1, np.nan, np.nan, 6, 3, np.nan, 2])
+    )
+
+
 @pytest.fixture
 def data_for_grouping(allow_in_pandas, dtype):
     """Data for factorization, grouping, and unique tests.

pandas/tests/extension/test_period.py

@@ -40,6 +40,12 @@ def data_missing_for_sorting(dtype):
     return PeriodArray([2018, iNaT, 2017], freq=dtype.freq)
 
 
+@pytest.fixture
+def data_multiple_nan(dtype):
+    return PeriodArray([2005, 2004, iNaT, 2001, iNaT, iNaT,
+                        2006, 2003, iNaT, 2002], freq=dtype.freq)
+
+
 @pytest.fixture
 def data_for_grouping(dtype):
     B = 2018

pandas/tests/extension/test_sparse.py

@@ -65,6 +65,12 @@ def data_missing_for_sorting(request):
     return SparseArray([2, np.nan, 1], fill_value=request.param)
 
 
+@pytest.fixture(params=[0, np.nan])
+def data_multiple_nan(request):
+    return SparseArray([5, 4, np.nan, 1, np.nan,
+                        np.nan, 6, 3, np.nan, 2], fill_value=request.param)
+
+
 @pytest.fixture
 def na_value():
     return np.nan