-
-
Notifications
You must be signed in to change notification settings - Fork 18.4k
De-duplicate masking/fallback logic in ops #19613
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
Changes from all commits
1c65d80
1e18968
d9e8518
0accde4
221f8a6
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -398,6 +398,79 @@ def _make_flex_doc(op_name, typ): | |
return doc | ||
|
||
|
||
# ----------------------------------------------------------------------------- | ||
# Masking NA values and fallbacks for operations numpy does not support | ||
|
||
def fill_binop(left, right, fill_value): | ||
""" | ||
If a non-None fill_value is given, replace null entries in left and right | ||
with this value, but only in positions where _one_ of left/right is null, | ||
not both. | ||
|
||
Parameters | ||
---------- | ||
left : array-like | ||
right : array-like | ||
fill_value : object | ||
|
||
Returns | ||
------- | ||
left : array-like | ||
right : array-like | ||
|
||
Notes | ||
----- | ||
Makes copies if fill_value is not None | ||
""" | ||
# TODO: can we make a no-copy implementation? | ||
if fill_value is not None: | ||
left_mask = isna(left) | ||
right_mask = isna(right) | ||
left = left.copy() | ||
right = right.copy() | ||
|
||
# one but not both | ||
mask = left_mask ^ right_mask | ||
left[left_mask & mask] = fill_value | ||
right[right_mask & mask] = fill_value | ||
return left, right | ||
|
||
|
||
def mask_cmp_op(x, y, op, allowed_types): | ||
""" | ||
Apply the function `op` to only non-null points in x and y. | ||
|
||
Parameters | ||
---------- | ||
x : array-like | ||
y : array-like | ||
op : binary operation | ||
allowed_types : class or tuple of classes | ||
|
||
Returns | ||
------- | ||
result : ndarray[bool] | ||
""" | ||
# TODO: Can we make the allowed_types arg unnecessary? | ||
xrav = x.ravel() | ||
result = np.empty(x.size, dtype=bool) | ||
if isinstance(y, allowed_types): | ||
yrav = y.ravel() | ||
mask = notna(xrav) & notna(yrav) | ||
result[mask] = op(np.array(list(xrav[mask])), | ||
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. Do you know why we build the intermediate 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. No idea. I can try removing it and see if that affects any tests. 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. Looks like a complex-dtype test would be affected. |
||
np.array(list(yrav[mask]))) | ||
else: | ||
mask = notna(xrav) | ||
result[mask] = op(np.array(list(xrav[mask])), y) | ||
|
||
if op == operator.ne: # pragma: no cover | ||
np.putmask(result, ~mask, True) | ||
else: | ||
np.putmask(result, ~mask, False) | ||
result = result.reshape(x.shape) | ||
return result | ||
|
||
|
||
# ----------------------------------------------------------------------------- | ||
# Functions that add arithmetic methods to objects, given arithmetic factory | ||
# methods | ||
|
@@ -1127,23 +1200,7 @@ def na_op(x, y): | |
with np.errstate(invalid='ignore'): | ||
result = op(x, y) | ||
except TypeError: | ||
xrav = x.ravel() | ||
result = np.empty(x.size, dtype=bool) | ||
if isinstance(y, (np.ndarray, ABCSeries)): | ||
yrav = y.ravel() | ||
mask = notna(xrav) & notna(yrav) | ||
result[mask] = op(np.array(list(xrav[mask])), | ||
np.array(list(yrav[mask]))) | ||
else: | ||
mask = notna(xrav) | ||
result[mask] = op(np.array(list(xrav[mask])), y) | ||
|
||
if op == operator.ne: # pragma: no cover | ||
np.putmask(result, ~mask, True) | ||
else: | ||
np.putmask(result, ~mask, False) | ||
result = result.reshape(x.shape) | ||
|
||
result = mask_cmp_op(x, y, op, (np.ndarray, ABCSeries)) | ||
return result | ||
|
||
@Appender('Wrapper for flexible comparison methods {name}' | ||
|
@@ -1221,23 +1278,7 @@ def na_op(x, y): | |
try: | ||
result = expressions.evaluate(op, str_rep, x, y) | ||
except TypeError: | ||
xrav = x.ravel() | ||
result = np.empty(x.size, dtype=bool) | ||
if isinstance(y, np.ndarray): | ||
yrav = y.ravel() | ||
mask = notna(xrav) & notna(yrav) | ||
result[mask] = op(np.array(list(xrav[mask])), | ||
np.array(list(yrav[mask]))) | ||
else: | ||
mask = notna(xrav) | ||
result[mask] = op(np.array(list(xrav[mask])), y) | ||
|
||
if op == operator.ne: # pragma: no cover | ||
np.putmask(result, ~mask, True) | ||
else: | ||
np.putmask(result, ~mask, False) | ||
result = result.reshape(x.shape) | ||
|
||
result = mask_cmp_op(x, y, op, np.ndarray) | ||
return result | ||
|
||
@Appender('Wrapper for comparison method {name}'.format(name=name)) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
you might be able to simplify this even more here (IOW remove _arith_op) and just in-line it
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This overlaps with #19611. This will probably get in before that does, so I'll take a look at this suggestion after rebasing there.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
aok