PERF, BENCH: Fix performance issue when indexing into non-unique DatetimeIndex/PeriodIndex. (#27136)

Author: qwhelan

asv_bench/benchmarks/categoricals.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import pandas.util.testing as tm
+import warnings
 
 try:
     from pandas.api.types import union_categoricals
@@ -122,11 +123,16 @@ def setup(self):
         ncats = 100
 
         self.s_str = pd.Series(tm.makeCategoricalIndex(N, ncats)).astype(str)
-        self.s_str_cat = self.s_str.astype("category")
-        self.s_str_cat_ordered = self.s_str_cat.cat.as_ordered()
+        self.s_str_cat = pd.Series(self.s_str, dtype="category")
+        with warnings.catch_warnings(record=True):
+            str_cat_type = pd.CategoricalDtype(set(self.s_str), ordered=True)
+            self.s_str_cat_ordered = self.s_str.astype(str_cat_type)
+
         self.s_int = pd.Series(np.random.randint(0, ncats, size=N))
-        self.s_int_cat = self.s_int.astype("category")
-        self.s_int_cat_ordered = self.s_int_cat.cat.as_ordered()
+        self.s_int_cat = pd.Series(self.s_int, dtype="category")
+        with warnings.catch_warnings(record=True):
+            int_cat_type = pd.CategoricalDtype(set(self.s_int), ordered=True)
+            self.s_int_cat_ordered = self.s_int.astype(int_cat_type)
 
     def time_rank_string(self):
         self.s_str.rank()

asv_bench/benchmarks/indexing.py

@@ -15,6 +15,7 @@
     concat,
     date_range,
     option_context,
+    period_range,
 )
 
 
@@ -93,22 +94,30 @@ def time_loc_slice(self, index, index_structure):
 class NonNumericSeriesIndexing:
 
     params = [
-        ("string", "datetime"),
-        ("unique_monotonic_inc", "nonunique_monotonic_inc"),
+        ("string", "datetime", "period"),
+        ("unique_monotonic_inc", "nonunique_monotonic_inc", "non_monotonic"),
     ]
     param_names = ["index_dtype", "index_structure"]
 
     def setup(self, index, index_structure):
         N = 10 ** 6
-        indexes = {
-            "string": tm.makeStringIndex(N),
-            "datetime": date_range("1900", periods=N, freq="s"),
-        }
-        index = indexes[index]
+        if index == "string":
+            index = tm.makeStringIndex(N)
+        elif index == "datetime":
+            index = date_range("1900", periods=N, freq="s")
+        elif index == "period":
+            index = period_range("1900", periods=N, freq="s")
+        index = index.sort_values()
+        assert index.is_unique and index.is_monotonic_increasing
         if index_structure == "nonunique_monotonic_inc":
             index = index.insert(item=index[2], loc=2)[:-1]
+        elif index_structure == "non_monotonic":
+            index = index[::2].append(index[1::2])
+            assert len(index) == N
         self.s = Series(np.random.rand(N), index=index)
         self.lbl = index[80000]
+        # warm up index mapping
+        self.s[self.lbl]
 
     def time_getitem_label_slice(self, index, index_structure):
         self.s[: self.lbl]

doc/source/whatsnew/v0.25.0.rst

@@ -938,6 +938,8 @@ Performance improvements
 - Improved performance when building :class:`MultiIndex` with at least one :class:`CategoricalIndex` level (:issue:`22044`)
 - Improved performance by removing the need for a garbage collect when checking for ``SettingWithCopyWarning`` (:issue:`27031`)
 - For :meth:`to_datetime` changed default value of cache parameter to ``True`` (:issue:`26043`)
+- Improved performance of :class:`DatetimeIndex` and :class:`PeriodIndex` slicing given non-unique, monotonic data (:issue:`27136`).
+
 
 .. _whatsnew_0250.bug_fixes:
 

pandas/core/indexes/base.py

