Add __array_ufunc__ to Series / Array
State: closed

This PR:

* adds a basic (but incomplete) `__array_ufunc__` implementation for `IntegerArray` (to be able to check it is correctly used from Series for ExtensionArrays)
* adds `Series.__array_ufunc__` (not yet for `Index`)

---

@TomAugspurger revived my branch, will try to work on it a bit further this afternoon

What this already does is a basic implementation of the protocol for IntegerArray for simple ufuncs (call) for all IntegerArrays, and Series dispatching to the underlying values.

One question is if we want to force EA authors to implement an `__array_ufunc__` (eg by having a default implementation returning NotImplemented).  

Labels:
name: API Design
name: Numeric Operations
description: Arithmetic, Comparison, and Logical operations

Comments:
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 2019-07-01 18:48:08 UTC
BTW, if you also already had some stuff on this, feel free to push the branch if I could use something of it
Considering the IntegerArray ``__array_ufunc__`` implementation, the question comes up to what extent do we want to integrate our `_create_arithmetic_method` with it. 
Because the `__array_ufunc__` needs to be able to do everything that  `_create_arithmetic_method` does. So do we want to use the dunder methods created in `_create_arithmetic_method` inside `__array_ufunc__`, or do we want to base the ops dunder methods on the `__array_ufunc__` (like the `NDArrayOperatorsMixin` does: https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html#numpy.lib.mixins.NDArrayOperatorsMixin)
@TomAugspurger the first option (use `__add__` inside `__array_ufunc__(np.add, ..)` instead of the other way around) is basically what you did for SparseArray:

https://github.com/pandas-dev/pandas/blob/5a0c91e5a5bc61c71da08d6968d973255829678d/pandas/core/arrays/sparse.py#L1450-L1487

# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/23293?src=pr&el=h1) Report
> Merging [#23293](https://codecov.io/gh/pandas-dev/pandas/pull/23293?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/e9555159afd8caabc88b45ca6c06a16efe2f6afc?src=pr&el=desc) will **decrease** coverage by `0.04%`.
> The diff coverage is `94.68%`.

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

```diff
@@            Coverage Diff             @@
##           master   #23293      +/-   ##
==========================================
- Coverage   92.04%   91.99%   -0.05%     
==========================================
  Files         180      180              
  Lines       50714    50848     +134     
==========================================
+ Hits        46679    46777      +98     
- Misses       4035     4071      +36
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.63% <94.68%> (-0.05%)` | :arrow_down: |
| #single | `41.86% <45.74%> (-0.1%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/23293?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/ops.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvb3BzLnB5) | `94.81% <100%> (+0.15%)` | :arrow_up: |
| [pandas/core/arrays/categorical.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2NhdGVnb3JpY2FsLnB5) | `95.95% <100%> (+0.03%)` | :arrow_up: |
| [pandas/core/arrays/sparse.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3NwYXJzZS5weQ==) | `94.15% <100%> (-0.05%)` | :arrow_down: |
| [pandas/core/arrays/integer.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2ludGVnZXIucHk=) | `95.44% <87.5%> (-0.87%)` | :arrow_down: |
| [pandas/core/series.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc2VyaWVzLnB5) | `93.72% <96.87%> (+0.05%)` | :arrow_up: |
| [pandas/io/gbq.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2lvL2dicS5weQ==) | `88.88% <0%> (-11.12%)` | :arrow_down: |
| [pandas/core/internals/managers.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL21hbmFnZXJzLnB5) | `96% <0%> (-0.87%)` | :arrow_down: |
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `94.38% <0%> (-0.77%)` | :arrow_down: |
| [pandas/core/dtypes/concat.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2NvbmNhdC5weQ==) | `96.58% <0%> (-0.46%)` | :arrow_down: |
| [pandas/core/internals/concat.py](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2NvbmNhdC5weQ==) | `96.48% <0%> (-0.37%)` | :arrow_down: |
| ... and [11 more](https://codecov.io/gh/pandas-dev/pandas/pull/23293/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/23293?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/23293?src=pr&el=footer). Last update [e955515...6e770e8](https://codecov.io/gh/pandas-dev/pandas/pull/23293?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

cc @shoyer if you have thoughts here as well.
@shoyer Thanks for the feedback!
Question: do we want ufuncs to work on IntegerArray for older numpy versions? (that don't support yet `__array_ufunc__` ? 

@TomAugspurger for sparse, you did this I suppose by defining `__array_wrap__` in terms of `__array_ufunc__`:

https://github.com/pandas-dev/pandas/blob/f928b46036009ac7c0e4c182fba51a0c448d606c/pandas/core/arrays/sparse.py#L1433-L1439

but that seems a bit like a hack? (and I also don't fully understand it; this will compute the ufunc twice, and once on the densified data before the result of that is passed to `__array_wrap__`?)
It is indeed a hack, since it converts the data to dense.

I'm not sure why the ufunc would be computed twice though.
Because it is the result of the ufunc that is passed to `__array_wrap__` ?
Ah, I though the input was passed... Will take a look.
cc @jreback @jbrockmendel 
> Question: do we want ufuncs to work on IntegerArray for older numpy versions? (that don't support yet `__array_ufunc__`?

You are of course welcome to try, but I don't think it's worth the effort. `__array_wrap__` is a pretty half-baked solution.
@jorisvandenbossche worth updating for the rc?
@jorisvandenbossche can you merge master and update
Pushed a largeish change to `Series.__array_ufunc__`. The basic flow is

1. maybe dispatch to dunder op
2. maybe align input series (handling dataframes is a TODO; I think we want to define a `DataFrame.__array_ufunc__` and defer to that)
3. extract the arrays from the input Series / Index
4. dtype check for whether we can handle the op. This assumes that the underlying array defining `__array_ufunc__` defines a `_HANDLED_TYPES` tuple. We'll require that for ExtensionArray subclasses, which is the world I care about. If we see an unhandled type, we return NotImplemented. *Unless* we're a `Series[object]` dtype. This exception should be revisited.
5. Apply the ufunc
6. Wrap the result(s).

probably some cleanup to do, certainly some docs to write, but this makes sense to me.
@shoyer I'm a bit confused about the errors on NumPy dev (actually, I think this happens on NumPy 1.16 as well; we just elevate warnings to errors for the NumPy dev build).

It comes from trying to compare a 0-d scalar `numpy.datetime64` to a `Series[datetime64[ns, tz]]`. We'd like pandas to raise a TypeError there.

The `Series == scalar` direction is fine

```pytb
In [12]:     dts = pd.Series(pd.date_range('2016', periods=2, tz='UTC'))
    ...:     other = np.datetime64('2016-01-01')
    ...:

In [13]: dts == other
Traceback (most recent call last):
  File "<ipython-input-13-1aa12ffb5b5f>", line 1, in <module>
    dts == other
  File "/Users/taugspurger/sandbox/pandas/pandas/core/ops.py", line 1855, in wrapper
    pd.DatetimeIndex)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/ops.py", line 1330, in dispatch_to_index_op
    result = op(left_idx, right)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/indexes/datetimelike.py", line 123, in wrapper
    result = op(self._data, maybe_unwrap_index(other))
  File "/Users/taugspurger/sandbox/pandas/pandas/core/arrays/datetimes.py", line 142, in wrapper
    self._assert_tzawareness_compat(other)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/arrays/datetimes.py", line 700, in _assert_tzawareness_compat
    raise TypeError('Cannot compare tz-naive and tz-aware '
TypeError: Cannot compare tz-naive and tz-aware datetime-like objects
```

The reverse is raising a couple issues because NumPy apparently tries to do the operation again, after casting the Series to an ndarray?

```pytb
In [14]: other == dts
/Users/taugspurger/Envs/pandas-dev/lib/python3.7/site-packages/IPython/core/interactiveshell.py:3296: FutureWarning: Converting timezone-aware DatetimeArray to timezone-naive ndarray with 'datetime64[ns]' dtype. In the future, this will return an ndarray with 'object' dtype where each element is a 'pandas.Timestamp' with the correct 'tz'.
        To accept the future behavior, pass 'dtype=object'.
        To keep the old behavior, pass 'dtype="datetime64[ns]"'.
  exec(code_obj, self.user_global_ns, self.user_ns)
/Users/taugspurger/.virtualenvs/pandas-dev/bin/ipython:1: DeprecationWarning: elementwise comparison failed; this will raise an error in the future.
  #!/Users/taugspurger/Envs/pandas-dev/bin/python
Traceback (most recent call last):
  File "<ipython-input-14-85ca43130dc6>", line 1, in <module>
    other == dts
  File "/Users/taugspurger/sandbox/pandas/pandas/core/ops.py", line 1855, in wrapper
    pd.DatetimeIndex)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/ops.py", line 1330, in dispatch_to_index_op
    result = op(left_idx, right)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/indexes/datetimelike.py", line 123, in wrapper
    result = op(self._data, maybe_unwrap_index(other))
  File "/Users/taugspurger/sandbox/pandas/pandas/core/arrays/datetimes.py", line 142, in wrapper
    self._assert_tzawareness_compat(other)
  File "/Users/taugspurger/sandbox/pandas/pandas/core/arrays/datetimes.py", line 700, in _assert_tzawareness_compat
    raise TypeError('Cannot compare tz-naive and tz-aware '
TypeError: Cannot compare tz-naive and tz-aware datetime-like objects
```

The first warning comes from pandas, doing `Series.__array__(dtype=None)` on a `datetime64[ns, tz]` Series.

The second is from NumPy.


@TomAugspurger I suspect the problem has to do with how you've implemented `__array_ufunc__` for `np.equal`, which I'm guessing ends up recursively calling `ndarray.__eq__` via `maybe_dispatch_ufunc_to_dunder_op`.

When interacting with others that define `__array_ufunc__`, NumPy defines `ndarray.__eq__` (which implements `==`) by calling corresponding NumPy ufuncs, something like:
```
class ndarray:
    def __eq__(self, other):
        # if other has a __array_ufunc__ method 
        return np.equal(self, other)
```
(`NDArrayOperatorsMixin` defines equivalent methods with the full details in pure Python, if you're curious)

Apologies if this was already clear to you. I'm still not sure exactly why you get this warning, but hopefully this helps put you on the right track to figuring it out...
Thanks @shoyer. Here's a simpler example, using NDArrayOperatorMixin

The only special part is a bit of debug printing and

```python
        if any(x.ndim == 0 for x in inputs):
            raise TypeError('bad!')
```

```python
import numbers
import numpy as np


i = 0

class ArrayLike(np.lib.mixins.NDArrayOperatorsMixin):
    def __init__(self, value):
        self.value = np.asarray(value)

    # One might also consider adding the built-in list type to this
    # list, to support operations like np.add(array_like, list)
    _HANDLED_TYPES = (np.ndarray, numbers.Number)

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        global i
        i += 1

        print(i, inputs)
        out = kwargs.get('out', ())
        for x in inputs + out:
            # Only support operations with instances of _HANDLED_TYPES.
            # Use ArrayLike instead of type(self) for isinstance to
            # allow subclasses that don't override __array_ufunc__ to
            # handle ArrayLike objects.
            if not isinstance(x, self._HANDLED_TYPES + (ArrayLike,)):
                return NotImplemented

        # Defer to the implementation of the ufunc on unwrapped values.
        inputs = tuple(x.value if isinstance(x, ArrayLike) else x
                       for x in inputs)

        if any(x.ndim == 0 for x in inputs):
            raise TypeError('bad!')

        if out:
            kwargs['out'] = tuple(
                x.value if isinstance(x, ArrayLike) else x
                for x in out)
        result = getattr(ufunc, method)(*inputs, **kwargs)

        if type(result) is tuple:
            # multiple return values
            return tuple(type(self)(x) for x in result)
        elif method == 'at':
            # no return value
            return None
        else:
            # one return value
            return type(self)(result)

    def __repr__(self):
        return '%s(%r)' % (type(self).__name__, self.value)


if __name__ == '__main__':
    a = ArrayLike([1, 2, 3])
    b = np.array(1)

    print(b == a)
```

Running that, I see

```pytb
1 (array(1), ArrayLike(array([1, 2, 3])))
/Users/taugspurger/sandbox/pandas/bar.py:60: DeprecationWarning: elementwise comparison failed; this will raise an error in the future.
  print(b == a)
2 (ArrayLike(array([1, 2, 3])), array(1))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/sandbox/pandas/bar.py in <module>
     58     b = np.array(1)
     59
---> 60     print(b == a)

~/Envs/pandas-dev/lib/python3.7/site-packages/numpy/lib/mixins.py in func(self, other)
     23         if _disables_array_ufunc(other):
     24             return NotImplemented
---> 25         return ufunc(self, other)
     26     func.__name__ = '__{}__'.format(name)
     27     return func

~/sandbox/pandas/bar.py in __array_ufunc__(self, ufunc, method, *inputs, **kwargs)
     32
     33         if any(x.ndim == 0 for x in inputs):
---> 34             raise TypeError('bad!')
     35
     36         if out:

TypeError: bad!
```

So it seems that someone (NumPy? Python?) falls back to calling `eq(a, b)` when `eq(b, a)` fails? This doesn't happen with other ops like `>`.
@TomAugspurger Thanks for the clear test case! I'm not sure if this is fixable, but I filed a NumPy bug nonetheless: https://github.com/numpy/numpy/issues/13841
Regarding the NA behaviour, there is an active discussion on masked arrays on the numpy dev mailing list. I didn't read the whole thing, but it might be relevant to see what they think there how a masked array should behave for ufuncs, because in the end the IntegerArray is similar to a MaskedArray (if it can have deviating behaviour regarding skipping NAs)
Yep, I'm following that discussion as well.
I implemented `DecimalArray.__array_ufunc__` as a basic test for an EA defining it. Things seem to be working well.
> Instead of using numbers.Number, you could explicitly check for a lack of __array_priority__ and __array_ufunc__ attributes to indicate that something is a scalar / can be safely coerced. That is basically what NumPy does.

@shoyer that would also solve what we discussed before about for object dtype not knowing what the handled types are? (currently we are skipping the type check for object dtype data)
I think this is good enough to put in the RC
> > Instead of using numbers.Number, you could explicitly check for a lack of **array_priority** and **array_ufunc** attributes to indicate that something is a scalar / can be safely coerced. That is basically what NumPy does.
> 
> @shoyer that would also solve what we discussed before about for object dtype not knowing what the handled types are? (currently we are skipping the type check for object dtype data)

Yes, I think this would be an elegant way to handle this. Here's what I wrote up last October:
https://github.com/numpy/numpy/issues/12258#issuecomment-432858949

It's up to you whether you want to check for `__array_priority__` *and* `__array_ufunc__` attributes. The former is just a legacy thing now from NumPy's perspective (it's impossible to map operator priorities to integers outside of a single project), but it's probably still worth checking because classes that set `__array_priority__` were explicitly opting out of NumPy's automatic arithmetic operations that treat everything as a scalar. This is a good indication that they shouldn't be treated as scalars by pandas, either.
Ah, thanks for the link to that issue. Forgot about that discussion.
I'll update to incorporate https://github.com/pandas-dev/pandas/pull/23293#issuecomment-507056773 later today. Fixing the CI for pandas and dask first.
also, is there any way we can pin numpy 1.17rc1 in one of the builds just to be sure (as saw some other breakages on this); even though we test master.
I tried to implement the array_priority / array_ufunc check. A bit rushed today so I may have messed something up.

The main wrinkle is that `Series.__array_ufunc__` *doesn't* want to return `NotImpelmented` for ExtensionArrays defining `__array_ufunc__`. We want them to defer to us, so that we get consistent behavior for binops between Series and EA (we take care of unboxing and calling the ufunc on the array).