API: ensure IntervalIndex.left/right are 64bit if numeric part II

Author: topper-123

pandas/core/arrays/interval.py

@@ -35,6 +35,7 @@
     Dtype,
     IntervalClosedType,
     NpDtype,
+    NumpyIndexT,
     PositionalIndexer,
     ScalarIndexer,
     SequenceIndexer,
@@ -55,7 +56,9 @@
     is_list_like,
     is_object_dtype,
     is_scalar,
+    is_signed_integer_dtype,
     is_string_dtype,
+    is_unsigned_integer_dtype,
     needs_i8_conversion,
     pandas_dtype,
 )
@@ -177,6 +180,21 @@
 """
 
 
+def maybe_convert_numeric_to_64bit(arr: NumpyIndexT) -> NumpyIndexT:
+    # IntervalTree only supports 64 bit numpy array
+    dtype = arr.dtype
+    if not np.issubclass_(dtype.type, np.number):
+        return arr
+    elif is_signed_integer_dtype(dtype) and dtype != np.int64:
+        return arr.astype(np.int64)
+    elif is_unsigned_integer_dtype(dtype) and dtype != np.uint64:
+        return arr.astype(np.uint64)
+    elif is_float_dtype(dtype) and dtype != np.float64:
+        return arr.astype(np.float64)
+    else:
+        return arr
+
+
 @Appender(
     _interval_shared_docs["class"]
     % {
@@ -248,6 +266,7 @@ def __new__(
 
             # might need to convert empty or purely na data
             data = _maybe_convert_platform_interval(data)
+            data = maybe_convert_numeric_to_64bit(data)
             left, right, infer_closed = intervals_to_interval_bounds(
                 data, validate_closed=closed is None
             )

pandas/core/indexes/interval.py

@@ -59,8 +59,6 @@
     is_number,
     is_object_dtype,
     is_scalar,
-    is_signed_integer_dtype,
-    is_unsigned_integer_dtype,
 )
 from pandas.core.dtypes.dtypes import IntervalDtype
 from pandas.core.dtypes.missing import is_valid_na_for_dtype
@@ -69,6 +67,7 @@
 from pandas.core.arrays.interval import (
     IntervalArray,
     _interval_shared_docs,
+    maybe_convert_numeric_to_64bit,
 )
 import pandas.core.common as com
 from pandas.core.indexers import is_valid_positional_slice
@@ -520,13 +519,12 @@ def _maybe_convert_i8(self, key):
             The original key if no conversion occurred, int if converted scalar,
             Int64Index if converted list-like.
         """
-        original = key
         if is_list_like(key):
             key = ensure_index(key)
-            key = self._maybe_convert_numeric_to_64bit(key)
+            key = maybe_convert_numeric_to_64bit(key)
 
         if not self._needs_i8_conversion(key):
-            return original
+            return key
 
         scalar = is_scalar(key)
         if is_interval_dtype(key) or isinstance(key, Interval):
@@ -569,20 +567,6 @@ def _maybe_convert_i8(self, key):
 
         return key_i8
 
-    def _maybe_convert_numeric_to_64bit(self, idx: Index) -> Index:
-        # IntervalTree only supports 64 bit numpy array
-        dtype = idx.dtype
-        if np.issubclass_(dtype.type, np.number):
-            return idx
-        elif is_signed_integer_dtype(dtype) and dtype != np.int64:
-            return idx.astype(np.int64)
-        elif is_unsigned_integer_dtype(dtype) and dtype != np.uint64:
-            return idx.astype(np.uint64)
-        elif is_float_dtype(dtype) and dtype != np.float64:
-            return idx.astype(np.float64)
-        else:
-            return idx
-
     def _searchsorted_monotonic(self, label, side: Literal["left", "right"] = "left"):
         if not self.is_non_overlapping_monotonic:
             raise KeyError(

pandas/tests/indexes/interval/test_interval.py

@@ -18,7 +18,10 @@
     timedelta_range,
 )
 import pandas._testing as tm
-from pandas.core.api import Float64Index
+from pandas.core.api import (
+    Float64Index,
+    NumericIndex,
+)
 import pandas.core.common as com
 
 
@@ -435,9 +438,12 @@ def test_maybe_convert_i8_numeric(self, breaks, make_key):
         index = IntervalIndex.from_breaks(breaks)
         key = make_key(breaks)
 
-        # no conversion occurs for numeric
         result = index._maybe_convert_i8(key)
-        assert result is key
+        if not isinstance(result, NumericIndex):
+            assert result is key
+        else:
+            expected = NumericIndex(key)
+            tm.assert_index_equal(result, expected)
 
     @pytest.mark.parametrize(
         "breaks1, breaks2",