Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Functions for arithmetic and comparison operations on NumPy arrays and
3ExtensionArrays.
4"""
6from __future__ import annotations
8import datetime
9from functools import partial
10import operator
11from typing import (
12 TYPE_CHECKING,
13 Any,
14)
16import numpy as np
18from pandas._libs import (
19 NaT,
20 Timedelta,
21 Timestamp,
22 lib,
23 ops as libops,
24)
25from pandas._libs.tslibs import (
26 BaseOffset,
27 get_supported_dtype,
28 is_supported_dtype,
29 is_unitless,
30)
32from pandas.core.dtypes.cast import (
33 construct_1d_object_array_from_listlike,
34 find_common_type,
35)
36from pandas.core.dtypes.common import (
37 ensure_object,
38 is_bool_dtype,
39 is_list_like,
40 is_numeric_v_string_like,
41 is_object_dtype,
42 is_scalar,
43)
44from pandas.core.dtypes.generic import (
45 ABCExtensionArray,
46 ABCIndex,
47 ABCSeries,
48)
49from pandas.core.dtypes.missing import (
50 isna,
51 notna,
52)
54from pandas.core import roperator
55from pandas.core.computation import expressions
56from pandas.core.construction import (
57 ensure_wrapped_if_datetimelike,
58 sanitize_array,
59)
60from pandas.core.ops import missing
61from pandas.core.ops.dispatch import should_extension_dispatch
62from pandas.core.ops.invalid import (
63 disallow_datetimelike_logical_op,
64 invalid_comparison,
65)
67if TYPE_CHECKING:
68 from pandas._typing import (
69 ArrayLike,
70 Shape,
71 )
73# -----------------------------------------------------------------------------
74# Masking NA values and fallbacks for operations numpy does not support
77def fill_binop(left, right, fill_value):
78 """
79 If a non-None fill_value is given, replace null entries in left and right
80 with this value, but only in positions where _one_ of left/right is null,
81 not both.
83 Parameters
84 ----------
85 left : array-like
86 right : array-like
87 fill_value : object
89 Returns
90 -------
91 left : array-like
92 right : array-like
94 Notes
95 -----
96 Makes copies if fill_value is not None and NAs are present.
97 """
98 if fill_value is not None:
99 left_mask = isna(left)
100 right_mask = isna(right)
102 # one but not both
103 mask = left_mask ^ right_mask
105 if left_mask.any():
106 # Avoid making a copy if we can
107 left = left.copy()
108 left[left_mask & mask] = fill_value
110 if right_mask.any():
111 # Avoid making a copy if we can
112 right = right.copy()
113 right[right_mask & mask] = fill_value
115 return left, right
118def comp_method_OBJECT_ARRAY(op, x, y):
119 if isinstance(y, list):
120 # e.g. test_tuple_categories
121 y = construct_1d_object_array_from_listlike(y)
123 if isinstance(y, (np.ndarray, ABCSeries, ABCIndex)):
124 if not is_object_dtype(y.dtype):
125 y = y.astype(np.object_)
127 if isinstance(y, (ABCSeries, ABCIndex)):
128 y = y._values
130 if x.shape != y.shape:
131 raise ValueError("Shapes must match", x.shape, y.shape)
132 result = libops.vec_compare(x.ravel(), y.ravel(), op)
133 else:
134 result = libops.scalar_compare(x.ravel(), y, op)
135 return result.reshape(x.shape)
138def _masked_arith_op(x: np.ndarray, y, op) -> np.ndarray:
139 """
140 If the given arithmetic operation fails, attempt it again on
141 only the non-null elements of the input array(s).
143 Parameters
144 ----------
145 x : np.ndarray
146 y : np.ndarray, Series, Index
147 op : binary operator
148 """
149 # For Series `x` is 1D so ravel() is a no-op; calling it anyway makes
150 # the logic valid for both Series and DataFrame ops.
151 xrav = x.ravel()
153 if isinstance(y, np.ndarray):
154 dtype = find_common_type([x.dtype, y.dtype])
155 result = np.empty(x.size, dtype=dtype)
157 if len(x) != len(y):
158 raise ValueError(x.shape, y.shape)
159 ymask = notna(y)
161 # NB: ravel() is only safe since y is ndarray; for e.g. PeriodIndex
162 # we would get int64 dtype, see GH#19956
163 yrav = y.ravel()
164 mask = notna(xrav) & ymask.ravel()
166 # See GH#5284, GH#5035, GH#19448 for historical reference
167 if mask.any():
168 result[mask] = op(xrav[mask], yrav[mask])
170 else:
171 if not is_scalar(y):
172 raise TypeError(
173 f"Cannot broadcast np.ndarray with operand of type {type(y)}"
174 )
176 # mask is only meaningful for x
177 result = np.empty(x.size, dtype=x.dtype)
178 mask = notna(xrav)
180 # 1 ** np.nan is 1. So we have to unmask those.
181 if op is pow:
182 mask = np.where(x == 1, False, mask)
183 elif op is roperator.rpow:
184 mask = np.where(y == 1, False, mask)
186 if mask.any():
187 result[mask] = op(xrav[mask], y)
189 np.putmask(result, ~mask, np.nan)
190 result = result.reshape(x.shape) # 2D compat
191 return result
194def _na_arithmetic_op(left: np.ndarray, right, op, is_cmp: bool = False):
195 """
196 Return the result of evaluating op on the passed in values.
198 If native types are not compatible, try coercion to object dtype.
200 Parameters
201 ----------
202 left : np.ndarray
203 right : np.ndarray or scalar
204 Excludes DataFrame, Series, Index, ExtensionArray.
205 is_cmp : bool, default False
206 If this is a comparison operation.
208 Returns
209 -------
210 array-like
212 Raises
213 ------
214 TypeError : invalid operation
215 """
216 if isinstance(right, str):
217 # can never use numexpr
218 func = op
219 else:
220 func = partial(expressions.evaluate, op)
222 try:
223 result = func(left, right)
224 except TypeError:
225 if not is_cmp and (
226 left.dtype == object or getattr(right, "dtype", None) == object
227 ):
228 # For object dtype, fallback to a masked operation (only operating
229 # on the non-missing values)
230 # Don't do this for comparisons, as that will handle complex numbers
231 # incorrectly, see GH#32047
232 result = _masked_arith_op(left, right, op)
233 else:
234 raise
236 if is_cmp and (is_scalar(result) or result is NotImplemented):
237 # numpy returned a scalar instead of operating element-wise
238 # e.g. numeric array vs str
239 # TODO: can remove this after dropping some future numpy version?
240 return invalid_comparison(left, right, op)
242 return missing.dispatch_fill_zeros(op, left, right, result)
245def arithmetic_op(left: ArrayLike, right: Any, op):
246 """
247 Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ...
249 Note: the caller is responsible for ensuring that numpy warnings are
250 suppressed (with np.errstate(all="ignore")) if needed.
252 Parameters
253 ----------
254 left : np.ndarray or ExtensionArray
255 right : object
256 Cannot be a DataFrame or Index. Series is *not* excluded.
257 op : {operator.add, operator.sub, ...}
258 Or one of the reversed variants from roperator.
260 Returns
261 -------
262 ndarray or ExtensionArray
263 Or a 2-tuple of these in the case of divmod or rdivmod.
264 """
265 # NB: We assume that extract_array and ensure_wrapped_if_datetimelike
266 # have already been called on `left` and `right`,
267 # and `maybe_prepare_scalar_for_op` has already been called on `right`
268 # We need to special-case datetime64/timedelta64 dtypes (e.g. because numpy
269 # casts integer dtypes to timedelta64 when operating with timedelta64 - GH#22390)
270 if isinstance(right, list):
271 # GH#62423
272 right = sanitize_array(right, None)
273 right = ensure_wrapped_if_datetimelike(right)
275 if (
276 should_extension_dispatch(left, right)
277 or isinstance(right, (Timedelta, BaseOffset, Timestamp))
278 or right is NaT
279 ):
280 # Timedelta/Timestamp and other custom scalars are included in the check
281 # because numexpr will fail on it, see GH#31457
282 res_values = op(left, right)
283 else:
284 # TODO we should handle EAs consistently and move this check before the if/else
285 # (https://github.com/pandas-dev/pandas/issues/41165)
286 # error: Argument 2 to "_bool_arith_check" has incompatible type
287 # "Union[ExtensionArray, ndarray[Any, Any]]"; expected "ndarray[Any, Any]"
288 _bool_arith_check(op, left, right) # type: ignore[arg-type]
290 # error: Argument 1 to "_na_arithmetic_op" has incompatible type
291 # "Union[ExtensionArray, ndarray[Any, Any]]"; expected "ndarray[Any, Any]"
292 res_values = _na_arithmetic_op(left, right, op) # type: ignore[arg-type]
294 return res_values
297def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike:
298 """
299 Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`.
301 Note: the caller is responsible for ensuring that numpy warnings are
302 suppressed (with np.errstate(all="ignore")) if needed.
304 Parameters
305 ----------
306 left : np.ndarray or ExtensionArray
307 right : object
308 Cannot be a DataFrame, Series, or Index.
309 op : {operator.eq, operator.ne, operator.gt, operator.ge, operator.lt, operator.le}
311 Returns
312 -------
313 ndarray or ExtensionArray
314 """
315 # NB: We assume extract_array has already been called on left and right
316 lvalues = ensure_wrapped_if_datetimelike(left)
317 rvalues = ensure_wrapped_if_datetimelike(right)
319 rvalues = lib.item_from_zerodim(rvalues)
321 # Special handling needed if rvalues is a zerodim np.ndarray subclass, see GH#63205
322 rvalues_is_zerodim: bool = getattr(rvalues, "ndim", None) == 0
324 if isinstance(rvalues, list):
325 # We don't catch tuple here bc we may be comparing e.g. MultiIndex
326 # to a tuple that represents a single entry, see test_compare_tuple_strs
327 rvalues = sanitize_array(rvalues, None)
328 rvalues = ensure_wrapped_if_datetimelike(rvalues)
330 if isinstance(rvalues, (np.ndarray, ABCExtensionArray)) and not rvalues_is_zerodim:
331 # TODO: make this treatment consistent across ops and classes.
332 # We are not catching all listlikes here (e.g. frozenset, tuple)
333 # The ambiguous case is object-dtype. See GH#27803
334 if len(lvalues) != len(rvalues):
335 raise ValueError(
336 "Lengths must match to compare", lvalues.shape, rvalues.shape
337 )
339 if should_extension_dispatch(lvalues, rvalues) or (
340 (isinstance(rvalues, (Timedelta, BaseOffset, Timestamp)) or right is NaT)
341 and lvalues.dtype != object
342 ):
343 # Call the method on lvalues
344 res_values = op(lvalues, rvalues)
346 # TODO: but not pd.NA?
347 elif (is_scalar(rvalues) or rvalues_is_zerodim) and isna(rvalues):
348 # numpy does not like comparisons vs None
349 if op is operator.ne:
350 res_values = np.ones(lvalues.shape, dtype=bool)
351 else:
352 res_values = np.zeros(lvalues.shape, dtype=bool)
354 elif is_numeric_v_string_like(lvalues, rvalues):
355 # GH#36377 going through the numexpr path would incorrectly raise
356 return invalid_comparison(lvalues, rvalues, op)
358 elif lvalues.dtype == object or isinstance(rvalues, str):
359 res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues)
361 else:
362 res_values = _na_arithmetic_op(lvalues, rvalues, op, is_cmp=True)
364 return res_values
367def na_logical_op(x: np.ndarray, y, op):
368 try:
369 # For exposition, write:
370 # yarr = isinstance(y, np.ndarray)
371 # yint = is_integer(y) or (yarr and y.dtype.kind == "i")
372 # ybool = is_bool(y) or (yarr and y.dtype.kind == "b")
373 # xint = x.dtype.kind == "i"
374 # xbool = x.dtype.kind == "b"
375 # Then Cases where this goes through without raising include:
376 # (xint or xbool) and (yint or bool)
377 result = op(x, y)
378 except TypeError:
379 if isinstance(y, np.ndarray):
380 # bool-bool dtype operations should be OK, should not get here
381 assert not (x.dtype.kind == "b" and y.dtype.kind == "b")
382 x = ensure_object(x)
383 y = ensure_object(y)
384 result = libops.vec_binop(x.ravel(), y.ravel(), op)
385 else:
386 # let null fall thru
387 assert lib.is_scalar(y)
388 if not isna(y):
389 y = bool(y)
390 try:
391 result = libops.scalar_binop(x, y, op)
392 except (
393 TypeError,
394 ValueError,
395 AttributeError,
396 OverflowError,
397 NotImplementedError,
398 ) as err:
399 typ = type(y).__name__
400 raise TypeError(
401 f"Cannot perform '{op.__name__}' with a dtyped [{x.dtype}] array "
402 f"and scalar of type [{typ}]"
403 ) from err
405 return result.reshape(x.shape)
408def logical_op(left: ArrayLike, right: Any, op) -> ArrayLike:
409 """
410 Evaluate a logical operation `|`, `&`, or `^`.
412 Parameters
413 ----------
414 left : np.ndarray or ExtensionArray
415 right : object
416 Cannot be a DataFrame, Series, or Index.
417 op : {operator.and_, operator.or_, operator.xor}
418 Or one of the reversed variants from roperator.
420 Returns
421 -------
422 ndarray or ExtensionArray
423 """
425 def fill_bool(x, left=None):
426 # if `left` is specifically not-boolean, we do not cast to bool
427 if x.dtype.kind in "cfO":
428 # dtypes that can hold NA
429 mask = isna(x)
430 if mask.any():
431 x = x.astype(object)
432 x[mask] = False
434 if left is None or left.dtype.kind == "b":
435 x = x.astype(bool)
436 return x
438 right = lib.item_from_zerodim(right)
439 if is_list_like(right) and not hasattr(right, "dtype"):
440 # e.g. list, tuple
441 raise TypeError(
442 # GH#52264
443 "Logical ops (and, or, xor) between Pandas objects and dtype-less "
444 "sequences (e.g. list, tuple) are no longer supported. "
445 "Wrap the object in a Series, Index, or np.array "
446 "before operating instead.",
447 )
449 # NB: We assume extract_array has already been called on left and right
450 lvalues = ensure_wrapped_if_datetimelike(left)
451 rvalues = right
453 disallow_datetimelike_logical_op(lvalues, rvalues, op)
455 if should_extension_dispatch(lvalues, rvalues):
456 # Call the method on lvalues
457 res_values = op(lvalues, rvalues)
459 else:
460 if isinstance(rvalues, np.ndarray):
461 is_other_int_dtype = rvalues.dtype.kind in "iu"
462 if not is_other_int_dtype:
463 rvalues = fill_bool(rvalues, lvalues)
465 else:
466 # i.e. scalar
467 is_other_int_dtype = lib.is_integer(rvalues)
469 res_values = na_logical_op(lvalues, rvalues, op)
471 # For int vs int `^`, `|`, `&` are bitwise operators and return
472 # integer dtypes. Otherwise these are boolean ops
473 if not (left.dtype.kind in "iu" and is_other_int_dtype):
474 res_values = fill_bool(res_values)
476 return res_values
479def get_array_op(op):
480 """
481 Return a binary array operation corresponding to the given operator op.
483 Parameters
484 ----------
485 op : function
486 Binary operator from operator or roperator module.
488 Returns
489 -------
490 functools.partial
491 """
492 if isinstance(op, partial):
493 # We get here via dispatch_to_series in DataFrame case
494 # e.g. test_rolling_consistency_var_debiasing_factors
495 return op
497 op_name = op.__name__.strip("_").lstrip("r")
498 if op_name == "arith_op":
499 # Reached via DataFrame._combine_frame i.e. flex methods
500 # e.g. test_df_add_flex_filled_mixed_dtypes
501 return op
503 if op_name in {"eq", "ne", "lt", "le", "gt", "ge"}:
504 return partial(comparison_op, op=op)
505 elif op_name in {"and", "or", "xor", "rand", "ror", "rxor"}:
506 return partial(logical_op, op=op)
507 elif op_name in {
508 "add",
509 "sub",
510 "mul",
511 "truediv",
512 "floordiv",
513 "mod",
514 "divmod",
515 "pow",
516 }:
517 return partial(arithmetic_op, op=op)
518 else:
519 raise NotImplementedError(op_name)
522def maybe_prepare_scalar_for_op(obj, shape: Shape):
523 """
524 Cast non-pandas objects to pandas types to unify behavior of arithmetic
525 and comparison operations.
527 Parameters
528 ----------
529 obj: object
530 shape : tuple[int]
532 Returns
533 -------
534 out : object
536 Notes
537 -----
538 Be careful to call this *after* determining the `name` attribute to be
539 attached to the result of the arithmetic operation.
540 """
541 if type(obj) is datetime.timedelta:
542 # GH#22390 cast up to Timedelta to rely on Timedelta
543 # implementation; otherwise operation against numeric-dtype
544 # raises TypeError
545 return Timedelta(obj)
546 elif type(obj) is datetime.datetime:
547 # cast up to Timestamp to rely on Timestamp implementation, see Timedelta above
548 return Timestamp(obj)
549 elif isinstance(obj, np.datetime64):
550 # GH#28080 numpy casts integer-dtype to datetime64 when doing
551 # array[int] + datetime64, which we do not allow
552 if isna(obj):
553 from pandas.core.arrays import DatetimeArray
555 # Avoid possible ambiguities with pd.NaT
556 # GH 52295
557 if is_unitless(obj.dtype):
558 # Use second resolution to ensure that the result of e.g.
559 # `left - np.datetime64("NaT")` retains the unit of left.unit
560 obj = obj.astype("datetime64[s]")
561 elif not is_supported_dtype(obj.dtype):
562 new_dtype = get_supported_dtype(obj.dtype)
563 obj = obj.astype(new_dtype)
564 right = np.broadcast_to(obj, shape)
565 return DatetimeArray._simple_new(right, dtype=right.dtype)
567 return Timestamp(obj)
569 elif isinstance(obj, np.timedelta64):
570 if isna(obj):
571 from pandas.core.arrays import TimedeltaArray
573 # wrapping timedelta64("NaT") in Timedelta returns NaT,
574 # which would incorrectly be treated as a datetime-NaT, so
575 # we broadcast and wrap in a TimedeltaArray
576 # GH 52295
577 if is_unitless(obj.dtype):
578 # Use second resolution to ensure that the result of e.g.
579 # `left + np.timedelta64("NaT")` retains the unit of left.unit
580 obj = obj.astype("timedelta64[s]")
581 elif not is_supported_dtype(obj.dtype):
582 new_dtype = get_supported_dtype(obj.dtype)
583 obj = obj.astype(new_dtype)
584 right = np.broadcast_to(obj, shape)
585 return TimedeltaArray._simple_new(right, dtype=right.dtype)
587 # In particular non-nanosecond timedelta64 needs to be cast to
588 # nanoseconds, or else we get undesired behavior like
589 # np.timedelta64(3, 'D') / 2 == np.timedelta64(1, 'D')
590 return Timedelta(obj)
592 # We want NumPy numeric scalars to behave like Python scalars
593 # post NEP 50
594 elif isinstance(obj, np.integer):
595 return int(obj)
597 elif isinstance(obj, np.floating):
598 return float(obj)
600 return obj
603_BOOL_OP_NOT_ALLOWED = {
604 operator.truediv,
605 roperator.rtruediv,
606 operator.floordiv,
607 roperator.rfloordiv,
608 operator.pow,
609 roperator.rpow,
610 divmod,
611 roperator.rdivmod,
612}
615def _bool_arith_check(op, a: np.ndarray, b) -> None:
616 """
617 In contrast to numpy, pandas raises an error for certain operations
618 with booleans.
619 """
620 if op in _BOOL_OP_NOT_ALLOWED:
621 if a.dtype.kind == "b" and (is_bool_dtype(b) or lib.is_bool(b)):
622 op_name = op.__name__.strip("_").lstrip("r")
623 raise NotImplementedError(
624 f"operator '{op_name}' not implemented for bool dtypes"
625 )