Support NDFrame.shift with EAs
State: closed

Uses take internally.

Closes https://github.com/pandas-dev/pandas/issues/22386

One question: do we want to add `shift` to the interface, and dispatch to that? This just does the right thing at the block level, by writing a `shift` that'll work with an ExtensionArray.

CategoricalBlock still overrides `.shift` to call `Categorical.shift`, which uses a slightly faster implementation based around the codes and `np.roll` (~1.5-2x faster based on microbenchmarks)`.

We could add move `ExtensionBlock.shift `to `ExtensionArray.shift`, and just call that from ExtensionBlock. Then we can delete CategoricalBlock.shift.

I'm 50/50 on whether it's worth adding to the interface.

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

Comments:
i think it’s worth adding to the interface itself
Wouldn't it be possible to do this without take, but with slicing?

Something like (very crude implementation for positive shifts):

```
def shift(a, p):
    a_sliced = a[:-p]
    a_NA = a._from_sequence([a.dtype.na_value]*p, dtype=a.dtype)
    return a._concat_same_type([a_NA, a_sliced])
```

Speed wise I would expect this to be very similar to the current `Categorical.shift` (which only uses the `np.roll` to basically do this instead of manually concatenating both parts)
# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/22387?src=pr&el=h1) Report
> Merging [#22387](https://codecov.io/gh/pandas-dev/pandas/pull/22387?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/68273a7a8b05049b311ea207f78e6d6f30c6eaf5?src=pr&el=desc) will **decrease** coverage by `<.01%`.
> The diff coverage is `91.66%`.

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

```diff
@@            Coverage Diff             @@
##           master   #22387      +/-   ##
==========================================
- Coverage   92.05%   92.05%   -0.01%     
==========================================
  Files         169      169              
  Lines       50733    50743      +10     
