PERF: axis=1 reductions with EA dtypes
State: closed

- [x] closes #52016 
- [x] closes #54389
- [x] [Tests added and passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#writing-tests) if fixing a bug or adding a new feature
- [x] All [code checks passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit).
- [x] Added [type annotations](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#type-hints) to new arguments/methods/functions.
- [x] Added an entry in the latest `doc/source/whatsnew/v2.1.0.rst` file if fixing a bug or adding a new feature.


cc @rhshadrach - this includes the test you wrote in #51923 (with a few edits) as it looks like axis=1 reductions with EA dtypes are not well tested.

This PR avoids special-casing the internal EA types (although special-casing might allow for further optimization). 

xref: #51474

Timings on a slow laptop:

```
import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randn(10000, 4), dtype="float64[pyarrow]")
%timeit df.sum(axis=1)

# 3.79 s ± 137 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)      -> main
# 1.35 s ± 40.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)     -> 2.0.3
# 250 ms ± 53 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)       -> 1.5.3
# 5.42 ms ± 306 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)  -> PR

df = pd.DataFrame(np.random.randn(10000, 4), dtype="Float64")
%timeit df.sum(axis=1)

# 1.85 s ± 44.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)     -> main
# 1.2 s ± 58.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)      -> 2.0.3
# 16.4 ms ± 745 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)  -> 1.5.3
# 5.15 ms ± 286 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)  -> PR
```


Labels:
name: Performance
description: Memory or execution speed performance
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.
name: Reduction Operations
description: sum, mean, min, max, etc.

Comments:
Wow - this is quite clever! On my `AMD Ryzen 9 5950X` machine, I get

```
df = pd.DataFrame(np.random.randn(10000, 4), dtype="float64[pyarrow]")
%timeit df.sum(axis=1)
698 ms ± 2.82 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <-- main
929 µs ± 4.33 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)  <-- PR

df = pd.DataFrame(np.random.randn(10000, 4), dtype="Float64")
%timeit df.sum(axis=1)
401 ms ± 1.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  <-- main
1.16 ms ± 2.52 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)  <-- PR
```

Also, for wide inputs, there appears to not be a slowdown

```
df = pd.DataFrame(np.random.randn(4, 10000), dtype="float64[pyarrow]")
%timeit df.sum(axis=1)
29 ms ± 157 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)  <-- main
28.2 ms ± 440 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)  <-- PR

df = pd.DataFrame(np.random.randn(4, 10000), dtype="Float64")
%timeit df.sum(axis=1)
22.9 ms ± 48.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) <-- main
22.4 ms ± 273 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) <-- PR
```

Edit:

I also ran ASVs for `frame_methods` and `stat_ops`. `frame_methods` showed no change, `stat_ops`:

```
       before           after         ratio
     [c93e8034]       [8335a189]
     <main>           <perf-axis1-ea-reductions>
-      4.25±0.03s       10.2±0.3ms     0.00  stat_ops.FrameOps.time_op('sum', 'Int64', 1)
-         4.24±0s       10.0±0.3ms     0.00  stat_ops.FrameOps.time_op('prod', 'Int64', 1)
-      6.85±0.06s       13.0±0.3ms     0.00  stat_ops.FrameOps.time_op('skew', 'Int64', 1)
-      7.66±0.03s       13.7±0.3ms     0.00  stat_ops.FrameOps.time_op('median', 'Int64', 1)
-      7.46±0.04s       10.6±0.3ms     0.00  stat_ops.FrameOps.time_op('var', 'Int64', 1)
-      7.61±0.04s       10.8±0.3ms     0.00  stat_ops.FrameOps.time_op('std', 'Int64', 1)
```
> I also ran ASVs for `frame_methods` and `stat_ops`. `frame_methods` showed no change

isn't that showing a huge change, seconds -> milliseconds?

Edit: nevermind, I see the "no change" comment was for `frame_methods`. I'm not sure if any of the benchmarks in `frame_methods` would be impacted by this but it also looks like `frame_methods` generally does not measure EA types.

Right - I wasn't sure where axis=1 ops might show up (haven't thoroughly looked through the benchmarks); these seemed to be the most likely places.
how would this compare to using reduce_axis_1 approach?  i think that would avoid copies/allocations.  agreed this is clever, but if theres a way to avoid depending on the groupby code id prefer it
> how would this compare to using reduce_axis_1 approach? i think that would avoid copies/allocations. agreed this is clever, but if theres a way to avoid depending on the groupby code id prefer it

I might be missing something but my read of `_reduce_axis1` is that it would only provide a benefit for 2D arrays whereas this specifically targets EA (1D) arrays. 
@lukemanley - the rough idea of `_reduce_axis_1` is to do `df['a'] + df['b']` instead of transpose and sum along axis 0. This will help EAs but not 2d blocks (which have a fast transpose).

@jbrockmendel - I think looking into `_reduce_axis_1` approach is a good idea, but it won't work for reductions that don't have an online algorithm - e.g. median.
The additional copy required to construct the 1D array would be nice to avoid. I took a quick look at using `_reduce_axis1` for `sum` and it looks promising - comparable performance to this PR and avoids the copy. As @rhshadrach points out it would only be an option for a subset of reduction types. Any thoughts on what to do with the other reductions?
 
Somewhat surprising is that with this current PR, `axis=1` reductions seem to be faster than `axis=0` reductions for EA dtypes:

```
import pandas as pd
import numpy as np

data = np.random.randn(1000, 1000)

df = pd.DataFrame(data, dtype="Float64")
%timeit df.sum(axis=0)
%timeit df.sum(axis=1)
# 150 ms ± 5.45 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# 62.7 ms ± 3.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

df = pd.DataFrame(data, dtype="float64[pyarrow]")
%timeit df.sum(axis=0)
%timeit df.sum(axis=1)
# 301 ms ± 38.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 63.6 ms ± 4.49 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```
> Any thoughts on what to do with the other reductions?

median comes to mind.  Are there others?

At some point for #53261 I'm hoping to implement chunked-friendly implementations of the groupby reductions.  I think that combined with the trick in this PR could get the best of both worlds avoiding the copy.


> > Any thoughts on what to do with the other reductions?
> 
> median comes to mind. Are there others?

It looks like just nunique and quantile.


> Somewhat surprising is that with this current PR, axis=1 reductions seem to be faster than axis=0 reductions for EA dtypes:

It appears as if axis=0 has more overhead; I don't believe the results hold up for size e.g. (10000, 10000).
> It appears as if axis=0 has more overhead; I don't believe the results hold up for size e.g. (10000, 10000).

That seems reasonable.  I'd be interested in seeing profiling results though!
> > It appears as if axis=0 has more overhead; I don't believe the results hold up for size e.g. (10000, 10000).
> 
> That seems reasonable. I'd be interested in seeing profiling results though!

On my slower laptop now:

```
data = np.random.randn(10000, 10000)

df = pd.DataFrame(data, dtype="Float64")
%timeit df.sum(axis=0)
%timeit df.sum(axis=1)
# 1.02 s ± 5.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 2.62 s ± 46.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

df = pd.DataFrame(data, dtype="float64[pyarrow]")
%timeit df.sum(axis=0)
%timeit df.sum(axis=1)
# 1.24 s ± 7.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# 2.55 s ± 5.26 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
I would be okay with saying this closes the last two of the three issues mentioned in the OP, leaving the `_reduce_axis_1` open for further improvements.
@lukemanley - there's a whatsnew conflict here
Thanks @lukemanley - very nice!
Yeah, very nice indeed 👍 .
might also close #28487?  and less likely but related #36123?
> might also close #28487? and less likely but related #36123?

doesn't look like it closes either of these