ENH: Add NumIndex for indexic of any numeric type

pandas-dev · jreback · Aug 5, 2021 · Apr 25, 2021 · Apr 25, 2021 · Apr 26, 2021
commit 9739f851600fea6f1b2de96b1b61c08a820ba95c
diff --git a/pandas/__init__.py b/pandas/__init__.py
@@ -75,6 +75,7 @@
     UInt64Index,
     RangeIndex,
     Float64Index,
+    NumIndex,
     MultiIndex,
     IntervalIndex,
     TimedeltaIndex,

diff --git a/pandas/_libs/join.pyx b/pandas/_libs/join.pyx
@@ -265,6 +265,9 @@ ctypedef fused join_t:
     int16_t
     int32_t
     int64_t
+    uint8_t
+    uint16_t
+    uint32_t
     uint64_t
 
 

diff --git a/pandas/_testing/__init__.py b/pandas/_testing/__init__.py
@@ -50,6 +50,7 @@
     Int64Index,
     IntervalIndex,
     MultiIndex,
+    NumIndex,
     RangeIndex,
     Series,
     UInt64Index,

diff --git a/pandas/_testing/asserters.py b/pandas/_testing/asserters.py
@@ -267,8 +267,9 @@ def assert_index_equal(
     right : Index
     exact : bool or {'equiv'}, default 'equiv'
         Whether to check the Index class, dtype and inferred_type
-        are identical. If 'equiv', then RangeIndex can be substituted for
-        Int64Index as well.
+        are identical. If 'equiv', RangeIndex can be substituted for
+        Int64Index and signed integer dtypes will be equivalent to each other, unsigned
+        integer to each other and float dtypes equivalent to each other.
     check_names : bool, default True
         Whether to check the names attribute.
     check_less_precise : bool or int, default False

diff --git a/pandas/conftest.py b/pandas/conftest.py
@@ -460,6 +460,16 @@ def _create_mi_with_dt64tz_level():
     "uint": tm.makeUIntIndex(100),
     "range": tm.makeRangeIndex(100),
     "float": tm.makeFloatIndex(100),
+    "num_int64": tm.makeNumIndex(100, dtype="int64"),
+    "num_int32": tm.makeNumIndex(100, dtype="int32"),
+    "num_int16": tm.makeNumIndex(100, dtype="int16"),
+    "num_int8": tm.makeNumIndex(100, dtype="int8"),
+    "num_uint64": tm.makeNumIndex(100, dtype="uint64"),
+    "num_uint32": tm.makeNumIndex(100, dtype="uint32"),
+    "num_uint16": tm.makeNumIndex(100, dtype="uint16"),
+    "num_uint8": tm.makeNumIndex(100, dtype="uint8"),
+    "num_float64": tm.makeNumIndex(100, dtype="float64"),
+    "num_float32": tm.makeNumIndex(100, dtype="float32"),
     "bool": tm.makeBoolIndex(10),
     "categorical": tm.makeCategoricalIndex(100),
     "interval": tm.makeIntervalIndex(100),
@@ -511,7 +521,10 @@ def index_flat(request):
     params=[
         key
         for key in indices_dict
-        if key not in ["int", "uint", "range", "empty", "repeats"]
+        if not (
+            key in ["int", "uint", "range", "empty", "repeats"]
+            or key.startswith("num_")
+        )
         and not isinstance(indices_dict[key], MultiIndex)
     ]
 )

diff --git a/pandas/core/api.py b/pandas/core/api.py
@@ -57,6 +57,7 @@
     Int64Index,
     IntervalIndex,
     MultiIndex,
+    NumIndex,
     PeriodIndex,
     RangeIndex,
     TimedeltaIndex,

diff --git a/pandas/core/dtypes/generic.py b/pandas/core/dtypes/generic.py
@@ -100,6 +100,7 @@ def _check(cls, inst) -> bool:
             "rangeindex",
             "float64index",
             "uint64index",
+            "numindex",
             "multiindex",
             "datetimeindex",
             "timedeltaindex",

diff --git a/pandas/core/indexes/api.py b/pandas/core/indexes/api.py
@@ -23,6 +23,7 @@
     Float64Index,
     Int64Index,
     NumericIndex,
