ENH: Add NumIndex for indexic of any numeric type

Author: topper-123

doc/source/whatsnew/v1.3.0.rst

@@ -23,6 +23,30 @@ including other versions of pandas.
 Enhancements
 ~~~~~~~~~~~~
 
+.. _whatsnew_130.numindex:
+
+NumIndex: New Index type than can hold all numpy numeric dtypes
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+We've added :class:`NumIndex`, a index type whose dtype can be any of int64/32/16/8,
+uint64/32/16/8 or float64/32 numpy dtypes. For example:
+
+.. ipython:: python
+
+    idx = pd.NumIndex(range(1_000), dtype="uint16")
+    idx
+    ser = pd.Series(range(1_000), index=idx)
+    ser
+
+To use ``NumIndex``, you need to instantiate it directly, as e.g.
+``pd.Index(range(1_000), dtype="uint16")`` will return a ``UInt64Index`` for reasons
+of backwards compatibility.
+
+``NumIndex`` will become the default index for numeric indexes in Pandas 2.0 and
+``Int64Index``, ``UInt64Index`` and ``Float64Index`` will be deprecated in a
+future version of Pandas and removed in version 2.0.
+
+
 .. _whatsnew_130.read_csv_json_http_headers:
 
 Custom HTTP(s) headers when reading csv or json files

pandas/__init__.py

@@ -78,6 +78,7 @@
     UInt64Index,
     RangeIndex,
     Float64Index,
+    NumIndex,
     MultiIndex,
     IntervalIndex,
     TimedeltaIndex,

pandas/_libs/join.pyx

@@ -263,6 +263,9 @@ ctypedef fused join_t:
     int16_t
     int32_t
     int64_t
+    uint8_t
+    uint16_t
+    uint32_t
     uint64_t
 
 

pandas/_testing/__init__.py

@@ -33,6 +33,7 @@
     is_period_dtype,
     is_sequence,
     is_timedelta64_dtype,
+    pandas_dtype,
 )
 
 import pandas as pd
@@ -44,6 +45,7 @@
     Index,
     IntervalIndex,
     MultiIndex,
+    NumIndex,
     RangeIndex,
     Series,
     bdate_range,
@@ -309,6 +311,20 @@ def makeFloatIndex(k=10, name=None):
     return Index(values * (10 ** np.random.randint(0, 9)), name=name)
 
 
+def makeNumIndex(k=10, name=None, *, dtype):
+    dtype = pandas_dtype(dtype)
+    if dtype.kind == "i":
+        return NumIndex(list(range(k)), name=name, dtype=dtype)
+    elif dtype.kind == "u":
+        start_num = 2 ** (dtype.itemsize * 8 - 1)
+        return NumIndex([start_num + i for i in range(k)], name=name, dtype=dtype)
+    elif dtype.kind == "f":
+        values = sorted(np.random.random_sample(k)) - np.random.random_sample(1)
+        return NumIndex(values, name=name, dtype=dtype)
+    else:
+        raise NotImplementedError()
+
+
 def makeDateIndex(k: int = 10, freq="B", name=None, **kwargs) -> DatetimeIndex:
     dt = datetime(2000, 1, 1)
     dr = bdate_range(dt, periods=k, freq=freq, name=name)

pandas/_testing/asserters.py

@@ -265,8 +265,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
@@ -313,7 +314,10 @@ def _check_types(left, right, obj="Index"):
             assert_class_equal(left, right, exact=exact, obj=obj)
 
             # Skip exact dtype checking when `check_categorical` is False
-            if check_categorical:
+            if check_categorical and "categorical" in (
+                left.inferred_type,
+                right.inferred_type,
+            ):
                 assert_attr_equal("dtype", left, right, obj=obj)
 
             # allow string-like to have different inferred_types

pandas/conftest.py

@@ -471,6 +471,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),
@@ -522,7 +532,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)
     ]
 )

pandas/core/algorithms.py

@@ -143,11 +143,11 @@ def _ensure_data(values: ArrayLike) -> tuple[np.ndarray, DtypeObj]:
             # until our algos support uint8 directly (see TODO)
             return np.asarray(values).astype("uint64"), np.dtype("bool")
         elif is_signed_integer_dtype(values):
-            return ensure_int64(values), np.dtype("int64")
+            return ensure_int64(values), values.dtype
         elif is_unsigned_integer_dtype(values):
-            return ensure_uint64(values), np.dtype("uint64")
+            return ensure_uint64(values), values.dtype
         elif is_float_dtype(values):
-            return ensure_float64(values), np.dtype("float64")
+            return ensure_float64(values), values.dtype
         elif is_complex_dtype(values):
 
             # ignore the fact that we are casting to float

pandas/core/api.py

@@ -57,6 +57,7 @@
     Int64Index,
     IntervalIndex,
     MultiIndex,
+    NumIndex,
     PeriodIndex,
     RangeIndex,
     TimedeltaIndex,

pandas/core/dtypes/generic.py

@@ -100,6 +100,7 @@ def _check(cls, inst) -> bool:
             "rangeindex",
             "float64index",
             "uint64index",
+            "numindex",
             "multiindex",
             "datetimeindex",
             "timedeltaindex",

pandas/core/indexes/api.py

@@ -25,6 +25,7 @@
     Float64Index,
     Int64Index,
     NumericIndex,
+    NumIndex,
     UInt64Index,
 )
 from pandas.core.indexes.period import PeriodIndex
@@ -46,6 +47,7 @@
 __all__ = [
     "Index",
     "MultiIndex",
+    "NumIndex",
     "NumericIndex",
     "Float64Index",
     "Int64Index",

pandas/core/indexes/base.py

@@ -390,6 +390,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)
@@ -441,6 +442,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):
@@ -2939,7 +2942,11 @@ def union(self, other, sort=None):
                 # float | [u]int -> float  (the special case)
                 # <T>   | <T>    -> T
                 # <T>   | <U>    -> object
-                if not (is_integer_dtype(self.dtype) and is_integer_dtype(other.dtype)):
+                if is_float_dtype(self.dtype) and is_float_dtype(other.dtype):
+                    dtype = np.dtype("object")
+                elif not (
+                    is_integer_dtype(self.dtype) and is_integer_dtype(other.dtype)
+                ):
                     dtype = np.dtype("float64")
                 else:
                     # one is int64 other is uint64
@@ -5465,8 +5472,9 @@ 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 = super()._map_values(mapper, na_action=na_action)
+        new_values = self._map_values(mapper, na_action=na_action)
 
         attributes = self._get_attributes_dict()
 
@@ -5485,6 +5493,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

pandas/core/indexes/category.py

@@ -25,6 +25,7 @@
     ensure_platform_int,
     is_categorical_dtype,
     is_scalar,
+    pandas_dtype,
 )
 from pandas.core.dtypes.missing import (
     is_valid_na_for_dtype,
@@ -277,6 +278,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.