Skip to content

Commit d9551c8

Browse files
jbrockmendeljreback
authored andcommitted
De-duplicate masking/fallback logic in ops (#19613)
1 parent d6fe194 commit d9551c8

File tree

3 files changed

+78
-58
lines changed

3 files changed

+78
-58
lines changed

pandas/core/frame.py

+1-11
Original file line numberDiff line numberDiff line change
@@ -3943,17 +3943,7 @@ def _combine_frame(self, other, func, fill_value=None, level=None):
39433943
new_index, new_columns = this.index, this.columns
39443944

39453945
def _arith_op(left, right):
3946-
if fill_value is not None:
3947-
left_mask = isna(left)
3948-
right_mask = isna(right)
3949-
left = left.copy()
3950-
right = right.copy()
3951-
3952-
# one but not both
3953-
mask = left_mask ^ right_mask
3954-
left[left_mask & mask] = fill_value
3955-
right[right_mask & mask] = fill_value
3956-
3946+
left, right = ops.fill_binop(left, right, fill_value)
39573947
return func(left, right)
39583948

39593949
if this._is_mixed_type or other._is_mixed_type:

pandas/core/ops.py

+75-34
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,79 @@ def _make_flex_doc(op_name, typ):
398398
return doc
399399

400400

401+
# -----------------------------------------------------------------------------
402+
# Masking NA values and fallbacks for operations numpy does not support
403+
404+
def fill_binop(left, right, fill_value):
405+
"""
406+
If a non-None fill_value is given, replace null entries in left and right
407+
with this value, but only in positions where _one_ of left/right is null,
408+
not both.
409+
410+
Parameters
411+
----------
412+
left : array-like
413+
right : array-like
414+
fill_value : object
415+
416+
Returns
417+
-------
418+
left : array-like
419+
right : array-like
420+
421+
Notes
422+
-----
423+
Makes copies if fill_value is not None
424+
"""
425+
# TODO: can we make a no-copy implementation?
426+
if fill_value is not None:
427+
left_mask = isna(left)
428+
right_mask = isna(right)
429+
left = left.copy()
430+
right = right.copy()
431+
432+
# one but not both
433+
mask = left_mask ^ right_mask
434+
left[left_mask & mask] = fill_value
435+
right[right_mask & mask] = fill_value
436+
return left, right
437+
438+
439+
def mask_cmp_op(x, y, op, allowed_types):
440+
"""
441+
Apply the function `op` to only non-null points in x and y.
442+
443+
Parameters
444+
----------
445+
x : array-like
446+
y : array-like
447+
op : binary operation
448+
allowed_types : class or tuple of classes
449+
450+
Returns
451+
-------
452+
result : ndarray[bool]
453+
"""
454+
# TODO: Can we make the allowed_types arg unnecessary?
455+
xrav = x.ravel()
456+
result = np.empty(x.size, dtype=bool)
457+
if isinstance(y, allowed_types):
458+
yrav = y.ravel()
459+
mask = notna(xrav) & notna(yrav)
460+
result[mask] = op(np.array(list(xrav[mask])),
461+
np.array(list(yrav[mask])))
462+
else:
463+
mask = notna(xrav)
464+
result[mask] = op(np.array(list(xrav[mask])), y)
465+
466+
if op == operator.ne: # pragma: no cover
467+
np.putmask(result, ~mask, True)
468+
else:
469+
np.putmask(result, ~mask, False)
470+
result = result.reshape(x.shape)
471+
return result
472+
473+
401474
# -----------------------------------------------------------------------------
402475
# Functions that add arithmetic methods to objects, given arithmetic factory
403476
# methods
@@ -1127,23 +1200,7 @@ def na_op(x, y):
11271200
with np.errstate(invalid='ignore'):
11281201
result = op(x, y)
11291202
except TypeError:
1130-
xrav = x.ravel()
1131-
result = np.empty(x.size, dtype=bool)
1132-
if isinstance(y, (np.ndarray, ABCSeries)):
1133-
yrav = y.ravel()
1134-
mask = notna(xrav) & notna(yrav)
1135-
result[mask] = op(np.array(list(xrav[mask])),
1136-
np.array(list(yrav[mask])))
1137-
else:
1138-
mask = notna(xrav)
1139-
result[mask] = op(np.array(list(xrav[mask])), y)
1140-
1141-
if op == operator.ne: # pragma: no cover
1142-
np.putmask(result, ~mask, True)
1143-
else:
1144-
np.putmask(result, ~mask, False)
1145-
result = result.reshape(x.shape)
1146-
1203+
result = mask_cmp_op(x, y, op, (np.ndarray, ABCSeries))
11471204
return result
11481205

11491206
@Appender('Wrapper for flexible comparison methods {name}'
@@ -1221,23 +1278,7 @@ def na_op(x, y):
12211278
try:
12221279
result = expressions.evaluate(op, str_rep, x, y)
12231280
except TypeError:
1224-
xrav = x.ravel()
1225-
result = np.empty(x.size, dtype=bool)
1226-
if isinstance(y, np.ndarray):
1227-
yrav = y.ravel()
1228-
mask = notna(xrav) & notna(yrav)
1229-
result[mask] = op(np.array(list(xrav[mask])),
1230-
np.array(list(yrav[mask])))
1231-
else:
1232-
mask = notna(xrav)
1233-
result[mask] = op(np.array(list(xrav[mask])), y)
1234-
1235-
if op == operator.ne: # pragma: no cover
1236-
np.putmask(result, ~mask, True)
1237-
else:
1238-
np.putmask(result, ~mask, False)
1239-
result = result.reshape(x.shape)
1240-
1281+
result = mask_cmp_op(x, y, op, np.ndarray)
12411282
return result
12421283

12431284
@Appender('Wrapper for comparison method {name}'.format(name=name))

pandas/core/series.py

+2-13
Original file line numberDiff line numberDiff line change
@@ -1725,19 +1725,8 @@ def _binop(self, other, func, level=None, fill_value=None):
17251725
copy=False)
17261726
new_index = this.index
17271727

1728-
this_vals = this.values
1729-
other_vals = other.values
1730-
1731-
if fill_value is not None:
1732-
this_mask = isna(this_vals)
1733-
other_mask = isna(other_vals)
1734-
this_vals = this_vals.copy()
1735-
other_vals = other_vals.copy()
1736-
1737-
# one but not both
1738-
mask = this_mask ^ other_mask
1739-
this_vals[this_mask & mask] = fill_value
1740-
other_vals[other_mask & mask] = fill_value
1728+
this_vals, other_vals = ops.fill_binop(this.values, other.values,
1729+
fill_value)
17411730

17421731
with np.errstate(all='ignore'):
17431732
result = func(this_vals, other_vals)

0 commit comments

Comments
 (0)