Repository: pandas-dev/pandas
Pull request: API: generalized check_array_indexer for validating array-like getitem indexers
Closes https://github.com/pandas-dev/pandas/issues/30738. Also fixes the performance issue for other arrays from https://github.com/pandas-dev/pandas/issues/30744, and related to https://github.com/pandas-dev/pandas/pull/30308#issuecomment-571146305

This generalizes the `check_bool_array_indexer` helper method that we added for 1.0.0 to not be specific to boolean arrays, but any array-like input, and ensures that the output is a proper numpy array that can be used to index into numpy arrays.

I think such a more general "check" is useful, to avoid that all (external+internal) EAs need to do what the test EAs are already doing (checking for integer arrays as well) at https://github.com/pandas-dev/pandas/blob/master/pandas/tests/extension/decimal/array.py#L118-L126, and also to fix the performance issue in a general way (as was now only done for Categorical in https://github.com/pandas-dev/pandas/pull/30747/files)

If we agree on the general idea, I still need to clean up this PR (eg remove the existing `check_bool_array_indexer`, update the extending docs, etc)

cc @TomAugspurger 

Problem statement:
API: positional indexing with IntegerArray
Currently, the following does not work (but probably should):

```
In [12]: arr1 = pd.array([1, 2, 3]) 

In [13]: arr2 = pd.array([0, 2])

In [14]: arr1[arr2]  
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-14-1646c66d4d26> in <module>
----> 1 arr1[arr2]

~/scipy/pandas/pandas/core/arrays/integer.py in __getitem__(self, item)
    375             item = check_bool_array_indexer(self, item)
    376 
--> 377         return type(self)(self._data[item], self._mask[item])
    378 
    379     def _coerce_to_ndarray(self, dtype=None, na_value=lib._no_default):

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices
```

So that raises the following questions:

