Skip to content

Fix Bug with NA value in Grouping for Groupby.nth #26152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
May 5, 2019
Merged

Fix Bug with NA value in Grouping for Groupby.nth #26152

Show file tree
Hide file tree
Changes from 14 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
e62a14f
Added failing test
WillAyd Apr 17, 2019
e113b55
Merge remote-tracking branch 'upstream/master' into groupby-nth-nan
WillAyd Apr 19, 2019
823225b
Added fix in code
WillAyd Apr 19, 2019
068936f
Expanded test coverage scope
WillAyd Apr 19, 2019
8716764
Whatsnew note
WillAyd Apr 19, 2019
3000d2b
Fixed implementation
WillAyd Apr 19, 2019
84ddd6d
Typing and doc fixups
WillAyd Apr 19, 2019
84d343a
Fixed dropna annotation
WillAyd Apr 19, 2019
6f40fbe
Replaced Collection reference with List in annotation
WillAyd Apr 19, 2019
e2d006d
Fixed cast expression
WillAyd Apr 19, 2019
69014a7
Merge remote-tracking branch 'upstream/master' into groupby-nth-nan
WillAyd Apr 22, 2019
0b7eb6c
Moved comment position
WillAyd Apr 22, 2019
48d90ee
Removed re-assignment of nth_values for typing
WillAyd Apr 22, 2019
7eea5e3
typing fixup
WillAyd Apr 22, 2019
4dec450
Merge remote-tracking branch 'upstream/master' into groupby-nth-nan
WillAyd Apr 23, 2019
ce0abcd
Merge remote-tracking branch 'upstream/master' into groupby-nth-nan
WillAyd Apr 29, 2019
c2a0b8e
dedented and added comment
WillAyd Apr 29, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.25.0.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,7 @@ Groupby/Resample/Rolling
- Bug in :meth:`pandas.core.window.Rolling.min` and :meth:`pandas.core.window.Rolling.max` that caused a memory leak (:issue:`25893`)
- Bug in :meth:`pandas.core.groupby.GroupBy.idxmax` and :meth:`pandas.core.groupby.GroupBy.idxmin` with datetime column would return incorrect dtype (:issue:`25444`, :issue:`15306`)
- Bug in :meth:`pandas.core.groupby.GroupBy.cumsum`, :meth:`pandas.core.groupby.GroupBy.cumprod`, :meth:`pandas.core.groupby.GroupBy.cummin` and :meth:`pandas.core.groupby.GroupBy.cummax` with categorical column having absent categories, would return incorrect result or segfault (:issue:`16771`)
- Bug in :meth:`pandas.core.groupby.GroupBy.nth` where NA values in the grouping would return incorrect results (:issue:`26011`)

Reshaping
^^^^^^^^^
Expand Down
159 changes: 85 additions & 74 deletions pandas/core/groupby/groupby.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ class providing the base-class of operations.
import datetime
from functools import partial, wraps
import types
from typing import FrozenSet, Optional, Tuple, Type
from typing import FrozenSet, List, Optional, Tuple, Type, Union
import warnings

import numpy as np
Expand Down Expand Up @@ -1546,15 +1546,16 @@ def backfill(self, limit=None):

@Substitution(name='groupby')
@Substitution(see_also=_common_see_also)
def nth(self, n, dropna=None):
def nth(self,
n: Union[int, List[int]],
dropna: Optional[str] = None) -> DataFrame:
"""
Take the nth row from each group if n is an int, or a subset of rows
if n is a list of ints.

If dropna, will take the nth non-null row, dropna is either
Truthy (if a Series) or 'all', 'any' (if a DataFrame);
this is equivalent to calling dropna(how=dropna) before the
groupby.
'all' or 'any'; this is equivalent to calling dropna(how=dropna)
before the groupby.

Parameters
----------
Expand Down Expand Up @@ -1617,92 +1618,102 @@ def nth(self, n, dropna=None):
4 2 5.0
"""

if isinstance(n, int):
nth_values = [n]
elif isinstance(n, (set, list, tuple)):
nth_values = list(set(n))
if dropna is not None:
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of explicitly checking if dropna is not None this condition was refactored to sit in a branch that follows the implicit condition of if dropna.

There is however a slight behavior difference between this PR and master, where now these are equivalent:

>>> df = pd.DataFrame([[0, 1], [0, 2]], columns=['foo', 'bar'])
>>> df.groupby('foo').nth([0], dropna=None)
     bar
foo
0      1
>>> df.groupby('foo').nth([0], dropna=False)
     bar
foo
0      1

Whereas on master these would not yield the same thing:

>>> df = pd.DataFrame([[0, 1], [0, 2]], columns=['foo', 'bar'])
>>> df.groupby('foo').nth([0], dropna=None)
     bar
foo
0      1
>>> df.groupby('foo').nth([0], dropna=False)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/williamayd/clones/pandas/pandas/core/groupby/groupby.py", line 1626, in nth
    "dropna option with a list of nth values is not supported")
ValueError: dropna option with a list of nth values is not supported

I think the new behavior is more consistent and preferred. Might be worth a whatsnew

raise ValueError(
"dropna option with a list of nth values is not supported")
else:
valid_containers = (set, list, tuple)
if not isinstance(n, (valid_containers, int)):
raise TypeError("n needs to be an int or a list/set/tuple of ints")

nth_values = np.array(nth_values, dtype=np.intp)
self._set_group_selection()

if not dropna:
mask_left = np.in1d(self._cumcount_array(), nth_values)

if isinstance(n, int):
nth_values = [n]
elif isinstance(n, valid_containers):
nth_values = list(set(n))

