PERF: DataFrame.where for EA dtype mask

[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 when the mask/cond is "boolean" (EA) dtype.

Ending up with an EA mask is common when working with EA-backed data:

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randn(10, 2), dtype="Float64")
mask = df > 0

print(mask.dtypes) 

# 0    boolean
# 1    boolean
# dtype: object

ASV added:

       before           after         ratio
     [ddae6946]       [48684e5c]
     <main>           <where-putmask-ea-boolean>
-      51.0±0.7ms       18.4±0.4ms     0.36  frame_methods.Where.time_where(False, 'Float64')
-      45.3±0.6ms       13.5±0.4ms     0.30  frame_methods.Where.time_where(True, 'Float64')
-      44.4±0.8ms       13.1±0.4ms     0.29  frame_methods.Where.time_where(False, 'float64[pyarrow]')
-      47.2±0.9ms       13.5±0.7ms     0.29  frame_methods.Where.time_where(True, 'float64[pyarrow]')

EA-backed frames are not widely covered by ASVs, but there are other methods it may help with, e.g.:

       before           after         ratio
     [ddae6946]       [48684e5c]
     <main>           <where-putmask-ea-boolean>
+      11.0±0.7ms       13.6±0.4ms     1.24  frame_methods.Clip.time_clip('float64')
-        95.9±4ms         41.0±2ms     0.43  frame_methods.Clip.time_clip('Float64')
-        82.0±2ms       30.8±0.2ms     0.38  frame_methods.Clip.time_clip('float64[pyarrow]')

Note a slight slowdown in the non-EA clip ASV. However, I think the simplification is worth it and it puts clip back to where it was a few days ago in terms of perf (prior to #51472). In addition, we might be able to improve perf via the discussion in #51547.

[jbrockmendel]

is there a viable option to avoid getting here with DataFrames?

[lukemanley]

I suspect something could be done in NDFrame._where to avoid most cases. And then maybe .clip could just call _where?

If you're already thinking about changes to ._where in #51547 I can close this if it makes sense to combine the two.

[lukemanley]

I've moved this to NDFrame._where and updated the timings in the OP.

[phofl]

I think we can do them separately, but I agree that this should be handled before getting to the manager level

[phofl]

This basically reverts the initial clip pr?

[lukemanley]

This basically reverts the initial clip pr?

Partially, but it still retains the new inplace behavior and the perf improvements for EAs.

[phofl]

Thx, forgot about the inplace fix

[phofl]

Can you fix conflicts?

[phofl]

thx @lukemanley