- This should probably simply work for the above case? (converting the IntegerArray to a numpy integer array, instead of object array, so numpy's indexing works) We might want to combine this with the boolean array checking?
- What if there are missing values? This should probably simply raise an error for now (which is what pandas 0.25 also does), although we could consider propagating an NA value as well, I think.

For Series, this seems to work partly. For `iloc` it works for the case without missing values. For `__getitem__` you get the same error as above.

Hints from issue discussion:

Commit message: API: generalized check_array_indexer for validating array-like indexers

Patch to review:
diff --git a/pandas/api/indexers/__init__.py b/pandas/api/indexers/__init__.py
index 10654eb0888ee..e0050326e791a 100644
--- a/pandas/api/indexers/__init__.py
+++ b/pandas/api/indexers/__init__.py
@@ -2,7 +2,7 @@
 Public API for Rolling Window Indexers.
 """
 
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer, check_bool_array_indexer
 from pandas.core.window.indexers import BaseIndexer
 
-__all__ = ["check_bool_array_indexer", "BaseIndexer"]
+__all__ = ["check_array_indexer", "check_bool_array_indexer", "BaseIndexer"]
diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py
index 9d7359dd9c614..3c717458797b5 100644
--- a/pandas/core/arrays/categorical.py
+++ b/pandas/core/arrays/categorical.py
@@ -39,7 +39,7 @@
 )
 from pandas.core.dtypes.dtypes import CategoricalDtype
 from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
-from pandas.core.dtypes.inference import is_array_like, is_hashable
+from pandas.core.dtypes.inference import is_hashable
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import ops
@@ -54,7 +54,7 @@
 from pandas.core.base import NoNewAttributesMixin, PandasObject, _shared_docs
 import pandas.core.common as com
 from pandas.core.construction import array, extract_array, sanitize_array
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer
 from pandas.core.missing import interpolate_2d
 from pandas.core.ops.common import unpack_zerodim_and_defer
 from pandas.core.sorting import nargsort
@@ -2001,11 +2001,8 @@ def __getitem__(self, key):
             else:
                 return self.categories[i]
 
-        if is_list_like(key) and not is_array_like(key):
-            key = np.asarray(key)
-
-        if com.is_bool_indexer(key):
-            key = check_bool_array_indexer(self, key)
+        if is_list_like(key):
+            key = check_array_indexer(self, key)
 
         result = self._codes[key]
         if result.ndim > 1:
diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py
index 70637026c278d..7fb3d31f34e8e 100644
--- a/pandas/core/arrays/datetimelike.py
+++ b/pandas/core/arrays/datetimelike.py
@@ -42,7 +42,7 @@
 from pandas.core.algorithms import checked_add_with_arr, take, unique1d, value_counts
 from pandas.core.arrays.base import ExtensionArray, ExtensionOpsMixin
 import pandas.core.common as com
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer
 from pandas.core.ops.common import unpack_zerodim_and_defer
 from pandas.core.ops.invalid import invalid_comparison, make_invalid_op
 
@@ -517,8 +517,10 @@ def __getitem__(self, key):
                 return self._box_func(val)
             return type(self)(val, dtype=self.dtype)
 
+        if is_list_like(key):
+            key = check_array_indexer(self, key)
+
         if com.is_bool_indexer(key):
-            key = check_bool_array_indexer(self, key)
             if key.all():
                 key = slice(0, None, None)
             else:
diff --git a/pandas/core/arrays/interval.py b/pandas/core/arrays/interval.py
index 37d2baed2c09e..23f2151bb593a 100644
--- a/pandas/core/arrays/interval.py
+++ b/pandas/core/arrays/interval.py
@@ -40,6 +40,7 @@
 from pandas.core.arrays.categorical import Categorical
 import pandas.core.common as com
 from pandas.core.construction import array
+from pandas.core.indexers import check_array_indexer
 from pandas.core.indexes.base import ensure_index
 
 _VALID_CLOSED = {"left", "right", "both", "neither"}
@@ -495,6 +496,8 @@ def __len__(self) -> int:
         return len(self.left)
 
     def __getitem__(self, value):
+        if is_list_like(value):
+            value = check_array_indexer(self, value)
         left = self.left[value]
         right = self.right[value]
 
diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py
index 47605413ff1a6..50beb58d79b0e 100644
--- a/pandas/core/arrays/masked.py
+++ b/pandas/core/arrays/masked.py
@@ -4,13 +4,17 @@
 
 from pandas._libs import lib, missing as libmissing
 
-from pandas.core.dtypes.common import is_integer, is_object_dtype, is_string_dtype
+from pandas.core.dtypes.common import (
+    is_integer,
+    is_list_like,
+    is_object_dtype,
+    is_string_dtype,
+)
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core.algorithms import take
 from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
-import pandas.core.common as com
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer
 
 if TYPE_CHECKING:
     from pandas._typing import Scalar
@@ -35,8 +39,8 @@ def __getitem__(self, item):
                 return self.dtype.na_value
             return self._data[item]
 
-        elif com.is_bool_indexer(item):
-            item = check_bool_array_indexer(self, item)
+        elif is_list_like(item):
+            item = check_array_indexer(self, item)
 
         return type(self)(self._data[item], self._mask[item])
 
diff --git a/pandas/core/arrays/numpy_.py b/pandas/core/arrays/numpy_.py
index 4db3d3010adaf..903b4ad851907 100644
--- a/pandas/core/arrays/numpy_.py
+++ b/pandas/core/arrays/numpy_.py
@@ -9,6 +9,7 @@
 from pandas.util._decorators import Appender
 from pandas.util._validators import validate_fillna_kwargs
 
+from pandas.core.dtypes.common import is_list_like
 from pandas.core.dtypes.dtypes import ExtensionDtype
 from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
 from pandas.core.dtypes.inference import is_array_like
@@ -18,9 +19,8 @@
 from pandas.core import nanops
 from pandas.core.algorithms import searchsorted, take, unique
 from pandas.core.arrays.base import ExtensionArray, ExtensionOpsMixin
-import pandas.core.common as com
 from pandas.core.construction import extract_array
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer
 from pandas.core.missing import backfill_1d, pad_1d
 
 
@@ -235,8 +235,8 @@ def __getitem__(self, item):
         if isinstance(item, type(self)):
             item = item._ndarray
 
-        elif com.is_bool_indexer(item):
-            item = check_bool_array_indexer(self, item)
+        elif is_list_like(item):
+            item = check_array_indexer(self, item)
 
         result = self._ndarray[item]
         if not lib.is_scalar(item):
diff --git a/pandas/core/arrays/sparse/array.py b/pandas/core/arrays/sparse/array.py
index e2562a375515d..0f21e4e03b8d2 100644
--- a/pandas/core/arrays/sparse/array.py
+++ b/pandas/core/arrays/sparse/array.py
@@ -29,6 +29,7 @@
     is_datetime64_any_dtype,
     is_dtype_equal,
     is_integer,
+    is_list_like,
     is_object_dtype,
     is_scalar,
     is_string_dtype,
@@ -43,6 +44,7 @@
 from pandas.core.base import PandasObject
 import pandas.core.common as com
 from pandas.core.construction import sanitize_array
+from pandas.core.indexers import check_array_indexer
 from pandas.core.missing import interpolate_2d
 import pandas.core.ops as ops
 from pandas.core.ops.common import unpack_zerodim_and_defer
@@ -768,6 +770,9 @@ def __getitem__(self, key):
                 else:
                     key = np.asarray(key)
 
+            if is_list_like(key):
+                key = check_array_indexer(self, key)
+
             if com.is_bool_indexer(key):
                 key = check_bool_indexer(self, key)
 
diff --git a/pandas/core/indexers.py b/pandas/core/indexers.py
index 4d45769d2fea9..c59f6536683b3 100644
--- a/pandas/core/indexers.py
+++ b/pandas/core/indexers.py
@@ -5,7 +5,12 @@
 
 from pandas._typing import AnyArrayLike
 
-from pandas.core.dtypes.common import is_list_like
+from pandas.core.dtypes.common import (
+    is_array_like,
+    is_bool_dtype,
+    is_integer_dtype,
+    is_list_like,
+)
 from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
 
 # -----------------------------------------------------------
@@ -307,3 +312,62 @@ def check_bool_array_indexer(array: AnyArrayLike, mask: AnyArrayLike) -> np.ndar
     if len(result) != len(array):
         raise IndexError(f"Item wrong length {len(result)} instead of {len(array)}.")
     return result
+
+
+def check_array_indexer(array, indexer) -> np.ndarray:
+    """
+    Check if `indexer` is a valid array indexer for `array`.
+
+    `array` and `indexer` are checked to have the same length, and the
+    dtype is validated. If it is an integer or boolean ExtensionArray, it is
+    checked if there are missing values present, and it is converted to
+    the appropriate numpy array.
+
+    .. versionadded:: 1.0.0
+
+    Parameters
+    ----------
+    array : array
+        The array that's being indexed (only used for the length).
+    indexer : array-like
+        The array-like that's used to index.
+
+    Returns
+    -------
+    numpy.ndarray
+        The validated indexer.
+
+    Raises
+    ------
+    IndexError
+        When the lengths don't match.
+    ValueError
+        When `indexer` cannot be converted to a numpy ndarray.
+
+    """
+    import pandas as pd
+
+    if not is_array_like(indexer):
+        indexer = pd.array(indexer)
+    dtype = indexer.dtype
+    if is_bool_dtype(dtype):
+        try:
+            indexer = np.asarray(indexer, dtype=bool)
+        except ValueError:
+            raise ValueError("Cannot mask with a boolean indexer containing NA values")
+
+        # GH26658
+        if len(indexer) != len(array):
+            raise IndexError(
+                f"Item wrong length {len(indexer)} instead of {len(array)}."
+            )
+
+    elif is_integer_dtype(dtype):
+        try:
+            indexer = np.asarray(indexer, dtype=int)
+        except ValueError:
+            raise ValueError(
+                "Cannot index with an integer indexer containing NA values"
+            )
+
+    return indexer
diff --git a/pandas/tests/extension/base/getitem.py b/pandas/tests/extension/base/getitem.py
index dc1f62c4c97c5..5c31d6f04335e 100644
--- a/pandas/tests/extension/base/getitem.py
+++ b/pandas/tests/extension/base/getitem.py
@@ -160,6 +160,38 @@ def test_getitem_boolean_array_mask_raises(self, data):
         with pytest.raises(ValueError):
             s[mask]
 
+    @pytest.mark.parametrize(
+        "idx",
+        [[0, 1, 2], pd.array([0, 1, 2], dtype="Int64"), np.array([0, 1, 2])],
+        ids=["list", "integer-array", "numpy-array"],
+    )
+    def test_getitem_integer_array(self, data, idx):
+        result = data[idx]
+        assert len(result) == 3
+        assert isinstance(result, type(data))
+        expected = data.take([0, 1, 2])
+        self.assert_extension_array_equal(result, expected)
+
+        expected = pd.Series(expected)
+        result = pd.Series(data)[idx]
+        self.assert_series_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "idx",
+        [[0, 1, 2, pd.NA], pd.array([0, 1, 2, pd.NA], dtype="Int64")],
+        ids=["list", "integer-array"],
+    )
+    def test_getitem_integer_with_missing_raises(self, data, idx):
+        msg = "Cannot index with an integer indexer containing NA values"
+        with pytest.raises(ValueError, match=msg):
+            data[idx]
+
+        # TODO this raises KeyError about labels not found (it tries label-based)
+        # import pandas._testing as tm
+        # s = pd.Series(data, index=[tm.rands(4) for _ in range(len(data))])
+        # with pytest.raises(ValueError, match=msg):
+        #    s[idx]
+
     def test_getitem_slice(self, data):
         # getitem[slice] should return an array
         result = data[slice(0)]  # empty
diff --git a/pandas/tests/extension/decimal/array.py b/pandas/tests/extension/decimal/array.py
index 85bd5f7a33fe1..12a906b5afbe1 100644
--- a/pandas/tests/extension/decimal/array.py
+++ b/pandas/tests/extension/decimal/array.py
@@ -117,13 +117,7 @@ def __getitem__(self, item):
         else:
             # array, slice.
             if pd.api.types.is_list_like(item):
-                if not pd.api.types.is_array_like(item):
-                    item = pd.array(item)
-                dtype = item.dtype
-                if pd.api.types.is_bool_dtype(dtype):
-                    item = pd.api.indexers.check_bool_array_indexer(self, item)
-                elif pd.api.types.is_integer_dtype(dtype):
-                    item = np.asarray(item, dtype="int")
+                item = pd.api.indexers.check_array_indexer(self, item)
             return type(self)(self._data[item])
 
     def take(self, indexer, allow_fill=False, fill_value=None):
diff --git a/pandas/tests/extension/json/array.py b/pandas/tests/extension/json/array.py
index 17bc2773aad19..fe83200152794 100644
--- a/pandas/tests/extension/json/array.py
+++ b/pandas/tests/extension/json/array.py
@@ -82,11 +82,8 @@ def __getitem__(self, item):
             # slice
             return type(self)(self.data[item])
         else:
-            if not pd.api.types.is_array_like(item):
-                item = pd.array(item)
-            dtype = item.dtype
-            if pd.api.types.is_bool_dtype(dtype):
-                item = pd.api.indexers.check_bool_array_indexer(self, item)
+            item = pd.api.indexers.check_array_indexer(self, item)
+            if pd.api.types.is_bool_dtype(item.dtype):
                 return self._from_sequence([x for x, m in zip(self, item) if m])
             # integer
             return type(self)([self.data[i] for i in item])


Code review comments:
diff_hunk: @@ -2,7 +2,7 @@
 Public API for Rolling Window Indexers.
 """
 
-from pandas.core.indexers import check_bool_array_indexer
+from pandas.core.indexers import check_array_indexer, check_bool_array_indexer
path: pandas/api/indexers/__init__.py
text: @user1:
Is there any value in keeping both?

@author:
> Is there any value in keeping both?

No, I don't think so. Just need to clean-up if we are agree on the path forward.

@user1:
Yeah, that seems fine.

diff_hunk: @@ -307,3 +312,62 @@ def check_bool_array_indexer(array: AnyArrayLike, mask: AnyArrayLike) -> np.ndar
     if len(result) != len(array):
         raise IndexError(f"Item wrong length {len(result)} instead of {len(array)}.")
     return result
+
+
+def check_array_indexer(array, indexer) -> np.ndarray:
+    """
+    Check if `indexer` is a valid array indexer for `array`.
+
+    `array` and `indexer` are checked to have the same length, and the
+    dtype is validated. If it is an integer or boolean ExtensionArray, it is
+    checked if there are missing values present, and it is converted to
+    the appropriate numpy array.
+
+    .. versionadded:: 1.0.0
+
+    Parameters
+    ----------
+    array : array
+        The array that's being indexed (only used for the length).
+    indexer : array-like
path: pandas/core/indexers.py
text: @user1:
In a few places above, you've done `is_list_like`, but here we require an array (with a dtype).

Thoughts on what we want? Requiring an array is certainly easier, so that we don't have to infer the types. But users may be passing arbitrary objects to `__getitem__`.

@author:
We actually don't require an array with a dtype. The first thing that this function does is:

```
    if not is_array_like(indexer):
        indexer = pd.array(indexer)
```

to deal with eg lists.

So I probably meant to update the array into "list-like" instead of "array-like"