Skip to content

BUG: transform with nunique should have dtype int64 #35152

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.

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 3 commits into from
Jul 10, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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/v1.1.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -1080,6 +1080,7 @@ Groupby/resample/rolling
- Bug in :meth:`DataFrame.groupby` lost index, when one of the ``agg`` keys referenced an empty list (:issue:`32580`)
- Bug in :meth:`Rolling.apply` where ``center=True`` was ignored when ``engine='numba'`` was specified (:issue:`34784`)
- Bug in :meth:`DataFrame.ewm.cov` was throwing ``AssertionError`` for :class:`MultiIndex` inputs (:issue:`34440`)
- Bug in :meth:`core.groupby.DataFrameGroupBy.transform` when ``func='nunique'`` and columns are of type ``datetime64``, the result would also be of type ``datetime64`` instead of ``int64`` (:issue:`35109`)

Reshaping
^^^^^^^^^
Expand Down
21 changes: 20 additions & 1 deletion pandas/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@
"""

from collections import abc, defaultdict
import contextlib
from datetime import datetime, timedelta
from functools import partial
import inspect
from typing import Any, Collection, Iterable, List, Union
from typing import Any, Collection, Iterable, Iterator, List, Union
import warnings

import numpy as np
Expand Down Expand Up @@ -502,3 +503,21 @@ def convert_to_list_like(
return list(values)

return [values]


@contextlib.contextmanager
def temp_setattr(obj, attr: str, value) -> Iterator[None]:
"""Temporarily set attribute on an object.

Args:
obj: Object whose attribute will be modified.
attr: Attribute to modify.
value: Value to temporarily set attribute to.

Yields:
obj with modified attribute.
"""
old_value = getattr(obj, attr)
setattr(obj, attr, value)
yield obj
setattr(obj, attr, old_value)
30 changes: 11 additions & 19 deletions pandas/core/groupby/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -485,8 +485,10 @@ def transform(self, func, *args, engine="cython", engine_kwargs=None, **kwargs):
# If func is a reduction, we need to broadcast the
# result to the whole group. Compute func result
# and deal with possible broadcasting below.
result = getattr(self, func)(*args, **kwargs)
return self._transform_fast(result, func)
# Temporarily set observed for dealing with categoricals.
with com.temp_setattr(self, "observed", True):
result = getattr(self, func)(*args, **kwargs)
return self._transform_fast(result)

def _transform_general(
self, func, *args, engine="cython", engine_kwargs=None, **kwargs
Expand Down Expand Up @@ -539,17 +541,14 @@ def _transform_general(
result.index = self._selected_obj.index
return result

def _transform_fast(self, result, func_nm: str) -> Series:
def _transform_fast(self, result) -> Series:
"""
fast version of transform, only applicable to
builtin/cythonizable functions
"""
ids, _, ngroup = self.grouper.group_info
result = result.reindex(self.grouper.result_index, copy=False)
cast = self._transform_should_cast(func_nm)
out = algorithms.take_1d(result._values, ids)
if cast:
out = maybe_cast_result(out, self.obj, how=func_nm)
return self.obj._constructor(out, index=self.obj.index, name=self.obj.name)

def filter(self, func, dropna=True, *args, **kwargs):
Expand Down Expand Up @@ -1467,25 +1466,23 @@ def transform(self, func, *args, engine="cython", engine_kwargs=None, **kwargs):
# If func is a reduction, we need to broadcast the
# result to the whole group. Compute func result
# and deal with possible broadcasting below.
result = getattr(self, func)(*args, **kwargs)
# Temporarily set observed for dealing with categoricals.
with com.temp_setattr(self, "observed", True):
result = getattr(self, func)(*args, **kwargs)

if isinstance(result, DataFrame) and result.columns.equals(
self._obj_with_exclusions.columns
):
return self._transform_fast(result, func)
return self._transform_fast(result)

return self._transform_general(
func, engine=engine, engine_kwargs=engine_kwargs, *args, **kwargs
)

def _transform_fast(self, result: DataFrame, func_nm: str) -> DataFrame:
def _transform_fast(self, result: DataFrame) -> DataFrame:
"""
Fast transform path for aggregations
"""
# if there were groups with no observations (Categorical only?)
# try casting data to original dtype
cast = self._transform_should_cast(func_nm)

obj = self._obj_with_exclusions

# for each col, reshape to to size of original frame
Expand All @@ -1494,12 +1491,7 @@ def _transform_fast(self, result: DataFrame, func_nm: str) -> DataFrame:
result = result.reindex(self.grouper.result_index, copy=False)
output = []
for i, _ in enumerate(result.columns):
res = algorithms.take_1d(result.iloc[:, i].values, ids)
# TODO: we have no test cases that get here with EA dtypes;
# maybe_cast_result may not be needed if EAs never get here
if cast:
res = maybe_cast_result(res, obj.iloc[:, i], how=func_nm)
output.append(res)
output.append(algorithms.take_1d(result.iloc[:, i].values, ids))

return self.obj._constructor._from_arrays(
output, columns=result.columns, index=obj.index
Expand Down
8 changes: 8 additions & 0 deletions pandas/tests/groupby/test_nunique.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,3 +167,11 @@ def test_nunique_preserves_column_level_names():
result = test.groupby([0, 0, 0]).nunique()
expected = pd.DataFrame([2], columns=test.columns)
tm.assert_frame_equal(result, expected)


def test_nunique_transform_with_datetime():
# GH 35109 - transform with nunique on datetimes results in integers
df = pd.DataFrame(date_range("2008-12-31", "2009-01-02"), columns=["date"])
result = df.groupby([0, 0, 1])["date"].transform("nunique")
expected = pd.Series([2, 2, 1], name="date")
tm.assert_series_equal(result, expected)