BUG/Perf: Support ExtensionArrays in where
We need some way to do `.where` on EA object for DatetimeArray. Adding it
to the interface is, I think, the easiest way.

Initially I started to write a version on ExtensionBlock, but it proved
to be unwieldy. to write a version that performed well for all types.
It *may* be possible to do using `_ndarray_values` but we'd need a few more
things around that (missing values, converting an arbitrary array to the
"same' ndarary_values, error handling, re-constructing). It seemed easier
to push this down to the array.

The implementation on ExtensionArray is readable, but likely slow since
it'll involve a conversion to object-dtype.

Closes #24077
Closes #24142
Closes #16983
Hello @TomAugspurger! Thanks for submitting the PR.

 - There are no PEP8 issues in the file [`pandas/core/arrays/base.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/arrays/base.py) !

 - There are no PEP8 issues in the file [`pandas/core/arrays/categorical.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/arrays/categorical.py) !

 - There are no PEP8 issues in the file [`pandas/core/arrays/interval.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/arrays/interval.py) !

 - There are no PEP8 issues in the file [`pandas/core/arrays/period.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/arrays/period.py) !

 - There are no PEP8 issues in the file [`pandas/core/arrays/sparse.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/arrays/sparse.py) !

 - There are no PEP8 issues in the file [`pandas/core/dtypes/base.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/dtypes/base.py) !

 - There are no PEP8 issues in the file [`pandas/core/indexes/category.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/indexes/category.py) !

 - There are no PEP8 issues in the file [`pandas/core/internals/blocks.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/core/internals/blocks.py) !

 - There are no PEP8 issues in the file [`pandas/tests/arrays/categorical/test_indexing.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/arrays/categorical/test_indexing.py) !

 - There are no PEP8 issues in the file [`pandas/tests/arrays/interval/test_interval.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/arrays/interval/test_interval.py) !

 - There are no PEP8 issues in the file [`pandas/tests/arrays/test_period.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/arrays/test_period.py) !

 - There are no PEP8 issues in the file [`pandas/tests/extension/base/methods.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/extension/base/methods.py) !

 - There are no PEP8 issues in the file [`pandas/tests/extension/json/test_json.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/extension/json/test_json.py) !

 - There are no PEP8 issues in the file [`pandas/tests/extension/test_sparse.py`](https://github.com/TomAugspurger/pandas/blob/c4604df534abb15a9127adb887066e35dc16a3cc/pandas/tests/extension/test_sparse.py) !

 categorical perf:

master


```
In [2]: import pandas as pd

In [3]: cser = pd.Series(pd.Categorical(['a', 'b', 'c'] * 10000))

In [4]: %timeit cser.where(cser == 'a', 'c')
1.18 ms ± 37.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
```

PR:


```
In [3]: %timeit cser.where(cser == 'a', 'c')
755 µs ± 23.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
``` # [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=h1) Report
> Merging [#24114](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/9f2c7164ec581a484ee167c04478bf9e4f81426b?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `94.44%`.

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

```diff
@@            Coverage Diff             @@
##           master   #24114      +/-   ##
==========================================
+ Coverage    92.2%    92.2%   +<.01%     
==========================================
  Files         162      162              
  Lines       51714    51782      +68     
==========================================
+ Hits        47682    47747      +65     
- Misses       4032     4035       +3
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.6% <94.44%> (ø)` | :arrow_up: |
| #single | `43% <12.5%> (-0.02%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/dtypes/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Jhc2UucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `96.81% <100%> (+0.04%)` | :arrow_up: |
| [pandas/core/indexes/category.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW5kZXhlcy9jYXRlZ29yeS5weQ==) | `97.88% <100%> (-0.02%)` | :arrow_down: |
| [pandas/core/arrays/interval.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2ludGVydmFsLnB5) | `93.16% <100%> (+0.17%)` | :arrow_up: |
| [pandas/core/arrays/sparse.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3NwYXJzZS5weQ==) | `91.92% <88.88%> (-0.04%)` | :arrow_down: |
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `93.71% <93.33%> (-0.01%)` | :arrow_down: |
| [pandas/core/arrays/categorical.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2NhdGVnb3JpY2FsLnB5) | `95.37% <94.11%> (-0.03%)` | :arrow_down: |
| [pandas/core/arrays/period.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3BlcmlvZC5weQ==) | `98.31% <94.44%> (-0.16%)` | :arrow_down: |
| [pandas/util/testing.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL3V0aWwvdGVzdGluZy5weQ==) | `87.51% <0%> (+0.09%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?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/24114?src=pr&el=footer). Last update [9f2c716...56470c3](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
 # [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=h1) Report
> Merging [#24114](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/9f2c7164ec581a484ee167c04478bf9e4f81426b?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `94.44%`.

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

```diff
@@            Coverage Diff             @@
##           master   #24114      +/-   ##
==========================================
+ Coverage    92.2%    92.2%   +<.01%     
==========================================
  Files         162      162              
  Lines       51714    51782      +68     
==========================================
+ Hits        47682    47747      +65     
- Misses       4032     4035       +3
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.6% <94.44%> (ø)` | :arrow_up: |
| #single | `43% <12.5%> (-0.02%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/dtypes/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Jhc2UucHk=) | `100% <100%> (ø)` | :arrow_up: |
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `96.81% <100%> (+0.04%)` | :arrow_up: |
| [pandas/core/indexes/category.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW5kZXhlcy9jYXRlZ29yeS5weQ==) | `97.88% <100%> (-0.02%)` | :arrow_down: |
| [pandas/core/arrays/interval.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2ludGVydmFsLnB5) | `93.16% <100%> (+0.17%)` | :arrow_up: |
| [pandas/core/arrays/sparse.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3NwYXJzZS5weQ==) | `91.92% <88.88%> (-0.04%)` | :arrow_down: |
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `93.71% <93.33%> (-0.01%)` | :arrow_down: |
| [pandas/core/arrays/categorical.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2NhdGVnb3JpY2FsLnB5) | `95.37% <94.11%> (-0.03%)` | :arrow_down: |
| [pandas/core/arrays/period.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3BlcmlvZC5weQ==) | `98.31% <94.44%> (-0.16%)` | :arrow_down: |
| [pandas/util/testing.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL3V0aWwvdGVzdGluZy5weQ==) | `87.51% <0%> (+0.09%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?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/24114?src=pr&el=footer). Last update [9f2c716...56470c3](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
 # [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=h1) Report
> Merging [#24114](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/029cde201367827f35ac03779bce476a13c4232e?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `97.43%`.

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

```diff
@@            Coverage Diff             @@
##           master   #24114      +/-   ##
==========================================
+ Coverage   92.21%   92.21%   +<.01%     
==========================================
  Files         162      162              
  Lines       51723    51761      +38     
==========================================
+ Hits        47694    47731      +37     
- Misses       4029     4030       +1
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.61% <97.43%> (ø)` | :arrow_up: |
| #single | `43% <7.69%> (-0.03%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `97.41% <ø> (ø)` | :arrow_up: |
| [pandas/core/indexes/category.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW5kZXhlcy9jYXRlZ29yeS5weQ==) | `97.9% <ø> (ø)` | :arrow_up: |
| [pandas/core/arrays/sparse.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL3NwYXJzZS5weQ==) | `92.08% <ø> (ø)` | :arrow_up: |
| [pandas/core/internals/blocks.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW50ZXJuYWxzL2Jsb2Nrcy5weQ==) | `93.81% <100%> (+0.13%)` | :arrow_up: |
| [pandas/core/arrays/categorical.py](https://codecov.io/gh/pandas-dev/pandas/pull/24114/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2NhdGVnb3JpY2FsLnB5) | `95.3% <83.33%> (-0.1%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/24114?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/24114?src=pr&el=footer). Last update [029cde2...539d3cb](https://codecov.io/gh/pandas-dev/pandas/pull/24114?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
 Hmm I don't like the return dtype depending on the values.

Perhaps we do this with a deprecation warning?

On Wed, Dec 5, 2018 at 1:30 PM Jeff Reback <notifications@github.com> wrote:

> *@jreback* commented on this pull request.
> ------------------------------
>
> In doc/source/whatsnew/v0.24.0.rst
> <https://github.com/pandas-dev/pandas/pull/24114#discussion_r239205989>:
>
> > @@ -1262,6 +1264,7 @@ Categorical
>  - In meth:`Series.unstack`, specifying a ``fill_value`` not present in the categories now raises a ``TypeError`` rather than ignoring the ``fill_value`` (:issue:`23284`)
>  - Bug when resampling :meth:`Dataframe.resample()` and aggregating on categorical data, the categorical dtype was getting lost. (:issue:`23227`)
>  - Bug in many methods of the ``.str``-accessor, which always failed on calling the ``CategoricalIndex.str`` constructor (:issue:`23555`, :issue:`23556`)
> +- Bug in :meth:`Series.where` losing the categorical dtype for categorical data (:issue:`24077`)
>
> if its in the categories then this should work and return categorical, if
> its not hen i think coercing to object is ok
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pandas-dev/pandas/pull/24114#discussion_r239205989>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ABQHItYzZZM79tE07IEHNarQnkEaGctNks5u2B6_gaJpZM4ZDPyj>
> .
>
 @jorisvandenbossche do you have any objections to adding `ExtensionArray.where` to the interface? Or do the behavior change on `Series[category].where`? Just for context: how is this different from `eaarray[cond] = other` ?

The behaviour change to keep the categorical dtype is certainly fine. I suppose that `np.where` would work on items that don't implement
`__setitem__`. But other than that they should be identical
for 1-d arrays, right?

On Thu, Dec 6, 2018 at 9:44 AM Joris Van den Bossche <
notifications@github.com> wrote:

> Just for context: how is this different from eaarray[cond] = other ?
>
> The behaviour change to keep the categorical dtype is certainly fine.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pandas-dev/pandas/pull/24114#issuecomment-444916642>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ABQHIkP1MIBozQDarm1XFzpLj6LMJeHJks5u2TtCgaJpZM4ZDPyj>
> .
>
 > But other than that they should be identical for 1-d arrays, right?

And since EAs are 1D, and our internal EAs support setitem, why is the new code needed? Or what in setitem is not working as it should right now? (maybe I am missing some context) This came out  of the DatetimeArray refactor. I'll have to take another
look at exactly what the failures where. They were pretty deep in the
internals.

On Thu, Dec 6, 2018 at 9:53 AM Joris Van den Bossche <
notifications@github.com> wrote:

> But other than that they should be identical for 1-d arrays, right?
>
> And since EAs are 1D, and our internal EAs support setitem, why is the new
> code needed? Or what in setitem is not working as it should right now?
> (maybe I am missing some context)
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/pandas-dev/pandas/pull/24114#issuecomment-444920085>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/ABQHInMDbZ9aQt7DQXAVqcteU2aBIjFXks5u2T1ngaJpZM4ZDPyj>
> .
>
 On further reflection, I realize that ndarrays don't have a `where` method, so I don't think we should add `ExtensionArray.where`.

I'll see if setitem on a copy is sufficient. Updated to use setitem...

Sorry, this kind of grew in scope. I'm going to rebase my datetimearray PR and see which subset is actually needed to fix DTA. > Or what in setitem is not working as it should right now? (maybe I am missing some context)

@jorisvandenbossche OK, here's some context :) The most immediate failure is mismatched block dimensions / shapes for in `DataFrame.where` for EA (not just DatetimeArray, but that was tested). This uses Categorical

```python
In [8]: df = pd.DataFrame({"A": pd.Categorical([1, 2, 3])})

In [9]: df.where(pd.DataFrame({"A": [True, False, True]}))
```

```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-9-56dcebf7e672> in <module>
----> 1 df.where(pd.DataFrame({"A": [True, False, True]}))

... 

~/sandbox/pandas/pandas/core/internals/blocks.py in __init__(self, values, placement, ndim)
     84             raise ValueError(
     85                 'Wrong number of items passed {val}, placement implies '
---> 86                 '{mgr}'.format(val=len(self.values), mgr=len(self.mgr_locs)))
     87
     88     def _check_ndim(self, values, ndim):

ValueError: Wrong number of items passed 3, placement implies 1
```

The broadcasting is all messed up since the shapes aren't right (we're using `Block.where`).

```python
ipdb> cond
array([[ True],
       [False],
       [ True]])
ipdb> values
[1, 2, 3]
Categories (3, int64): [1, 2, 3]
ipdb> other
nan
```

A hacky, but shorter fix is to use the following (this is in Block.where)

```diff
diff --git a/pandas/core/internals/blocks.py b/pandas/core/internals/blocks.py
index 618b9eb12..2356a226d 100644
--- a/pandas/core/internals/blocks.py
+++ b/pandas/core/internals/blocks.py
@@ -1319,12 +1319,20 @@ class Block(PandasObject):
 
         values = self.values
         orig_other = other
+        if not self._can_consolidate:
+            transpose = False
+
         if transpose:
             values = values.T
 
         other = getattr(other, '_values', getattr(other, 'values', other))
         cond = getattr(cond, 'values', cond)
 
+        if not self._can_consolidate:
+            if cond.ndim == 2:
+                assert cond.shape[-1] == 1
+                cond = cond.ravel()
+
         # If the default broadcasting would go in the wrong direction, then
         # explicitly reshape other instead
         if getattr(other, 'ndim', 0) >= 1:
```

That fixes most of the issues I was having on the DTA branch. Still running the tests to see if any were re-broken.
(edit: it's slightly more complicated, have to handle reshaping `other` as well, so ~15 more LOC).

---

So, in summary

1. we need to support EAs in DataFrame.where
2. The diff just above is a hacky way of achieving just that (no more).
3. This PR will still be useful for avoiding the conversion to object (DatetimeTZBlock will avoid it even without this PR, thanks to `_try_coerce_args` coercing datetimes to ints)

 Updated. Main outstanding point is whether or not we should create a SparseBlock just for this. I don't have a preference. Cleaned things up a bit I think. All green. All green. thanks!