Skip to content

ENH: use correct dtype in groupby cython ops when it is known (without try/except) #38291

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 8 commits into from
Dec 8, 2020
Merged

ENH: use correct dtype in groupby cython ops when it is known (without try/except) #38291

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
279c4d1
REF: groupby op casting without try/except
jbrockmendel Dec 4, 2020
e6dc529
Float64
jbrockmendel Dec 4, 2020
ea79027
add tests for expected dtype of cython agg ops with nullable dtypes
jorisvandenbossche Dec 5, 2020
97fcd22
fix casting to float numpy array for FloatingArray
jorisvandenbossche Dec 5, 2020
b04d91f
fix tests
jorisvandenbossche Dec 5, 2020
202bee8
update rules of known result dtypes
jorisvandenbossche Dec 5, 2020
e888b3e
Merge branch 'master' of https://github.com/pandas-dev/pandas into gb…
jbrockmendel Dec 8, 2020
2566ec4
retain dtype
jbrockmendel Dec 8, 2020
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
3 changes: 3 additions & 0 deletions pandas/core/dtypes/cast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,12 +357,15 @@ def maybe_cast_result_dtype(dtype: DtypeObj, how: str) -> DtypeObj:
The desired dtype of the result.
"""
from pandas.core.arrays.boolean import BooleanDtype
from pandas.core.arrays.floating import Float64Dtype
from pandas.core.arrays.integer import Int64Dtype

if how in ["add", "cumsum", "sum"] and (dtype == np.dtype(bool)):
return np.dtype(np.int64)
elif how in ["add", "cumsum", "sum"] and isinstance(dtype, BooleanDtype):
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
elif how in ["add", "cumsum", "sum"] and isinstance(dtype, BooleanDtype):
elif how in ["add", "cumsum", "sum"] and isinstance(dtype, (BooleanDtype, IntegerDtype)):

A bit more off-topic here, but: for int dtypes with lower precision, we actually want int64 for those (eg sum of int8 gives int64)

Copy link
Member Author

Choose a reason for hiding this comment

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

to the extent that we can separate fixing of these from the avoiding the try/except goal, I'd like to do that

return Int64Dtype()
elif how in ["mean", "median", "var"] and isinstance(dtype, Int64Dtype):
return Float64Dtype()
return dtype


Expand Down
15 changes: 14 additions & 1 deletion pandas/core/groupby/ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
is_datetime64_any_dtype,
is_datetime64tz_dtype,
is_extension_array_dtype,
is_float_dtype,
is_integer_dtype,
is_numeric_dtype,
is_period_dtype,
Expand Down Expand Up @@ -507,7 +508,19 @@ def _ea_wrap_cython_operation(
res_values = self._cython_operation(
kind, values, how, axis, min_count, **kwargs
)
result = maybe_cast_result(result=res_values, obj=orig_values, how=how)
dtype = maybe_cast_result_dtype(orig_values.dtype, how)
if is_extension_array_dtype(dtype):
cls = dtype.construct_array_type()
return cls._from_sequence(res_values)
Copy link
Member

Choose a reason for hiding this comment

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

I think you need to pass the dtype here as well (to ensure lower precision gets preserved)

return res_values

elif is_float_dtype(values.dtype):
Copy link
Contributor

Choose a reason for hiding this comment

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

so would really like to move this entire wrapping to a method on EA / generic casting. We do this in multiple places (e.g. also on _reduce operatiosn), and this is likely leading to missing functionaility in various places.

# FloatingArray
values = values.to_numpy(na_value=np.nan)
res_values = self._cython_operation(
kind, values, how, axis, min_count, **kwargs
)
result = type(orig_values)._from_sequence(res_values)
return result

raise NotImplementedError(values.dtype)
Expand Down
5 changes: 4 additions & 1 deletion pandas/tests/arrays/integer/test_arithmetic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,10 @@ def test_reduce_to_float(op):
result = getattr(df.groupby("A"), op)()

expected = pd.DataFrame(
{"B": np.array([1.0, 3.0]), "C": integer_array([1, 3], dtype="Int64")},
{
"B": np.array([1.0, 3.0]),
"C": pd.array([1, 3], dtype="Float64"),
},
index=pd.Index(["a", "b"], name="A"),
)
tm.assert_frame_equal(result, expected)
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/groupby/test_function.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1072,7 +1072,7 @@ def test_apply_to_nullable_integer_returns_float(values, function):
output = 0.5 if function == "var" else 1.5
arr = np.array([output] * 3, dtype=float)
idx = Index([1, 2, 3], dtype=object, name="a")
expected = DataFrame({"b": arr}, index=idx)
expected = DataFrame({"b": arr}, index=idx).astype("Float64")

groups = DataFrame(values, dtype="Int64").groupby("a")

Expand Down
4 changes: 3 additions & 1 deletion pandas/tests/resample/test_datetime_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,9 @@ def test_resample_integerarray():

result = ts.resample("3T").mean()
expected = Series(
[1, 4, 7], index=pd.date_range("1/1/2000", periods=3, freq="3T"), dtype="Int64"
[1, 4, 7],
index=pd.date_range("1/1/2000", periods=3, freq="3T"),
dtype="Float64",
)
tm.assert_series_equal(result, expected)

Expand Down