PERF: MultiIndex._engine.get_loc() handles non-unique fine

Author: toobaz

doc/source/whatsnew/v0.24.0.txt

@@ -501,6 +501,7 @@ Performance Improvements
 - Improved performance of :func:`Series.describe` in case of numeric dtpyes (:issue:`21274`)
 - Improved performance of :func:`pandas.core.groupby.GroupBy.rank` when dealing with tied rankings (:issue:`21237`)
 - Improved performance of :func:`DataFrame.set_index` with columns consisting of :class:`Period` objects (:issue:`21582`,:issue:`21606`)
+- Improved performance of :func:`MultiIndex.get_loc` for non-unique indexes, which as a consequence does not emit a ``PerformanceWarning`` any more (:issue:`19464`)
 - Improved performance of membership checks in :class:`Categorical` and :class:`CategoricalIndex`
   (i.e. ``x in cat``-style checks are much faster). :meth:`CategoricalIndex.contains`
   is likewise much faster (:issue:`21369`, :issue:`21508`)

pandas/core/indexes/multi.py

@@ -16,6 +16,7 @@
 from pandas.core.dtypes.common import (
     ensure_int64,
     ensure_platform_int,
+    is_integer,
     is_categorical_dtype,
     is_object_dtype,
     is_hashable,
@@ -2197,10 +2198,16 @@ def _maybe_to_slice(loc):
             raise KeyError('Key length ({0}) exceeds index depth ({1})'
                            ''.format(keylen, self.nlevels))
 
-        if keylen == self.nlevels and self.is_unique:
-            return self._engine.get_loc(key)
+        # If the index is monotonic, the code for partial selection or
+        # non-unique index (below) is more efficient than the following:
+        if keylen == self.nlevels and not self.is_monotonic:
+            loc = self._engine.get_loc(key)
+            if not self.is_unique and is_integer(loc):
+                # Indexers expect a slice from indexing a non-unique index
+                loc = slice(loc, loc + 1)
+            return loc
 
-        # -- partial selection or non-unique index
+        # -- partial selection or non-unique index or monotonic index
         # break the key into 2 parts based on the lexsort_depth of the index;
         # the first part returns a continuous slice of the index; the 2nd part
         # needs linear search within the slice
@@ -2213,6 +2220,10 @@ def _maybe_to_slice(loc):
             raise KeyError(key)
 
         if not follow_key:
+            # Indexers expect an integer from indexing a key in a unique index
+            if self.is_unique:
+                # Breaks if we pass a np.int64. TODO: investigate why
+                return int(start)
             return slice(start, stop)
 
         warnings.warn('indexing past lexsort depth may impact performance.',

pandas/tests/frame/test_sort_values_level_as_str.py

@@ -2,7 +2,6 @@
 import pytest
 
 from pandas import DataFrame, Index
-from pandas.errors import PerformanceWarning
 from pandas.util import testing as tm
 from pandas.util.testing import assert_frame_equal
 
@@ -85,14 +84,7 @@ def test_sort_column_level_and_index_label(
                                   ascending=ascending,
                                   axis=1)
 
-    if len(levels) > 1:
-        # Accessing multi-level columns that are not lexsorted raises a
-        # performance warning
-        with tm.assert_produces_warning(PerformanceWarning,
-                                        check_stacklevel=False):
-            assert_frame_equal(result, expected)
-    else:
-        assert_frame_equal(result, expected)
+    assert_frame_equal(result, expected)
 
 
 def test_sort_values_column_index_level_precedence():

pandas/tests/indexing/test_ix.py

@@ -11,7 +11,6 @@
 from pandas.compat import lrange
 from pandas import Series, DataFrame, option_context, MultiIndex
 from pandas.util import testing as tm
-from pandas.errors import PerformanceWarning
 
 
 class TestIX(object):
@@ -187,9 +186,7 @@ def test_ix_general(self):
         df = DataFrame(data).set_index(keys=['col', 'year'])
         key = 4.0, 2012
 
-        # emits a PerformanceWarning, ok
-        with tm.assert_produces_warning(PerformanceWarning):
-            tm.assert_frame_equal(df.loc[key], df.iloc[2:])
+        tm.assert_frame_equal(df.loc[key], df.iloc[2:])
 
         # this is ok
         df.sort_index(inplace=True)

pandas/tests/indexing/test_multiindex.py

@@ -366,10 +366,6 @@ def test_multiindex_perf_warn(self):
                         'joe': ['x', 'x', 'z', 'y'],
                         'jolie': np.random.rand(4)}).set_index(['jim', 'joe'])
 
-        with tm.assert_produces_warning(PerformanceWarning,
-                                        clear=[pd.core.index]):
-            df.loc[(1, 'z')]
-
         df = df.iloc[[2, 1, 3, 0]]
         with tm.assert_produces_warning(PerformanceWarning):
             df.loc[(0, )]