PERF: group by manipulation is slower with new arrow engine

[wegamekinglc]

Pandas version checks

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

• I have confirmed this issue exists on the latest version of pandas.

• I have confirmed this issue exists on the main branch of pandas.

Reproducible Example

Following codes will re-produce my issue:

import numpy as np
import pandas as pd

# generate a bunch sample data for later use
n = 5000000

s_samples = [f"s_{i}" for i in range(1, 101)]
i_samples = [f"i_{i}" for i in range(1, 201)]
bool_samples = [True, False]

ssamples = np.random.choice(s_samples, n)
isamples = np.random.choice(i_samples, n)
d_values = np.random.randn(3, n)
b_values = np.random.choice(bool_samples, n)

df = pd.DataFrame(
    dict(s=ssamples, i=isamples, v1=d_values[0], v2=d_values[1], v3=d_values[2], f1=b_values, f2=b_values)
)

df.to_csv("sample.csv", index=None)

# read in data with different engine
df_new = pd.read_csv("sample.csv", engine="pyarrow", dtype_backend="pyarrow")
df_old = pd.read_csv("sample.csv")

# do the bechmark
%timeit df_new.groupby("s")["v1"].sum()
# >> 660 ms ± 20.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df_old.groupby("s")["v1"].sum()
# >> 311 ms ± 13.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

The new engine is 2x slower than the old engine.

Installed Versions

INSTALLED VERSIONS

commit : c2a7f1a
python : 3.9.16.final.0
python-bits : 64
OS : Windows
OS-release : 10
Version : 10.0.19045
machine : AMD64
processor : AMD64 Family 23 Model 1 Stepping 1, AuthenticAMD
byteorder : little
LC_ALL : None
LANG : None
LOCALE : Chinese (Simplified)_China.936

pandas : 2.0.0rc1
numpy : 1.23.5
pytz : 2022.7
dateutil : 2.8.2
setuptools : 65.6.3
pip : 23.0.1
Cython : 0.29.33
pytest : 7.1.2
hypothesis : None
sphinx : 5.0.2
blosc : None
feather : None
xlsxwriter : 3.0.3
lxml.etree : 4.9.1
html5lib : None
pymysql : 1.0.2
psycopg2 : 2.9.3
jinja2 : 3.1.2
IPython : 8.10.0
pandas_datareader: None
bs4 : 4.11.1
bottleneck : 1.3.5
brotli :
fastparquet : None
fsspec : 2022.11.0
gcsfs : None
matplotlib : 3.7.1
numba : 0.56.4
numexpr : 2.8.4
odfpy : None
openpyxl : 3.0.10
pandas_gbq : None
pyarrow : 11.0.0
pyreadstat : None
pyxlsb : None
s3fs : None
scipy : 1.10.0
snappy :
sqlalchemy : 1.4.39
tables : 3.7.0
tabulate : 0.8.10
xarray : 2022.11.0
xlrd : 2.0.1
zstandard : 0.19.0
tzdata : None
qtpy : 2.2.0
pyqt5 : None

Prior Performance

No response

[phofl]

This is expected. GroupBy isn't implemented for arrow yet

[wegamekinglc]

@phofl thanks for your comments

[topper-123]

Is there a explanation about the current limitations when using arrow? There have been several blog posts talking up benefits of having arrow in Pandas and I think it could be a good idea laying out the current performance limitations of using arrow in Pandas.

[jbrockmendel]

I'm working on a patch here and having trouble making a performant conversion from the ArrowArray to MaskedArray. The non-working method looks like:

def _to_masked(self):
           pa_dtype = self._pa_array.type
        if pa.types.is_floating(pa_dtype):
            nbits = pa_dtype.bit_width
            dtype = f"Float{nbits}"
            np_dtype = dtype.lower()
            from pandas.core.arrays import FloatingArray as arr_cls
        elif pa.types.is_unsigned_integer(pa_dtype):
            nbits = pa_dtype.bit_width
            dtype = f"UInt{nbits}"
            np_dtype = dtype.lower()
            from pandas.core.arrays import IntegerArray as arr_cls
        elif pa.types.is_signed_integer(pa_dtype):
            nbits = pa_dtype.bit_width
            dtype = f"Int{nbits}"
            np_dtype = dtype.lower()
            from pandas.core.arrays import IntegerArray as arr_cls
        elif pa.types.is_boolean(pa_dtype):
            dtype = "boolean"
            np_dtype = "bool"
            from pandas.core.arrays import BooleanArray as arr_cls
        else:
            raise NotImplementedError

        data = self._pa_array.combine_chunks()
        buffs = data.buffers()
        assert len(buffs) == 2
        mask = self.isna()
        arr = np.array(buffs[1], dtype=np_dtype)
        return arr_cls(arr, mask)

But it looks like this is not the correct way to get arr. @jorisvandenbossche how do I get the correct ndarray here?

[phofl]

Did you try to_numpy on the ArrowExtensionArray and setting the na_value to 1?

[jbrockmendel]

That seems to work, thanks. Branch is ready for once the EA._groupby_op PR is merged.

[jorisvandenbossche]

Did you try to_numpy on the ArrowExtensionArray and setting the na_value to 1?

That doesn't give you a masked array, though (if that's what is needed). And will make a copy if there are missing values.

We normally already have this functionality in the __from_arrow__ methods on the MaskedArray classes. The easiest way to directly reuse this might be self._pa_array.to_pandas(types_mapper=pd.io._util._arrow_dtype_mapping().get). But the underlying functionality could maybe be refactored as a more general utility to use internally (pyarrow will returns a Series here, which we don't need)

[phofl]

@jorisvandenbossche Regarding the copy: I am aware that we create a copy when we have missing values, but wouldn't arrow do the same?