Skip to content

Commit

Permalink
Move offsets downcasting to MultiIndex._engine
Browse files Browse the repository at this point in the history
  • Loading branch information
GianlucaFicarelli committed Apr 25, 2024
1 parent 3ab8601 commit 180098a
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 6 deletions.
8 changes: 3 additions & 5 deletions pandas/_libs/index.pyx
Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion pandas/core/indexes/multi.py
Expand Up @@ -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:
Expand Down

0 comments on commit 180098a

Please sign in to comment.