From f721ab921c50d67863f775c2f7a54cb90956ed6b Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Sat, 4 Jan 2020 15:21:01 -0800 Subject: [PATCH] Implement ExtensionIndex --- pandas/core/indexes/category.py | 22 +++++----------------- pandas/core/indexes/datetimelike.py | 18 ++++++++++-------- pandas/core/indexes/extension.py | 24 ++++++++++++++++++++++++ pandas/core/indexes/interval.py | 4 ++-- 4 files changed, 41 insertions(+), 27 deletions(-) diff --git a/pandas/core/indexes/category.py b/pandas/core/indexes/category.py index b44b83cec7b71..f60f2140553ed 100644 --- a/pandas/core/indexes/category.py +++ b/pandas/core/indexes/category.py @@ -9,7 +9,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 ( @@ -29,7 +28,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 @@ -66,7 +65,7 @@ typ="method", overwrite=True, ) -class CategoricalIndex(Index, accessor.PandasDelegate): +class CategoricalIndex(ExtensionIndex, accessor.PandasDelegate): """ Index based on an underlying :class:`Categorical`. @@ -722,20 +721,9 @@ 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) - - take_nd = take + def take_nd(self, *args, **kwargs): + """Alias for `take`""" + return self.take(*args, **kwargs) @Appender(_index_shared_docs["_maybe_cast_slice_bound"]) def _maybe_cast_slice_bound(self, label, side, kind): diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index 7001b32c66204..4b0dd750a2f03 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -40,7 +40,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) @@ -80,7 +85,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. """ @@ -245,16 +250,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 diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 3c98d31e34b7d..6f581c4ebb594 100644 --- a/pandas/core/indexes/extension.py +++ b/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) diff --git a/pandas/core/indexes/interval.py b/pandas/core/indexes/interval.py index 35cdf840a55b2..2677ed4d95158 100644 --- a/pandas/core/indexes/interval.py +++ b/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"]