Skip to content

REF/PERF: MultiIndex.get_locs to use boolean arrays internally #46330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Mar 18, 2022
Merged

REF/PERF: MultiIndex.get_locs to use boolean arrays internally #46330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
9412158
MultiIndex.get_locs to use boolean arrays internally
lukemanley Mar 11, 2022
4797c0a
whatsnew
lukemanley Mar 11, 2022
2f70e0d
mypy
lukemanley Mar 11, 2022
eb142b4
Merge branch 'main' into multiindex-get-locs-bool-arrays
lukemanley Mar 12, 2022
7ee139e
update comments
lukemanley Mar 12, 2022
a8da325
Merge branch 'main' into multiindex-get-locs-bool-arrays
lukemanley Mar 12, 2022
7ae1992
update whatsnew
lukemanley Mar 12, 2022
83c2291
Merge remote-tracking branch 'upstream/main' into multiindex-get-locs…
lukemanley Mar 16, 2022
3381e13
Merge remote-tracking branch 'upstream/main' into multiindex-get-locs…
lukemanley Mar 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.5.0.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ Performance improvements
- Performance improvement in :meth:`.GroupBy.diff` (:issue:`16706`)
- Performance improvement in :meth:`.GroupBy.transform` when broadcasting values for user-defined functions (:issue:`45708`)
- Performance improvement in :meth:`.GroupBy.transform` for user-defined functions when only a single group exists (:issue:`44977`)
- Performance improvement in :meth:`MultiIndex.get_locs` (:issue:`45681`, :issue:`46040`)
- Performance improvement in :meth:`DataFrame.loc` and :meth:`Series.loc` for tuple-based indexing of a :class:`MultiIndex` (:issue:`45681`, :issue:`46040`, :issue:`46330`)
- Performance improvement in :attr:`MultiIndex.values` when the MultiIndex contains levels of type DatetimeIndex, TimedeltaIndex or ExtensionDtypes (:issue:`46288`)
- Performance improvement in :func:`merge` when left and/or right are empty (:issue:`45838`)
- Performance improvement in :meth:`DataFrame.join` when left and/or right are empty (:issue:`46015`)
Expand Down
190 changes: 68 additions & 122 deletions pandas/core/indexes/multi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,6 @@
get_unanimous_names,
)
from pandas.core.indexes.frozen import FrozenList
from pandas.core.indexes.numeric import Int64Index
from pandas.core.ops.invalid import make_invalid_op
from pandas.core.sorting import (
get_group_index,
Expand Down Expand Up @@ -3160,47 +3159,37 @@ def maybe_mi_droplevels(indexer, levels):
return indexer, result_index

def _get_level_indexer(
self, key, level: int = 0, indexer: Int64Index | None = None
self, key, level: int = 0, indexer: npt.NDArray[np.bool_] | None = None
):
# `level` kwarg is _always_ positional, never name
# return an indexer, boolean array or a slice showing where the key is
# return a boolean array or slice showing where the key is
# in the totality of values
# if the indexer is provided, then use this

level_index = self.levels[level]
level_codes = self.codes[level]

def convert_indexer(start, stop, step, indexer=indexer, codes=level_codes):
# given the inputs and the codes/indexer, compute an indexer set
# if we have a provided indexer, then this need not consider
# the entire labels set
r = np.arange(start, stop, step)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add an explanation (similar to the below) say around L3160, e.g. for a future reader to understand what this algorithm is doing.


if indexer is not None and len(indexer) != len(codes):

# we have an indexer which maps the locations in the labels
# that we have already selected (and is not an indexer for the
# entire set) otherwise this is wasteful so we only need to
# examine locations that are in this set the only magic here is
# that the result are the mappings to the set that we have
# selected
from pandas import Series

mapper = Series(indexer)
indexer = codes.take(ensure_platform_int(indexer))
result = Series(Index(indexer).isin(r).nonzero()[0])
m = result.map(mapper)
# error: Incompatible types in assignment (expression has type
# "ndarray", variable has type "Series")
m = np.asarray(m) # type: ignore[assignment]

# Compute a bool indexer to identify the positions to take.
# If we have an existing indexer, we only need to examine the
# subset of positions where the existing indexer is True.
if indexer is not None:
# we only need to look at the subset of codes where the
# existing indexer equals True
codes = codes[indexer]

if step is None or step == 1:
new_indexer = (codes >= start) & (codes < stop)
else:
# error: Incompatible types in assignment (expression has type
# "ndarray", variable has type "Series")
m = np.zeros(len(codes), dtype=bool) # type: ignore[assignment]
m[np.in1d(codes, r, assume_unique=Index(codes).is_unique)] = True
r = np.arange(start, stop, step, dtype=codes.dtype)
new_indexer = algos.isin(codes, r)

if indexer is None:
return new_indexer

return m
indexer = indexer.copy()
indexer[indexer] = new_indexer
return indexer

if isinstance(key, slice):
# handle a slice, returning a slice if we can
Expand Down Expand Up @@ -3327,62 +3316,41 @@ def get_locs(self, seq):
f"on levels {true_slices}, lexsort depth {self._lexsort_depth}"
)

