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 invalid_comparison
64if TYPE_CHECKING:
65 from pandas._typing import (
66 ArrayLike,
67 Shape,
68 )
70# -----------------------------------------------------------------------------
71# Masking NA values and fallbacks for operations numpy does not support
74def fill_binop(left, right, fill_value):
75 """
76 If a non-None fill_value is given, replace null entries in left and right
77 with this value, but only in positions where _one_ of left/right is null,
78 not both.
80 Parameters
81 ----------
82 left : array-like
83 right : array-like
84 fill_value : object
86 Returns
87 -------
88 left : array-like
89 right : array-like
91 Notes
92 -----
93 Makes copies if fill_value is not None and NAs are present.
94 """
95 if fill_value is not None:
96 left_mask = isna(left)
97 right_mask = isna(right)
99 # one but not both
100 mask = left_mask ^ right_mask
102 if left_mask.any():
103 # Avoid making a copy if we can
104 left = left.copy()
105 left[left_mask & mask] = fill_value
107 if right_mask.any():
108 # Avoid making a copy if we can
109 right = right.copy()
110 right[right_mask & mask] = fill_value
112 return left, right
115def comp_method_OBJECT_ARRAY(op, x, y):
116 if isinstance(y, list):
117 # e.g. test_tuple_categories
118 y = construct_1d_object_array_from_listlike(y)
120 if isinstance(y, (np.ndarray, ABCSeries, ABCIndex)):
121 if not is_object_dtype(y.dtype):
122 y = y.astype(np.object_)
124 if isinstance(y, (ABCSeries, ABCIndex)):
125 y = y._values
127 if x.shape != y.shape:
128 raise ValueError("Shapes must match", x.shape, y.shape)
129 result = libops.vec_compare(x.ravel(), y.ravel(), op)
130 else:
131 result = libops.scalar_compare(x.ravel(), y, op)
132 return result.reshape(x.shape)
135def _masked_arith_op(x: np.ndarray, y, op) -> np.ndarray:
136 """
137 If the given arithmetic operation fails, attempt it again on
138 only the non-null elements of the input array(s).
140 Parameters
141 ----------
142 x : np.ndarray
143 y : np.ndarray, Series, Index
144 op : binary operator
145 """
146 # For Series `x` is 1D so ravel() is a no-op; calling it anyway makes
147 # the logic valid for both Series and DataFrame ops.
148 xrav = x.ravel()
150 if isinstance(y, np.ndarray):
151 dtype = find_common_type([x.dtype, y.dtype])
152 result = np.empty(x.size, dtype=dtype)
154 if len(x) != len(y):
155 raise ValueError(x.shape, y.shape)
156 ymask = notna(y)
158 # NB: ravel() is only safe since y is ndarray; for e.g. PeriodIndex
159 # we would get int64 dtype, see GH#19956
160 yrav = y.ravel()
161 mask = notna(xrav) & ymask.ravel()
163 # See GH#5284, GH#5035, GH#19448 for historical reference
164 if mask.any():
165 result[mask] = op(xrav[mask], yrav[mask])
167 else:
168 if not is_scalar(y):
169 raise TypeError(
170 f"Cannot broadcast np.ndarray with operand of type {type(y)}"
171 )
173 # mask is only meaningful for x
174 result = np.empty(x.size, dtype=x.dtype)
175 mask = notna(xrav)
177 # 1 ** np.nan is 1. So we have to unmask those.
178 if op is pow:
179 mask = np.where(x == 1, False, mask)
180 elif op is roperator.rpow:
181 mask = np.where(y == 1, False, mask)
183 if mask.any():
184 result[mask] = op(xrav[mask], y)
186 np.putmask(result, ~mask, np.nan)
187 result = result.reshape(x.shape) # 2D compat
188 return result
191def _na_arithmetic_op(left: np.ndarray, right, op, is_cmp: bool = False):
192 """
193 Return the result of evaluating op on the passed in values.
195 If native types are not compatible, try coercion to object dtype.
197 Parameters
198 ----------
199 left : np.ndarray
200 right : np.ndarray or scalar
201 Excludes DataFrame, Series, Index, ExtensionArray.
202 is_cmp : bool, default False
203 If this a comparison operation.
205 Returns
206 -------
207 array-like
209 Raises
210 ------
211 TypeError : invalid operation
212 """
213 if isinstance(right, str):
214 # can never use numexpr
215 func = op
216 else:
217 func = partial(expressions.evaluate, op)
219 try:
220 result = func(left, right)
221 except TypeError:
222 if not is_cmp and (
223 left.dtype == object or getattr(right, "dtype", None) == object
224 ):
225 # For object dtype, fallback to a masked operation (only operating
226 # on the non-missing values)
227 # Don't do this for comparisons, as that will handle complex numbers
228 # incorrectly, see GH#32047
229 result = _masked_arith_op(left, right, op)
230 else:
231 raise
233 if is_cmp and (is_scalar(result) or result is NotImplemented):
234 # numpy returned a scalar instead of operating element-wise
235 # e.g. numeric array vs str
236 # TODO: can remove this after dropping some future numpy version?
237 return invalid_comparison(left, right, op)
239 return missing.dispatch_fill_zeros(op, left, right, result)
242def arithmetic_op(left: ArrayLike, right: Any, op):
243 """
244 Evaluate an arithmetic operation `+`, `-`, `*`, `/`, `//`, `%`, `**`, ...
246 Note: the caller is responsible for ensuring that numpy warnings are
247 suppressed (with np.errstate(all="ignore")) if needed.
249 Parameters
250 ----------
251 left : np.ndarray or ExtensionArray
252 right : object
253 Cannot be a DataFrame or Index. Series is *not* excluded.
254 op : {operator.add, operator.sub, ...}
255 Or one of the reversed variants from roperator.
257 Returns
258 -------
259 ndarray or ExtensionArray
260 Or a 2-tuple of these in the case of divmod or rdivmod.
261 """
262 # NB: We assume that extract_array and ensure_wrapped_if_datetimelike
263 # have already been called on `left` and `right`,
264 # and `maybe_prepare_scalar_for_op` has already been called on `right`
265 # We need to special-case datetime64/timedelta64 dtypes (e.g. because numpy
266 # casts integer dtypes to timedelta64 when operating with timedelta64 - GH#22390)
267 if isinstance(right, list):
268 # GH#62423
269 right = sanitize_array(right, None)
270 right = ensure_wrapped_if_datetimelike(right)
272 if (
273 should_extension_dispatch(left, right)
274 or isinstance(right, (Timedelta, BaseOffset, Timestamp))
275 or right is NaT
276 ):
277 # Timedelta/Timestamp and other custom scalars are included in the check
278 # because numexpr will fail on it, see GH#31457
279 res_values = op(left, right)
280 else:
281 # TODO we should handle EAs consistently and move this check before the if/else
282 # (https://github.com/pandas-dev/pandas/issues/41165)
283 # error: Argument 2 to "_bool_arith_check" has incompatible type
284 # "Union[ExtensionArray, ndarray[Any, Any]]"; expected "ndarray[Any, Any]"
285 _bool_arith_check(op, left, right) # type: ignore[arg-type]
287 # error: Argument 1 to "_na_arithmetic_op" has incompatible type
288 # "Union[ExtensionArray, ndarray[Any, Any]]"; expected "ndarray[Any, Any]"
289 res_values = _na_arithmetic_op(left, right, op) # type: ignore[arg-type]
291 return res_values
294def comparison_op(left: ArrayLike, right: Any, op) -> ArrayLike:
295 """
296 Evaluate a comparison operation `=`, `!=`, `>=`, `>`, `<=`, or `<`.
298 Note: the caller is responsible for ensuring that numpy warnings are
299 suppressed (with np.errstate(all="ignore")) if needed.
301 Parameters
302 ----------
303 left : np.ndarray or ExtensionArray
304 right : object
305 Cannot be a DataFrame, Series, or Index.
306 op : {operator.eq, operator.ne, operator.gt, operator.ge, operator.lt, operator.le}
308 Returns
309 -------
310 ndarray or ExtensionArray
311 """
312 # NB: We assume extract_array has already been called on left and right
313 lvalues = ensure_wrapped_if_datetimelike(left)
314 rvalues = ensure_wrapped_if_datetimelike(right)
316 rvalues = lib.item_from_zerodim(rvalues)
318 # Special handling needed if rvalues is a zerodim np.ndarray subclass, see GH#63205
319 rvalues_is_zerodim: bool = getattr(rvalues, "ndim", None) == 0
321 if isinstance(rvalues, list):
322 # We don't catch tuple here bc we may be comparing e.g. MultiIndex
323 # to a tuple that represents a single entry, see test_compare_tuple_strs
324 rvalues = sanitize_array(rvalues, None)
325 rvalues = ensure_wrapped_if_datetimelike(rvalues)
327 if isinstance(rvalues, (np.ndarray, ABCExtensionArray)) and not rvalues_is_zerodim:
328 # TODO: make this treatment consistent across ops and classes.
329 # We are not catching all listlikes here (e.g. frozenset, tuple)
330 # The ambiguous case is object-dtype. See GH#27803
331 if len(lvalues) != len(rvalues):
332 raise ValueError(
333 "Lengths must match to compare", lvalues.shape, rvalues.shape
334 )
336 if should_extension_dispatch(lvalues, rvalues) or (
337 (isinstance(rvalues, (Timedelta, BaseOffset, Timestamp)) or right is NaT)
338 and lvalues.dtype != object
339 ):
340 # Call the method on lvalues
341 res_values = op(lvalues, rvalues)
343 # TODO: but not pd.NA?
344 elif (is_scalar(rvalues) or rvalues_is_zerodim) and isna(rvalues):
345 # numpy does not like comparisons vs None
346 if op is operator.ne:
347 res_values = np.ones(lvalues.shape, dtype=bool)
348 else:
349 res_values = np.zeros(lvalues.shape, dtype=bool)
351 elif is_numeric_v_string_like(lvalues, rvalues):
352 # GH#36377 going through the numexpr path would incorrectly raise
353 return invalid_comparison(lvalues, rvalues, op)
355 elif lvalues.dtype == object or isinstance(rvalues, str):
356 res_values = comp_method_OBJECT_ARRAY(op, lvalues, rvalues)
358 else:
359 res_values = _na_arithmetic_op(lvalues, rvalues, op, is_cmp=True)
361 return res_values
364def na_logical_op(x: np.ndarray, y, op):
365 try:
366 # For exposition, write:
367 # yarr = isinstance(y, np.ndarray)
368 # yint = is_integer(y) or (yarr and y.dtype.kind == "i")
369 # ybool = is_bool(y) or (yarr and y.dtype.kind == "b")
370 # xint = x.dtype.kind == "i"
371 # xbool = x.dtype.kind == "b"
372 # Then Cases where this goes through without raising include:
373 # (xint or xbool) and (yint or bool)
374 result = op(x, y)
375 except TypeError:
376 if isinstance(y, np.ndarray):
377 # bool-bool dtype operations should be OK, should not get here
378 assert not (x.dtype.kind == "b" and y.dtype.kind == "b")
379 x = ensure_object(x)
380 y = ensure_object(y)
381 result = libops.vec_binop(x.ravel(), y.ravel(), op)
382 else:
383 # let null fall thru
384 assert lib.is_scalar(y)
385 if not isna(y):
386 y = bool(y)
387 try:
388 result = libops.scalar_binop(x, y, op)
389 except (
390 TypeError,
391 ValueError,
392 AttributeError,
393 OverflowError,
394 NotImplementedError,
395 ) as err:
396 typ = type(y).__name__
397 raise TypeError(
398 f"Cannot perform '{op.__name__}' with a dtyped [{x.dtype}] array "
399 f"and scalar of type [{typ}]"
400 ) from err
402 return result.reshape(x.shape)
405def logical_op(left: ArrayLike, right: Any, op) -> ArrayLike:
406 """
407 Evaluate a logical operation `|`, `&`, or `^`.
409 Parameters
410 ----------
411 left : np.ndarray or ExtensionArray
412 right : object
413 Cannot be a DataFrame, Series, or Index.
414 op : {operator.and_, operator.or_, operator.xor}
415 Or one of the reversed variants from roperator.
417 Returns
418 -------
419 ndarray or ExtensionArray
420 """
422 def fill_bool(x, left=None):
423 # if `left` is specifically not-boolean, we do not cast to bool
424 if x.dtype.kind in "cfO":
425 # dtypes that can hold NA
426 mask = isna(x)
427 if mask.any():
428 x = x.astype(object)
429 x[mask] = False
431 if left is None or left.dtype.kind == "b":
432 x = x.astype(bool)
433 return x
435 right = lib.item_from_zerodim(right)
436 if is_list_like(right) and not hasattr(right, "dtype"):
437 # e.g. list, tuple
438 raise TypeError(
439 # GH#52264
440 "Logical ops (and, or, xor) between Pandas objects and dtype-less "
441 "sequences (e.g. list, tuple) are no longer supported. "
442 "Wrap the object in a Series, Index, or np.array "
443 "before operating instead.",
444 )
446 # NB: We assume extract_array has already been called on left and right
447 lvalues = ensure_wrapped_if_datetimelike(left)
448 rvalues = right
450 if should_extension_dispatch(lvalues, rvalues):
451 # Call the method on lvalues
452 res_values = op(lvalues, rvalues)
454 else:
455 if isinstance(rvalues, np.ndarray):
456 is_other_int_dtype = rvalues.dtype.kind in "iu"
457 if not is_other_int_dtype:
458 rvalues = fill_bool(rvalues, lvalues)
460 else:
461 # i.e. scalar
462 is_other_int_dtype = lib.is_integer(rvalues)
464 res_values = na_logical_op(lvalues, rvalues, op)
466 # For int vs int `^`, `|`, `&` are bitwise operators and return
467 # integer dtypes. Otherwise these are boolean ops
468 if not (left.dtype.kind in "iu" and is_other_int_dtype):
469 res_values = fill_bool(res_values)
471 return res_values
474def get_array_op(op):
475 """
476 Return a binary array operation corresponding to the given operator op.
478 Parameters
479 ----------
480 op : function
481 Binary operator from operator or roperator module.
483 Returns
484 -------
485 functools.partial
486 """
487 if isinstance(op, partial):
488 # We get here via dispatch_to_series in DataFrame case
489 # e.g. test_rolling_consistency_var_debiasing_factors
490 return op
492 op_name = op.__name__.strip("_").lstrip("r")
493 if op_name == "arith_op":
494 # Reached via DataFrame._combine_frame i.e. flex methods
495 # e.g. test_df_add_flex_filled_mixed_dtypes
496 return op
498 if op_name in {"eq", "ne", "lt", "le", "gt", "ge"}:
499 return partial(comparison_op, op=op)
500 elif op_name in {"and", "or", "xor", "rand", "ror", "rxor"}:
501 return partial(logical_op, op=op)
502 elif op_name in {
503 "add",
504 "sub",
505 "mul",
506 "truediv",
507 "floordiv",
508 "mod",
509 "divmod",
510 "pow",
511 }:
512 return partial(arithmetic_op, op=op)
513 else:
514 raise NotImplementedError(op_name)
517def maybe_prepare_scalar_for_op(obj, shape: Shape):
518 """
519 Cast non-pandas objects to pandas types to unify behavior of arithmetic
520 and comparison operations.
522 Parameters
523 ----------
524 obj: object
525 shape : tuple[int]
527 Returns
528 -------
529 out : object
531 Notes
532 -----
533 Be careful to call this *after* determining the `name` attribute to be
534 attached to the result of the arithmetic operation.
535 """
536 if type(obj) is datetime.timedelta:
537 # GH#22390 cast up to Timedelta to rely on Timedelta
538 # implementation; otherwise operation against numeric-dtype
539 # raises TypeError
540 return Timedelta(obj)
541 elif type(obj) is datetime.datetime:
542 # cast up to Timestamp to rely on Timestamp implementation, see Timedelta above
543 return Timestamp(obj)
544 elif isinstance(obj, np.datetime64):
545 # GH#28080 numpy casts integer-dtype to datetime64 when doing
546 # array[int] + datetime64, which we do not allow
547 if isna(obj):
548 from pandas.core.arrays import DatetimeArray
550 # Avoid possible ambiguities with pd.NaT
551 # GH 52295
552 if is_unitless(obj.dtype):
553 # Use second resolution to ensure that the result of e.g.
554 # `left - np.datetime64("NaT")` retains the unit of left.unit
555 obj = obj.astype("datetime64[s]")
556 elif not is_supported_dtype(obj.dtype):
557 new_dtype = get_supported_dtype(obj.dtype)
558 obj = obj.astype(new_dtype)
559 right = np.broadcast_to(obj, shape)
560 return DatetimeArray._simple_new(right, dtype=right.dtype)
562 return Timestamp(obj)
564 elif isinstance(obj, np.timedelta64):
565 if isna(obj):
566 from pandas.core.arrays import TimedeltaArray
568 # wrapping timedelta64("NaT") in Timedelta returns NaT,
569 # which would incorrectly be treated as a datetime-NaT, so
570 # we broadcast and wrap in a TimedeltaArray
571 # GH 52295
572 if is_unitless(obj.dtype):
573 # Use second resolution to ensure that the result of e.g.
574 # `left + np.timedelta64("NaT")` retains the unit of left.unit
575 obj = obj.astype("timedelta64[s]")
576 elif not is_supported_dtype(obj.dtype):
577 new_dtype = get_supported_dtype(obj.dtype)
578 obj = obj.astype(new_dtype)
579 right = np.broadcast_to(obj, shape)
580 return TimedeltaArray._simple_new(right, dtype=right.dtype)
582 # In particular non-nanosecond timedelta64 needs to be cast to
583 # nanoseconds, or else we get undesired behavior like
584 # np.timedelta64(3, 'D') / 2 == np.timedelta64(1, 'D')
585 return Timedelta(obj)
587 # We want NumPy numeric scalars to behave like Python scalars
588 # post NEP 50
589 elif isinstance(obj, np.integer):
590 return int(obj)
592 elif isinstance(obj, np.floating):
593 return float(obj)
595 return obj
598_BOOL_OP_NOT_ALLOWED = {
599 operator.truediv,
600 roperator.rtruediv,
601 operator.floordiv,
602 roperator.rfloordiv,
603 operator.pow,
604 roperator.rpow,
605 divmod,
606 roperator.rdivmod,
607}
610def _bool_arith_check(op, a: np.ndarray, b) -> None:
611 """
612 In contrast to numpy, pandas raises an error for certain operations
613 with booleans.
614 """
615 if op in _BOOL_OP_NOT_ALLOWED:
616 if a.dtype.kind == "b" and (is_bool_dtype(b) or lib.is_bool(b)):
617 op_name = op.__name__.strip("_").lstrip("r")
618 raise NotImplementedError(
619 f"operator '{op_name}' not implemented for bool dtypes"
620 )