ENH: Support operators for ExtensionArray
State: closed

- [x] closes #20659, #19577, #20825
- [x] tests added / passed
  - pandas/tests/extension/decimal/test_decimal.py
    - Tests for operators using existing operator tests for fundamental types
  - pandas/tests/extension/json/test_json.py
    - Create an `__add__` method for json and test it, make sure others fail
  - pandas/tests/extension/category/test_categorical.py
    - Adds a `test_combine` method
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry
  - Not included since I'm not sure if it is v0.23 or v0.24, and I wasn't about to start a v0.24 whatsnew

### Some Notes

- @jorisvandenbossche suggested that I put up an initial implementation for review, so I expect there will be lots of suggestions for doing it more, better, or different.
- There is a hook in there that is untested, where if an implementer (like me) wants to have the logical operators not return a `Series` with `bool`, they can do that.  I'm testing it with my library, but once we agree on this basic implementation, I will add tests for that specific use case.
- The idea is that you first see if the `ExtensionArray` class has implemented the operator.  If not, then you just go element by element.

Labels:
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.

Comments:
# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=h1) Report
> :exclamation: No coverage uploaded for pull request base (`master@93e7123`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference#section-missing-base-commit).
> The diff coverage is `93.81%`.

[![Impacted file tree graph](https://codecov.io/gh/pandas-dev/pandas/pull/20889/graphs/tree.svg?token=eZ4WkYLtcO&src=pr&width=650&height=150)](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=tree)

```diff
@@            Coverage Diff            @@
##             master   #20889   +/-   ##
=========================================
  Coverage          ?   91.78%           
=========================================
  Files             ?      153           
  Lines             ?    49423           
  Branches          ?        0           
=========================================
  Hits              ?    45362           
  Misses            ?     4061           
  Partials          ?        0
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.17% <93.81%> (?)` | |
| #single | `41.89% <10.3%> (?)` | |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/util/testing.py](https://codecov.io/gh/pandas-dev/pandas/pull/20889/diff?src=pr&el=tree#diff-cGFuZGFzL3V0aWwvdGVzdGluZy5weQ==) | `84.43% <100%> (ø)` | |
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/20889/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `84.7% <100%> (ø)` | |
| [pandas/core/indexes/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/20889/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW5kZXhlcy9iYXNlLnB5) | `96.59% <91.66%> (ø)` | |
| [pandas/core/series.py](https://codecov.io/gh/pandas-dev/pandas/pull/20889/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc2VyaWVzLnB5) | `93.87% <91.66%> (ø)` | |
| [pandas/core/ops.py](https://codecov.io/gh/pandas-dev/pandas/pull/20889/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvb3BzLnB5) | `96.19% <94.11%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=footer). Last update [93e7123...8aea5cf](https://codecov.io/gh/pandas-dev/pandas/pull/20889?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

just an FYI, rather than a single PR, this needs to be a series of smaller, self-consistent PR's. much easier to review and revise, otherwise this will get bogged down.
@jreback  Right now, I had to include some code from an early version of #20885.  Once that is approved, that will shorten this one a little bit. I can then first submit something that deals with #20825 . And once that is approved, I can submit the things related to ops.

Let me know if this plan works for you.
> just an FYI, rather than a single PR, this needs to be a series of smaller, self-consistent PR's. much easier to review and revise, otherwise this will get bogged down.

@Dr-Irv There are a bunch of planned/discussed steps between the status quo and the functionality in this PR.  If you're up to help with the whole transition that'd be great.  If you're mostly interested in getting this specific feature in place ASAP, we can talk about workarounds that you can apply in the interim.

@jreback LMK if I'm overstepping by offering to shepherd this process.  It's up my alley.
@jbrockmendel I'd like both.  I need the operator functionality for `ExtensionArray` .  So if you guys don't want to accept this PR, then I'd like to see the workaround so I can get my use case working, and since I have put a lot of thought into getting this PR to work, I can also help with the transition.
> @jreback LMK if I'm overstepping by offering to shepherd this process. It's up my alley.

I'd appreciate your expertise here @jbrockmendel.

FWIW, I'll have some availability this week, but very little next week.
> then I'd like to see the workaround so I can get my use case working

The workaround I have in mind assumes you are mostly interested in patching Series arithmetic and logical operators.  Is that accurate?

> and since I have put a lot of thought into getting this PR to work, I can also help with the transition

@Dr-Irv sounds good.  The first thing I'd suggest is figuring out if there is anything in this PR that is not strictly necessary for the feature(s) you need or can otherwise be separated into independent PRs.  Based on a very quick glance it looks like the tests for `get` or `combine` might be independent of the rest of the PR.

Are the changes to boolean ops logically independent of the changes to arithmetic ops?  If so, I suggest focusing on the boolean ops first, since they aren't _nearly_ so wound up with Index vs Series shenanigans.  There might be something in #19795 worth resurrecting.

@jbrockmendel 
> The workaround I have in mind assumes you are mostly interested in patching Series arithmetic and logical operators. Is that accurate?

Yes, that is correct.  I also need the results of arithmetic and logical operators to be Series of objects.

This PR has 3 things in it. 

1. Code from #20885 that was still in review mode when I submitted the present PR.  The tests for `get` that you mention fall into that category.  I can remove that to make this PR more up-to-date.
2. Fixes to support `combine` correctly.  This addresses #20825 .  I can break that out into a separate PR.  But that is needed for testing the operator stuff.
3. The code to support the arithmetic and logical operators for `ExtensionArray`.  It is the same code that does both.

So, if you want, I can create a new PR that only works on `combine`.  Then this one we are discussing here will only have the operator support, the meat of which is in the new internal method `dispatch_to_extension_ops`.

FYI I have this almost all implemented for integer na support: https://github.com/jreback/pandas/commit/6fc19f9c6f01f0fd42d0551aa1defd6f20bc3fb9 (this is a bit stale), will be updated soon.
I think it is still worthwhile to push forward this PR to implement this separately from the integer-array PR (which might take more discussion, and it also a very big one right now). 

@jreback can you give more concrete feedback on the way @Dr-Irv implemented things here, based on what you did in the other PR? (for the changes in `ops.py`)
(once the unrelated changes are removed (for `combine` and `get`), I don't think this PR needs to be splitted)

> There are a bunch of planned/discussed steps between the status quo and the functionality in this PR. 

@jbrockmendel can you elaborate on this a bit? (regarding this PR, what would you do instead)
>> There are a bunch of planned/discussed steps between the status quo and the functionality in this PR.

> can you elaborate on this a bit? (regarding this PR, what would you do instead)

@jorisvandenbossche 
- I'm pretty sure there is some test code here that is orthogonal to the rest of the PR; I'd separate that out to make review easier.
- The place to define the element-by-element fallback (essentially object-dtype behavior) is in the downstream EA subclass.
- Eventually _all_ of the arith/comparison Series ops are intended to dispatch to their underlying EA objects (possibly via `self._data.blocks[0])`.  (This will render dispatch_to_index_op unnecessary)
- Some steps towards making that happen include:
  - Something like #21160 (and analogues for float, bool, object, datetimelike (#19902), ... dtypes)
  - Align remaining mismatches between Series vs Index arithmetic/comparison behavior #19322, #18824, #19369.  Note 19322 is the most relevant one for the Integer EA case.
  - Categorical EA (and possibly the base class) should implement the relevant ops, even if they just raise.  They also need to figure out some inconsistencies between Categorical vs CategoricalIndex behavior (#19513)
 - As mentioned in a comment above, the Series bool ops need some serious work, and I'd much prefer to have that squared away before trying to make this transition. #19792, #19759

@jbrockmendel but is it needed to fix all of this (make sure all of Series ops are dispatched correctly etc, resolve duplication between index and series, ...) before getting basic ops working for ExtensionArrays? 
It's true there is a lot of work to do the full internal transition, but that doesn't mean we can already support basic ops for external extension arrays.
echo the above comments, this has a lot of othrogonal changes going on. let's strip out things to the minimal changes.
In response to @jbrockmendel:

> The place to define the element-by-element fallback (essentially object-dtype behavior) is in the downstream EA subclass.

I really disagree with this.  If you already have the operators defined for the `ExtensionDtype`, why should you have to implement the operators for `ExtensionArray`?  My implementation gives the person implementing an `ExtensionArray` subclass the option of defining the operators at the EA level, or letting the operators defined (if they are defined) for the `ExtensionDtype` do all of the work.
@jreback Since you wanted something smaller, I have split off the code that fixes #20825 into #21183 .  This is to handle the issue with `combine`.  I suggest that people reviewing this PR look at #21183 first.  Once it is accepted, I will then update this PR to only have the ops functionality.

Note that the `combine` bug has to be fixed in order for the tests in this PR to work right.  So hopefully you will accept #21183, and then this PR for ops will end up being smaller.

Why do the tests need the combine fix?
@jorisvandenbossche To test the operator functionality, I am leveraging the code that tests operators for `Series` that is in `pandas/tests/series/test_operators.py:TestSeriesOperators`, and those tests use `combine` to verify the results.  By fixing the bug in `combine`, then the tests for operators are leveraging existing test code.
> I really disagree with this. If you already have the operators defined for the ExtensionDtype, why should you have to implement the operators for ExtensionArray? My implementation gives the person implementing an ExtensionArray subclass the option of defining the operators at the EA level, or letting the operators defined (if they are defined) for the ExtensionDtype do all of the work.

Implementing it on the DType is fine, I guess.  It should be _really easy_ to go from there to implementing it on the EA subclass.  Once that is done, the changes to `core.ops` (and I hope/think `Series._binop`) can pretty much all be skipped and just call `dispatch_to_index_op` with the appropriate EA subclass (this is already done with pd.Categorical on L1193)

In most cases returning `NotImplemented` means "try the reversed operation".  I'd like to avoid giving that return value a double meaning if possible.
@jbrockmendel You wrote:

> Implementing it on the DType is fine, I guess. It should be _really easy_ to go from there to implementing it on the EA subclass. 

Are you suggesting that for someone implementing their own EA subclass, it should be really easy for _them_ to implement the operators on the EA subclass?  Or are you suggesting that pandas provide default operators for `ExtensionArray` that are defined based on the operators on the `ExtensionDtype`? If the former, I don't see how we can do that.  If the latter, the implementation in this PR does that.

The changes in this PR that are in `Series._binop` are now moved to #21183 in order to make `combine` work.

I can investigate whether `dispatch_to_index_ops` can do what `dispatch_to_extension_ops` does, but I'm doubtful, because my goal was to allow supporting the operators whether they were defined by the user at the EA subclass level or at the `ExtensionDtype` subclass level.


To try to take a step back from the discussion:

- We want that it is easy for ExtensionArray ops to just fall back to the scalar ops (the `ExtensionDtype.type` objects might already define ops, so use that). 
  - For this reason, @Dr-Irv implemented this fall back inside `ops.py` in case the ExtensionArray has no method for the specific `op`.
- On the other hand, we can also argue that it might be easier that the ops are always defined on the ExtensionArray. That would make the implementation inside `ops.py` a lot easier, as the fallback to the `ExtensionDtype.type` scalars could be removed.
- I would like to add that for my use case, I want that it is possible for the ExtensionArray to indicate that *no* such fallback should be tried (as this can be expensive). Defining the op and letting it raise might be enough for that (however, currently `TypeError` is catched and then still the fallback is tried. I think we should not do that).

Given the above, I think we can try to harmonize the goal of making this fallback both easy as explicit (1 and 3), and keep the ops defined on the ExtensionArray and `ops.py` simple (2). 
What if we **provide a mixin that defines all those basic fallbacks?** In that case, if the extension author decides it would like to enable the fallback, it can simply inherit from the mixin. Also, if then the behaviour of certain methods (like the booleans) are not exactly what he/she wants, just overwrite that specific method again in your array class.





Regarding the elementwise fallback to `ExtensionDtype.type`, isn't that relatively easy to define on the EA subclass?

```python
# in MyEASubclass
def __lt__(self, other):
    return np.array([a < b for a, b in zip(self, other)])
```

likewise for the other methods.

> I would like to add that for my use case, I want that it is possible for the ExtensionArray to indicate that no such fallback should be tried (as this can be expensive).

For cyberpandas, the scalar fall back would be incorrect. We use IPAddress(0) as the NA value, so comparisions with it should be NA to match pandas behavior.

```python
In [3]: arr = IPArray([0, 1, 2, 3])

In [4]: arr <= arr
Out[4]: array([False,  True,  True,  True])

In [5]: arr.astype(object) <= arr.astype(object)
Out[5]: array([ True,  True,  True,  True])
```

A mixin seems like a reasonable compromise.
> Regarding the elementwise fallback to ExtensionDtype.type, isn't that relatively easy to define on the EA subclass?

It's indeed relatively straightforward, although a bit more complicated as you small example if you want to cover all cases (scalar vs array `other` etc). Therefore providing a default implementation available as mixin would make this easier.
@TomAugspurger You wrote:

> Regarding the elementwise fallback to ExtensionDtype.type, isn't that relatively easy to define on the EA subclass?

Yes, but having to do it for *all* of the operators is painful.  What I implemented avoids that.  Take the example of `DecimalArray`.  You'd like to have to avoid implementing every single operator supported by `decimal.Decimal` in `DecimalArray`.

I like the idea of using a mixin as proposed by @jorisvandenbossche.  However, it seems that I ought to wait until @jreback finishes #21191 and #21160 and it is merged in, and then I can see where things are and make modifications as needed for my use case.
> Yes, but having to do it for all of the operators is painful.

It should be possible to write a function that generates all the different operators (which is how we often do it inside pandas). But we can use this approach for the mixin.

> However, it seems that I ought to wait until @jreback finishes #21191 and #21160 and it is merged in, and then I can see where things are and make modifications as needed for my use case.

I don't think you need to wait with eg writing the mixin with fallback (as that use case will not be covered in #21191).
@Dr-Irv 
> Are you suggesting that for someone implementing their own EA subclass, it should be really easy for them to implement the operators on the EA subclass?

Yes, this.

> Or are you suggesting that pandas provide default operators for ExtensionArray that are defined based on the operators on the ExtensionDtype?

The element-by-element implementation in the PR is effectively the object-dtype behavior, and I'd be fine with that being a default at the BaseArray level (@TomAugspurger thoughts?  It would also be reasonable for BaseArray to implement these ops to raise `com.AbtractMethodError`).

> The changes in this PR that are in Series._binop are now moved to #21183 in order to make combine work.

I'll take a closer look at #21883, thanks.

BTW, thanks for pushing forward on this.  You've clearly put some thought+effort into this and it is very much appreciated.

@jorisvandenbossche 
> but is it needed to fix all of this (make sure all of Series ops are dispatched correctly etc, resolve duplication between index and series, ...) before getting basic ops working for ExtensionArrays?

Not at all, the steps I mentioned were very much for the "what would you do instead" question.  I _do_ think that implementing ops at the array level will minimize duplicated effort in `core.ops` (a lot of effort went in to un-tangling older special-casing in core.ops, a process which I'm not eager to repeat).

> What if we provide a mixin that defines all those basic fallbacks? In that case, if the extension author decides it would like to enable the fallback, it can simply inherit from the mixin. Also, if then the behaviour of certain methods (like the booleans) are not exactly what he/she wants, just overwrite that specific method again in your array class.

+1.
My plan right now is to look at the mixin approach on May 29 or later.

When I created this PR, the whole point was to inspire discussion on the design and implementation, so I think this discussion has been healthy in terms of us moving forward.

Closing this, as there is a new proposed implementation in #21261 .