use lexsort to sort values and mask, so that +/- inf values are placed before nan

Author: peterpanmj

pandas/_libs/algos_rank_helper.pxi.in

@@ -44,26 +44,24 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
 
  {{if dtype == 'object'}}
  ndarray sorted_data, values
- ndarray[np.uint8_t, cast=True] sorted_namask
- {{elif dtype != 'uint64'}}
- ndarray[{{ctype}}] sorted_data, values
- ndarray[np.uint8_t, cast=True] sorted_namask
  {{else}}
  ndarray[{{ctype}}] sorted_data, values
  {{endif}}
 
  ndarray[float64_t] ranks
  ndarray[int64_t] argsorted
+ ndarray[np.uint8_t, cast=True] sorted_mask
 
  {{if dtype == 'uint64'}}
  {{ctype}} val
  {{else}}
- {{ctype}} val, nan_value, isnan
+ {{ctype}} val, nan_value
  {{endif}}
 
  float64_t sum_ranks = 0
  int tiebreak = 0
  bint keep_na = 0
+ bint isnan
  float count = 0.0
  tiebreak = tiebreakers[ties_method]
 
@@ -95,14 +93,16 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
  {{endif}}
 
  np.putmask(values, mask, nan_value)
+ {{else}}
+ mask = np.zeros(shape=len(values), dtype=bool)
  {{endif}}
 
  n = len(values)
  ranks = np.empty(n, dtype='f8')
 
  {{if dtype == 'object'}}
  try:
- _as = values.argsort()
+ _as = np.lexsort(keys=(mask, values))
  except TypeError:
  if not retry:
  raise
@@ -116,40 +116,37 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
  {{else}}
  if tiebreak == TIEBREAK_FIRST:
  # need to use a stable sort here
- _as = values.argsort(kind='mergesort')
+ _as = np.lexsort(keys=(mask, values))
  if not ascending:
  tiebreak = TIEBREAK_FIRST_DESCENDING
  else:
- _as = values.argsort()
+ _as = np.lexsort(keys=(mask, values))
  {{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
- {{if dtype != 'uint64'}}
- sorted_namask = mask.take(_as)
- sorted_namask = sorted_namask.astype(np.bool)
- {{endif}}
+ sorted_mask = mask.take(_as)
  argsorted = _as.astype('i8')
 
  {{if dtype == 'object'}}
  for i in range(n):
  sum_ranks += i + 1
  dups += 1
- isnan = sorted_namask[i]
+ isnan = sorted_mask[i]
  val = util.get_value_at(sorted_data, i)
 
  if isnan and keep_na:
  ranks[argsorted[i]] = nan
- sum_ranks = dups = 0
  continue
 
  count += 1.0
 
  if (i == n - 1 or
- are_diff(util.get_value_at(sorted_data, i + 1), val)):
+ are_diff(util.get_value_at(sorted_data, i + 1), val) or
+ sorted_mask[i + 1]):
  if tiebreak == TIEBREAK_AVERAGE:
  for j in range(i - dups + 1, i + 1):
  ranks[argsorted[j]] = sum_ranks / dups
@@ -178,16 +175,15 @@ def rank_1d_{{dtype}}(object in_arr, ties_method='average', ascending=True,
  val = sorted_data[i]
 
  {{if dtype != 'uint64'}}
- isnan = sorted_namask[i]
- if  isnan and keep_na:
+ isnan = sorted_mask[i]
+ if isnan and keep_na:
  ranks[argsorted[i]] = nan
- sum_ranks = dups = 0
  continue
  {{endif}}
 
  count += 1.0
 
- if i == n - 1 or sorted_data[i + 1] != val:
+ if i == n - 1 or sorted_data[i + 1] != val or sorted_mask[i + 1]:
  if tiebreak == TIEBREAK_AVERAGE:
  for j in range(i - dups + 1, i + 1):
  ranks[argsorted[j]] = sum_ranks / dups