n = len(self)
# indexer is the list of all positions that we want to take; it
# is created on the first entry in seq and narrowed down as we
# look at remaining entries
indexer = None

if any(x is Ellipsis for x in seq):
raise NotImplementedError(
"MultiIndex does not support indexing with Ellipsis"
)

def _convert_to_indexer(r) -> Int64Index:
# return an indexer
if isinstance(r, slice):
m = np.zeros(n, dtype=bool)
m[r] = True
r = m.nonzero()[0]
elif com.is_bool_indexer(r):
if len(r) != n:
raise ValueError(
"cannot index with a boolean indexer "
"that is not the same length as the "
"index"
)
r = r.nonzero()[0]
return Int64Index(r)
n = len(self)

def _update_indexer(idxr: Index, indexer: Index | None) -> Index:
if indexer is None:
return idxr
indexer_intersection = indexer.intersection(idxr)
if indexer_intersection.empty and not idxr.empty and not indexer.empty:
raise KeyError(seq)
return indexer_intersection
def _to_bool_indexer(indexer) -> npt.NDArray[np.bool_]:
if isinstance(indexer, slice):
new_indexer = np.zeros(n, dtype=np.bool_)
new_indexer[indexer] = True
return new_indexer
return indexer

# a bool indexer for the positions we want to take
indexer: npt.NDArray[np.bool_] | None = None

for i, k in enumerate(seq):

lvl_indexer: npt.NDArray[np.bool_] | slice | None = None

if com.is_bool_indexer(k):
# a boolean indexer, must be the same length!
k = np.asarray(k)
lvl_indexer = _convert_to_indexer(k)
indexer = _update_indexer(lvl_indexer, indexer=indexer)
if len(k) != n:
raise ValueError(
"cannot index with a boolean indexer that "
"is not the same length as the index"
)
lvl_indexer = np.asarray(k)

elif is_list_like(k):
# a collection of labels to include from this level (these
# are or'd)

indexers: Int64Index | None = None
# a collection of labels to include from this level (these are or'd)