@@ -4792,7 +4792,6 @@ def get_indexer_non_unique(self, target):
             return pself.get_indexer_non_unique(ptarget)
 
         if self.is_all_dates:
-            self = Index(self.asi8)
             tgt_values = target.asi8
         else:
             tgt_values = target._ndarray_values

pandas/core/indexes/period.py

@@ -9,6 +9,7 @@
 from pandas.util._decorators import Appender, Substitution, cache_readonly
 
 from pandas.core.dtypes.common import (
+    ensure_platform_int,
     is_bool_dtype,
     is_datetime64_any_dtype,
     is_float,
@@ -618,7 +619,7 @@ def get_value(self, series, key):
                 elif grp == freqn:
                     key = Period(asdt, freq=self.freq).ordinal
                     return com.maybe_box(
-                        self, self._engine.get_value(s, key), series, key
+                        self, self._int64index.get_value(s, key), series, key
                     )
                 else:
                     raise KeyError(key)
@@ -627,7 +628,7 @@ def get_value(self, series, key):
 
             period = Period(key, self.freq)
             key = period.value if isna(period) else period.ordinal
-            return com.maybe_box(self, self._engine.get_value(s, key), series, key)
+            return com.maybe_box(self, self._int64index.get_value(s, key), series, key)
 
     @Appender(_index_shared_docs["get_indexer"] % _index_doc_kwargs)
     def get_indexer(self, target, method=None, limit=None, tolerance=None):
@@ -648,6 +649,23 @@ def get_indexer(self, target, method=None, limit=None, tolerance=None):
             tolerance = self._convert_tolerance(tolerance, target)
         return Index.get_indexer(self._int64index, target, method, limit, tolerance)
 
+    @Appender(_index_shared_docs["get_indexer_non_unique"] % _index_doc_kwargs)
+    def get_indexer_non_unique(self, target):
+        target = ensure_index(target)
+
+        if isinstance(target, PeriodIndex):
+            target = target.asi8
+            if hasattr(target, "freq") and target.freq != self.freq:
+                msg = DIFFERENT_FREQ.format(
+                    cls=type(self).__name__,
+                    own_freq=self.freqstr,
+                    other_freq=target.freqstr,
+                )
+                raise IncompatibleFrequency(msg)
+
+        indexer, missing = self._int64index.get_indexer_non_unique(target)
+        return ensure_platform_int(indexer), missing
+
     def _get_unique_index(self, dropna=False):
         """
         wrap Index._get_unique_index to handle NaT
@@ -954,6 +972,12 @@ def base(self):
         )
         return np.asarray(self._data)
 
+    def memory_usage(self, deep=False):
+        result = super().memory_usage(deep=deep)
+        if hasattr(self, "_cache") and "_int64index" in self._cache:
+            result += self._int64index.memory_usage(deep=deep)
+        return result
+
 
 PeriodIndex._add_comparison_ops()
 PeriodIndex._add_numeric_methods_disabled()

pandas/tests/series/indexing/test_datetime.py

@@ -604,6 +604,30 @@ def test_indexing_over_size_cutoff():
         _index._SIZE_CUTOFF = old_cutoff
 
 
+def test_indexing_over_size_cutoff_period_index():
+    # GH 27136
+
+    old_cutoff = _index._SIZE_CUTOFF
+    try:
+        _index._SIZE_CUTOFF = 1000
+
+        n = 1100
+        idx = pd.period_range("1/1/2000", freq="T", periods=n)
+        assert idx._engine.over_size_threshold
+
+        s = pd.Series(np.random.randn(len(idx)), index=idx)
+
+        pos = n - 1
+        timestamp = idx[pos]
+        assert timestamp in s.index
+
+        # it works!
+        s[timestamp]
+        assert len(s.loc[[timestamp]]) > 0
+    finally:
+        _index._SIZE_CUTOFF = old_cutoff
+
+
 def test_indexing_unordered():
     # GH 2437
     rng = date_range(start="2011-01-01", end="2011-01-15")