+    NumIndex,
     UInt64Index,
 )
 from pandas.core.indexes.period import PeriodIndex
@@ -44,6 +45,7 @@
 __all__ = [
     "Index",
     "MultiIndex",
+    "NumIndex",
     "NumericIndex",
     "Float64Index",
     "Int64Index",

diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -385,6 +385,7 @@ def __new__(
             )
 
         from pandas.core.arrays import PandasArray
+        from pandas.core.indexes.numeric import NumIndex
         from pandas.core.indexes.range import RangeIndex
 
         name = maybe_extract_name(name, data, cls)
@@ -436,6 +437,8 @@ def __new__(
             return Index._simple_new(data, name=name)
 
         # index-like
+        elif isinstance(data, NumIndex) and dtype is None:
+            return NumIndex(data, name=name, copy=copy)
         elif isinstance(data, (np.ndarray, Index, ABCSeries)):
 
             if isinstance(data, ABCMultiIndex):
@@ -5699,6 +5702,7 @@ def map(self, mapper, na_action=None):
             a MultiIndex will be returned.
         """
         from pandas.core.indexes.multi import MultiIndex
+        from pandas.core.indexes.numeric import NumIndex
 
         new_values = self._map_values(mapper, na_action=na_action)
 
@@ -5719,6 +5723,9 @@ def map(self, mapper, na_action=None):
             # empty
             attributes["dtype"] = self.dtype
 
+        if isinstance(self, NumIndex):
+            return NumIndex(new_values, **attributes)
+
         return Index(new_values, **attributes)
 
     # TODO: De-duplicate with map, xref GH#32349

diff --git a/pandas/core/indexes/category.py b/pandas/core/indexes/category.py
@@ -21,6 +21,7 @@
 from pandas.core.dtypes.common import (
     is_categorical_dtype,
     is_scalar,
+    pandas_dtype,
 )
 from pandas.core.dtypes.missing import (
     is_valid_na_for_dtype,
@@ -280,6 +281,26 @@ def _is_dtype_compat(self, other) -> Categorical:
 
         return other
 
+    @doc(Index.astype)
+    def astype(self, dtype, copy: bool = True) -> Index:
+        from pandas import NumIndex
+
+        dtype = pandas_dtype(dtype)
+
+        cat = self.categories
+        if isinstance(cat, NumIndex):
+            try:
+                cat._validate_dtype(dtype)
+            except ValueError:
+                pass
+            else:
+                new_values = self._data.astype(dtype, copy=copy)
+                # pass copy=False because any copying has been done in the
+                #  _data.astype call above
+                return NumIndex(new_values, name=self.name, copy=False)
+
+        return super().astype(dtype, copy=copy)
+
     def equals(self, other: object) -> bool:
         """
         Determine if two CategoricalIndex objects contain the same elements.

diff --git a/pandas/core/indexes/numeric.py b/pandas/core/indexes/numeric.py
@@ -23,6 +23,7 @@
 
 from pandas.core.dtypes.cast import astype_nansafe
 from pandas.core.dtypes.common import (
+    is_categorical_dtype,
     is_dtype_equal,
     is_extension_array_dtype,
     is_float,

diff --git a/pandas/tests/indexes/common.py b/pandas/tests/indexes/common.py
@@ -13,7 +13,7 @@
 from pandas.core.dtypes.dtypes import CategoricalDtype
 
 import pandas as pd
-from pandas import (
+from pandas import (  # noqa
     CategoricalIndex,
     DatetimeIndex,
     Float64Index,
@@ -533,7 +533,7 @@ def test_hasnans_isnans(self, index_flat):
             return
         elif isinstance(index, DatetimeIndexOpsMixin):
             values[1] = iNaT
-        elif isinstance(index, (Int64Index, UInt64Index, RangeIndex)):
+        elif issubclass(index.dtype.type, np.integer):
             return
         else:
             values[1] = np.nan
@@ -572,7 +572,7 @@ def test_fillna(self, index):
 
             if isinstance(index, DatetimeIndexOpsMixin):
                 values[1] = iNaT
-            elif isinstance(index, (Int64Index, UInt64Index, RangeIndex)):
+            elif issubclass(index.dtype.type, np.integer):
                 return
             else:
                 values[1] = np.nan

diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py
@@ -24,6 +24,7 @@
     Float64Index,
     Int64Index,
     IntervalIndex,
+    NumIndex,
     PeriodIndex,
     RangeIndex,
     Series,
@@ -713,6 +714,8 @@ def test_map_dictlike(self, index, mapper):
         if index.empty:
             # to match proper result coercion for uints
             expected = Index([])
+        elif isinstance(index, NumIndex):
+            expected = NumIndex(np.arange(len(index), 0, -1), dtype=index.dtype)
         else:
             expected = Index(np.arange(len(index), 0, -1))
 

diff --git a/pandas/tests/indexes/test_common.py b/pandas/tests/indexes/test_common.py
@@ -21,6 +21,7 @@
     CategoricalIndex,
     DatetimeIndex,
     MultiIndex,
+    NumIndex,
     PeriodIndex,
     RangeIndex,
     TimedeltaIndex,
@@ -261,7 +262,8 @@ def test_drop_duplicates(self, index_flat, keep):
         # make unique index
         holder = type(index)
         unique_values = list(set(index))
-        unique_idx = holder(unique_values)
+        dtype = index.dtype if isinstance(index, NumIndex) else None
+        unique_idx = holder(unique_values, dtype=dtype)
 
         # make duplicated index
         n = len(unique_idx)
@@ -289,7 +291,8 @@ def test_drop_duplicates_no_duplicates(self, index_flat):
         else:
             holder = type(index)
             unique_values = list(set(index))
-            unique_idx = holder(unique_values)
+            dtype = index.dtype if isinstance(index, NumIndex) else None
+            unique_idx = holder(unique_values, dtype=dtype)
 
         # check on unique index
         expected_duplicated = np.array([False] * len(unique_idx), dtype="bool")

diff --git a/pandas/tests/indexes/test_numpy_compat.py b/pandas/tests/indexes/test_numpy_compat.py
@@ -5,14 +5,12 @@
     DatetimeIndex,
     Float64Index,
     Index,
-    Int64Index,
     PeriodIndex,
-    RangeIndex,
     TimedeltaIndex,
-    UInt64Index,
 )
 import pandas._testing as tm
 from pandas.core.indexes.datetimelike import DatetimeIndexOpsMixin
+from pandas.core.indexes.numeric import NumericIndex
 
 
 @pytest.mark.parametrize(
@@ -51,7 +49,7 @@ def test_numpy_ufuncs_basic(index, func):
         with tm.external_error_raised((TypeError, AttributeError)):
             with np.errstate(all="ignore"):
                 func(index)
-    elif isinstance(index, (Float64Index, Int64Index, UInt64Index, RangeIndex)):
+    elif isinstance(index, NumericIndex):
         # coerces to float (e.g. np.sin)
         with np.errstate(all="ignore"):
             result = func(index)
@@ -96,7 +94,7 @@ def test_numpy_ufuncs_other(index, func, request):
         with tm.external_error_raised(TypeError):
             func(index)
 
-    elif isinstance(index, (Float64Index, Int64Index, UInt64Index, RangeIndex)):
+    elif isinstance(index, NumericIndex):
         # Results in bool array
         result = func(index)
         assert isinstance(result, np.ndarray)

diff --git a/pandas/tests/indexes/test_setops.py b/pandas/tests/indexes/test_setops.py
@@ -31,7 +31,21 @@
 
 COMPATIBLE_INCONSISTENT_PAIRS = [
     (np.float64, np.int64),
+    (np.float64, np.int32),
+    (np.float64, np.int16),
+    (np.float64, np.int8),
     (np.float64, np.uint64),
+    (np.float64, np.uint32),
+    (np.float64, np.uint16),
+    (np.float64, np.uint8),
+    (np.float32, np.int64),
+    (np.float32, np.int32),
+    (np.float32, np.int16),
+    (np.float32, np.int8),
+    (np.float32, np.uint64),
+    (np.float32, np.uint32),
+    (np.float32, np.uint16),
+    (np.float32, np.uint8),
 ]
-Original file line number
+Diff line change
@@ Expand Up / @@ -265,6 +265,9 @@ ctypedef fused join_t: @@
         int16_t
         int32_t
         int64_t
+        uint8_t
+        uint16_t
+        uint32_t
         uint64_t
@@ Expand Down @@