ENH/PERF: enable column-wise reductions for EA-backed columns
State: closed

Currently, for reductions on a DataFrame, we convert the full DataFrame to a single "interleaved" array and then perform the operation. That's the default, but when `numeric_only=True` is specified, it is done block-wise.

Enabling column-wise reductions (or block-wise for EAs):

- Gives better performance in common cases / no need to create a new 2D array (which goes through object dtype for nullable ints)
- Ensures to use the reduction implementation of the EA itself, which can be more correct / more efficient than converting to an ndarray and using our generic nanops.

For illustration purposes, I added a `column_wise` keyword in this PR (not meant to keep this, just for testing), so we can compare a few cases:

```
In [9]: df_wide = pd.DataFrame(np.random.randint(1000, size=(1000,100))).astype("Int64").copy() 

In [15]: %timeit df_wide.mean()   
9.68 ms ± 100 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [16]: %timeit df_wide.mean(numeric_only=True)
10.1 ms ± 345 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [17]: %timeit df_wide.mean(column_wise=True)  
5.22 ms ± 29.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [18]: df_long = pd.DataFrame(np.random.randint(1000, size=(10000,10))).astype("Int64").copy()  

In [19]: %timeit df_long.mean()
7.77 ms ± 115 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [20]: %timeit df_long.mean(numeric_only=True)         
2.07 ms ± 31.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [21]: %timeit df_long.mean(column_wise=True) 
1.04 ms ± 4.54 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

So I experimented with two approaches:

- First by iterating through the columns and calling the `_reduce` of the underlying EA (this path gets taken by using the temporary keyword `column_wise=True`)
- Fixing the block-wise case for extension blocks (triggered by `numeric_only=True`) by changing that to also use `_reduce` of the EA (currently this was failing by calling nanops functions on the EA)

The first gives better performance (it is simpler in implementation by not involding the blocks), but requires some more new code (it uses less the existing machinery).

Ideally, for EA columns, we should always use their own reduction implementation (thus call `EA._reduce`), I think. So for both approaches, the question will be how to trigger this behaviour.


Closes #32651,  closes #34520

Labels:
name: Performance
description: Memory or execution speed performance
name: Numeric Operations
description: Arithmetic, Comparison, and Logical operations
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.
name: Stale

Comments:
cc @jbrockmendel 
Come to think of it, the place where this dispatch belongs may be in the relevant nanops functions
> Come to think of it, the place where this dispatch belongs may be in the relevant nanops functions

Personally, I prefer the nanops to be about ops on numpy arrays, and not deal with extension arrays
So the bigger question is: how can we get to use this by default (at least for EAs). Some ideas:

- The `column_wise` route could for example be triggered by checking if all columns are ExtensionBlock columns (we already have a `BlockManager.any_extension_types`, could also have an `all_extension_types`). 
  This way, this would only be triggered in rare cases though (certainly as long we don't yet have a float extension type). We could also trigger it with `any_extension_types`: for the new nullable dtype that would be possible (since those are new, they can introduces changes in behaviour), but not for the older extension dtypes (categorical, datetimelikes)

- Currently, with `numeric_only is None`, we first try on the full values (`DataFrame.values`), and then fall back to `frame_apply` with `ignore_failures=True`. 
  I *suppose* we could also do this block-wise, just there ignore failures in case of `numeric_only is None`, and in such a case skip those blocks to assemble the output (but will need to check what all fails by doing this ..). 
  If that works, we could get rid of the full `numeric_only is None`-block in `_reduce`.





> Currently, with numeric_only is None, we first try on the full values (DataFrame.values), and then fall back to frame_apply with ignore_failures=True.

I've got a branch that identifies the cases where frame_apply is used and does that before the .values call.  Trying to 1) avoid a .values call and 2) de-nest _reduce.  I'll move that branch up the priority-list.


> Personally, I prefer the nanops to be about ops on numpy arrays, and not deal with extension arrays

Depends on the public-ness of those functions.  ATM their docstrings have examples with Series, not sure if we have other docs or tests with those.

> I've got a branch that identifies the cases where frame_apply is used and does that before the .values call. Trying to 1) avoid a .values call and 2) de-nest _reduce. I'll move that branch up the priority-list.

Note that my suggestion was to eliminate this part entirely (by using block-wise for all). So let's ensure we don't do duplicate / conflicting work. Does the branch already have something? (can you maybe push it to your fork?)

> Depends on the public-ness of those functions. ATM their docstrings have examples with Series, not sure if we have other docs or tests with those.

nanops are not public (regardless of their docstrings). What I meanly meant is that (in my head) they are meant to work on numpy arrays (whether extracted from a Series first or not)
https://github.com/jbrockmendel/pandas/tree/cln-reduce

> Note that my suggestion was to eliminate this part entirely (by using block-wise for all).

This would require making `BlockManager._reduce` handle `numeric_only=None` similar to how `frame_apply` handles `allow_failures=True`, right?  I've been reticent to do that (because I dont like the ignoring-exceptions pattern, xref #28900), but it may end up being the best option (off the top of my head, this would make it easy to solve #28773).
Thanks!

> because I dont like the ignoring-exceptions pattern,

Yes, but as long as we keep the `numeric_only=None` behaviour, that's the whole point of it .. (to ignore exceptions). So it would still be nice to try to unify the two code paths a bit more than what we have now. 
Will try to take a look at it next week.
hmm block-wise and column-wise with ignore_failures wont necessarily be equivalent for object-dtype
> hmm block-wise and column-wise with ignore_failures wont necessarily be equivalent for object-dtype

Ah, that's a good point. So for ObjectBlock, we would still need to do it column-wise, if we want to get rid of the fallback in general. Will take a further look one of the coming days if that looks feasable
Hello @jorisvandenbossche! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:









There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2021-02-11 01:34:29 UTC
@jbrockmendel did some further tests with trying to always use the block-wise path. Tests are far from passing, but I already included some special cases (like object block still needing to fall back to column-wise)

This will make the the block-wise path of course a *lot* more complex (given all those special cases, and with try/except blocks ...). So would like to get some feedback on this possible path forward before I put more time in it fixing all corner cases. 
I *think* it will make some special cases more explicit in the code, though, and the original goal was to get use the EA-native reduce by default, which needs either column or block wise. So some extra complexity will be inevitable for achieving that goal, I think (when preserving all special behaviours of `numeric_only=None`).
Alternative of course could be to not try to put everything in the block-wise case, but just go with column-wise instead of block-wise. But that would mean a perf degradation for 2D blocks. Unless we would do the column-wise for `numeric_only=None` and the block-wise for `numeric_only=True/False`, but that would give some duplication.
> Unless we would do the column-wise for numeric_only=None and the block-wise for numeric_only=True/False, but that would give some duplication.

I think this makes a lot of sense.  ATM we only do block-wise for `numeric_only is not None and axis in [0, 1]`.  It sounds like there is room to optimize the re-comining-results step, but the code for this is pretty clean.

We also already have frame_apply that supports both column-wise and allow_failures.  ATM we only use frame_apply as a fall-back.  Can we pick up the performance gains you have in mind by using it more aggressively?

Going back to the OP:

> Ideally, for EA columns, we should always use their own reduction implementation (thus call EA._reduce), I think. So for both approaches, the question will be how to trigger this behaviour.

Can we do this independent of the rest of this PR (and the difficult decisions it entails) with a small edit to `blk_func`?
> Can we do this independent of the rest of this PR (and the difficult decisions it entails) with a small edit to blk_func?

No, because that is (in the current state of master) only used for `numeric_only=True/False`, and I want EA._reduce be used by default. 

(the small edit to blk_func is still useful for the  `numeric_only=True/False`, for sure)
> We also already have frame_apply that supports both column-wise and allow_failures. ATM we only use frame_apply as a fall-back. Can we pick up the performance gains you have in mind by using it more aggressively?

Yeah, that might be a good idea. Will need to take a look to see if that is possible.
> (the small edit to blk_func is still useful for the numeric_only=True/False, for sure)

Yah, thats what i was suggesting
@jorisvandenbossche is this still active?
Yes, just need to find the time to investigate the `frame_apply` option (I already looked at it briefly, and my first impression is that it is doing way too much things that are all not needed if you know you are doing a simple reduction, but need to check the impact of this)
gentle encouragement; i think this could make a significant improvement
I would hold off reviewing for a moment, it is in a totally messy state, I am cleaning it up right now (I only rebased up to now)
Updated / cleaned up this branch. I now implemented column (array) wise for `axis=0` when the block-wise approach is not used. 
To limit potentially unforeseen backwards incompatible changes, we could also only take this route if there is at least one extension array column (or more restrictive: any extension array column with nullable dtype).

Remaining failures are all dubious cases (for which the solution might also be to fix the test):

- Any / all on timedelta/datetime dtypes -> https://github.com/pandas-dev/pandas/issues/34479
- Reduction on empty object-dtype DataFrame, currently returns float for DataFrame but integer for Series:

    ```
    In [6]: df_empty = pd.DataFrame(columns=['a', 'b', 'c']) 

    In [8]: df_empty.dtypes 
    Out[8]: 
    a    object
    b    object
    c    object
    dtype: object

    In [9]: df_empty.sum()  
    Out[9]: 
    a    0.0
    b    0.0
    c    0.0
    dtype: float64

    In [10]: s_empty = pd.Series(dtype=object)  

    In [11]: s_empty.sum() 
    Out[11]: 0

    In [12]: type(_)   
    Out[12]: int
    ```
  while on this PR the DataFrame behaviour follows Series to return integer series.

> Remaining failures are all dubious cases (for which the solution might also be to fix the test):

Agreed those cases look sketchy.  Especially the one with DataFrame not matching Series.

> We might want to limit the new code above to if axis == 0 and self._mgr.any_extension_types:, and then keep the block here for backwards compatibility

Did you try the previously-discussed approach of going through frame_apply to operate blockwise?  
> Did you try the previously-discussed approach of going through frame_apply to operate blockwise?

Yes, I did, but forgot to add a note about this. I am planning to open some issues for this, but the main reason I didn't take that route of `frame_apply` now: it's rather complex code (it's more general / doing a lot more than just the reductions I need here), it's also not yet working for ExtensionArrays (which is of course something that needs to be fixed in itself), and those two things together mean that it would be more involved to properly fix `frame_apply` to serve the purpose for `DataFrame._reduce`.  In addition, `frame_apply` also doesn't give the same performance increase as the current code in the PR (in part because it is doing a lot more).

We should certainly improve `frame_apply` as well (and at that point, the code I am adding here might be movable to `frame_apply` almost in its entirety), but since properly supporting EAs there is a larger endeavor, I went for now to fix it only here for `DataFrame._reduce`, as it would be nice to at least already start using the EA reductions in the DataFrame reductions.
AFAICT there are two perf-based motivations: 1) avoiding realizing `frame.values` (mentioned in the "interleaving" part of the OP) and 2) overhead associated with constructing a Series from block-wise op.  Do you have a clear picture of the relative importance of these?
Does this close https://github.com/pandas-dev/pandas/issues/34520?
Yes, but might be good to add a test for it (will do that here)
> AFAICT there are two perf-based motivations: 1) avoiding realizing frame.values (mentioned in the "interleaving" part of the OP) and 2) overhead associated with constructing a Series from block-wise op. Do you have a clear picture of the relative importance of these?

Regarding 1), avoiding working on `frame.values` is not only for performance, but also for ensuring we use the actual EA._reduce (which might be faster, but also simply more correct). 

Regarding 2), the overhead from constructing a Series from block-wise ops is real but also not huge (I don't remember the actual number, but blockwise is certainly already a significant improvement compared to master performance wise), but the main reason that I gave up on handling everything inside BlockManager._reduce (the attempt is still somewhere in the git history of this branch) was that it became quite complex. One of the reasons for this is that for ObjectBlock, we then need to do it column-wise anyway there, because different columns of an ObjectBlock might hold different types and have differently. 
In this commit I removed it again:  https://github.com/pandas-dev/pandas/pull/32867/commits/7088cfcbbf37cff20fe06432dd5062bcabfe1f09 (see the diff in managers.py)

Of course, now is the complexity (to handle potential failures etc) inside frame.py, and since it needs to live somewhere, it can certainly also be somewhere else.
> which might be faster, but also simply more correct

+1 on "more correct" being the priority.

> the overhead from constructing a Series from block-wise ops is real but also not huge

This seems to contradict a [previous comment](https://github.com/pandas-dev/pandas/pull/32867#discussion_r436390017).  If the `if numeric_only is not None and axis in [0, 1]:` case is correct as-is and there is not a compelling perf reason to change it, I would prefer to keep it as is, as it significantly narrows down the set of cases to be considered in the rest of the method.

> Of course, now is the complexity (to handle potential failures etc) inside frame.py, and since it needs to live somewhere, it can certainly also be somewhere else.

longer-term we should try to re-use the existing machinery in `frame_apply` (i think `apply_series_generator` with `allow_failures=True`), but if that is burdensome short term, I'll take your word for it.

Two other medium-term options are #34714 and #28900.