BUG: Fix problems in group rank when both nans and infinity are present (#20561)

Author: peterpanmj

pandas/_libs/groupby_helper.pxi.in

@@ -508,7 +508,8 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
 
             # if keep_na, check for missing values and assign back
             # to the result where appropriate
-            if keep_na and masked_vals[_as[i]] == nan_fill_val:
+
+            if keep_na and mask[_as[i]]:
                 grp_na_count += 1
                 out[_as[i], 0] = nan
             else:
@@ -548,9 +549,9 @@ def group_rank_{{name}}(ndarray[float64_t, ndim=2] out,
             # reset the dups and sum_ranks, knowing that a new value is 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]]) and not
-                    (mask[_as[i]] and mask[_as[i+1]]))):
+            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
                 val_start = i
                 grp_vals_seen += 1

pandas/tests/groupby/test_groupby.py

@@ -1965,6 +1965,55 @@ def test_rank_args(self, grps, vals, ties_method, ascending, pct, exp):
         exp_df = DataFrame(exp * len(grps), columns=['val'])
         assert_frame_equal(result, exp_df)
 
+    @pytest.mark.parametrize("grps", [
+        ['qux'], ['qux', 'quux']])
+    @pytest.mark.parametrize("vals", [
+        [-np.inf, -np.inf, np.nan, 1., np.nan, np.inf, np.inf],
+    ])
+    @pytest.mark.parametrize("ties_method,ascending,na_option,exp", [
+        ('average', True, 'keep', [1.5, 1.5, np.nan, 3, np.nan, 4.5, 4.5]),
+        ('average', True, 'top', [3.5, 3.5, 1.5, 5., 1.5, 6.5, 6.5]),
+        ('average', True, 'bottom', [1.5, 1.5, 6.5, 3., 6.5, 4.5, 4.5]),
+        ('average', False, 'keep', [1.5, 1.5, np.nan, 3, np.nan, 4.5, 4.5
+                                    ][::-1]),
+        ('average', False, 'top', [3.5, 3.5, 1.5, 5., 1.5, 6.5, 6.5][::-1]),
+        ('average', False, 'bottom', [1.5, 1.5, 6.5, 3., 6.5, 4.5, 4.5][::-1]),
+        ('min', True, 'keep', [1., 1., np.nan, 3., np.nan, 4., 4.]),
+        ('min', True, 'top', [3., 3., 1., 5., 1., 6., 6.]),
+        ('min', True, 'bottom', [1., 1., 6., 3., 6., 4., 4.]),
+        ('min', False, 'keep', [1., 1., np.nan, 3., np.nan, 4., 4.][::-1]),
+        ('min', False, 'top', [3., 3., 1., 5., 1., 6., 6.][::-1]),
+        ('min', False, 'bottom', [1., 1., 6., 3., 6., 4., 4.][::-1]),
+        ('max', True, 'keep', [2., 2., np.nan, 3., np.nan, 5., 5.]),
+        ('max', True, 'top', [4., 4., 2., 5., 2., 7., 7.]),
+        ('max', True, 'bottom', [2., 2., 7., 3., 7., 5., 5.]),
+        ('max', False, 'keep', [2., 2., np.nan, 3., np.nan, 5., 5.][::-1]),
+        ('max', False, 'top', [4., 4., 2., 5., 2., 7., 7.][::-1]),
+        ('max', False, 'bottom', [2., 2., 7., 3., 7., 5., 5.][::-1]),
+        ('first', True, 'keep', [1., 2., np.nan, 3., np.nan, 4., 5.]),
+        ('first', True, 'top', [3., 4., 1., 5., 2., 6., 7.]),
+        ('first', True, 'bottom', [1., 2., 6., 3., 7., 4., 5.]),
+        ('first', False, 'keep', [4., 5., np.nan, 3., np.nan, 1., 2.]),
+        ('first', False, 'top', [6., 7., 1., 5., 2., 3., 4.]),
+        ('first', False, 'bottom', [4., 5., 6., 3., 7., 1., 2.]),
+        ('dense', True, 'keep', [1., 1., np.nan, 2., np.nan, 3., 3.]),
+        ('dense', True, 'top', [2., 2., 1., 3., 1., 4., 4.]),
+        ('dense', True, 'bottom', [1., 1., 4., 2., 4., 3., 3.]),
+        ('dense', False, 'keep', [3., 3., np.nan, 2., np.nan, 1., 1.]),
+        ('dense', False, 'top', [4., 4., 1., 3., 1., 2., 2.]),
+        ('dense', False, 'bottom', [3., 3., 4., 2., 4., 1., 1.]),
+    ])
+    def test_infs_n_nans(self, grps, vals, ties_method, ascending, na_option,
+                         exp):
+        key = np.repeat(grps, len(vals))
+        vals = vals * len(grps)
+        df = DataFrame({'key': key, 'val': vals})
+        result = df.groupby('key').rank(method=ties_method,
+                                        ascending=ascending,
+                                        na_option=na_option)
+        exp_df = DataFrame(exp * len(grps), columns=['val'])
+        assert_frame_equal(result, exp_df)
+
     @pytest.mark.parametrize("grps", [
         ['qux'], ['qux', 'quux']])
     @pytest.mark.parametrize("vals", [