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,
@@ -48,6 +49,7 @@
 
 from pandas.core.dtypes.cast import LossySetitemError
 from pandas.core.dtypes.common import (
+    is_array_like,
     is_categorical_dtype,
     is_dtype_equal,
     is_float_dtype,
@@ -56,7 +58,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,
 )
@@ -180,6 +184,24 @@
 """
 
 
+def maybe_convert_numeric_to_64bit(arr: NumpyIndexT) -> NumpyIndexT:
+    # IntervalTree only supports 64 bit numpy array
+
+    if not is_array_like(arr):
+        return arr
+    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"]
     % {
@@ -304,7 +326,10 @@ def _ensure_simple_new_inputs(
         from pandas.core.indexes.base import ensure_index
 
         left = ensure_index(left, copy=copy)
+        left = maybe_convert_numeric_to_64bit(left)
+
         right = ensure_index(right, copy=copy)
+        right = maybe_convert_numeric_to_64bit(right)
 
         if closed is None and isinstance(dtype, IntervalDtype):
             closed = dtype.closed

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",