change the lexsort logic, first by mask/~mask then by values to ensure

to keep nans at front or at the end

Author: peterpanmj

pandas/_libs/algos.pyx

@@ -68,15 +68,17 @@ class Infinity(object):
     __le__ = lambda self, other: isinstance(other, Infinity)
     __eq__ = lambda self, other: isinstance(other, Infinity)
     __ne__ = lambda self, other: not isinstance(other, Infinity)
-    __gt__ = lambda self, other: not isinstance(other, Infinity)
-    __ge__ = lambda self, other: True
+    __gt__ = lambda self, other: (not isinstance(other, Infinity) and
+                                  not lib.checknull(other))
+    __ge__ = lambda self, other: not lib.checknull(other)
 
 
 class NegInfinity(object):
     """ provide a negative Infinity comparision method for ranking """
 
-    __lt__ = lambda self, other: not isinstance(other, NegInfinity)
-    __le__ = lambda self, other: True
+    __lt__ = lambda self, other: (not isinstance(other, NegInfinity) and
+                                  not lib.checknull(other))
+    __le__ = lambda self, other: not lib.checknull(other)
     __eq__ = lambda self, other: isinstance(other, NegInfinity)
     __ne__ = lambda self, other: not isinstance(other, NegInfinity)
     __gt__ = lambda self, other: False

pandas/_libs/algos_rank_helper.pxi.in

