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

Author: lukemanley

doc/source/whatsnew/v1.5.0.rst

@@ -318,7 +318,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`)

pandas/core/indexes/multi.py

@@ -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,
@@ -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)
-
-            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
@@ -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
@@ -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
@@ -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)
 
         # 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