REF: Back IntervalArray by array instead of Index (pandas-dev#36310)

Author: jbrockmendel

doc/source/whatsnew/v1.2.0.rst

@@ -282,6 +282,7 @@ Performance improvements
 - Performance improvement in :meth:`GroupBy.transform` with the ``numba`` engine (:issue:`36240`)
 - ``Styler`` uuid method altered to compress data transmission over web whilst maintaining reasonably low table collision probability (:issue:`36345`)
 - Performance improvement in :meth:`pd.to_datetime` with non-ns time unit for ``float`` ``dtype`` columns (:issue:`20445`)
+- Performance improvement in setting values on a :class:`IntervalArray` (:issue:`36310`)
 
 .. ---------------------------------------------------------------------------
 

pandas/_testing.py

@@ -977,8 +977,14 @@ def assert_interval_array_equal(left, right, exact="equiv", obj="IntervalArray")
     """
     _check_isinstance(left, right, IntervalArray)
 
-    assert_index_equal(left.left, right.left, exact=exact, obj=f"{obj}.left")
-    assert_index_equal(left.right, right.right, exact=exact, obj=f"{obj}.left")
+    kwargs = {}
+    if left._left.dtype.kind in ["m", "M"]:
+        # We have a DatetimeArray or TimedeltaArray
+        kwargs["check_freq"] = False
+
+    assert_equal(left._left, right._left, obj=f"{obj}.left", **kwargs)
+    assert_equal(left._right, right._right, obj=f"{obj}.left", **kwargs)
+
     assert_attr_equal("closed", left, right, obj=obj)
 
 
@@ -989,20 +995,22 @@ def assert_period_array_equal(left, right, obj="PeriodArray"):
     assert_attr_equal("freq", left, right, obj=obj)
 
 
-def assert_datetime_array_equal(left, right, obj="DatetimeArray"):
+def assert_datetime_array_equal(left, right, obj="DatetimeArray", check_freq=True):
     __tracebackhide__ = True
     _check_isinstance(left, right, DatetimeArray)
 
     assert_numpy_array_equal(left._data, right._data, obj=f"{obj}._data")
-    assert_attr_equal("freq", left, right, obj=obj)
+    if check_freq:
+        assert_attr_equal("freq", left, right, obj=obj)
     assert_attr_equal("tz", left, right, obj=obj)
 
 
-def assert_timedelta_array_equal(left, right, obj="TimedeltaArray"):
+def assert_timedelta_array_equal(left, right, obj="TimedeltaArray", check_freq=True):
     __tracebackhide__ = True
     _check_isinstance(left, right, TimedeltaArray)
     assert_numpy_array_equal(left._data, right._data, obj=f"{obj}._data")
-    assert_attr_equal("freq", left, right, obj=obj)
+    if check_freq:
+        assert_attr_equal("freq", left, right, obj=obj)
 
 
 def raise_assert_detail(obj, message, left, right, diff=None, index_values=None):

pandas/core/arrays/interval.py

@@ -31,7 +31,6 @@
 from pandas.core.dtypes.dtypes import IntervalDtype
 from pandas.core.dtypes.generic import (
     ABCDatetimeIndex,
-    ABCIndexClass,
     ABCIntervalIndex,
     ABCPeriodIndex,
     ABCSeries,
@@ -42,7 +41,7 @@
 from pandas.core.arrays.base import ExtensionArray, _extension_array_shared_docs
 from pandas.core.arrays.categorical import Categorical
 import pandas.core.common as com
-from pandas.core.construction import array
+from pandas.core.construction import array, extract_array
 from pandas.core.indexers import check_array_indexer
 from pandas.core.indexes.base import ensure_index
 
@@ -161,12 +160,14 @@ def __new__(
         verify_integrity: bool = True,
     ):
 
-        if isinstance(data, ABCSeries) and is_interval_dtype(data.dtype):
-            data = data._values
+        if isinstance(data, (ABCSeries, ABCIntervalIndex)) and is_interval_dtype(
+            data.dtype
+        ):
+            data = data._values  # TODO: extract_array?
 
-        if isinstance(data, (cls, ABCIntervalIndex)):
-            left = data.left
-            right = data.right
+        if isinstance(data, cls):
+            left = data._left
+            right = data._right
             closed = closed or data.closed
         else:
 
@@ -243,6 +244,20 @@ def _simple_new(
             )
             raise ValueError(msg)
 
+        # For dt64/td64 we want DatetimeArray/TimedeltaArray instead of ndarray
+        from pandas.core.ops.array_ops import maybe_upcast_datetimelike_array
+
+        left = maybe_upcast_datetimelike_array(left)
+        left = extract_array(left, extract_numpy=True)
+        right = maybe_upcast_datetimelike_array(right)
+        right = extract_array(right, extract_numpy=True)
+
+        lbase = getattr(left, "_ndarray", left).base
+        rbase = getattr(right, "_ndarray", right).base
+        if lbase is not None and lbase is rbase:
+            # If these share data, then setitem could corrupt our IA
+            right = right.copy()
+
         result._left = left
         result._right = right
         result._closed = closed
@@ -476,18 +491,18 @@ def _validate(self):
         if self.closed not in VALID_CLOSED:
             msg = f"invalid option for 'closed': {self.closed}"
             raise ValueError(msg)
-        if len(self.left) != len(self.right):
+        if len(self._left) != len(self._right):
             msg = "left and right must have the same length"
             raise ValueError(msg)
-        left_mask = notna(self.left)
-        right_mask = notna(self.right)
+        left_mask = notna(self._left)
+        right_mask = notna(self._right)
         if not (left_mask == right_mask).all():
             msg = (
                 "missing values must be missing in the same "
                 "location both left and right sides"
             )
             raise ValueError(msg)
-        if not (self.left[left_mask] <= self.right[left_mask]).all():
+        if not (self._left[left_mask] <= self._right[left_mask]).all():
             msg = "left side of interval must be <= right side"
             raise ValueError(msg)
 
@@ -527,37 +542,29 @@ def __iter__(self):
         return iter(np.asarray(self))
 
     def __len__(self) -> int:
-        return len(self.left)
+        return len(self._left)
 
     def __getitem__(self, value):
         value = check_array_indexer(self, value)
-        left = self.left[value]
-        right = self.right[value]
+        left = self._left[value]
+        right = self._right[value]
 
-        # scalar
-        if not isinstance(left, ABCIndexClass):
+        if not isinstance(left, (np.ndarray, ExtensionArray)):
+            # scalar
             if is_scalar(left) and isna(left):
                 return self._fill_value
-            if np.ndim(left) > 1:
-                # GH#30588 multi-dimensional indexer disallowed
-                raise ValueError("multi-dimensional indexing not allowed")
             return Interval(left, right, self.closed)
-
+        if np.ndim(left) > 1:
+            # GH#30588 multi-dimensional indexer disallowed
+            raise ValueError("multi-dimensional indexing not allowed")
         return self._shallow_copy(left, right)
 
     def __setitem__(self, key, value):
         value_left, value_right = self._validate_setitem_value(value)
         key = check_array_indexer(self, key)
 
-        # Need to ensure that left and right are updated atomically, so we're
-        # forced to copy, update the copy, and swap in the new values.
-        left = self.left.copy(deep=True)
-        left._values[key] = value_left
-        self._left = left
-
-        right = self.right.copy(deep=True)
-        right._values[key] = value_right
-        self._right = right
+        self._left[key] = value_left
+        self._right[key] = value_right
 
     def __eq__(self, other):
         # ensure pandas array for list-like and eliminate non-interval scalars
@@ -588,7 +595,7 @@ def __eq__(self, other):
         if is_interval_dtype(other_dtype):
             if self.closed != other.closed:
                 return np.zeros(len(self), dtype=bool)
-            return (self.left == other.left) & (self.right == other.right)
+            return (self._left == other.left) & (self._right == other.right)
 
         # non-interval/non-object dtype -> no matches
         if not is_object_dtype(other_dtype):
@@ -601,8 +608,8 @@ def __eq__(self, other):
             if (
                 isinstance(obj, Interval)
                 and self.closed == obj.closed
-                and self.left[i] == obj.left
-                and self.right[i] == obj.right
+                and self._left[i] == obj.left
+                and self._right[i] == obj.right
             ):
                 result[i] = True
 
@@ -665,6 +672,7 @@ def astype(self, dtype, copy=True):
         array : ExtensionArray or ndarray
             ExtensionArray or NumPy ndarray with 'dtype' for its dtype.
         """
+        from pandas import Index
         from pandas.core.arrays.string_ import StringDtype
 
         if dtype is not None:
@@ -676,8 +684,10 @@ def astype(self, dtype, copy=True):
 
             # need to cast to different subtype
             try:
-                new_left = self.left.astype(dtype.subtype)
-                new_right = self.right.astype(dtype.subtype)
+                # We need to use Index rules for astype to prevent casting
+                #  np.nan entries to int subtypes
+                new_left = Index(self._left, copy=False).astype(dtype.subtype)
+                new_right = Index(self._right, copy=False).astype(dtype.subtype)
             except TypeError as err:
                 msg = (
                     f"Cannot convert {self.dtype} to {dtype}; subtypes are incompatible"
@@ -726,14 +736,14 @@ def copy(self):
         -------
         IntervalArray
         """
-        left = self.left.copy(deep=True)
-        right = self.right.copy(deep=True)
+        left = self._left.copy()
+        right = self._right.copy()
         closed = self.closed
         # TODO: Could skip verify_integrity here.
         return type(self).from_arrays(left, right, closed=closed)
 
-    def isna(self):
-        return isna(self.left)
+    def isna(self) -> np.ndarray:
+        return isna(self._left)
 
     def shift(self, periods: int = 1, fill_value: object = None) -> "IntervalArray":
         if not len(self) or periods == 0:
@@ -749,7 +759,9 @@ def shift(self, periods: int = 1, fill_value: object = None) -> "IntervalArray":
 
         empty_len = min(abs(periods), len(self))
         if isna(fill_value):
-            fill_value = self.left._na_value
+            from pandas import Index
+
+            fill_value = Index(self._left, copy=False)._na_value
             empty = IntervalArray.from_breaks([fill_value] * (empty_len + 1))
         else:
             empty = self._from_sequence([fill_value] * empty_len)
@@ -815,10 +827,10 @@ def take(self, indices, allow_fill=False, fill_value=None, axis=None, **kwargs):
             fill_left, fill_right = self._validate_fill_value(fill_value)
 
         left_take = take(
-            self.left, indices, allow_fill=allow_fill, fill_value=fill_left
+            self._left, indices, allow_fill=allow_fill, fill_value=fill_left
         )
         right_take = take(
-            self.right, indices, allow_fill=allow_fill, fill_value=fill_right
+            self._right, indices, allow_fill=allow_fill, fill_value=fill_right
         )
 
         return self._shallow_copy(left_take, right_take)
@@ -977,15 +989,19 @@ def left(self):
         Return the left endpoints of each Interval in the IntervalArray as
         an Index.
         """
-        return self._left
+        from pandas import Index
+
+        return Index(self._left, copy=False)
 
     @property
     def right(self):
         """
         Return the right endpoints of each Interval in the IntervalArray as
         an Index.
         """
-        return self._right
+        from pandas import Index
+
+        return Index(self._right, copy=False)
 
     @property
     def length(self):
@@ -1146,7 +1162,7 @@ def set_closed(self, closed):
             raise ValueError(msg)
 
         return type(self)._simple_new(
-            left=self.left, right=self.right, closed=closed, verify_integrity=False
+            left=self._left, right=self._right, closed=closed, verify_integrity=False
         )
 
     _interval_shared_docs[
@@ -1172,15 +1188,15 @@ def is_non_overlapping_monotonic(self):
         # at a point when both sides of intervals are included
         if self.closed == "both":
             return bool(
-                (self.right[:-1] < self.left[1:]).all()
-                or (self.left[:-1] > self.right[1:]).all()
+                (self._right[:-1] < self._left[1:]).all()
+                or (self._left[:-1] > self._right[1:]).all()
             )
 
         # non-strict inequality when closed != 'both'; at least one side is
         # not included in the intervals, so equality does not imply overlapping
         return bool(
-            (self.right[:-1] <= self.left[1:]).all()
-            or (self.left[:-1] >= self.right[1:]).all()
+            (self._right[:-1] <= self._left[1:]).all()
+            or (self._left[:-1] >= self._right[1:]).all()
         )
 
     # ---------------------------------------------------------------------
@@ -1191,8 +1207,8 @@ def __array__(self, dtype=None) -> np.ndarray:
         Return the IntervalArray's data as a numpy array of Interval
         objects (with dtype='object')
         """
-        left = self.left
-        right = self.right
+        left = self._left
+        right = self._right
         mask = self.isna()
         closed = self._closed
 
@@ -1222,8 +1238,8 @@ def __arrow_array__(self, type=None):
         interval_type = ArrowIntervalType(subtype, self.closed)
         storage_array = pyarrow.StructArray.from_arrays(
             [
-                pyarrow.array(self.left, type=subtype, from_pandas=True),
-                pyarrow.array(self.right, type=subtype, from_pandas=True),
+                pyarrow.array(self._left, type=subtype, from_pandas=True),
+                pyarrow.array(self._right, type=subtype, from_pandas=True),
             ],
             names=["left", "right"],
         )
@@ -1277,7 +1293,7 @@ def __arrow_array__(self, type=None):
         _interval_shared_docs["to_tuples"] % dict(return_type="ndarray", examples="")
     )
     def to_tuples(self, na_tuple=True):
-        tuples = com.asarray_tuplesafe(zip(self.left, self.right))
+        tuples = com.asarray_tuplesafe(zip(self._left, self._right))
         if not na_tuple:
             # GH 18756
             tuples = np.where(~self.isna(), tuples, np.nan)
@@ -1343,8 +1359,8 @@ def contains(self, other):
         if isinstance(other, Interval):
             raise NotImplementedError("contains not implemented for two intervals")
 
-        return (self.left < other if self.open_left else self.left <= other) & (
-            other < self.right if self.open_right else other <= self.right
+        return (self._left < other if self.open_left else self._left <= other) & (
+            other < self._right if self.open_right else other <= self._right
         )