@@ -27,7 +27,7 @@ dtypes = [('object', 'object', 'Infinity()', 'NegInfinity()'),
 {{if dtype == 'object'}}
 
 
-def rank_1d_{{dtype}}(object in_arr, bint retry=1, ties_method='average',
+def rank_1d_{{dtype}}(object in_arr, ties_method='average',
                       ascending=True, na_option='keep', pct=False):
 {{else}}
 
@@ -40,7 +40,7 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
     """
 
     cdef:
-        Py_ssize_t i, j, n, dups = 0, total_tie_count = 0
+        Py_ssize_t i, j, n, dups = 0, total_tie_count = 0, non_na_idx = 0
 
         {{if dtype == 'object'}}
         ndarray sorted_data, values
@@ -78,12 +78,6 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
 
     keep_na = na_option == 'keep'
 
-    {{if dtype != 'uint64'}}
-    if ascending ^ (na_option == 'top'):
-        nan_value = {{pos_nan_value}}
-    else:
-        nan_value = {{neg_nan_value}}
-
     {{if dtype == 'object'}}
     mask = lib.isnaobj(values)
     {{elif dtype == 'float64'}}
@@ -92,49 +86,52 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
     mask = values == iNaT
     {{endif}}
 
+    # doulbe sort first by mask and then by values to ensure nan values are
+    # either at the beginning or the end. mask/(~mask) controls padding at
+    # tail or the head
+    {{if dtype != 'uint64'}}
+    if ascending ^ (na_option == 'top'):
+        nan_value = {{pos_nan_value}}
+        order = (values, mask)
+    else:
+        nan_value = {{neg_nan_value}}
+        order = (values, ~mask)
     np.putmask(values, mask, nan_value)
     {{else}}
     mask = np.zeros(shape=len(values), dtype=bool)
+    order = (values, mask)
     {{endif}}
 
     n = len(values)
     ranks = np.empty(n, dtype='f8')
 
     {{if dtype == 'object'}}
+
     try:
-        # lexsort on object array will raise TypeError for numpy version 
-        # earlier than 1.11.0
-        _as = np.lexsort(keys=(mask, values))
+        _as = np.lexsort(keys=order)
     except TypeError:
-        try:
-            _dt = [('values', 'O'), ('mask', '?')]
-            _values = np.asarray(list(zip(values, mask)), dtype=_dt)
-            _as = np.argsort(_values, kind='mergesort', order=('values', 'mask'))
-        except TypeError:
-            if not retry:
-                raise
-            valid_locs = (~mask).nonzero()[0]
-            ranks.put(valid_locs, rank_1d_object(values.take(valid_locs), 0,
-                                                 ties_method=ties_method,
-                                                 ascending=ascending))
-            np.putmask(ranks, mask, np.nan)
-            return ranks
+        # lexsort on object array will raise TypeError for numpy version
+        # earlier than 1.11.0. Use argsort with order argument instead.
+        _dt = [('values', 'O'), ('mask', '?')]
+        _values = np.asarray(list(zip(order[0], order[1])), dtype=_dt)
+        _as = np.argsort(_values, kind='mergesort', order=('mask', 'values'))
     {{else}}
     if tiebreak == TIEBREAK_FIRST:
         # need to use a stable sort here
-        _as = np.lexsort(keys=(mask, values))
+        _as = np.lexsort(keys=order)
         if not ascending:
             tiebreak = TIEBREAK_FIRST_DESCENDING
     else:
-        _as = np.lexsort(keys=(mask, values))
+        _as = np.lexsort(keys=order)
     {{endif}}
 
     if not ascending:
         _as = _as[::-1]
 
     sorted_data = values.take(_as)
-    # need to distinguish between pos/neg nan and real nan when keep_na is true
     sorted_mask = mask.take(_as)
+    _indices = order[1].take(_as).nonzero()[0]
+    non_na_idx = _indices[0] if len(_indices) > 0 else -1
     argsorted = _as.astype('i8')
 
     {{if dtype == 'object'}}
@@ -147,12 +144,11 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
         if isnan and keep_na:
             ranks[argsorted[i]] = nan
             continue
-
         count += 1.0
 
         if (i == n - 1 or
                 are_diff(util.get_value_at(sorted_data, i + 1), val) or
-                sorted_mask[i + 1]):
+                i == non_na_idx - 1):
             if tiebreak == TIEBREAK_AVERAGE:
                 for j in range(i - dups + 1, i + 1):
                     ranks[argsorted[j]] = sum_ranks / dups
@@ -177,7 +173,6 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
         for i in range(n):
             sum_ranks += i + 1
             dups += 1
-
             val = sorted_data[i]
 
             {{if dtype != 'uint64'}}
@@ -189,7 +184,8 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
 
             count += 1.0
 
-            if i == n - 1 or sorted_data[i + 1] != val or sorted_mask[i + 1]:
+            if (i == n - 1 or sorted_data[i + 1] != val or
+                i == non_na_idx - 1):
                 if tiebreak == TIEBREAK_AVERAGE:
                     for j in range(i - dups + 1, i + 1):
                         ranks[argsorted[j]] = sum_ranks / dups

pandas/tests/series/test_rank.py

@@ -224,13 +224,18 @@ def test_rank_inf(self, contents, dtype):
             'int64': iNaT,
             'object': None
         }
+        # Insert nans at random positions if underlying dtype has missing
+        # value. Then adjust the expected order by adding nans accordingly
+        # This is for testing whether rank calculation is affected
+        # when values are interwined with nan values.
         values = np.array(contents, dtype=dtype)
-        nan_indices = np.random.choice(range(len(values)), 3)
         exp_order = np.array(range(len(values)), dtype='float64') + 1.0
         if dtype in dtype_na_map:
             na_value = dtype_na_map[dtype]
+            nan_indices = np.random.choice(range(len(values)), 5)
             values = np.insert(values, nan_indices, na_value)
             exp_order = np.insert(exp_order, nan_indices, np.nan)
+        # shuffle the testing array and expected results in the same way
         random_order = np.random.permutation(len(values))
         iseries = Series(values[random_order])
         exp = Series(exp_order[random_order], dtype='float64')
@@ -254,6 +259,39 @@ def _check(s, expected, method='average'):
             series = s if dtype is None else s.astype(dtype)
             _check(series, results[method], method=method)
 
+    def test_rank_tie_methods_on_infs_nans(self):
+        dtypes = [('object', None, Infinity(), NegInfinity()),
+                  ('float64', np.nan, np.inf, -np.inf)]
+        chunk = 3
+        disabled = set([('object', 'first')])
+
+        def _check(s, expected, method='average', na_option='keep'):
+            result = s.rank(method=method, na_option=na_option)
+            tm.assert_series_equal(result, Series(expected, dtype='float64'))
+
+        exp_ranks = {
+            'average': ([2, 2, 2], [5, 5, 5], [8, 8, 8]),
+            'min': ([1, 1, 1], [4, 4, 4], [7, 7, 7]),
+            'max': ([3, 3, 3], [6, 6, 6], [9, 9, 9]),
+            'first': ([1, 2, 3], [4, 5, 6], [7, 8, 9]),
+            'dense': ([1, 1, 1], [2, 2, 2], [3, 3, 3])
+        }
+        na_options = ('top', 'bottom', 'keep')
+        for dtype, na_value, pos_inf, neg_inf in dtypes:
+            in_arr = [neg_inf] * chunk + [na_value] * chunk + [pos_inf] * chunk
+            iseries = Series(in_arr, dtype=dtype)
+            for method, na_opt in product(exp_ranks.keys(), na_options):
+                ranks = exp_ranks[method]
+                if (dtype, method) in disabled:
+                    continue
+                if na_opt == 'top':
+                    order = ranks[1] + ranks[0] + ranks[2]
+                elif na_opt == 'bottom':
+                    order = ranks[0] + ranks[2] + ranks[1]
+                else:
+                    order = ranks[0] + [np.nan] * chunk + ranks[1]
+                _check(iseries, order, method, na_opt)
+
     def test_rank_methods_series(self):
         pytest.importorskip('scipy.stats.special')
         rankdata = pytest.importorskip('scipy.stats.rankdata')

pandas/tests/test_algos.py

@@ -1271,44 +1271,48 @@ def test_infinity_sort():
     NegInf = libalgos.NegInfinity()
 
     ref_nums = [NegInf, float("-inf"), -1e100, 0, 1e100, float("inf"), Inf]
-    ref_objects = [NegInf, 'A', 'Ac', 'B', 'C', 'D', 'a', 'z', Inf]
 
     assert all(Inf >= x for x in ref_nums)
     assert all(Inf > x or x is Inf for x in ref_nums)
     assert Inf >= Inf and Inf == Inf
-    assert Inf >= libalgos.Infinity() and Inf == libalgos.Infinity()
     assert not Inf < Inf and not Inf > Inf
-    assert not Inf < libalgos.Infinity() and not Inf > libalgos.Infinity()
     assert libalgos.Infinity() == libalgos.Infinity()
     assert not libalgos.Infinity() != libalgos.Infinity()
 
     assert all(NegInf <= x for x in ref_nums)
     assert all(NegInf < x or x is NegInf for x in ref_nums)
     assert NegInf <= NegInf and NegInf == NegInf
-    assert (NegInf <= libalgos.NegInfinity() and
-            NegInf == libalgos.NegInfinity())
     assert not NegInf < NegInf and not NegInf > NegInf
-    assert (not NegInf < libalgos.NegInfinity() and not
-            NegInf > libalgos.NegInfinity())
     assert libalgos.NegInfinity() == libalgos.NegInfinity()
     assert not libalgos.NegInfinity() != libalgos.NegInfinity()
 
-    assert Inf > np.nan and Inf >= np.nan and Inf != np.nan
-    assert not Inf < np.nan and not Inf <= np.nan
-    assert NegInf < np.nan and NegInf <= np.nan and NegInf != np.nan
-    assert (not NegInf > np.nan and not NegInf >= np.nan and not
-            NegInf == np.nan)
-
     for perm in permutations(ref_nums):
         assert sorted(perm) == ref_nums
-    for perm in permutations(ref_objects):
-        assert sorted(perm) == ref_objects
 
     # smoke tests
     np.array([libalgos.Infinity()] * 32).argsort()
     np.array([libalgos.NegInfinity()] * 32).argsort()
 
 
+def test_infinity_against_nan():
+    Inf = libalgos.Infinity()
+    NegInf = libalgos.NegInfinity()
+
+    assert not Inf > np.nan
+    assert not Inf >= np.nan
+    assert not Inf < np.nan
+    assert not Inf <= np.nan
+    assert not Inf == np.nan
+    assert Inf != np.nan
+
+    assert not NegInf > np.nan
+    assert not NegInf >= np.nan
+    assert not NegInf < np.nan
+    assert not NegInf <= np.nan
+    assert not NegInf == np.nan
+    assert NegInf != np.nan
+
+
 def test_ensure_platform_int():
     arr = np.arange(100, dtype=np.intp)