Simplify handling of nan values in groupby rank

Author: peterpanmj

doc/source/whatsnew/v0.24.0.txt

@@ -63,7 +63,7 @@ Removal of prior version deprecations/changes
 Performance Improvements
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
-- Improved performance of :func:`pandas.core.groupby.GroupBy.rank` (:issue:`21237`)
+- Improved performance of :func:`pandas.core.groupby.GroupBy.rank` when dealing with tied rankings (:issue:`21237`)
 -
 -
 

pandas/_libs/groupby_helper.pxi.in

@@ -514,30 +514,6 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
             # Used to calculate tiebreakers
             dups += 1
             sum_ranks += i - grp_start + 1
-            # if keep_na, check for missing values and assign back
-            # to the result where appropriate
-            if keep_na and mask[_as[i]]:
-                grp_na_count += 1
-                out[_as[i], 0] = nan
-                grp_vals_seen = i + 1
-                # when label transition happens, update grp_size to keep track
-                # of number of nans encountered and increment grp_tie_count
-                if (i== N-1) or (labels[_as[i]] != labels[_as[i+1]]):
-                    grp_tie_count +=1
-                    if tiebreak != TIEBREAK_DENSE:
-                        for j in range(grp_start, i + 1):
-                            grp_sizes[_as[j], 0] = (i - grp_start + 1 -
-                                                    grp_na_count)
-                    else:
-                        for j in range(grp_start, i + 1):
-                            grp_sizes[_as[j], 0] = (grp_tie_count -
-                                                    (grp_na_count > 0))
-                    dups = sum_ranks = 0
-                    grp_na_count = 0
-                    grp_tie_count = 0
-                    grp_start = i + 1
-                    grp_vals_seen = 1
-                continue
 
             # Update out only when there is a transition of values or labels.
             # When a new value or group is encountered, go back #dups steps(
@@ -548,7 +524,13 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
                     (masked_vals[_as[i]] != masked_vals[_as[i+1]]) or
                     (mask[_as[i]] ^ mask[_as[i+1]]) or
                     (labels[_as[i]] != labels[_as[i+1]])):
-                if tiebreak == TIEBREAK_AVERAGE:
+                # if keep_na, check for missing values and assign back
+                # to the result where appropriate
+                if keep_na and mask[_as[i]]:
+                    for j in range(i - dups + 1, i + 1):
+                        out[_as[j], 0] = nan
+                        grp_na_count = dups
+                elif tiebreak == TIEBREAK_AVERAGE:
                     for j in range(i - dups + 1, i + 1):
                         out[_as[j], 0] = sum_ranks / <float64_t>dups
                 elif tiebreak == TIEBREAK_MIN:
@@ -570,14 +552,14 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
                 # look forward to the next value (using the sorting in _as)
                 # if the value does not equal the current value then we need to
                 # reset the dups and sum_ranks, knowing that a new value is
-                # coming up. the condition also needs to handle nan equality
+                # coming up. the conditional also needs to handle nan equality
                 # and the end of iteration
                 if (i == N - 1 or
                         (masked_vals[_as[i]] != masked_vals[_as[i+1]]) or
                         (mask[_as[i]] ^ mask[_as[i+1]])):
                     dups = sum_ranks = 0
                     grp_vals_seen += 1
-                    grp_tie_count +=1
+                    grp_tie_count += 1
 
                 # Similar to the previous conditional, check now if we are
                 # moving to a new group. If so, keep track of the index where