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"""
2Methods that can be shared by many array-like classes or subclasses:
3 Series
4 Index
5 ExtensionArray
6"""
8from __future__ import annotations
10import operator
11from typing import Any
13import numpy as np
15from pandas._libs import lib
16from pandas._libs.ops_dispatch import maybe_dispatch_ufunc_to_dunder_op
18from pandas.core.dtypes.cast import maybe_unbox_numpy_scalar
19from pandas.core.dtypes.generic import ABCNDFrame
21from pandas.core import roperator
22from pandas.core.construction import extract_array
23from pandas.core.ops.common import unpack_zerodim_and_defer
25REDUCTION_ALIASES = {
26 "maximum": "max",
27 "minimum": "min",
28 "add": "sum",
29 "multiply": "prod",
30}
33class OpsMixin:
34 # -------------------------------------------------------------
35 # Comparisons
37 def _cmp_method(self, other, op):
38 return NotImplemented
40 @unpack_zerodim_and_defer("__eq__")
41 def __eq__(self, other):
42 return self._cmp_method(other, operator.eq)
44 @unpack_zerodim_and_defer("__ne__")
45 def __ne__(self, other):
46 return self._cmp_method(other, operator.ne)
48 @unpack_zerodim_and_defer("__lt__")
49 def __lt__(self, other):
50 return self._cmp_method(other, operator.lt)
52 @unpack_zerodim_and_defer("__le__")
53 def __le__(self, other):
54 return self._cmp_method(other, operator.le)
56 @unpack_zerodim_and_defer("__gt__")
57 def __gt__(self, other):
58 return self._cmp_method(other, operator.gt)
60 @unpack_zerodim_and_defer("__ge__")
61 def __ge__(self, other):
62 return self._cmp_method(other, operator.ge)
64 # -------------------------------------------------------------
65 # Logical Methods
67 def _logical_method(self, other, op):
68 return NotImplemented
70 @unpack_zerodim_and_defer("__and__")
71 def __and__(self, other):
72 return self._logical_method(other, operator.and_)
74 @unpack_zerodim_and_defer("__rand__")
75 def __rand__(self, other):
76 return self._logical_method(other, roperator.rand_)
78 @unpack_zerodim_and_defer("__or__")
79 def __or__(self, other):
80 return self._logical_method(other, operator.or_)
82 @unpack_zerodim_and_defer("__ror__")
83 def __ror__(self, other):
84 return self._logical_method(other, roperator.ror_)
86 @unpack_zerodim_and_defer("__xor__")
87 def __xor__(self, other):
88 return self._logical_method(other, operator.xor)
90 @unpack_zerodim_and_defer("__rxor__")
91 def __rxor__(self, other):
92 return self._logical_method(other, roperator.rxor)
94 # -------------------------------------------------------------
95 # Arithmetic Methods
97 def _arith_method(self, other, op):
98 return NotImplemented
100 @unpack_zerodim_and_defer("__add__")
101 def __add__(self, other):
102 """
103 Get Addition of DataFrame and other, column-wise.
105 Equivalent to ``DataFrame.add(other)``.
107 Parameters
108 ----------
109 other : scalar, sequence, Series, dict or DataFrame
110 Object to be added to the DataFrame.
112 Returns
113 -------
114 DataFrame
115 The result of adding ``other`` to DataFrame.
117 See Also
118 --------
119 DataFrame.add : Add a DataFrame and another object, with option for index-
120 or column-oriented addition.
122 Examples
123 --------
124 >>> df = pd.DataFrame(
125 ... {"height": [1.5, 2.6], "weight": [500, 800]}, index=["elk", "moose"]
126 ... )
127 >>> df
128 height weight
129 elk 1.5 500
130 moose 2.6 800
132 Adding a scalar affects all rows and columns.
134 >>> df[["height", "weight"]] + 1.5
135 height weight
136 elk 3.0 501.5
137 moose 4.1 801.5
139 Each element of a list is added to a column of the DataFrame, in order.
141 >>> df[["height", "weight"]] + [0.5, 1.5]
142 height weight
143 elk 2.0 501.5
144 moose 3.1 801.5
146 Keys of a dictionary are aligned to the DataFrame, based on column names;
147 each value in the dictionary is added to the corresponding column.
149 >>> df[["height", "weight"]] + {"height": 0.5, "weight": 1.5}
150 height weight
151 elk 2.0 501.5
152 moose 3.1 801.5
154 When `other` is a :class:`Series`, the index of `other` is aligned with the
155 columns of the DataFrame.
157 >>> s1 = pd.Series([0.5, 1.5], index=["weight", "height"])
158 >>> df[["height", "weight"]] + s1
159 height weight
160 elk 3.0 500.5
161 moose 4.1 800.5
163 Even when the index of `other` is the same as the index of the DataFrame,
164 the :class:`Series` will not be reoriented. If index-wise alignment is desired,
165 :meth:`DataFrame.add` should be used with `axis='index'`.
167 >>> s2 = pd.Series([0.5, 1.5], index=["elk", "moose"])
168 >>> df[["height", "weight"]] + s2
169 elk height moose weight
170 elk NaN NaN NaN NaN
171 moose NaN NaN NaN NaN
173 >>> df[["height", "weight"]].add(s2, axis="index")
174 height weight
175 elk 2.0 500.5
176 moose 4.1 801.5
178 When `other` is a :class:`DataFrame`, both columns names and the
179 index are aligned.
181 >>> other = pd.DataFrame(
182 ... {"height": [0.2, 0.4, 0.6]}, index=["elk", "moose", "deer"]
183 ... )
184 >>> df[["height", "weight"]] + other
185 height weight
186 deer NaN NaN
187 elk 1.7 NaN
188 moose 3.0 NaN
189 """
190 return self._arith_method(other, operator.add)
192 @unpack_zerodim_and_defer("__radd__")
193 def __radd__(self, other):
194 return self._arith_method(other, roperator.radd)
196 @unpack_zerodim_and_defer("__sub__")
197 def __sub__(self, other):
198 return self._arith_method(other, operator.sub)
200 @unpack_zerodim_and_defer("__rsub__")
201 def __rsub__(self, other):
202 return self._arith_method(other, roperator.rsub)
204 @unpack_zerodim_and_defer("__mul__")
205 def __mul__(self, other):
206 return self._arith_method(other, operator.mul)
208 @unpack_zerodim_and_defer("__rmul__")
209 def __rmul__(self, other):
210 return self._arith_method(other, roperator.rmul)
212 @unpack_zerodim_and_defer("__truediv__")
213 def __truediv__(self, other):
214 return self._arith_method(other, operator.truediv)
216 @unpack_zerodim_and_defer("__rtruediv__")
217 def __rtruediv__(self, other):
218 return self._arith_method(other, roperator.rtruediv)
220 @unpack_zerodim_and_defer("__floordiv__")
221 def __floordiv__(self, other):
222 return self._arith_method(other, operator.floordiv)
224 @unpack_zerodim_and_defer("__rfloordiv")
225 def __rfloordiv__(self, other):
226 return self._arith_method(other, roperator.rfloordiv)
228 @unpack_zerodim_and_defer("__mod__")
229 def __mod__(self, other):
230 return self._arith_method(other, operator.mod)
232 @unpack_zerodim_and_defer("__rmod__")
233 def __rmod__(self, other):
234 return self._arith_method(other, roperator.rmod)
236 @unpack_zerodim_and_defer("__divmod__")
237 def __divmod__(self, other):
238 return self._arith_method(other, divmod)
240 @unpack_zerodim_and_defer("__rdivmod__")
241 def __rdivmod__(self, other):
242 return self._arith_method(other, roperator.rdivmod)
244 @unpack_zerodim_and_defer("__pow__")
245 def __pow__(self, other):
246 return self._arith_method(other, operator.pow)
248 @unpack_zerodim_and_defer("__rpow__")
249 def __rpow__(self, other):
250 return self._arith_method(other, roperator.rpow)
253# -----------------------------------------------------------------------------
254# Helpers to implement __array_ufunc__
257def array_ufunc(self, ufunc: np.ufunc, method: str, *inputs: Any, **kwargs: Any):
258 """
259 Compatibility with numpy ufuncs.
261 See also
262 --------
263 numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__
264 """
265 from pandas.core.frame import (
266 DataFrame,
267 Series,
268 )
269 from pandas.core.generic import NDFrame
270 from pandas.core.internals import BlockManager
272 cls = type(self)
274 kwargs = _standardize_out_kwarg(**kwargs)
276 # for binary ops, use our custom dunder methods
277 result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs)
278 if result is not NotImplemented:
279 return result
281 # Determine if we should defer.
282 no_defer = (
283 np.ndarray.__array_ufunc__,
284 cls.__array_ufunc__,
285 )
287 for item in inputs:
288 higher_priority = (
289 hasattr(item, "__array_priority__")
290 and item.__array_priority__ > self.__array_priority__
291 )
292 has_array_ufunc = (
293 hasattr(item, "__array_ufunc__")
294 and type(item).__array_ufunc__ not in no_defer
295 and not isinstance(item, self._HANDLED_TYPES)
296 )
297 if higher_priority or has_array_ufunc:
298 return NotImplemented
300 # align all the inputs.
301 types = tuple(type(x) for x in inputs)
302 alignable = [
303 x for x, t in zip(inputs, types, strict=True) if issubclass(t, NDFrame)
304 ]
306 if len(alignable) > 1:
307 # This triggers alignment.
308 # At the moment, there aren't any ufuncs with more than two inputs
309 # so this ends up just being x1.index | x2.index, but we write
310 # it to handle *args.
311 set_types = set(types)
312 if len(set_types) > 1 and {DataFrame, Series}.issubset(set_types):
313 # We currently don't handle ufunc(DataFrame, Series)
314 # well. Previously this raised an internal ValueError. We might
315 # support it someday, so raise a NotImplementedError.
316 raise NotImplementedError(
317 f"Cannot apply ufunc {ufunc} to mixed DataFrame and Series inputs."
318 )
319 axes = self.axes
320 for obj in alignable[1:]:
321 # this relies on the fact that we aren't handling mixed
322 # series / frame ufuncs.
323 for i, (ax1, ax2) in enumerate(zip(axes, obj.axes, strict=True)):
324 axes[i] = ax1.union(ax2)
326 reconstruct_axes = dict(zip(self._AXIS_ORDERS, axes, strict=True))
327 inputs = tuple(
328 x.reindex(**reconstruct_axes) if issubclass(t, NDFrame) else x
329 for x, t in zip(inputs, types, strict=True)
330 )
331 else:
332 reconstruct_axes = dict(zip(self._AXIS_ORDERS, self.axes, strict=True))
334 if self.ndim == 1:
335 names = {x.name for x in inputs if hasattr(x, "name")}
336 name = names.pop() if len(names) == 1 else None
337 reconstruct_kwargs = {"name": name}
338 else:
339 reconstruct_kwargs = {}
341 def reconstruct(result):
342 if ufunc.nout > 1:
343 # np.modf, np.frexp, np.divmod
344 return tuple(_reconstruct(x) for x in result)
346 return _reconstruct(result)
348 def _reconstruct(result):
349 if lib.is_scalar(result):
350 return result
352 if result.ndim != self.ndim:
353 if method == "outer":
354 raise NotImplementedError
355 return result
356 if isinstance(result, BlockManager):
357 # we went through BlockManager.apply e.g. np.sqrt
358 result = self._constructor_from_mgr(result, axes=result.axes)
359 else:
360 # we converted an array, lost our axes
361 result = self._constructor(
362 result, **reconstruct_axes, **reconstruct_kwargs, copy=False
363 )
364 # TODO: When we support multiple values in __finalize__, this
365 # should pass alignable to `__finalize__` instead of self.
366 # Then `np.add(a, b)` would consider attrs from both a and b
367 # when a and b are NDFrames.
368 if len(alignable) == 1:
369 result = result.__finalize__(self)
370 return result
372 if "out" in kwargs:
373 # e.g. test_multiindex_get_loc
374 result = dispatch_ufunc_with_out(self, ufunc, method, *inputs, **kwargs)
375 return reconstruct(result)
377 if method == "reduce":
378 # e.g. test.series.test_ufunc.test_reduce
379 result = dispatch_reduction_ufunc(self, ufunc, method, *inputs, **kwargs)
380 if result is not NotImplemented:
381 return result
383 # We still get here with kwargs `axis` for e.g. np.maximum.accumulate
384 # and `dtype` and `keepdims` for np.ptp
386 if self.ndim > 1 and (len(inputs) > 1 or ufunc.nout > 1):
387 # Just give up on preserving types in the complex case.
388 # In theory we could preserve them for them.
389 # * nout>1 is doable if BlockManager.apply took nout and
390 # returned a Tuple[BlockManager].
391 # * len(inputs) > 1 is doable when we know that we have
392 # aligned blocks / dtypes.
394 # e.g. my_ufunc, modf, logaddexp, heaviside, subtract, add
395 inputs = tuple(np.asarray(x) for x in inputs)
396 # Note: we can't use default_array_ufunc here bc reindexing means
397 # that `self` may not be among `inputs`
398 result = getattr(ufunc, method)(*inputs, **kwargs)
399 elif self.ndim == 1:
400 # ufunc(series, ...)
401 inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs)
402 result = getattr(ufunc, method)(*inputs, **kwargs)
403 # ufunc(dataframe)
404 elif method == "__call__" and not kwargs:
405 # for np.<ufunc>(..) calls
406 # kwargs cannot necessarily be handled block-by-block, so only
407 # take this path if there are no kwargs
408 mgr = inputs[0]._mgr # pyright: ignore[reportGeneralTypeIssues]
409 result = mgr.apply(getattr(ufunc, method))
410 else:
411 # otherwise specific ufunc methods (eg np.<ufunc>.accumulate(..))
412 # Those can have an axis keyword and thus can't be called block-by-block
413 result = default_array_ufunc(inputs[0], ufunc, method, *inputs, **kwargs) # pyright: ignore[reportGeneralTypeIssues]
414 # e.g. np.negative (only one reached), with "where" and "out" in kwargs
416 result = reconstruct(result)
417 return result
420def _standardize_out_kwarg(**kwargs) -> dict:
421 """
422 If kwargs contain "out1" and "out2", replace that with a tuple "out"
424 np.divmod, np.modf, np.frexp can have either `out=(out1, out2)` or
425 `out1=out1, out2=out2)`
426 """
427 if "out" not in kwargs and "out1" in kwargs and "out2" in kwargs:
428 out1 = kwargs.pop("out1")
429 out2 = kwargs.pop("out2")
430 out = (out1, out2)
431 kwargs["out"] = out
432 return kwargs
435def dispatch_ufunc_with_out(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
436 """
437 If we have an `out` keyword, then call the ufunc without `out` and then
438 set the result into the given `out`.
439 """
441 # Note: we assume _standardize_out_kwarg has already been called.
442 out = kwargs.pop("out")
443 where = kwargs.pop("where", None)
445 result = getattr(ufunc, method)(*inputs, **kwargs)
447 if result is NotImplemented:
448 return NotImplemented
450 if isinstance(result, tuple):
451 # i.e. np.divmod, np.modf, np.frexp
452 if not isinstance(out, tuple) or len(out) != len(result):
453 raise NotImplementedError
455 for arr, res in zip(out, result, strict=True):
456 _assign_where(arr, res, where)
458 return out
460 if isinstance(out, tuple):
461 if len(out) == 1:
462 out = out[0]
463 else:
464 raise NotImplementedError
466 _assign_where(out, result, where)
467 return out
470def _assign_where(out, result, where) -> None:
471 """
472 Set a ufunc result into 'out', masking with a 'where' argument if necessary.
473 """
474 if where is None:
475 # no 'where' arg passed to ufunc
476 out[:] = result
477 else:
478 np.putmask(out, where, result)
481def default_array_ufunc(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
482 """
483 Fallback to the behavior we would get if we did not define __array_ufunc__.
485 Notes
486 -----
487 We are assuming that `self` is among `inputs`.
488 """
489 if not any(x is self for x in inputs):
490 raise NotImplementedError
492 new_inputs = [x if x is not self else np.asarray(x) for x in inputs]
494 return getattr(ufunc, method)(*new_inputs, **kwargs)
497def dispatch_reduction_ufunc(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
498 """
499 Dispatch ufunc reductions to self's reduction methods.
500 """
501 assert method == "reduce"
503 if len(inputs) != 1 or inputs[0] is not self:
504 return NotImplemented
506 if ufunc.__name__ not in REDUCTION_ALIASES:
507 return NotImplemented
509 method_name = REDUCTION_ALIASES[ufunc.__name__]
511 # NB: we are assuming that min/max represent minimum/maximum methods,
512 # which would not be accurate for e.g. Timestamp.min
513 if not hasattr(self, method_name):
514 return NotImplemented
516 if self.ndim > 1:
517 if isinstance(self, ABCNDFrame):
518 # TODO: test cases where this doesn't hold, i.e. 2D DTA/TDA
519 kwargs["numeric_only"] = False
521 if "axis" not in kwargs:
522 # For DataFrame reductions we don't want the default axis=0
523 # Note: np.min is not a ufunc, but uses array_function_dispatch,
524 # so calls DataFrame.min (without ever getting here) with the np.min
525 # default of axis=None, which DataFrame.min catches and changes to axis=0.
526 # np.minimum.reduce(df) gets here bc axis is not in kwargs,
527 # so we set axis=0 to match the behavior of np.minimum.reduce(df.values)
528 kwargs["axis"] = 0
530 # By default, numpy's reductions do not skip NaNs, so we have to
531 # pass skipna=False
532 result = getattr(self, method_name)(skipna=False, **kwargs)
533 result = maybe_unbox_numpy_scalar(result)
534 return result