API: implement array_function for ExtensionArray by simonjayhawkins · Pull Request #35032 · pandas-dev/pandas (original) (raw)

while adding dispatch from numpy to concat_compat, it was discovered that the axis argument is ignored for EAs. changing this fixes another bug, so the relevant change has been broken off into a separate PR #35038 to be reviewed first and reduce the diff here to just the changes required to get tests passing with __array_function__

on master

>>> arr = pd.array([1, 2, 3, None], dtype="Int64")
>>> arr
<IntegerArray>
[1, 2, 3, <NA>]
Length: 4, dtype: Int64
>>>
>>> from pandas.core.dtypes.concat import concat_compat
>>>
>>> concat_compat([arr, arr])
<IntegerArray>
[1, 2, 3, <NA>, 1, 2, 3, <NA>]
Length: 8, dtype: Int64
>>>
>>> concat_compat([arr, arr], axis=1)
<IntegerArray>
[1, 2, 3, <NA>, 1, 2, 3, <NA>]
Length: 8, dtype: Int64
>>>
>>> np.concatenate([arr, arr])
array([1, 2, 3, <NA>, 1, 2, 3, <NA>], dtype=object)
>>>
>>> np.concatenate([arr, arr], axis=1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in concatenate
numpy.AxisError: axis 1 is out of bounds for array of dimension 1
>>>