Skip to content

Commit a9c32c6

Browse files
committed
BUG: implement new engine for codes-based MultiIndex indexing
closes pandas-dev#18519 closes pandas-dev#18818 closes pandas-dev#18520 closes pandas-dev#18485 closes pandas-dev#15994 closes pandas-dev#19086
1 parent 3aece1d commit a9c32c6

File tree

2 files changed

+108
-10
lines changed

2 files changed

+108
-10
lines changed

doc/source/whatsnew/v0.23.0.txt

+4
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,10 @@ Indexing
318318
- Bug in :func:`MultiIndex.get_level_values` which would return an invalid index on level of ints with missing values (:issue:`17924`)
319319
- Bug in :func:`MultiIndex.remove_unused_levels` which would fill nan values (:issue:`18417`)
320320
- Bug in :func:`MultiIndex.from_tuples`` which would fail to take zipped tuples in python3 (:issue:`18434`)
321+
- Bug in :func:`MultiIndex.get_loc`` which would fail to automatically cast values between float and int (:issue:`18818`, :issue:`15994`)
322+
- Bug in :func:`MultiIndex.get_loc`` which would cast boolean to integer labels (:issue:`19086`)
323+
- Bug in :func:`MultiIndex.get_loc`` which would fail to locate keys containing NaN (:issue:`18485`)
324+
- Bug in :func:`MultiIndex.get_loc`` in large :class:`MultiIndex` which would fail when levels had different dtypes (:issue:`18520`)
321325
- Bug in :class:`Index` construction from list of mixed type tuples (:issue:`18505`)
322326
- Bug in :func:`Index.drop` when passing a list of both tuples and non-tuples (:issue:`18304`)
323327
- Bug in :class:`IntervalIndex` where empty and purely NA data was constructed inconsistently depending on the construction method (:issue:`18421`)

pandas/core/indexes/multi.py

+104-10
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,101 @@
5050
target_klass='MultiIndex or list of tuples'))
5151

5252

53+
class BaseMultiIndexCodesEngine(object):
54+
def __init__(self, levels, labels, offsets, **kwargs):
55+
self._levels = levels
56+
self._offsets = offsets
57+
58+
# Map each combination to an integer
59+
codes = (np.array(labels, dtype='int64').T + 1).astype('uint64')
60+
lab_ints = self._codes_to_ints(codes)
61+
62+
# Initialize underlying index
63+
self._base.__init__(self, lambda: lab_ints, len(lab_ints), **kwargs)
64+
65+
def get_indexer(self, target, method=None, limit=None):
66+
level_codes = [self._levels[lev].get_indexer(codes, method=method) + 1
67+
for lev, codes in enumerate(zip(*target))]
68+
69+
keys_int = self._codes_to_ints(np.array(level_codes, dtype='uint64').T)
70+
71+
if method is not None:
72+
# keys must be sorted - the engine already is
73+
order = np.argsort(keys_int)
74+
keys_int = keys_int[order]
75+
sup_meth = getattr(self._base, 'get_{}_indexer'.format(method))
76+
indexer = sup_meth(self, keys_int, limit=limit)
77+
indexer = indexer[order]
78+
else:
79+
indexer = self._base.get_indexer(self, keys_int)
80+
81+
return indexer
82+
83+
def get_loc(self, key):
84+
if libindex.is_definitely_invalid_key(key):
85+
raise TypeError("'{key}' is an invalid key".format(key=key))
86+
if not isinstance(key, tuple):
87+
raise KeyError(key)
88+
try:
89+
idces = [0 if isna(val) else self._levels[lev].get_loc(val) + 1
90+
for lev, val in enumerate(key)]
91+
except KeyError:
92+
raise KeyError(key)
93+
idces = np.array(idces, ndmin=2, dtype='uint64')
94+
95+
key_int = self._codes_to_ints(idces)[0]
96+
97+
return self._base.get_loc(self, key_int)
98+
99+
def get_indexer_non_unique(self, target):
100+
# This needs to be overridden just because the default one works on
101+
# target._values, and target can be itself a MultiIndex.
102+
103+
level_codes = [self._levels[lev].get_indexer(codes) + 1
104+
for lev, codes in enumerate(zip(*target))]
105+
codes = np.array(level_codes, dtype='uint64').T
106+
keys_int = self._codes_to_ints(codes)
107+
108+
indexer = self._base.get_indexer_non_unique(self, keys_int)
109+
110+
return indexer
111+
112+
def __contains__(self, val):
113+
try:
114+
self.get_loc(val)
115+
return True
116+
except (KeyError, TypeError, ValueError):
117+
return False
118+
119+
120+
class MultiIndexUIntEngine(BaseMultiIndexCodesEngine, libindex.UInt64Engine):
121+
"""
122+
Manage a MultiIndex by mapping label combinations to positive integers.
123+
"""
124+
_base = libindex.UInt64Engine
125+
126+
def _codes_to_ints(self, codes):
127+
# Shift:
128+
codes <<= self._offsets
129+
# Now sum and OR are in fact interchangeable:
130+
return np.bitwise_or.reduce(codes, axis=1)
131+
132+
133+
class MultiIndexPyIntEngine(BaseMultiIndexCodesEngine, libindex.ObjectEngine):
134+
"""
135+
In those (extreme) cases in which the number of possible label combinations
136+
overflows the 64 bits integers, use an ObjectEngine containing Python
137+
integers.
138+
"""
139+
_base = libindex.ObjectEngine
140+
141+
def _codes_to_ints(self, codes):
142+
# Shift:
143+
codes = codes.astype('object') << self._offsets
144+
# Now sum and OR are in fact interchangeable:
145+
return np.bitwise_or.reduce(codes, axis=1)
146+
147+
53148
class MultiIndex(Index):
54149
"""
55150
A multi-level, or hierarchical, index object for pandas objects
@@ -691,16 +786,15 @@ def _get_level_number(self, level):
691786

692787
@cache_readonly
693788
def _engine(self):
789+
# Find powers of 2 which dominate level sizes - including -1 for NaN:
790+
lev_bits = np.cumsum(np.ceil(np.log2([len(l) + 1 for l in
791+
self.levels[::-1]])))[::-1]
792+
offsets = np.concatenate([lev_bits[1:], [0]]).astype('uint')
694793

695-
# choose our engine based on our size
696-
# the hashing based MultiIndex for larger
697-
# sizes, and the MultiIndexOjbect for smaller
698-
# xref: https://github.com/pandas-dev/pandas/pull/16324
699-
l = len(self)
700-
if l > 10000:
701-
return libindex.MultiIndexHashEngine(lambda: self, l)
702-
703-
return libindex.MultiIndexObjectEngine(lambda: self.values, l)
794+
if lev_bits[0] > 64:
795+
# The levels would overflow a 64 bit integer - use Python integers:
796+
return MultiIndexPyIntEngine(self.levels, self.labels, offsets)
797+
return MultiIndexUIntEngine(self.levels, self.labels, offsets)
704798

705799
@property
706800
def values(self):
@@ -1889,7 +1983,7 @@ def get_indexer(self, target, method=None, limit=None, tolerance=None):
18891983
if tolerance is not None:
18901984
raise NotImplementedError("tolerance not implemented yet "
18911985
'for MultiIndex')
1892-
indexer = self._get_fill_indexer(target, method, limit)
1986+
indexer = self._engine.get_indexer(target, method, limit)
18931987
elif method == 'nearest':
18941988
raise NotImplementedError("method='nearest' not implemented yet "
18951989
'for MultiIndex; see GitHub issue 9365')

0 commit comments

Comments
 (0)