Implement ExtensionIndex (#30703)

Author: jbrockmendel

pandas/core/indexes/category.py

@@ -10,7 +10,6 @@
 from pandas._libs.hashtable import duplicated_int64
 from pandas._typing import AnyArrayLike
 import pandas.compat as compat
-from pandas.compat.numpy import function as nv
 from pandas.util._decorators import Appender, cache_readonly
 
 from pandas.core.dtypes.common import (
@@ -30,7 +29,7 @@
 import pandas.core.common as com
 import pandas.core.indexes.base as ibase
 from pandas.core.indexes.base import Index, _index_shared_docs, maybe_extract_name
-from pandas.core.indexes.extension import make_wrapped_comparison_op
+from pandas.core.indexes.extension import ExtensionIndex, make_wrapped_comparison_op
 import pandas.core.missing as missing
 from pandas.core.ops import get_op_result_name
 
@@ -67,7 +66,7 @@
     typ="method",
     overwrite=True,
 )
-class CategoricalIndex(Index, accessor.PandasDelegate):
+class CategoricalIndex(ExtensionIndex, accessor.PandasDelegate):
     """
     Index based on an underlying :class:`Categorical`.
 
@@ -723,20 +722,8 @@ def _convert_arr_indexer(self, keyarr):
     def _convert_index_indexer(self, keyarr):
         return self._shallow_copy(keyarr)
 
-    @Appender(_index_shared_docs["take"] % _index_doc_kwargs)
-    def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
-        nv.validate_take(tuple(), kwargs)
-        indices = ensure_platform_int(indices)
-        taken = self._assert_take_fillable(
-            self._data,
-            indices,
-            allow_fill=allow_fill,
-            fill_value=fill_value,
-            na_value=self._data.dtype.na_value,
-        )
-        return self._shallow_copy(taken)
-
     def take_nd(self, *args, **kwargs):
+        """Alias for `take`"""
         warnings.warn(
             "CategoricalIndex.take_nd is deprecated, use CategoricalIndex.take instead",
             FutureWarning,

pandas/core/indexes/datetimelike.py

@@ -41,7 +41,12 @@
 
 from pandas.tseries.frequencies import DateOffset, to_offset
 
-from .extension import inherit_names, make_wrapped_arith_op, make_wrapped_comparison_op
+from .extension import (
+    ExtensionIndex,
+    inherit_names,
+    make_wrapped_arith_op,
+    make_wrapped_comparison_op,
+)
 
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
 
@@ -81,7 +86,7 @@ def wrapper(left, right):
     ["__iter__", "mean", "freq", "freqstr", "_ndarray_values", "asi8", "_box_values"],
     DatetimeLikeArrayMixin,
 )
-class DatetimeIndexOpsMixin(ExtensionOpsMixin):
+class DatetimeIndexOpsMixin(ExtensionIndex, ExtensionOpsMixin):
     """
     Common ops mixin to support a unified interface datetimelike Index.
     """
@@ -246,16 +251,13 @@ def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
         if isinstance(maybe_slice, slice):
             return self[maybe_slice]
 
-        taken = self._assert_take_fillable(
-            self._data,
-            indices,
-            allow_fill=allow_fill,
-            fill_value=fill_value,
-            na_value=NaT,
+        taken = ExtensionIndex.take(
+            self, indices, axis, allow_fill, fill_value, **kwargs
         )
 
         # keep freq in PeriodArray/Index, reset otherwise
         freq = self.freq if is_period_dtype(self) else None
+        assert taken.freq == freq, (taken.freq, freq, taken)
         return self._shallow_copy(taken, freq=freq)
 
     _can_hold_na = True

pandas/core/indexes/extension.py

@@ -3,10 +3,13 @@
 """
 from typing import List
 
+from pandas.compat.numpy import function as nv
 from pandas.util._decorators import cache_readonly
 
+from pandas.core.dtypes.common import ensure_platform_int
 from pandas.core.dtypes.generic import ABCSeries
 
+from pandas.core.arrays import ExtensionArray
 from pandas.core.ops import get_op_result_name
 
 from .base import Index
@@ -152,3 +155,24 @@ def _maybe_unwrap_index(obj):
     if isinstance(obj, Index):
         return obj._data
     return obj
+
+
+class ExtensionIndex(Index):
+    """
+    Index subclass for indexes backed by ExtensionArray.
+    """
+
+    _data: ExtensionArray
+
+    def take(self, indices, axis=0, allow_fill=True, fill_value=None, **kwargs):
+        nv.validate_take(tuple(), kwargs)
+        indices = ensure_platform_int(indices)
+
+        taken = self._assert_take_fillable(
+            self._data,
+            indices,
+            allow_fill=allow_fill,
+            fill_value=fill_value,
+            na_value=self._na_value,
+        )
+        return type(self)(taken, name=self.name)

pandas/core/indexes/interval.py

@@ -58,7 +58,7 @@
 from pandas.tseries.frequencies import to_offset
 from pandas.tseries.offsets import DateOffset
 
-from .extension import inherit_names
+from .extension import ExtensionIndex, inherit_names
 
 _VALID_CLOSED = {"left", "right", "both", "neither"}
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
@@ -213,7 +213,7 @@ def func(intvidx_self, other, sort=False):
     overwrite=True,
 )
 @inherit_names(["is_non_overlapping_monotonic", "mid"], IntervalArray, cache=True)
-class IntervalIndex(IntervalMixin, Index, accessor.PandasDelegate):
+class IntervalIndex(IntervalMixin, ExtensionIndex, accessor.PandasDelegate):
     _typ = "intervalindex"
     _comparables = ["name"]
     _attributes = ["name", "closed"]