ENH: NDArrayBackedExtensionArray.__array_function__ #38068
Conversation
|
@TomAugspurger meant to cc you in the OP but made a typo |
| @@ -304,6 +304,68 @@ def __repr__(self) -> str: | |||
| # ------------------------------------------------------------------------ | |||
| # __array_function__ methods | |||
|
|
|||
| def __array_function__(self, func, types, args, kwargs): | |||
There was a problem hiding this comment.
array_ufunc now exists in pandas/core/arraylike.py we should share as much as possible here (might mean breaking up the existing code).
There was a problem hiding this comment.
I think the only shared code will be the checking of the types. We've standardized on a pattern like
class Foo:
_HANDLED_TYPES = (...,)
And then ensuring that the set of types in types is a subset of _HANDLED_TYPES + (type(self),).
| @@ -62,9 +62,25 @@ def np_array_datetime64_compat(arr, *args, **kwargs): | |||
| return np.array(arr, *args, **kwargs) | |||
|
|
|||
|
|
|||
| def _is_nep18_active(): | |||
There was a problem hiding this comment.
We require NumPy>1.16.0, so this can maybe be simplified to
IS_NEP18_ACTIVE = not os.environ.get("NUMPY_EXPERIMENTAL_ARRAY_FUNCTION", "1") == "0"
See https://numpy.org/neps/nep-0018-array-function-protocol.html#implementation.
| @@ -304,6 +304,68 @@ def __repr__(self) -> str: | |||
| # ------------------------------------------------------------------------ | |||
| # __array_function__ methods | |||
|
|
|||
| def __array_function__(self, func, types, args, kwargs): | |||
There was a problem hiding this comment.
I think the only shared code will be the checking of the types. We've standardized on a pattern like
class Foo:
_HANDLED_TYPES = (...,)
And then ensuring that the set of types in types is a subset of _HANDLED_TYPES + (type(self),).
| res_data = func(self._ndarray, *args[1:], **kwargs) | ||
| return self._from_backing_data(res_data) |
There was a problem hiding this comment.
How do we know that there's 1 output argument from func here? Do we have a func.nout or something like that to check?
There was a problem hiding this comment.
we don't, thats why we're implementing only a specific handful of functions here
| if not issubclass(x, (np.ndarray, NDArrayBackedExtensionArray)): | ||
| return NotImplemented | ||
|
|
||
| if not args: |
There was a problem hiding this comment.
Can you give a code example of this path?
There was a problem hiding this comment.
np.delete(arr=np.arange(5), obj=4)
|
Closing to clear the queue; won't be super-useful until our min numpy version is 1.17 |
jbrockmendel
added
the
Mothballed
jbrockmendel
removed
the
Mothballed
motivated by getting np.delete and np.repeat working, had to implement a few others to get the tests passing.
If we go down this path, I plan to incrementally add support for others, some of which (e.g. vstack) will avoid some object-dtype casting.
cc @TomAugspurger @shoyer @seberg any suggestions to clean the implementation, particularly the fallback and the unknown-args comments?