nth_array = np.array(nth_values, dtype=np.intp)
self._set_group_selection()

mask_left = np.in1d(self._cumcount_array(), nth_array)
mask_right = np.in1d(self._cumcount_array(ascending=False) + 1,
-nth_values)
-nth_array)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note I had to change this because nth_values was inferred as List[int] upon initial assignment. This prevents reassignment from potentially obfuscating the type checker and code intent

mask = mask_left | mask_right

ids, _, _ = self.grouper.group_info

# Drop NA values in grouping
mask = mask & (ids != -1)

out = self._selected_obj[mask]
if not self.as_index:
return out

ids, _, _ = self.grouper.group_info
out.index = self.grouper.result_index[ids[mask]]

return out.sort_index() if self.sort else out

if dropna not in ['any', 'all']:
if isinstance(self._selected_obj, Series) and dropna is True:
warnings.warn("the dropna={dropna} keyword is deprecated,"
"use dropna='all' instead. "
"For a Series groupby, dropna must be "
"either None, 'any' or 'all'.".format(
dropna=dropna),
FutureWarning,
stacklevel=2)
dropna = 'all'
else:
# Note: when agg-ing picker doesn't raise this,
# just returns NaN
raise ValueError("For a DataFrame groupby, dropna must be "
"either None, 'any' or 'all', "
"(was passed {dropna}).".format(
dropna=dropna))

# old behaviour, but with all and any support for DataFrames.
# modified in GH 7559 to have better perf
max_len = n if n >= 0 else - 1 - n
dropped = self.obj.dropna(how=dropna, axis=self.axis)

# get a new grouper for our dropped obj
if self.keys is None and self.level is None:

# we don't have the grouper info available
# (e.g. we have selected out
# a column that is not in the current object)
axis = self.grouper.axis
grouper = axis[axis.isin(dropped.index)]

else:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you de-dent this here and remove the else which is unecessary; add a comment though of what this branch is

if isinstance(n, valid_containers):
raise ValueError(
"dropna option with a list of nth values is not supported")

# create a grouper with the original parameters, but on the dropped
# object
from pandas.core.groupby.grouper import _get_grouper
grouper, _, _ = _get_grouper(dropped, key=self.keys,
axis=self.axis, level=self.level,
sort=self.sort,
mutated=self.mutated)

grb = dropped.groupby(grouper, as_index=self.as_index, sort=self.sort)
sizes, result = grb.size(), grb.nth(n)
mask = (sizes < max_len).values

# set the results which don't meet the criteria
if len(result) and mask.any():
result.loc[mask] = np.nan

# reset/reindex to the original groups
if (len(self.obj) == len(dropped) or
len(result) == len(self.grouper.result_index)):
result.index = self.grouper.result_index
else:
result = result.reindex(self.grouper.result_index)
if dropna not in ['any', 'all']:
if isinstance(self._selected_obj, Series) and dropna is True:
warnings.warn("the dropna={dropna} keyword is deprecated,"
"use dropna='all' instead. "
"For a Series groupby, dropna must be "
"either None, 'any' or 'all'.".format(
dropna=dropna),
FutureWarning,
stacklevel=2)
dropna = 'all'
else:
# Note: when agg-ing picker doesn't raise this,
# just returns NaN
raise ValueError("For a DataFrame groupby, dropna must be "
"either None, 'any' or 'all', "
"(was passed {dropna}).".format(
dropna=dropna))

# old behaviour, but with all and any support for DataFrames.
# modified in GH 7559 to have better perf
max_len = n if n >= 0 else - 1 - n
dropped = self.obj.dropna(how=dropna, axis=self.axis)

# get a new grouper for our dropped obj
if self.keys is None and self.level is None:

# we don't have the grouper info available
# (e.g. we have selected out
# a column that is not in the current object)
axis = self.grouper.axis
grouper = axis[axis.isin(dropped.index)]

return result
else:

# create a grouper with the original parameters, but on dropped
# object
from pandas.core.groupby.grouper import _get_grouper
grouper, _, _ = _get_grouper(dropped, key=self.keys,
axis=self.axis, level=self.level,
sort=self.sort,
mutated=self.mutated)

grb = dropped.groupby(
grouper, as_index=self.as_index, sort=self.sort)
sizes, result = grb.size(), grb.nth(n)
mask = (sizes < max_len).values

# set the results which don't meet the criteria
if len(result) and mask.any():
result.loc[mask] = np.nan

# reset/reindex to the original groups
if (len(self.obj) == len(dropped) or
len(result) == len(self.grouper.result_index)):
result.index = self.grouper.result_index
else:
result = result.reindex(self.grouper.result_index)

return result

def quantile(self, q=0.5, interpolation='linear'):
"""
Expand Down
17 changes: 17 additions & 0 deletions pandas/tests/groupby/test_nth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -434,3 +434,20 @@ def test_nth_column_order():
columns=['C', 'B'],
index=Index([1, 2], name='A'))
assert_frame_equal(result, expected)


@pytest.mark.parametrize("dropna", [None, 'any', 'all'])
def test_nth_nan_in_grouper(dropna):
# GH 26011
df = DataFrame([
[np.nan, 0, 1],
['abc', 2, 3],
[np.nan, 4, 5],
['def', 6, 7],
[np.nan, 8, 9],
], columns=list('abc'))
result = df.groupby('a').nth(0, dropna=dropna)
expected = pd.DataFrame([[2, 3], [6, 7]], columns=list('bc'),
index=Index(['abc', 'def'], name='a'))

assert_frame_equal(result, expected)