# GH#27591 check if this is a single tuple key in the level
try:
# Argument "indexer" to "_get_level_indexer" of "MultiIndex"
# has incompatible type "Index"; expected "Optional[Int64Index]"
lev_loc = self._get_level_indexer(
k, level=i, indexer=indexer # type: ignore[arg-type]
)
lvl_indexer = self._get_level_indexer(k, level=i, indexer=indexer)
except (InvalidIndexError, TypeError, KeyError) as err:
# InvalidIndexError e.g. non-hashable, fall back to treating
# this as a sequence of labels
Expand All @@ -3394,11 +3362,8 @@ def _update_indexer(idxr: Index, indexer: Index | None) -> Index:
# e.g. slice
raise err
try:
# Argument "indexer" to "_get_level_indexer" of "MultiIndex"
# has incompatible type "Index"; expected
# "Optional[Int64Index]"
item_lvl_indexer = self._get_level_indexer(
x, level=i, indexer=indexer # type: ignore[arg-type]
item_indexer = self._get_level_indexer(
x, level=i, indexer=indexer
)
except KeyError:
# ignore not founds; see discussion in GH#39424
Expand All @@ -3418,82 +3383,63 @@ def _update_indexer(idxr: Index, indexer: Index | None) -> Index:
)
continue
else:
idxrs = _convert_to_indexer(item_lvl_indexer)

if indexers is None:
indexers = idxrs
if lvl_indexer is None:
lvl_indexer = _to_bool_indexer(item_indexer)
elif isinstance(item_indexer, slice):
lvl_indexer[item_indexer] = True # type: ignore[index]
else:
indexers = indexers.union(idxrs, sort=False)

else:
idxrs = _convert_to_indexer(lev_loc)
if indexers is None:
indexers = idxrs
else:
indexers = indexers.union(idxrs, sort=False)
lvl_indexer |= item_indexer

if indexers is not None:
indexer = _update_indexer(indexers, indexer=indexer)
else:
if lvl_indexer is None:
# no matches we are done
# test_loc_getitem_duplicates_multiindex_empty_indexer
return np.array([], dtype=np.intp)

elif com.is_null_slice(k):
# empty slice
if indexer is None:
indexer = Index(np.arange(n))
if indexer is None and i == len(seq) - 1:
return np.arange(n, dtype=np.intp)
continue

elif isinstance(k, slice):
else:
# a slice or a single label
lvl_indexer = self._get_level_indexer(k, level=i, indexer=indexer)

# a slice, include BOTH of the labels
# Argument "indexer" to "_get_level_indexer" of "MultiIndex" has
# incompatible type "Index"; expected "Optional[Int64Index]"
lvl_indexer = self._get_level_indexer(
k,
level=i,
indexer=indexer, # type: ignore[arg-type]
)
indexer = _update_indexer(
_convert_to_indexer(lvl_indexer),
indexer=indexer,
)
# update indexer
lvl_indexer = _to_bool_indexer(lvl_indexer)
if indexer is None:
indexer = lvl_indexer
else:
# a single label
lvl_indexer = self._get_loc_level(k, level=i)[0]
indexer = _update_indexer(
_convert_to_indexer(lvl_indexer),
indexer=indexer,
)
indexer &= lvl_indexer
if not np.any(indexer) and np.any(lvl_indexer):
raise KeyError(seq)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this hit by tests?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, covered by test_loc.py > test_missing_key_combination


# empty indexer
if indexer is None:
return np.array([], dtype=np.intp)

assert isinstance(indexer, Int64Index), type(indexer)
indexer = self._reorder_indexer(seq, indexer)

return indexer._values.astype(np.intp, copy=False)
pos_indexer = indexer.nonzero()[0]
return self._reorder_indexer(seq, pos_indexer)

# --------------------------------------------------------------------

def _reorder_indexer(
self,
seq: tuple[Scalar | Iterable | AnyArrayLike, ...],
indexer: Int64Index,
) -> Int64Index:
indexer: npt.NDArray[np.intp],
) -> npt.NDArray[np.intp]:
"""
Reorder an indexer of a MultiIndex (self) so that the label are in the
Reorder an indexer of a MultiIndex (self) so that the labels are in the
same order as given in seq

Parameters
----------
seq : label/slice/list/mask or a sequence of such
indexer: an Int64Index indexer of self
indexer: a position indexer of self

Returns
-------
indexer : a sorted Int64Index indexer of self ordered as seq
indexer : a sorted position indexer of self ordered as seq
"""
# If the index is lexsorted and the list_like label in seq are sorted
# then we do not need to sort
Expand Down