Skip to content

Commit cf8c7d5

Browse files
authored
ENH: NDArrayBackedExtensionArray swapaxes, delete-with-axis (#39963)
1 parent d84c57c commit cf8c7d5

File tree

2 files changed

+30
-2
lines changed

2 files changed

+30
-2
lines changed

pandas/core/arrays/_mixins.py

+10-2
Original file line numberDiff line numberDiff line change
@@ -369,8 +369,16 @@ def where(
369369
res_values = np.where(mask, self._ndarray, value)
370370
return self._from_backing_data(res_values)
371371

372-
def delete(self: NDArrayBackedExtensionArrayT, loc) -> NDArrayBackedExtensionArrayT:
373-
res_values = np.delete(self._ndarray, loc)
372+
def delete(
373+
self: NDArrayBackedExtensionArrayT, loc, axis: int = 0
374+
) -> NDArrayBackedExtensionArrayT:
375+
res_values = np.delete(self._ndarray, loc, axis=axis)
376+
return self._from_backing_data(res_values)
377+
378+
def swapaxes(
379+
self: NDArrayBackedExtensionArrayT, axis1, axis2
380+
) -> NDArrayBackedExtensionArrayT:
381+
res_values = self._ndarray.swapaxes(axis1, axis2)
374382
return self._from_backing_data(res_values)
375383

376384
# ------------------------------------------------------------------------

pandas/tests/extension/base/dim2.py

+20
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,26 @@ def maybe_xfail_masked_reductions(arr, request):
2525

2626

2727
class Dim2CompatTests(BaseExtensionTests):
28+
def test_swapaxes(self, data):
29+
arr2d = data.repeat(2).reshape(-1, 2)
30+
31+
result = arr2d.swapaxes(0, 1)
32+
expected = arr2d.T
33+
self.assert_extension_array_equal(result, expected)
34+
35+
def test_delete_2d(self, data):
36+
arr2d = data.repeat(3).reshape(-1, 3)
37+
38+
# axis = 0
39+
result = arr2d.delete(1, axis=0)
40+
expected = data.delete(1).repeat(3).reshape(-1, 3)
41+
self.assert_extension_array_equal(result, expected)
42+
43+
# axis = 1
44+
result = arr2d.delete(1, axis=1)
45+
expected = data.repeat(2).reshape(-1, 2)
46+
self.assert_extension_array_equal(result, expected)
47+
2848
def test_take_2d(self, data):
2949
arr2d = data.reshape(-1, 1)
3050

0 commit comments

Comments
 (0)