BUG: Poor GroupBy Performance with ArrowDtype(...) wrapped types

[kzvezdarov]

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

import pandas as pd

df = pd.DataFrame({"key": range(100000), "val": "test"})
%timeit df.groupby(["key"]).first();

pa_df = df.convert_dtypes(dtype_backend="pyarrow")
%timeit pa_df.groupby(["key"]).first();

pa_df = pa_df.astype({"val": pd.StringDtype("pyarrow")})
%timeit pa_df.groupby(["key"]).first();

Issue Description

Grouping by and then aggregating on a dataframe that contains ArrowDtype(pyarrow.string()) columns is orders of magnitude slower than performing the same operations on an equivalent dataframe whose corresponding string column is of any other acceptable string type (e.g. string, StringDtype("python"), StringDtype("pyarrow")). This is surprising in particular because StringDtype("pyarrow") does not exhibit the same problem.

Note that in the bug reproduction example, DataFrame.convert_dtypes with dtype_backend="pyarrow" converts string columns to ArrowDtype(pyarrow.string()) rather than StringDtype("pyarrow").

Finally, here's a sample run, with dtypes printed out for clarity; I've reproduced this on both OS X and OpenSuse Tumbleweed for the listed pandas and pyarrow versions (as well as current main):

In [7]: import pandas as pd

In [8]: df = pd.DataFrame({"key": range(100000), "val": "test"})

In [9]: df["val"].dtype
Out[9]: dtype('O')

In [10]: %timeit df.groupby(["key"]).first();
8.37 ms ± 599 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [11]: pa_df = df.convert_dtypes(dtype_backend="pyarrow")

In [13]: type(pa_df["val"].dtype)
Out[13]: pandas.core.dtypes.dtypes.ArrowDtype

In [14]: %timeit pa_df.groupby(["key"]).first();
2.39 s ± 142 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [15]: pa_df = pa_df.astype({"val": pd.StringDtype("pyarrow")})
    ...:

In [16]: type(pa_df["val"].dtype)
Out[16]: pandas.core.arrays.string_.StringDtype

In [17]: %timeit pa_df.groupby(["key"]).first();
12.9 ms ± 306 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Expected Behavior

Aggregation performance on ArrowDtype(pyarrow.string()) columns should be comparable to aggregation performance on StringDtype("pyarrow"), string typed columns.

Installed Versions

INSTALLED VERSIONS

commit : 0691c5c
python : 3.13.1
python-bits : 64
OS : Darwin
OS-release : 24.3.0
Version : Darwin Kernel Version 24.3.0: Thu Jan 2 20:24:16 PST 2025; root:xnu-11215.81.4~3/RELEASE_ARM64_T6000
machine : arm64
processor : arm
byteorder : little
LC_ALL : en_CA.UTF-8
LANG : None
LOCALE : en_CA.UTF-8

pandas : 2.2.3
numpy : 2.2.2
pytz : 2025.1
dateutil : 2.9.0.post0
pip : 24.3.1
Cython : None
sphinx : None
IPython : 8.32.0
adbc-driver-postgresql: None
adbc-driver-sqlite : None
bs4 : None
blosc : None
bottleneck : None
dataframe-api-compat : None
fastparquet : None
fsspec : 2025.2.0
html5lib : None
hypothesis : None
gcsfs : None
jinja2 : 3.1.5
lxml.etree : None
matplotlib : None
numba : None
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
psycopg2 : None
pymysql : None
pyarrow : 19.0.0
pyreadstat : None
pytest : None
python-calamine : None
pyxlsb : None
s3fs : None
scipy : None
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
xlsxwriter : None
zstandard : None
tzdata : 2025.1
qtpy : None
pyqt5 : None

[rhshadrach]

Thanks for the report!

Aggregation performance on ArrowDtype(pyarrow.string()) columns should be comparable to aggregation performance on StringDtype("pyarrow"), string typed columns.

