REF: call ensure_wrapped_if_datetimelike before the array arithmetic_op 
State: closed

~This is on top of https://github.com/pandas-dev/pandas/pull/39772, so only the last commit needs to be looked at for now: https://github.com/pandas-dev/pandas/pull/39820/commits/bd50388e019c61559c310638186aeea57d2d2d3e~

The actual goal here is to have a version of [`array_ops.py::arithmetic_op`](https://github.com/pandas-dev/pandas/blob/9f792cd90373177bffa0e9f85a14d5d5b30fb289/pandas/core/ops/array_ops.py#L159-L191) that does *not* call `ensure_wrapped_if_datetimelike` on its arrays. Reason: when performing ops column-wise, this becomes unnecessary overhead, and so the goal is to ensure this on a higher level, before calling the `array_op`.

So the experiment in this PR is to actually remove `ensure_wrapped_if_datetimelike` from `arithmetic_op`, and then add it in the places that will call `arithmetic_op`, to see how much changes this needs. This would ensure that we can keep a single shared `arithmetic_op` for use by both ArrayManager and BlockManager. The alternative would be to make a variant of `arithmetic_op` without the `ensure_wrapped_if_datetimelike`. That would require fewer code changes, but would give some duplication around `arithmetic_op` content (but it's not that it is a long function). Thoughts on which approach to prefer?

cc @jbrockmendel 

Labels:
name: Refactor
description: Internal refactoring of code
name: Numeric Operations
description: Arithmetic, Comparison, and Logical operations
name: Internals
description: Related to non-user accessible pandas implementation

Comments:
Any idea how big the perf impact is (both AM and non-AM)? needs to be worth the added complexity

Could trim down some of this by making DatetimeBlock/TimedeltaBlock backed directly by DTA/TDA

I am currently profiling `df + df` for an all numeric dataframe of shape (1000, 1000) for ArrayManager. And for this case the `ensure_wrapped_if_datetimelike` takes around 1-2 %. So not much at all (but with a relatively wider dataframe, this percentage will increase. And as other parts get optimized, this will also increase). 

So it's certainly not worth it on its own. But I think if we want to get the columnar frame ops closer to the blockwise ops (in performance), it will be thanks to multiple few-percentage improvements.  
Short term, I am planning to do a few other possible changes (moving `_maybe_upcast_for_op` out of `arithmetic_op`, moving the `np.seterr` out of `arithmetic_op`, etc). Hopefully that will inform us a bit more about the best way forward.

The main complication here is that I needed to basically re-implement BlockManager.apply to be able to call `ensure_wrapped_if_datetimelike` there (so indeed as you mention the cause for that is DatetimeBlock/TimedeltaBlock not storing EAs directly, but still using the EA in operations -> any idea how far off we are from being able to store the actual EA?)

To also repeat from above: an alternative is to have a variant of `arithmetic_op` that does less checks and is only used in the AraryManager for now (`arithmetic_op` is not a long function, so IMO that would also be acceptable).
> any idea how far off we are from being able to store the actual EA?

ive got a branch just about ready.  The sticking point ATM is arrow tests which I'll email you about shortly.

> Short term, I am planning to do a few other possible changes (moving _maybe_upcast_for_op out of arithmetic_op, moving the np.seterr out of arithmetic_op, etc). Hopefully that will inform us a bit more about the best way forward.

sounds promising


This pull request is stale because it has been open for thirty days with no activity. Please update or respond to this comment if you're still interested in working on this.
Simplified this PR: it's no longer based on top of #39772 (since the datetimelike blocks now store their values as EAs, I don't need to override `mgr.apply` for "operate_scalar", as I did originally, and for which I needed the changes in #39772)
@jbrockmendel any further comments here?
> A lot has changed since the OP. Any updated estimates on perf impact?

Eg in the `FrameWithFrameWide` benchmark (widest case, with operator.add), the `ensure_wrapped_if_datetimelike` calls take around 12% (without this PR).
thanks @jorisvandenbossche