Skip to content

Commit

Permalink
BUG: Fix problems in group rank when both nans and infinity are prese…
Browse files Browse the repository at this point in the history
  • Loading branch information
peterpanmj committed Apr 13, 2018
1 parent 2794474 commit f841ef8
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 4 deletions.
9 changes: 5 additions & 4 deletions pandas/_libs/groupby_helper.pxi.in
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
49 changes: 49 additions & 0 deletions pandas/tests/groupby/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", [
Expand Down

0 comments on commit f841ef8

Please sign in to comment.