ENH: allow storing ExtensionArrays in Index (#43930)

Author: jbrockmendel

doc/source/whatsnew/v1.4.0.rst

@@ -90,6 +90,43 @@ be removed in the future, see :ref:`here <whatsnew_140.deprecations.int64_uint64
 
 See :ref:`here <advanced.numericindex>` for more about :class:`NumericIndex`.
 
+
+.. _whatsnew_140.enhancements.ExtensionIndex:
+
+Index can hold arbitrary ExtensionArrays
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Until now, passing a custom :class:`ExtensionArray` to ``pd.Index`` would cast the
+array to ``object`` dtype. Now :class:`Index` can directly hold arbitrary ExtensionArrays (:issue:`43930`).
+
+*Previous behavior*:
+
+.. ipython:: python
+
+   arr = pd.array([1, 2, pd.NA])
+   idx = pd.Index(arr)
+
+In the old behavior, ``idx`` would be object-dtype:
+
+*Previous behavior*:
+
+.. code-block:: ipython
+
+   In [1]: idx
+   Out[1]: Index([1, 2, <NA>], dtype='object')
+
+With the new behavior, we keep the original dtype:
+
+*New behavior*:
+
+.. ipython:: python
+
+   idx
+
+One exception to this is ``SparseArray``, which will continue to cast to numpy
+dtype until pandas 2.0. At that point it will retain its dtype like other
+ExtensionArrays.
+
 .. _whatsnew_140.enhancements.styler:
 
 Styler

pandas/_libs/index.pyx

@@ -33,7 +33,9 @@ from pandas._libs import (
     hashtable as _hash,
 )
 
+from pandas._libs.lib cimport eq_NA_compat
 from pandas._libs.missing cimport (
+    C_NA as NA,
     checknull,
     is_matching_na,
 )
@@ -62,7 +64,7 @@ cdef ndarray _get_bool_indexer(ndarray values, object val):
     if values.descr.type_num == cnp.NPY_OBJECT:
         # i.e. values.dtype == object
         if not checknull(val):
-            indexer = values == val
+            indexer = eq_NA_compat(values, val)
 
         else:
             # We need to check for _matching_ NA values

pandas/_libs/lib.pxd

@@ -1 +1,6 @@
+from numpy cimport ndarray
+
+
 cdef bint c_is_list_like(object, bint) except -1
+
+cpdef ndarray eq_NA_compat(ndarray[object] arr, object key)

pandas/_libs/lib.pyx

@@ -3050,6 +3050,27 @@ def is_bool_list(obj: list) -> bool:
     return True
 
 
+cpdef ndarray eq_NA_compat(ndarray[object] arr, object key):
+    """
+    Check for `arr == key`, treating all values as not-equal to pd.NA.
+
+    key is assumed to have `not isna(key)`
+    """
+    cdef:
+        ndarray[uint8_t, cast=True] result = np.empty(len(arr), dtype=bool)
+        Py_ssize_t i
+        object item
+
+    for i in range(len(arr)):
+        item = arr[i]
+        if item is C_NA:
+            result[i] = False
+        else:
+            result[i] = item == key
+
+    return result
+
+
 def dtypes_all_equal(list types not None) -> bool:
     """
     Faster version for:

pandas/_testing/asserters.py

@@ -404,9 +404,9 @@ def _get_ilevel_values(index, level):
     # skip exact index checking when `check_categorical` is False
     if check_exact and check_categorical:
         if not left.equals(right):
-            diff = (
-                np.sum((left._values != right._values).astype(int)) * 100.0 / len(left)
-            )
+            mismatch = left._values != right._values
+
+            diff = np.sum(mismatch.astype(int)) * 100.0 / len(left)
             msg = f"{obj} values are different ({np.round(diff, 5)} %)"
             raise_assert_detail(obj, msg, left, right)
     else:

pandas/conftest.py

@@ -67,6 +67,14 @@
     MultiIndex,
 )
 
+try:
+    import pyarrow as pa
+except ImportError:
+    has_pyarrow = False
+else:
+    del pa
+    has_pyarrow = True
+
 # Until https://github.com/numpy/numpy/issues/19078 is sorted out, just suppress
 suppress_npdev_promotion_warning = pytest.mark.filterwarnings(
     "ignore:Promotion of numbers and bools:FutureWarning"
@@ -549,7 +557,15 @@ def _create_mi_with_dt64tz_level():
     "mi-with-dt64tz-level": _create_mi_with_dt64tz_level(),
     "multi": _create_multiindex(),
     "repeats": Index([0, 0, 1, 1, 2, 2]),
+    "nullable_int": Index(np.arange(100), dtype="Int64"),
+    "nullable_uint": Index(np.arange(100), dtype="UInt16"),
+    "nullable_float": Index(np.arange(100), dtype="Float32"),
+    "nullable_bool": Index(np.arange(100).astype(bool), dtype="boolean"),
+    "string-python": Index(pd.array(tm.makeStringIndex(100), dtype="string[python]")),
 }
+if has_pyarrow:
+    idx = Index(pd.array(tm.makeStringIndex(100), dtype="string[pyarrow]"))
+    indices_dict["string-pyarrow"] = idx
 
 
 @pytest.fixture(params=indices_dict.keys())

pandas/core/arrays/masked.py

@@ -714,10 +714,7 @@ def value_counts(self, dropna: bool = True) -> Series:
         data = self._data[~self._mask]
         value_counts = Index(data).value_counts()
 
-        # TODO(ExtensionIndex)
-        # if we have allow Index to hold an ExtensionArray
-        # this is easier
-        index = value_counts.index._values.astype(object)
+        index = value_counts.index
 
         # if we want nans, count the mask
         if dropna:
@@ -727,10 +724,9 @@ def value_counts(self, dropna: bool = True) -> Series:
             counts[:-1] = value_counts
             counts[-1] = self._mask.sum()
 
-            index = Index(
-                np.concatenate([index, np.array([self.dtype.na_value], dtype=object)]),
-                dtype=object,
-            )
+            index = index.insert(len(index), self.dtype.na_value)
+
+        index = index.astype(self.dtype)
 
         mask = np.zeros(len(counts), dtype="bool")
         counts = IntegerArray(counts, mask)

pandas/core/arrays/string_.py

@@ -470,7 +470,9 @@ def max(self, axis=None, skipna: bool = True, **kwargs) -> Scalar:
     def value_counts(self, dropna: bool = True):
         from pandas import value_counts
 
-        return value_counts(self._ndarray, dropna=dropna).astype("Int64")
+        result = value_counts(self._ndarray, dropna=dropna).astype("Int64")
+        result.index = result.index.astype(self.dtype)
+        return result
 
     def memory_usage(self, deep: bool = False) -> int:
         result = self._ndarray.nbytes

pandas/core/arrays/string_arrow.py

@@ -313,6 +313,13 @@ def __getitem__(
         elif isinstance(item, tuple):
             item = unpack_tuple_and_ellipses(item)
 
+        # error: Non-overlapping identity check (left operand type:
+        # "Union[Union[int, integer[Any]], Union[slice, List[int],
+        # ndarray[Any, Any]]]", right operand type: "ellipsis")
+        if item is Ellipsis:  # type: ignore[comparison-overlap]
+            # TODO: should be handled by pyarrow?
+            item = slice(None)
+
         if is_scalar(item) and not is_integer(item):
             # e.g. "foo" or 2.5
             # exception message copied from numpy
@@ -615,8 +622,7 @@ def value_counts(self, dropna: bool = True) -> Series:
         # No missing values so we can adhere to the interface and return a numpy array.
         counts = np.array(counts)
 
-        # Index cannot hold ExtensionArrays yet
-        index = Index(type(self)(values)).astype(object)
+        index = Index(type(self)(values))
 
         return Series(counts, index=index).astype("Int64")
 

pandas/core/dtypes/common.py

@@ -1325,12 +1325,8 @@ def is_bool_dtype(arr_or_dtype) -> bool:
         # now we use the special definition for Index
 
     if isinstance(arr_or_dtype, ABCIndex):
-
-        # TODO(jreback)
-        # we don't have a boolean Index class
-        # so its object, we need to infer to
-        # guess this
-        return arr_or_dtype.is_object() and arr_or_dtype.inferred_type == "boolean"
+        # Allow Index[object] that is all-bools or Index["boolean"]
+        return arr_or_dtype.inferred_type == "boolean"
     elif isinstance(dtype, ExtensionDtype):
         return getattr(dtype, "_is_boolean", False)