-
-
Notifications
You must be signed in to change notification settings - Fork 18.5k
Implement DataFrame.__array_ufunc__ #36955
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
20 commits
Select commit
Hold shift + click to select a range
0d725e8
Implement DataFrame.__array_ufunc__
TomAugspurger c4c1470
remove unnecessary decorator
TomAugspurger 4fcb1a4
Fixup
TomAugspurger 971659e
fixup finalize
TomAugspurger 6bd73dc
whatsnew
TomAugspurger 1085be4
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger 0afdf49
fixup
TomAugspurger 2260c83
fixup
TomAugspurger b3239e2
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger b1d93f5
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger 9cfcba1
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger 919ebb5
Move to arraylike
TomAugspurger c73ab12
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger 99acb86
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger acfe434
union
TomAugspurger 9a35023
Merge remote-tracking branch 'upstream/master' into frame-array-ufunc
TomAugspurger 2371499
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel d283446
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel 816c6dc
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel a6b120a
docstring, typo fixup
jbrockmendel File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -5,8 +5,15 @@ | |
ExtensionArray | ||
""" | ||
import operator | ||
from typing import Any, Callable | ||
import warnings | ||
|
||
from pandas.core.ops import roperator | ||
import numpy as np | ||
|
||
from pandas._libs import lib | ||
|
||
from pandas.core.construction import extract_array | ||
from pandas.core.ops import maybe_dispatch_ufunc_to_dunder_op, roperator | ||
from pandas.core.ops.common import unpack_zerodim_and_defer | ||
|
||
|
||
|
@@ -140,3 +147,138 @@ def __pow__(self, other): | |
@unpack_zerodim_and_defer("__rpow__") | ||
def __rpow__(self, other): | ||
return self._arith_method(other, roperator.rpow) | ||
|
||
|
||
def array_ufunc(self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any): | ||
""" | ||
Compatibility with numpy ufuncs. | ||
|
||
See also | ||
-------- | ||
numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__ | ||
""" | ||
from pandas.core.generic import NDFrame | ||
jreback marked this conversation as resolved.
Show resolved
Hide resolved
|
||
from pandas.core.internals import BlockManager | ||
|
||
cls = type(self) | ||
|
||
# for binary ops, use our custom dunder methods | ||
result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs) | ||
if result is not NotImplemented: | ||
return result | ||
|
||
# Determine if we should defer. | ||
no_defer = (np.ndarray.__array_ufunc__, cls.__array_ufunc__) | ||
|
||
for item in inputs: | ||
higher_priority = ( | ||
hasattr(item, "__array_priority__") | ||
and item.__array_priority__ > self.__array_priority__ | ||
) | ||
has_array_ufunc = ( | ||
hasattr(item, "__array_ufunc__") | ||
and type(item).__array_ufunc__ not in no_defer | ||
and not isinstance(item, self._HANDLED_TYPES) | ||
) | ||
if higher_priority or has_array_ufunc: | ||
return NotImplemented | ||
|
||
# align all the inputs. | ||
types = tuple(type(x) for x in inputs) | ||
alignable = [x for x, t in zip(inputs, types) if issubclass(t, NDFrame)] | ||
|
||
if len(alignable) > 1: | ||
# This triggers alignment. | ||
# At the moment, there aren't any ufuncs with more than two inputs | ||
# so this ends up just being x1.index | x2.index, but we write | ||
# it to handle *args. | ||
|
||
if len(set(types)) > 1: | ||
# We currently don't handle ufunc(DataFrame, Series) | ||
# well. Previously this raised an internal ValueError. We might | ||
# support it someday, so raise a NotImplementedError. | ||
raise NotImplementedError( | ||
"Cannot apply ufunc {} to mixed DataFrame and Series " | ||
"inputs.".format(ufunc) | ||
) | ||
axes = self.axes | ||
for obj in alignable[1:]: | ||
# this relies on the fact that we aren't handling mixed | ||
# series / frame ufuncs. | ||
for i, (ax1, ax2) in enumerate(zip(axes, obj.axes)): | ||
axes[i] = ax1.union(ax2) | ||
|
||
reconstruct_axes = dict(zip(self._AXIS_ORDERS, axes)) | ||
inputs = tuple( | ||
x.reindex(**reconstruct_axes) if issubclass(t, NDFrame) else x | ||
for x, t in zip(inputs, types) | ||
) | ||
else: | ||
reconstruct_axes = dict(zip(self._AXIS_ORDERS, self.axes)) | ||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ideally can you split this function up a bit (if easy) |
||
if self.ndim == 1: | ||
names = [getattr(x, "name") for x in inputs if hasattr(x, "name")] | ||
name = names[0] if len(set(names)) == 1 else None | ||
reconstruct_kwargs = {"name": name} | ||
else: | ||
reconstruct_kwargs = {} | ||
|
||
def reconstruct(result): | ||
if lib.is_scalar(result): | ||
return result | ||
if result.ndim != self.ndim: | ||
if method == "outer": | ||
if self.ndim == 2: | ||
# we already deprecated for Series | ||
msg = ( | ||
"outer method for ufunc {} is not implemented on " | ||
"pandas objects. Returning an ndarray, but in the " | ||
"future this will raise a 'NotImplementedError'. " | ||
"Consider explicitly converting the DataFrame " | ||
"to an array with '.to_numpy()' first." | ||
) | ||
warnings.warn(msg.format(ufunc), FutureWarning, stacklevel=4) | ||
return result | ||
raise NotImplementedError | ||
return result | ||
if isinstance(result, BlockManager): | ||
# we went through BlockManager.apply | ||
result = self._constructor(result, **reconstruct_kwargs, copy=False) | ||
else: | ||
# we converted an array, lost our axes | ||
result = self._constructor( | ||
result, **reconstruct_axes, **reconstruct_kwargs, copy=False | ||
) | ||
# TODO: When we support multiple values in __finalize__, this | ||
# should pass alignable to `__fianlize__` instead of self. | ||
# Then `np.add(a, b)` would consider attrs from both a and b | ||
# when a and b are NDFrames. | ||
if len(alignable) == 1: | ||
result = result.__finalize__(self) | ||
return result | ||
|
||
if self.ndim > 1 and ( | ||
len(inputs) > 1 or ufunc.nout > 1 # type: ignore[attr-defined] | ||
): | ||
# Just give up on preserving types in the complex case. | ||
# In theory we could preserve them for them. | ||
# * nout>1 is doable if BlockManager.apply took nout and | ||
# returned a Tuple[BlockManager]. | ||
# * len(inputs) > 1 is doable when we know that we have | ||
# aligned blocks / dtypes. | ||
inputs = tuple(np.asarray(x) for x in inputs) | ||
result = getattr(ufunc, method)(*inputs) | ||
elif self.ndim == 1: | ||
# ufunc(series, ...) | ||
inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs) | ||
result = getattr(ufunc, method)(*inputs, **kwargs) | ||
else: | ||
# ufunc(dataframe) | ||
mgr = inputs[0]._mgr | ||
result = mgr.apply(getattr(ufunc, method)) | ||
|
||
if ufunc.nout > 1: # type: ignore[attr-defined] | ||
result = tuple(reconstruct(x) for x in result) | ||
else: | ||
result = reconstruct(result) | ||
return result |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.