Flexible and powerful data analysis / manipulation library for Python, providing labeled data structures similar to R data.frame objects, statistical functions, and much more
[X] I have checked that this issue has not already been reported.
[X] I have confirmed this bug exists on the latest version of pandas.
[ ] I have confirmed this bug exists on the main branch of pandas.
Reproducible Example
import numpy as np
import pandas as pd
import numba as nb
from numba import njit
@njit(nb.int64(nb.int64[:], nb.int64[:]))
def return_one(values, index):
return np.int64(1)
df = pd.DataFrame({"group": [1, 1, 2, 2, 2], "y": [1, 1, 2, 3, 2]})
noagg = return_one(df["y"].values, df.index.values)
aggs1 = df.groupby("group").agg(lambda s: return_one(s.values, s.index.values))
aggs2 = df.groupby("group").agg(return_one, engine="numba")
print(type(noagg)) # int
print(aggs1["y"].dtype) # np.int64
print(aggs2["y"].dtype) # np.float64
Issue Description
When executing the numba-compiled function standalone or with groupby.agg() without the numba engine, ints are returned as expected. However, when using Pandas' numba-compiled agg function, ints are coerced into floats.
Ideally the numba engine behavior would match Pandas' and use the dtype of the aggregate function's result, but I'm not sure that's possible. A more practical solution would be to accept an optional result_dtype argument that defaults to np.float (making that default more explicit) and pass that to the result array constructor. It would also be helpful for this nuance to be documented.
rhshadrach
commented
4 months ago
Pandas version checks
[X] I have checked that this issue has not already been reported.
[X] I have confirmed this bug exists on the latest version of pandas.
[ ] I have confirmed this bug exists on the main branch of pandas.
Reproducible Example
Issue Description
When executing the numba-compiled function standalone or with groupby.agg() without the numba engine, ints are returned as expected. However, when using Pandas' numba-compiled agg function, ints are coerced into floats.
Best I can tell the issue is here https://github.com/pandas-dev/pandas/blob/b162331554d7c7f6fd46ddde1ff3908f2dc8bcce/pandas/core/groupby/numba_.py#L114, which seems to always initialize the array to float. When I add a
dtype=values.dtypeto that constructor, it works as expected.Expected Behavior
Ideally the numba engine behavior would match Pandas' and use the dtype of the aggregate function's result, but I'm not sure that's possible. A more practical solution would be to accept an optional
result_dtypeargument that defaults to np.float (making that default more explicit) and pass that to the result array constructor. It would also be helpful for this nuance to be documented.Installed Versions