PERF: Do not init cache in RangeIndex.take

Improve performance when passing an array to RangeIndex.take, DataFrame.loc, or DataFrame.iloc and the DataFrame is using a RangeIndex

Author: GianlucaFicarelli

doc/source/whatsnew/v2.0.2.rst

@@ -46,6 +46,7 @@ Bug fixes
 
 Other
 ~~~~~
+- Improved performance when passing an array to :meth:`RangeIndex.take`, :meth:`DataFrame.loc`, or :meth:`DataFrame.iloc` and the DataFrame is using a RangeIndex (:issue:`53387`)
 - Raised a better error message when calling :func:`Series.dt.to_pydatetime` with :class:`ArrowDtype` with ``pyarrow.date32`` or ``pyarrow.date64`` type (:issue:`52812`)
 
 .. ---------------------------------------------------------------------------

pandas/core/indexes/base.py

@@ -1157,7 +1157,6 @@ def take(
             taken = values.take(
                 indices, allow_fill=allow_fill, fill_value=self._na_value
             )
-        # _constructor so RangeIndex-> Index with an int64 dtype
         return self._constructor._simple_new(taken, name=self.name)
 
     @final
@@ -5543,6 +5542,10 @@ def equals(self, other: Any) -> bool:
         if not isinstance(other, Index):
             return False
 
+        if len(self) != len(other):
+            # quickly return if the lengths are different
+            return False
+
         if is_object_dtype(self.dtype) and not is_object_dtype(other.dtype):
             # if other is not object, use other's logic for coercion
             return other.equals(self)

pandas/core/indexes/range.py

@@ -49,6 +49,7 @@
 
 if TYPE_CHECKING:
     from pandas._typing import (
+        Axis,
         Dtype,
         NaPosition,
         Self,
@@ -1102,3 +1103,43 @@ def _arith_method(self, other, op):
         except (ValueError, TypeError, ZeroDivisionError):
             # test_arithmetic_explicit_conversions
             return super()._arith_method(other, op)
+
+    def take(
+        self,
+        indices,
+        axis: Axis = 0,
+        allow_fill: bool = True,
+        fill_value=None,
+        **kwargs,
+    ):
+        if kwargs:
+            nv.validate_take((), kwargs)
+        if is_scalar(indices):
+            raise TypeError("Expected indices to be array-like")
+        indices = ensure_platform_int(indices)
+        allow_fill = self._maybe_disallow_fill(allow_fill, fill_value, indices)
+        assert allow_fill is False, "allow_fill isn't supported by RangeIndex"
+
+        if len(indices) == 0:
+            taken = np.array([], dtype=self.dtype)
+        else:
+            ind_max = indices.max()
+            if ind_max >= len(self):
+                raise IndexError(
+                    f"index {ind_max} is out of bounds for axis 0 with size {len(self)}"
+                )
+            ind_min = indices.min()
+            if ind_min < -len(self):
+                raise IndexError(
+                    f"index {ind_min} is out of bounds for axis 0 with size {len(self)}"
+                )
+            taken = indices.astype(self.dtype, casting="safe")
+            if ind_min < 0:
+                taken %= len(self)
+            if self.step != 1:
+                taken *= self.step
+            if self.start != 0:
+                taken += self.start
+
+        # _constructor so RangeIndex-> Index with an int64 dtype
+        return self._constructor._simple_new(taken, name=self.name)

pandas/tests/indexes/ranges/test_indexing.py

@@ -76,10 +76,46 @@ def test_take_fill_value(self):
         with pytest.raises(ValueError, match=msg):
             idx.take(np.array([1, 0, -5]), fill_value=True)
 
+    def test_take_raises_index_error(self):
+        idx = RangeIndex(1, 4, name="xxx")
+
         msg = "index -5 is out of bounds for (axis 0 with )?size 3"
         with pytest.raises(IndexError, match=msg):
             idx.take(np.array([1, -5]))
 
+        msg = "index -4 is out of bounds for (axis 0 with )?size 3"
+        with pytest.raises(IndexError, match=msg):
+            idx.take(np.array([1, -4]))
+
+        # no errors
+        result = idx.take(np.array([1, -3]))
+        expected = Index([2, 1], dtype=np.int64, name="xxx")
+        tm.assert_index_equal(result, expected)
+
+    def test_take_accepts_empty_array(self):
+        idx = RangeIndex(1, 4, name="foo")
+        result = idx.take(np.array([]))
+        expected = Index([], dtype=np.int64, name="foo")
+        tm.assert_index_equal(result, expected)
+
+        # empty index
+        idx = RangeIndex(0, name="foo")
+        result = idx.take(np.array([]))
+        expected = Index([], dtype=np.int64, name="foo")
+        tm.assert_index_equal(result, expected)
+
+    def test_take_accepts_non_int64_array(self):
+        idx = RangeIndex(1, 4, name="foo")
+        result = idx.take(np.array([2, 1], dtype=np.uint32))
+        expected = Index([3, 2], dtype=np.int64, name="foo")
+        tm.assert_index_equal(result, expected)
+
+    def test_take_when_index_has_step(self):
+        idx = RangeIndex(1, 11, 3, name="foo")  # [1, 4, 7, 10]
+        result = idx.take(np.array([1, 0, -1, -4]))
+        expected = Index([4, 1, 10, 1], dtype=np.int64, name="foo")
+        tm.assert_index_equal(result, expected)
+
 
 class TestWhere:
     def test_where_putmask_range_cast(self):

pandas/tests/indexes/ranges/test_range.py

@@ -272,6 +272,24 @@ def test_cache(self):
         expected = np.arange(0, 100, 10, dtype="int64")
         tm.assert_numpy_array_equal(idx._cache["_data"], expected)
 
+    def test_cache_after_calling_loc_with_array(self):
+        # GH 53387
+        # the cache will contain a _constructor key, so it should be tested separately
+        idx = RangeIndex(0, 100, 10)
+        df = pd.DataFrame({"a": range(10)}, index=idx)
+
+        assert "_data" not in idx._cache
+
+        # take is internally called by loc, but it's also tested explicitly
+        idx.take([3, 0, 1])
+        assert "_data" not in idx._cache
+
+        df.loc[[50]]
+        assert "_data" not in idx._cache
+
+        df.iloc[[5, 6, 7, 8, 9]]
+        assert "_data" not in idx._cache
+
     def test_is_monotonic(self):
         index = RangeIndex(0, 20, 2)
         assert index.is_monotonic_increasing is True