|
50 | 50 | is_datetime64tz_dtype,
|
51 | 51 | is_datetime_or_timedelta_dtype,
|
52 | 52 | is_dtype_equal,
|
| 53 | + is_extension_array_dtype, |
53 | 54 | is_float,
|
54 | 55 | is_float_dtype,
|
55 | 56 | is_integer,
|
|
59 | 60 | is_number,
|
60 | 61 | is_object_dtype,
|
61 | 62 | is_scalar,
|
| 63 | + is_signed_integer_dtype, |
| 64 | + is_unsigned_integer_dtype, |
62 | 65 | )
|
63 | 66 | from pandas.core.dtypes.dtypes import IntervalDtype
|
64 | 67 | from pandas.core.dtypes.missing import is_valid_na_for_dtype
|
@@ -521,6 +524,7 @@ def _maybe_convert_i8(self, key):
|
521 | 524 | original = key
|
522 | 525 | if is_list_like(key):
|
523 | 526 | key = ensure_index(key)
|
| 527 | + key = self._maybe_convert_numeric_to_64bit(key) |
524 | 528 |
|
525 | 529 | if not self._needs_i8_conversion(key):
|
526 | 530 | return original
|
@@ -566,6 +570,20 @@ def _maybe_convert_i8(self, key):
|
566 | 570 |
|
567 | 571 | return key_i8
|
568 | 572 |
|
| 573 | + def _maybe_convert_numeric_to_64bit(self, idx: Index) -> Index: |
| 574 | + # IntervalTree only supports 64 bit numpy array |
| 575 | + dtype = idx.dtype |
| 576 | + if np.issubclass_(idx.dtype.type, np.number): |
| 577 | + return idx |
| 578 | + elif is_signed_integer_dtype(dtype) and dtype != np.int64: |
| 579 | + return idx.astype(np.int64) |
| 580 | + elif is_unsigned_integer_dtype(dtype) and dtype != np.uint64: |
| 581 | + return idx.astype(np.uint64) |
| 582 | + elif is_float_dtype(dtype) and dtype != np.float64: |
| 583 | + return idx.astype(np.float64) |
| 584 | + else: |
| 585 | + return idx |
| 586 | + |
569 | 587 | def _searchsorted_monotonic(self, label, side: Literal["left", "right"] = "left"):
|
570 | 588 | if not self.is_non_overlapping_monotonic:
|
571 | 589 | raise KeyError(
|
|
0 commit comments