BUG: DataFrame.groupby.count with arrow dtypes do not return arrow dtypes

[wirable23]

Pandas version checks

• I have checked that this issue has not already been reported.

• 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

>>> df = pd.DataFrame({"A": pd.Series([True, False, True, False], dtype="bool[pyarrow]"), "B": pd.Series([1,2,3,4], dtype="uint64[pyarrow]")})
>>> df.groupby("A").count().dtypes
B    int64
dtype: object
>>> df.groupby("A").std().dtypes
B    float64
dtype: object
>>>

Issue Description

Numpy types were returned when arrow types were provided to groupby.std()/count()

Expected Behavior

I would expect this to return "int64[pyarrow]" and "float64[pyarrow]". Other vectorized aggs such as var, sum, max, min return arrow dtypes when input is arrow backed.

Installed Versions

INSTALLED VERSIONS

commit : 965ceca
python : 3.11.2.final.0
python-bits : 64
OS : Windows
OS-release : 10
Version : 10.0.22621
machine : AMD64
processor : Intel64 Family 6 Model 140 Stepping 1, GenuineIntel
byteorder : little
LC_ALL : None
LANG : None
LOCALE : English_United States.1252

pandas : 2.0.2
numpy : 1.24.3
pytz : 2023.3
dateutil : 2.8.2
setuptools : 65.5.0
pip : 22.3.1
Cython : None
pytest : None
hypothesis : None
sphinx : None
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : None
IPython : None
pandas_datareader: None
bs4 : None
bottleneck : None
brotli : None
fastparquet : None
fsspec : None
gcsfs : None
matplotlib : None
numba : None
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : 12.0.1
pyreadstat : None
pyxlsb : None
s3fs : None
scipy : None
snappy : None
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
zstandard : None
tzdata : 2023.3
qtpy : None
pyqt5 : None

[wirable23]

Another weird example:

>>> pd.DataFrame({"A": pd.Series([True, False, True, False], dtype="bool[pyarrow]"), "B": pd.Series([17, 23, 106, 129], dtype="int64[pyarrow]")}).groupby("A").var().dtypes
B    float64
dtype: object
>>> pd.DataFrame({"A": pd.Series([True, False, True, False], dtype="bool[pyarrow]"), "B": pd.Series([17, 23, 107, 129], dtype="int64[pyarrow]")}).groupby("A").var().dtypes
B    int64[pyarrow]
dtype: object
>>>

Changing 3rd value in column B from 106 to 107 causes dtype of var to change from int64[pyarrow] to float64. Var should probably always return a floating type, in this case, "float64[pyarrow]"

[rhshadrach]

Thanks for the report. All instances except count return pyarrow dtypes on main. size also has the same issue.

[wirable23]

Thanks! Count would be expected to return "int64[pyarrow]"?

[mroeschke]

This is also an issue for the numpy masked types too

In [4]: >>> df = pd.DataFrame({"A": pd.Series([True, False, True, False], dtype="bool[pyarrow]"), "B": pd.Series([1,2,3,4], dtype="Int64")})

In [5]: df
Out[5]: 
       A  B
0   True  1
1  False  2
2   True  3
3  False  4

In [6]: df.groupby("A").count().dtypes
Out[6]: 
B    int64
dtype: object

[Charlie-XIAO]

Seems that count_level_2d is explicitly using dtype="i8", and similarly size is using dtype=int64 that lead to the issue. What is the recommended way to convert to the correct dtype backend based on the original data type?