I don't think this expectation is correct. ArrowExtensionArray (corresponding to ArrowDtype) is essentially a wrapper around general PyArrow data, whereas ArrowStringArray (corresponding to StringDtype) is a data container specially crafted to support PyArrow strings. I do not believe we aim for neither feature nor performance parity between the two. Users should prefer StringDtype.

@jorisvandenbossche @WillAyd - should

df.convert_dtypes(dtype_backend="pyarrow")

be converting to StringDtype here?

[WillAyd]

Hmm I don't think so - the StringDtype is associated with the "numpy_nullable" backend. The pyarrow backend has always provided back the ArrowExtensionArray wrapper types.

Do we know where the bottleneck is performance-wise? I agree with the OP that the performance should be equivalent between the two types, if not slightly faster for the ArrowExtensionArray

[rhshadrach]

Hmm I don't think so - the StringDtype is associated with the "numpy_nullable" backend.

pandas/pandas/core/arrays/string_arrow.py

Line 128 in 70edaa0

_dtype: StringDtype # type: ignore[assignment]

[snitish]

@WillAyd looks like the bottleneck is due to ArrowExtensionArray._groupby_op not properly handling the string[pyarrow] dtype. We have special handling for StringDtype but for string[pyarrow], it hits to_masked() which again does not handle this dtype. This causes the groupby operation to be performed in pure python causing the slowness.

[WillAyd]

Thanks for that insight @snitish . Yea anything we can do to remove the layers of indirection and particularly Arrow <> NumPy copies would be very beneficial

[WillAyd]

FWIW I realize the OP is talking about strings, but this likely applies to the wrapped Arrow types in general. You can see the same performance issues using another Arrow type like decimal:

In [108]: df = pd.DataFrame({
     ...:     "key": pd.Series(range(100_000), dtype=pd.ArrowDtype(pa.int64())),
     ...:     "val": pd.Series(["3.14"] * 100_000, dtype=pd.ArrowDtype(pa.decimal128(10, 4)))
     ...: })
In [109]: %timeit df.groupby("key").sum()
2.83 s ± 131 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [110]: %timeit pa.TableGroupBy(pa.Table.from_pandas(df), "key").aggregate([("val", "sum")])
8.18 ms ± 589 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[snitish]

@WillAyd also noticed that .astype("string[pyarrow]") converts the dtype to StringDtype and the aggregation is much faster, while .convert_dtypes(dtype_backend="pyarrow") does not. Both dtypes show as "string["pyarrow"], but only one is a StringDtype. Is this expected?

In [2]: df = pd.DataFrame({"key": range(100000), "val": "test"}).astype({"val": "string[pyarrow]"})

In [3]: df["val"].dtype
Out[3]: string[pyarrow]

In [4]: isinstance(df["val"].dtype, pd.StringDtype)
Out[4]: True

In [6]: %timeit df.groupby("key").first()
16.5 ms ± 234 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [7]: pa_df = df.convert_dtypes(dtype_backend="pyarrow")

In [8]: pa_df["val"].dtype
Out[8]: string[pyarrow]

In [9]: isinstance(pa_df["val"].dtype, pd.StringDtype)
Out[9]: False

In [10]: %timeit pa_df.groupby("key").first()
4.73 s ± 44.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[kzvezdarov]

I meant to open a separate issue about convert_dtypes but got carried away and forgot to, sorry about that. At least per the pandas.DataFrame.convert_dtypes doc, it seems that when convert_string is set to True, string objects should be converted to StringDtype, but in practice they appear to get converted to the wrapper type ArrowDtype(pyarrow.string()) instead. In both case the dtype's name property resolves to string[pyarrow].

In [3]: df = pd.DataFrame({"val": ["test"]}).convert_dtypes(dtype_backend="pyarrow")

In [4]: type(df.dtypes["val"])
Out[4]: pandas.core.dtypes.dtypes.ArrowDtype