==========================================
+ Hits        46702    46711       +9     
- Misses       4031     4032       +1
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.46% <91.66%> (-0.01%)` | :arrow_down: |
| #single | `42.24% <16.66%> (-0.01%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/22387?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/22387/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `93.84% <100%> (ø)` | :arrow_up: |
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/22387/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `88% <90%> (+0.14%)` | :arrow_up: |

------

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

> Speed wise I would expect this to be very similar to the current Categorical.shift

Categorical was apparently a very bad example as there it is like 5x slower than the existing `Categorical.shift`, but that is more because generating the missing values (`_from_sequence)` and `_concat_same_type` are very slow for Categorical (and could certainly be optimized), for IntegerArray it was as fast as I expected.



That seems preferable to me, as it avoids allocating the `indexer` array.
Just as a note, I couldn't quite use this `shift` implementation for `SparseArray`.

When `ExtensionArray.shift`creates the empty SparseArray, it ends up with NaN  for the fill value. When `shift`ing a SparseArray with, say, 0 as the fill value, you end up concatenating the empty (with fill of NaN) and the shifted (with a fill of 0) and that ends up with a fill of NaN.

The root issue is that SparseArray stores fill_value on the array, rather than on the dtype. I haven't looked into whether it *should* go on the dtype or not, it's not clear to me.

Regardless, I think this is good to go, if we're happy with the slice and concat approach.
From that it seems to makes sense to have this on the dtype then (but not that familiar with the rest of the code)
Indeed... On the other hand, it complicates things like equality a bit...

I'll push it onto the dtype and see what breaks.
Any objections to merging?
> shouldn't we be able to change the impl of interval / categorical shift for this?

As I mentioned above (https://github.com/pandas-dev/pandas/pull/22387#issuecomment-413658486), if we implement Categorical.shift as is done here (eg by simply inheriting this base method), it gets much slower because the creating of the NaNs and concatting same typed-categoricals is very slow (at least slower as it should be). 
`IntervalArray` did not yet have any shift implementation. 
It seems that `IntervalArray` may need some updates,

```python
In [22]: idx = IntervalArray.from_breaks(pd.date_range('2017', periods=10))

In [23]: idx.shift()
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-23-1b2c2192e1e6> in <module>()
----> 1 idx.shift()

~/sandbox/pandas/pandas/core/arrays/base.py in shift(self, periods)
    422             return self.copy()
    423         empty = self._from_sequence([self.dtype.na_value] * abs(periods),
--> 424                                     dtype=self.dtype)
    425         if periods > 0:
    426             a = empty

~/sandbox/pandas/pandas/core/arrays/interval.py in _from_sequence(cls, scalars, dtype, copy)
    193     @classmethod
    194     def _from_sequence(cls, scalars, dtype=None, copy=False):
--> 195         return cls(scalars, dtype=dtype, copy=copy)
    196
    197     @classmethod

~/sandbox/pandas/pandas/core/arrays/interval.py in __new__(cls, data, closed, dtype, copy, fastpath, verify_integrity)
    138
    139         return cls._simple_new(left, right, closed, copy=copy, dtype=dtype,
--> 140                                verify_integrity=verify_integrity)
    141
    142     @classmethod

~/sandbox/pandas/pandas/core/arrays/interval.py in _simple_new(cls, left, right, closed, copy, dtype, verify_integrity)
    156                 raise TypeError(msg.format(dtype=dtype))
    157             elif dtype.subtype is not None:
--> 158                 left = left.astype(dtype.subtype)
    159                 right = right.astype(dtype.subtype)
    160

~/sandbox/pandas/pandas/core/indexes/numeric.py in astype(self, dtype, copy)
    313             msg = ('Cannot convert Float64Index to dtype {dtype}; integer '
    314                    'values are required for conversion').format(dtype=dtype)
--> 315             raise TypeError(msg)
    316         elif is_integer_dtype(dtype) and self.hasnans:
    317             # GH 13149

TypeError: Cannot convert Float64Index to dtype datetime64[ns]; integer values are required for conversion

In [24]: idx = IntervalArray.from_breaks(range(10))

In [25]: idx.shift()
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-25-1b2c2192e1e6> in <module>()
----> 1 idx.shift()

~/sandbox/pandas/pandas/core/arrays/base.py in shift(self, periods)
    422             return self.copy()
    423         empty = self._from_sequence([self.dtype.na_value] * abs(periods),
--> 424                                     dtype=self.dtype)
    425         if periods > 0:
    426             a = empty

~/sandbox/pandas/pandas/core/arrays/interval.py in _from_sequence(cls, scalars, dtype, copy)
    193     @classmethod
    194     def _from_sequence(cls, scalars, dtype=None, copy=False):
--> 195         return cls(scalars, dtype=dtype, copy=copy)
    196
    197     @classmethod

~/sandbox/pandas/pandas/core/arrays/interval.py in __new__(cls, data, closed, dtype, copy, fastpath, verify_integrity)
    138
    139         return cls._simple_new(left, right, closed, copy=copy, dtype=dtype,
--> 140                                verify_integrity=verify_integrity)
    141
    142     @classmethod

~/sandbox/pandas/pandas/core/arrays/interval.py in _simple_new(cls, left, right, closed, copy, dtype, verify_integrity)
    156                 raise TypeError(msg.format(dtype=dtype))
    157             elif dtype.subtype is not None:
--> 158                 left = left.astype(dtype.subtype)
    159                 right = right.astype(dtype.subtype)
    160

~/sandbox/pandas/pandas/core/indexes/numeric.py in astype(self, dtype, copy)
    316         elif is_integer_dtype(dtype) and self.hasnans:
    317             # GH 13149
--> 318             raise ValueError('Cannot convert NA to integer')
    319         return super(Float64Index, self).astype(dtype, copy=copy)
    320

ValueError: Cannot convert NA to integer
```

The basic issue is 

```
    423         empty = self._from_sequence([self.dtype.na_value] * abs(periods),
--> 424                                     dtype=self.dtype)
 ```

The container (ndarray[self.dtype]) can't actually hold `self.dtype.na_value` (nan).

Does that change our opinion on this implementation?
I think that is rather a bug in IntervalArray? (or maybe better a current limitation) 

Anyway, if the point is that it (currently) cannot hold NaNs, it can also not implement `shift`, as it requires NaNs (so don't know if that changes anything about the shift implementation?)
Well, it can implement `.shift`, as long as it's allowed to change dtypes (upcast from int to float). The current implementation assumes / requires that the dtype not change.
Basically, the problem is that for IntervalArray, the data is stored in `left` and `right` arrays, and so it depends on the dtype of those arrays whether IntervalArray can have NaNs or not? 
So eg float and datetime64 can hold missing values, ints not.

For the datetime error above, that is actually solvable with the current model, because if the data is not all-NaN it works:

```
In [75]: IntervalArray([pd.Interval(pd.Timestamp('2017-01-01'), pd.Timestamp('2017-01-02'))], dtype='interval[datetime64[ns]]')
Out[75]: 
IntervalArray([(2017-01-01, 2017-01-02]],
              closed='right',
              dtype='interval[datetime64[ns]]')

In [76]: IntervalArray([pd.Interval(pd.Timestamp('2017-01-01'), pd.Timestamp('2017-01-02')), np.nan], dtype='interval[datetime64[ns]]')
Out[76]: 
IntervalArray([(2017-01-01, 2017-01-02], nan],
              closed='right',
              dtype='interval[datetime64[ns]]')

In [78]: IntervalArray([np.nan], dtype='interval[datetime64[ns]]').left
...
TypeError: Cannot convert Float64Index to dtype datetime64[ns]; integer values are required for conversion
```



> Well, it can implement .shift, as long as it's allowed to change dtypes

Ah, yes, the same happens with `Series([int]).shift()` of course.

I think ideally, we should improve the missing values handling on the IntervalArray level (since it's a composite array, we have the ability to do it properly, not requiring such upcasts), but that's certainly out of scope for this PR :-)
So using the `take` implementation, that does already do the upcasting for Interval, so the `shift` would then work for intervals. 

I don't really care much to whether to switch back to that or leave the implementation now as is (with the consequence shift doesn't work for Interval; although given the `take` precedent, upcasting the dtype seems somewhat consistent. 
But I do think we should discuss the missing value handling for IntervalArray independently from this (eg using a mask instead).


> I think ideally, we should improve the missing values handling on the IntervalArray level

That's my preference. Opened https://github.com/pandas-dev/pandas/issues/22428 for that.
CI is passing now (last one was a random failure).
thanks @TomAugspurger