ENH: implement EA.delete (#39405)

Author: jbrockmendel

pandas/core/arrays/base.py

@@ -997,12 +997,12 @@ def repeat(self, repeats, axis=None):
     # ------------------------------------------------------------------------
 
     def take(
-        self,
+        self: ExtensionArrayT,
         indices: Sequence[int],
         *,
         allow_fill: bool = False,
         fill_value: Any = None,
-    ) -> ExtensionArray:
+    ) -> ExtensionArrayT:
         """
         Take elements from an array.
 
@@ -1261,6 +1261,13 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs):
     def __hash__(self):
         raise TypeError(f"unhashable type: {repr(type(self).__name__)}")
 
+    # ------------------------------------------------------------------------
+    # Non-Optimized Default Methods
+
+    def delete(self: ExtensionArrayT, loc) -> ExtensionArrayT:
+        indexer = np.delete(np.arange(len(self)), loc)
+        return self.take(indexer)
+
 
 class ExtensionOpsMixin:
     """

pandas/core/arrays/interval.py

@@ -1413,6 +1413,15 @@ def to_tuples(self, na_tuple=True):
 
     # ---------------------------------------------------------------------
 
+    def delete(self: IntervalArrayT, loc) -> IntervalArrayT:
+        if isinstance(self._left, np.ndarray):
+            new_left = np.delete(self._left, loc)
+            new_right = np.delete(self._right, loc)
+        else:
+            new_left = self._left.delete(loc)
+            new_right = self._right.delete(loc)
+        return self._shallow_copy(left=new_left, right=new_right)
+
     @Appender(_extension_array_shared_docs["repeat"] % _shared_docs_kwargs)
     def repeat(self, repeats, axis=None):
         nv.validate_repeat((), {"axis": axis})

pandas/core/arrays/string_arrow.py

@@ -507,7 +507,7 @@ def __setitem__(self, key: Union[int, np.ndarray], value: Any) -> None:
 
     def take(
         self, indices: Sequence[int], allow_fill: bool = False, fill_value: Any = None
-    ) -> ExtensionArray:
+    ):
         """
         Take elements from an array.
 

pandas/core/indexes/extension.py

@@ -257,6 +257,17 @@ def searchsorted(self, value, side="left", sorter=None) -> np.ndarray:
     def _get_engine_target(self) -> np.ndarray:
         return np.asarray(self._data)
 
+    def delete(self, loc):
+        """
+        Make new Index with passed location(-s) deleted
+
+        Returns
+        -------
+        new_index : Index
+        """
+        arr = self._data.delete(loc)
+        return type(self)._simple_new(arr, name=self.name)
+
     def repeat(self, repeats, axis=None):
         nv.validate_repeat((), {"axis": axis})
         result = self._data.repeat(repeats, axis=axis)
@@ -333,17 +344,6 @@ class NDArrayBackedExtensionIndex(ExtensionIndex):
     def _get_engine_target(self) -> np.ndarray:
         return self._data._ndarray
 
-    def delete(self: _T, loc) -> _T:
-        """
-        Make new Index with passed location(-s) deleted
-
-        Returns
-        -------
-        new_index : Index
-        """
-        arr = self._data.delete(loc)
-        return type(self)._simple_new(arr, name=self.name)
-
     def insert(self: _T, loc: int, item) -> _T:
         """
         Make new Index inserting new item at location. Follows

pandas/core/indexes/interval.py

@@ -822,19 +822,6 @@ def where(self, cond, other=None):
         result = IntervalArray(values)
         return type(self)._simple_new(result, name=self.name)
 
-    def delete(self, loc):
-        """
-        Return a new IntervalIndex with passed location(-s) deleted
-
-        Returns
-        -------
-        IntervalIndex
-        """
-        new_left = self.left.delete(loc)
-        new_right = self.right.delete(loc)
-        result = self._data._shallow_copy(new_left, new_right)
-        return type(self)._simple_new(result, name=self.name)
-
     def insert(self, loc, item):
         """
         Return a new IntervalIndex inserting new item at location. Follows

pandas/tests/extension/base/methods.py

@@ -491,6 +491,15 @@ def test_repeat_raises(self, data, repeats, kwargs, error, msg, use_numpy):
             else:
                 data.repeat(repeats, **kwargs)
 
+    def test_delete(self, data):
+        result = data.delete(0)
+        expected = data[1:]
+        self.assert_extension_array_equal(result, expected)
+
+        result = data.delete([1, 3])
+        expected = data._concat_same_type([data[[0]], data[[2]], data[4:]])
+        self.assert_extension_array_equal(result, expected)
+
     @pytest.mark.parametrize("box", [pd.array, pd.Series, pd.DataFrame])
     def test_equals(self, data, na_value, as_series, box):
         data2 = type(data)._from_sequence([data[0]] * len(data), dtype=data.dtype)