Move offsets downcasting to MultiIndex._engine

Author: GianlucaFicarelli

pandas/_libs/index.pyx

@@ -699,21 +699,19 @@ cdef class BaseMultiIndexCodesEngine:
     Keys are located by first locating each component against the respective
     level, then locating (the integer representation of) codes.
     """
-    def __init__(self, object levels, object labels,
-                 ndarray[uint64_t, ndim=1] offsets):
+    def __init__(self, object levels, object labels, ndarray offsets):
         """
         Parameters
         ----------
         levels : list-like of numpy arrays
             Levels of the MultiIndex.
         labels : list-like of numpy arrays of integer dtype
             Labels of the MultiIndex.
-        offsets : numpy array of uint64 dtype
+        offsets : numpy array of int dtype
             Pre-calculated offsets, one for each level of the index.
         """
         self.levels = levels
-        # Downcast the type if possible, to prevent upcasting when shifting codes:
-        self.offsets = offsets.astype(np.min_scalar_type(offsets[0]), copy=False)
+        self.offsets = offsets
 
         # Transform labels in a single array, and add 2 so that we are working
         # with positive integers (-1 for NaN becomes 1). This enables us to

pandas/core/indexes/multi.py

@@ -1201,7 +1201,9 @@ def _engine(self):
         # equivalent to sorting lexicographically the codes themselves. Notice
         # that each level needs to be shifted by the number of bits needed to
         # represent the _previous_ ones:
-        offsets = np.concatenate([lev_bits[1:], [0]]).astype("uint64")
+        offsets = np.concatenate([lev_bits[1:], [0]])
+        # Downcast the type if possible, to prevent upcasting when shifting codes:
+        offsets = offsets.astype(np.min_scalar_type(int(offsets[0])))
 
         # Check the total number of bits needed for our representation:
         if lev_bits[0] > 64: