ExtensionArray.take default implementation
State: closed

Implements a take interface that's compatible with NumPy and optionally pandas'
NA semantics.

```python
In [1]: import pandas as pd

In [2]: from pandas.tests.extension.decimal.array import *

In [3]: arr = DecimalArray(['1.1', '1.2', '1.3'])

In [4]: arr.take([0, 1, -1])
Out[4]: DecimalArray(array(['1.1', '1.2', '1.3'], dtype=object))

In [5]: arr.take([0, 1, -1], fill_value=float('nan'))
Out[5]: DecimalArray(array(['1.1', '1.2', Decimal('NaN')], dtype=object))
```

Closes https://github.com/pandas-dev/pandas/issues/20640

- [ ] closes #xxxx
- [ ] tests added / passed
- [ ] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [ ] whatsnew entry


Labels:
name: Indexing
description: Related to indexing on series/frames, not to indexes themselves
name: ExtensionArray
description: Extending pandas with custom dtypes or arrays.

Comments:
Hello @TomAugspurger! Thanks for updating the PR.







































Cheers ! There are no PEP8 issues in this Pull Request. :beers: 

##### Comment last updated on April 27, 2018 at 11:03 Hours UTC
cc @jorisvandenbossche 
Oh fun


```python
In [3]: np.array([{'a': 10}])
Out[3]: array([{'a': 10}], dtype=object)

In [4]: np.array([UserDict({'a': 10})])
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-4-400145814467> in <module>()
----> 1 np.array([UserDict({'a': 10})])

/usr/local/Cellar/python/3.6.5/Frameworks/Python.framework/Versions/3.6/lib/python3.6/collections/__init__.py in __getitem__(self, key)
    989         if hasattr(self.__class__, "__missing__"):
    990             return self.__class__.__missing__(self, key)
--> 991         raise KeyError(key)
    992     def __setitem__(self, key, item): self.data[key] = item
    993     def __delitem__(self, key): del self.data[key]

KeyError: 0
```

(we use `UserDict` since pandas had some troubles with just regular `dict`s.
Quick general comment: why make a EA specific take method instead of just a general one that works for both ndarrays and EAs? (and for EA it would just dispatch to the EA). 

I am also not sure if it is worth to add the `_values_for_take`. If we have a general purpose `take` function, an `ExtensionArray.take` implementation could simply be (assuming `self.data` is the underlying numpy array):

```
           def take(self, indexer, fill_value=None):
               ... somehow handle default fill_value ...
               result = take(self.data, indexer, fill_value=fill_value)
               return self.__class__(result)
```

I think it is not too much asked for an ExtensionArray author to implement this themselves.

There is also a problem with what a public `na_value` on ExtensionDtype/ExtensionArray means. Is it the value that is used "publicly"? Or is it the value that is used in the underlying array? Eg for geopandas the one is `None`, the other is `0` (but the 0 nobody should care about, as that is an implementation detail)

> Quick general comment: why make a EA specific take method instead of just a general one that works for both ndarrays and EAs?


I think the way `take_ea` is written, it works on either.

```python
In [1]: from pandas.core.arrays.base import take_ea

In [2]: import numpy as np

In [3]: arr = np.array([1, 2, 3])

In [4]: take_ea(arr, [0, 1, -1])
Out[4]: array([1, 2, 3])
```

Suggestions for a better name?

> I think it is not too much asked for an ExtensionArray author to implement this themselves.

That seems fine.

> There is also a problem with what a public na_value on ExtensionDtype/ExtensionArray means.

This can be private.
> I think the way take_ea is written, it works on either.

The `_from_sequence` in there is only for EAs
How I had it in my head (but which doesn't mean it is better!), was that such a `take` method that you now wrote in algorithms.py would actually only work on numpy arrays (and on EAs as well, but for those we would dispatch), and then EA.take can call that function again but now with a numpy array of data.
Then the `take` method in algorithms.py would be "cleaner" in the sense that it doesn't need to be aware of EA's semantics. 
But trying to look now in more detail to what you did to better see advantages/disadvantages.

We should be able to have a clean implementation of `take_ea` (or `take` or whatever).

The tricky part is fixing the callers like `Block.reindex_indexer` / `Block.get_reindexed_values` / `NDFrame.align`, to call things correctly. These need to call `algos.take(thing, fill_value=fill_value)` with the right fill value for that dtype. Things should work, but it's a bit tricky.
Pushed another WIP if anyone has some thoughts on the general approach. The extension tests all pass but I'm still breaking things elsewhere, and I think things can be simplified.

The main idea is to have a sentinel `_default_na_value` that's used as the default for all reindexing operations (align, etc.). Then in internals, once we have a single-dtype block, we can figure out what default value to actually use based on the block type. In principle, this should work, but I'm struggling with various things, especially around upcasting. Hoping to clean it up today.
Regarding `_default_fill_value`, I have the feeling (again, from only taking a quick look) that we are mixing two different concepts: for `take` it means "no filling", for `reindex` it means "default fill value"
> that we are mixing two different concepts

That feeling is correct. Is that not what we want? I can re-add the allow_fill argument.
> That feeling is correct. Is that not what we want? I can re-add the allow_fill argument.

I am going a bit back and forth on it, not sure ...

Regarding what I said above about the dubious meaning of "na_value" (the physical value that is stored, or the public one), I think this is actually not really a problem. It should just be the public one, so the value that you get if you get an item of an EA that is missing (type of `arr[0]` if the first item is missing). So it is fine to keep this public.
Trying to get my head around the current logic. Some pseudo code flow path in case of a Series reindex operation with EA values:

```
s.reindex(self, ..., fill_value=_default_fill_value)
-> ... -> BlockManager.reindex_indexer(..., fill_value=fill_value)
-> Block.take_nd(..., 
        fill_tuple=(fill_value if fill_value is not _default_fill_value else Block.fill_value,)
-> algos.take_nd(values, ... allow_fill=True, fill_value=fill_tuple[1])
-> [in case of EA] values.take(..., fill_value=fill_value)
   -> if fill_value == _default_fill_value -> no filling
   -> if fill_value == Block.fill_value -> default filling behaviour -> algos.take with passed fill_value
```

(but in case of a reindex/aligment operation, the final `fill_value` actively passed to `values.take(..)` will never be `_default_fill_value` I think?)

So what I find complex about this: a) as said above different meaning of `_default_fill_value` throughout the code base in reindex and take, b) the fact that Block is responsible of passing an appropriate filling value to EA.take, while the EA could perfectly know and do that itself and c) that an EA.take implementation needs to be aware of the ``_default_fill_value`` object (which means we need to make that public).







