PERF: DataFrame.iloc[int] for EA dtypes

[lukemanley]

• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/v2.1.0.rst file if fixing a bug or adding a new feature.

Perf improvement in DataFrame.iloc when input is an integer and the dataframe is EA-backed. Most visible on wide frames.

import pandas as pd
import numpy as np

data = np.random.randn(4, 10_000)

df_wide = pd.DataFrame(data, dtype="float64[pyarrow]")
%timeit df_wide.iloc[1]

# 1.33 s ± 31.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)    <- main
# 98.6 ms ± 3.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)  <- PR


df_wide = pd.DataFrame(data, dtype="Float64")
%timeit df_wide.iloc[1]

# 97.9 ms ± 5.07 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)  <- main
# 51.2 ms ± 2.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)  <- PR

Also visible with DataFrame reductions of EA dtypes:

import pandas as pd
import numpy as np

data = np.random.randn(4, 10_000)

df_wide = pd.DataFrame(data, dtype="float64[pyarrow]")
%timeit df_wide.sum(axis=0)

# 3.16 s ± 57.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <- main
# 1.66 s ± 38.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <- PR

df_wide = pd.DataFrame(data, dtype="Float64")
%timeit df_wide.sum(axis=0)

# 1.49 s ± 28.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <- main
# 1.17 s ± 23.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <- PR

[jbrockmendel]

i expect this will be bad for e.g. DatetimeTZDtype

[lukemanley]

I just tried it for DatetimeTZDtype("ns", "UTC") and it seems to be ok - and more performant in that case as well.

I think it's ok since each element is pulled out individually which ensures wrapping in Timestamp.

[jbrockmendel]

Is the issue that the relevant EA.__setitem__ methods are non-performant?

[lukemanley]

Is the issue that the relevant EA.__setitem__ methods are non-performant?

Yes. In the case pyarrow, really non-performant to iteratively set each element.

[mroeschke]

Thanks @lukemanley

[jbrockmendel]

I'm not wild about this. Seems to be papering over a hacky __setitem__ implementation for the ArrowEA, which really should just be immutable.

[lukemanley]

I'm not wild about this. Seems to be papering over a hacky __setitem__ implementation for the ArrowEA, which really should just be immutable.

Fair enough. The __setitem__ performance for ArrrowEA was pretty rough here. It was copying the new array for every element set. We could revert, but it would reintroduce that behavior. Is there some change you would like to see in the near term? (revert, TODO note, special-case ArrowEA, something else?)

[jbrockmendel]

Is there some change you would like to see in the near term? (revert, TODO note, special-case ArrowEA, something else?)

Maybe a TODO note pointing back at the relevant part of this thread?