API: take interface for (Extension)Array-likes

[jorisvandenbossche]

Triggered by #20582, I was looking at the take implementation in ExtensionArray and Categorical (which is already an ExtensionArray subclass) and in the rest of pandas:

• ExtensionArray.take currently uses the "internal pandas"-like behaviour for take: -1 is an indicator for missing value (the behaviour we need for reindexing etc)
• Series.take actually uses the numpy behaviour, where negative values (including -1) start counting from the end of the array-like.

To illustrate the difference with a small example:

In [9]: pd.Categorical(['a', 'b', 'c']).take([0, -1])
Out[9]: 
[a, NaN]
Categories (3, object): [a, b, c]

In [10]: pd.Series(['a', 'b', 'c']).take([0, -1])
Out[10]: 
0    a
1    c
dtype: object

This difference is a bit unfortunate IMO. If ExtensionArray.take is a public method (which it is right now), it would be nice if it has consistent behaviour with Series.take.
If we agree on that, I was thinking about following options:

• make ExtensionArray.take private for now (eg require a _take method for the interface) and keep the "internal pandas"-like behaviour
• make ExtensionArray.take default behaviour consistent with Series.take, but still have the allow_fill/fill_value arguments so that when they are specified it has the "internal pandas"-like behavour (so that internal code that expects this behaviour which already passes those keywords keeps working)

[jreback]

.take() was originally a numpy method so -1 indicates a position rather than a missing value. I don't think we can easily change this, so would make EA consistent here.

[jorisvandenbossche]

@jreback argued (for good reason) in #20582 that having to deal with bounds checking / correct behaviour for empty arrays in each ExtensionArray is rather error prone and complex code in each array implementation.

So it would be good that we have in pandas a take implementation that does this for you (and can also be used in our own ExtensionArrays, to eg fix #20664).
The current take_1d however does currently no bounds checks whatsoever.

take_1d is internally used in many cases where the indexer passed to take_1d is the result of eg get_indexer, so this ensures that the indexer is correct, and for those no out-of-bounds checks are needed.

If we want to expose our take implementation to extension array authors, I would personally create a new take functions which just calls the existing take_1d, but additionally does first bounds checking.
That way, we keep on using the faster take_1d functions internally where we know bounds checks are not needed, and we can make the new take method public for external usage as well.

[TomAugspurger]

I would personally create a new take functions which just calls the existing take_1d, but additionally does first bounds checking.

+1.

[TomAugspurger]

Only other thing to add is that

• I don't want to break API on Categorical.take
• I want consistency between Categorical.take / _take and ExtensionArray.take / _take. I care less about consistency with Series than I do within EAs, though it'd be nice to have consistency here.

If we need, then we can make ExtensionArray._take the official method that's part of the interface.

[jorisvandenbossche]

I don't want to break API on Categorical.take

Do you think there are many people relying on the fact that -1 returns NaN for Categorical.take?
We can deprecate instead of just breaking it.
We in any case need to somehow fix Series.take, because now it dispatches to ExtensionArray.take, which gives wrong results for Series (because it has this difference in behaviour for -1).

[TomAugspurger]

Do you think there are many people relying on the fact that -1 returns NaN for Categorical.take?

Probably not :)

We in any case need to somehow fix Series.take, because now it dispatches to ExtensionArray.take

Hmm, so you're saying that (on master) we've already changed Series[Categorical].take? I didn't realize that... OK, that makes me less certain on what do do here (making a _take vs. just changing ExtensionArray/Categorical.take. I'd support either I think.).

[jorisvandenbossche]

Hmm, so you're saying that (on master) we've already changed Series[Categorical].take

No, we didn't change this on master. It's an already existing bug (but only from 0.22, at least in 0.20 it was still working correctly). It's just that if we keep the difference in behaviour between Array.take and Series.take regarding -1, we should not just dispatch from Series.take to Array.take, but deal with this difference in behaviour.
But for me that's a reason to actually align the default behaviour between Series and Array (Array.take still needs the filling behaviour with -1 which Series.take does not have, but it does not need to be the default).

[jorisvandenbossche]

@jreback @TomAugspurger more comments on this one?

What I am now thinking might be best:

• Make all our take implementation follow numpy semantics by default (meaning -1 gets converted to a positive indexer and does not indicate missingness). Consequences:
  • changes the ExtensionArray.take API to follow this, still provide allow_fill/fill_values keywords to optionally use -1 as missing indicator
  • change Categorical.take as well, but do this with deprecation path
  • This should fix BUG: Categorical.take and Series([Categorical]).take is inconsistent with other dtypes #20664
• Provide a top-level take function
  • This would be a wrapper around take_1d, but does conversion of negative to positive indices by default, does bounds checking, ..
  • This would be exposed as a general public take function (similar to factorize and unique) and can be used by extension array authors to apply it on their underlying values if possible.

[jorisvandenbossche]

The current example implementation (in the docstring) of:

           def take(self, indexer, allow_fill=True, fill_value=None):
               indexer = np.asarray(indexer)
               mask = indexer == -1

               # take on empty array not handled as desired by numpy
               # in case of -1 (all missing take)
               if not len(self) and mask.all():
                   return type(self)([np.nan] * len(indexer))

               result = self.data.take(indexer)
               result[mask] = np.nan  # NA for this type
               return type(self)(result)

would then become something like (not fully sure about the fill_value handling, need to check that):

           def take(self, indexer, allow_fill=True, fill_value=None):
               if fill_value is None:
                   fill_value = np.nan  # NA for this type
               result = pd.take(self.data, indexer, allow_fill=allow_fill, fill_value=fill_value)
               return type(self)(result)

[jreback]

I don't think we need to elevate .take to a top-level. Its just another way to do indexing, which we already have many, and its not settable, so not terrible useful. Why have another advertised method of doing the same thing that (.iloc) can do.

Also a bit -1 (pun) on changing the semantics for this; honestly I don't think the negative indexing is useful and the filling missing values IS useful (but numpy has trouble with this as it will never change dtypes; we just make this work).

Instead I like your 2nd option, but expose API things that we want EA authors (and internal routines) to use that are not quite first-class top-level things (e.g. .factorize() and .unique() are useful generally in there own right).

maybe in: pandas.core.arrays.api ?