> EA.take implementation needs to be aware of the _default_fill_value

I think if we use `allow_fill`, then we can just switch to `fill_value=None`. The issue with `fill_value=None` as the default is that an EA may want to use `None` as their NA value, and so wouldn't be able to tell if it's explicit or the default.

This should be able to be simplified greatly. Let's see...
Just brainstorming on an alternative:

```
s.reindex(self, ..., fill_value=_default_fill_value)
-> ... -> BlockManager.reindex_indexer(..., fill_value=fill_value)
-> Block.take_nd(..., fill_tuple=(fill_value,)
-> algos.take_nd(values, ... allow_fill=True, fill_value=fill_tuple[1])
  -> if fill_value == _default_fill_value -> values.take(..., allow_fill=True) (let EA use its default)
  -> else -> values.take(..., allow_fill=True, fill_value=fill_value) (override EA default)
-> values.take(..., allow_fill=False, fill_value=None) # can use their own default for fill_value
   -> if allow_fill == False -> no filling
   -> if allow_fill == True -> filling behaviour 
      -> if fill_value == None -> fill_value = EA.dtype.na_value
```

(I don't know if this scheme is anywhere near clear for somebody that did not write it ...?)

This introduces again the second keyword, but eliminates double usecase of ``_default_fill_value``, and EA.take does not need to be aware of it (the EA author can still use a similar trick if they want to have None as a valid fill_value instead of denoting the default, but that is then up to the EA author).

But of course this is only a single code path, I don't know how it interacts with all other usages of the take machinery.

> I think if we use allow_fill, then we can just switch to fill_value=None. The issue with fill_value=None as the default is that an EA may want to use None as their NA value, and so wouldn't be able to tell if it's explicit or the default.

Actually, I don't think that's true. There would still be the ambiguity. Anyway, we can leave that up to the EA author I think...

Your proposal looks sensible. I'll try to implement it.
Getting there...

This will still fail on `DataFrame.merge(..., how='outer')` that requires reindexing the dataframes. Somewhere in the join units stuff we're determining that 

```python
(Pdb) join_units
[JoinUnit(ExtensionBlock: slice(0, 1, 1), 1 x 3, dtype: decimal, {1: array([ 0,  0,  1,  2, -1])})]
```

should have `np.nan` for it's `upcasted_na`.
Alrighty, I think this is ready for another look. Overall, I think this is better. Thanks @jorisvandenbossche.

Some API questions:

1. Name for the new `take` method?
2. Remove `_values_for_take` and make `take` abstract again? With the idea that the new `take` method makes it straightforward?
3. Are we OK with the names `fill_value` and `allow_fill` for the parameter in `EA.take` / `take`?
4. Public `ExtensionDtype.na_value`?

probably others.l

@TomAugspurger I think the `ExtensionDtype.na_value` is needed and required so that the only NA values inside of an `ExtensionArray` are that specific value.  This relates to making it so that `_can_hold_na=True` for all `ExtensionArray` types (given the discussion at #20815
> 2. Remove _values_for_take and make take abstract again? With the idea that the new take method makes it straightforward?

I would personally remove `_values_for_take` and just have a `take`, to keep the interface simpler. I think most actual ExtensionArrays will not be able to use the `_values_for_take` anyhow, and will need to override `take` itself. This will be the case for geopandas, and I think also for cyberpandas? (as you will want to fill the underlying array with a native numpy scalar, and not with the public `na_value` (ipaddress.IPv4Address(0))?) 
If that is the case, there is even less reason to have it I think.
FYI, there's an occasional failure in `test_merge` for JSONArray. This happens when all of the arrays are the same length (issue is deep inside assert_frame_equal, which calls get_values, which goes to an ndarray, which fails on the UserDicts. I have a fix, but may do it separately).


69e7fe76b540807bb973496ee5a2e3e636b90287 had some notable changes:

1. Reversed order of take keywords
2. Added to extensions API
3. Removed default implementation


Thanks. I think the biggest question is how to the change to `get_empty_dtype_and_na` at https://github.com/pandas-dev/pandas/pull/20814/files#diff-e705e723b2d6e7c0e2a0443f80916abfR5454 can fit in better with the rest of what's going on there. I'm investigating that now.
Oh, another question: do we or NumPy have any code that'll do bounds checking for us? Or should I just do `if (indexer > len(arr)).any(): raise IndexError`
@TomAugspurger So my (now very small) changes that remove `_can_hold_na` as something people would set in `ExtensionArray` are all green in #20819 , so if they pass muster, you might want to merge them in and use it here, since there is a bit of a relation to your PR for `take`
@TomAugspurger see ``maybe_convert_indices`` (in generic.py/_take)
I wasn't really able to share code with `maybe_convert_indices`. It has a couple differences

1. modifies `indicies` in place, we just want bounds checking
2. allows for indices < -1, we want to raise for anything other than that
3. handles `take([], [-1, -1])` differently (we want to allow it)
I think the failure is a random one due to https://github.com/pandas-dev/pandas/pull/20827. Will re-run once that's in.
# [Codecov](https://codecov.io/gh/pandas-dev/pandas/pull/20814?src=pr&el=h1) Report
> Merging [#20814](https://codecov.io/gh/pandas-dev/pandas/pull/20814?src=pr&el=desc) into [master](https://codecov.io/gh/pandas-dev/pandas/commit/6cacdde5630c593999059833b516e1fec60aaf72?src=pr&el=desc) will **increase** coverage by `<.01%`.
> The diff coverage is `100%`.

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

```diff
@@            Coverage Diff             @@
##           master   #20814      +/-   ##
==========================================
+ Coverage   91.77%   91.78%   +<.01%     
==========================================
  Files         153      153              
  Lines       49280    49313      +33     
==========================================
+ Hits        45229    45261      +32     
- Misses       4051     4052       +1
```

| Flag | Coverage Δ | |
|---|---|---|
| #multiple | `90.17% <100%> (ø)` | :arrow_up: |
| #single | `41.88% <42.1%> (ø)` | :arrow_up: |

| [Impacted Files](https://codecov.io/gh/pandas-dev/pandas/pull/20814?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pandas/core/frame.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZnJhbWUucHk=) | `97.22% <ø> (ø)` | :arrow_up: |
| [pandas/core/dtypes/cast.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Nhc3QucHk=) | `88.06% <100%> (+0.2%)` | :arrow_up: |
| [pandas/core/indexing.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvaW5kZXhpbmcucHk=) | `93.14% <100%> (+0.05%)` | :arrow_up: |
| [pandas/core/dtypes/missing.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL21pc3NpbmcucHk=) | `92.94% <100%> (+0.08%)` | :arrow_up: |
| [pandas/core/algorithms.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYWxnb3JpdGhtcy5weQ==) | `94.48% <100%> (+0.08%)` | :arrow_up: |
| [pandas/core/arrays/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvYXJyYXlzL2Jhc2UucHk=) | `83.95% <100%> (ø)` | :arrow_up: |
| [pandas/core/series.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvc2VyaWVzLnB5) | `93.99% <100%> (ø)` | :arrow_up: |
| [pandas/core/dtypes/base.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZHR5cGVzL2Jhc2UucHk=) | `92.1% <100%> (+0.21%)` | :arrow_up: |
| [pandas/core/generic.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2NvcmUvZ2VuZXJpYy5weQ==) | `95.94% <100%> (ø)` | :arrow_up: |
| [pandas/api/extensions/\_\_init\_\_.py](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree#diff-cGFuZGFzL2FwaS9leHRlbnNpb25zL19faW5pdF9fLnB5) | `100% <100%> (ø)` | :arrow_up: |
| ... and [2 more](https://codecov.io/gh/pandas-dev/pandas/pull/20814/diff?src=pr&el=tree-more) | |

------

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

> can you add there a test case with a different fill_value (not None or nan), to make sure it is passed through correctly to EA.take

I don't think we can do that generically, since we don't know what a valid na_value is.

I'll add one in decimal though.
> I don't think we can do that generically, since we don't know what a valid na_value is.

If you fill it with an existing scalar (like `pd.Series(data[:2]).reindex([1, 2], fill_value=data[2])`), then we can expect that to work I think.

But it is actually a good question what should happen when you pass a value that is not accepted by the ExtensionArray. I suppose it is up to the EA to decide on that? (raise an error, or coerce to object)