Skip to content

Commit 0b7a1c8

Browse files
jbrockmendelkrajatclRajat Bishnoifangchenlimicahjsmith
authored and
Kevin D Smith
committed
REF/TYP: use OpsMixin for DataFrame (pandas-dev#37044)
* REF/TYP: use OpsMixin for arithmetic methods * REF: separate arith_method_FRAME from flex_arith_method_FRAME * whatsnew * REF/TYP: use OpsMixin for logical methods (pandas-dev#36964) * TST: insert 'match' to bare pytest raises in pandas/tests/tools/test_to_datetime.py (pandas-dev#37027) * TST: insert 'match' to bare pytest raises in pandas/tests/test_flags.py (pandas-dev#37026) Co-authored-by: Rajat Bishnoi <[email protected]> * TYP: generic, series, frame (pandas-dev#36989) * CI: pin pymysql pandas-dev#36465 (pandas-dev#36847) * CI: unpin sql to verify the bugs pandas-dev#36465 * CI: pin sqlalchemy * CI: pin pymsql * CI: pin sqlalchemy * CI: pin pymysql * CI: pin pymysql * CI: add note * CLN/REF: de-duplicate DatetimeTZBlock.setitem (pandas-dev#37019) * REF/TYP: define NDFrame numeric methods non-dynamically (pandas-dev#37017) * CLN: require td64 in TimedeltaBlock (pandas-dev#37018) * BUG: Raise ValueError instead of bare Exception in sanitize_array (pandas-dev#35769) * CLN: collected cleanups, warning suppression in tests (pandas-dev#37021) * REF/TYP: use OpsMixin for DataFrame * CLN: remove get_op_name * mypy fixup * de-privatize Co-authored-by: krajatcl <[email protected]> Co-authored-by: Rajat Bishnoi <[email protected]> Co-authored-by: Fangchen Li <[email protected]> Co-authored-by: Micah Smith <[email protected]>
1 parent 9958b5c commit 0b7a1c8

File tree

3 files changed

+161
-276
lines changed

3 files changed

+161
-276
lines changed

pandas/core/frame.py

+88-4
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@
124124
relabel_result,
125125
transform,
126126
)
127+
from pandas.core.arraylike import OpsMixin
127128
from pandas.core.arrays import Categorical, ExtensionArray
128129
from pandas.core.arrays.sparse import SparseFrameAccessor
129130
from pandas.core.construction import extract_array
@@ -336,7 +337,7 @@
336337
# DataFrame class
337338

338339

339-
class DataFrame(NDFrame):
340+
class DataFrame(NDFrame, OpsMixin):
340341
"""
341342
Two-dimensional, size-mutable, potentially heterogeneous tabular data.
342343
@@ -5838,7 +5839,87 @@ def reorder_levels(self, order, axis=0) -> DataFrame:
58385839
return result
58395840

58405841
# ----------------------------------------------------------------------
5841-
# Arithmetic / combination related
5842+
# Arithmetic Methods
5843+
5844+
def _cmp_method(self, other, op):
5845+
axis = 1 # only relevant for Series other case
5846+
5847+
self, other = ops.align_method_FRAME(self, other, axis, flex=False, level=None)
5848+
5849+
# See GH#4537 for discussion of scalar op behavior
5850+
new_data = self._dispatch_frame_op(other, op, axis=axis)
5851+
return self._construct_result(new_data)
5852+
5853+
def _arith_method(self, other, op):
5854+
if ops.should_reindex_frame_op(self, other, op, 1, 1, None, None):
5855+
return ops.frame_arith_method_with_reindex(self, other, op)
5856+
5857+
axis = 1 # only relevant for Series other case
5858+
5859+
self, other = ops.align_method_FRAME(self, other, axis, flex=True, level=None)
5860+
5861+
new_data = self._dispatch_frame_op(other, op, axis=axis)
5862+
return self._construct_result(new_data)
5863+
5864+
_logical_method = _arith_method
5865+
5866+
def _dispatch_frame_op(self, right, func, axis: Optional[int] = None):
5867+
"""
5868+
Evaluate the frame operation func(left, right) by evaluating
5869+
column-by-column, dispatching to the Series implementation.
5870+
5871+
Parameters
5872+
----------
5873+
right : scalar, Series, or DataFrame
5874+
func : arithmetic or comparison operator
5875+
axis : {None, 0, 1}
5876+
5877+
Returns
5878+
-------
5879+
DataFrame
5880+
"""
5881+
# Get the appropriate array-op to apply to each column/block's values.
5882+
array_op = ops.get_array_op(func)
5883+
5884+
right = lib.item_from_zerodim(right)
5885+
if not is_list_like(right):
5886+
# i.e. scalar, faster than checking np.ndim(right) == 0
5887+
bm = self._mgr.apply(array_op, right=right)
5888+
return type(self)(bm)
5889+
5890+
elif isinstance(right, DataFrame):
5891+
assert self.index.equals(right.index)
5892+
assert self.columns.equals(right.columns)
5893+
# TODO: The previous assertion `assert right._indexed_same(self)`
5894+
# fails in cases with empty columns reached via
5895+
# _frame_arith_method_with_reindex
5896+
5897+
bm = self._mgr.operate_blockwise(right._mgr, array_op)
5898+
return type(self)(bm)
5899+
5900+
elif isinstance(right, Series) and axis == 1:
5901+
# axis=1 means we want to operate row-by-row
5902+
assert right.index.equals(self.columns)
5903+
5904+
right = right._values
5905+
# maybe_align_as_frame ensures we do not have an ndarray here
5906+
assert not isinstance(right, np.ndarray)
5907+
5908+
arrays = [array_op(l, r) for l, r in zip(self._iter_column_arrays(), right)]
5909+
5910+
elif isinstance(right, Series):
5911+
assert right.index.equals(self.index) # Handle other cases later
5912+
right = right._values
5913+
5914+
arrays = [array_op(l, right) for l in self._iter_column_arrays()]
5915+
5916+
else:
5917+
# Remaining cases have less-obvious dispatch rules
5918+
raise NotImplementedError(right)
5919+
5920+
return type(self)._from_arrays(
5921+
arrays, self.columns, self.index, verify_integrity=False
5922+
)
58425923

58435924
def _combine_frame(self, other: DataFrame, func, fill_value=None):
58445925
# at this point we have `self._indexed_same(other)`
@@ -5857,7 +5938,7 @@ def _arith_op(left, right):
58575938
left, right = ops.fill_binop(left, right, fill_value)
58585939
return func(left, right)
58595940

5860-
new_data = ops.dispatch_to_series(self, other, _arith_op)
5941+
new_data = self._dispatch_frame_op(other, _arith_op)
58615942
return new_data
58625943

58635944
def _construct_result(self, result) -> DataFrame:
@@ -5879,6 +5960,9 @@ def _construct_result(self, result) -> DataFrame:
58795960
out.index = self.index
58805961
return out
58815962

5963+
# ----------------------------------------------------------------------
5964+
# Combination-Related
5965+
58825966
@Appender(
58835967
"""
58845968
Returns
@@ -7254,7 +7338,7 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> DataFrame:
72547338
bm_axis = self._get_block_manager_axis(axis)
72557339

72567340
if bm_axis == 0 and periods != 0:
7257-
return self - self.shift(periods, axis=axis) # type: ignore[operator]
7341+
return self - self.shift(periods, axis=axis)
72587342

72597343
new_data = self._mgr.diff(n=periods, axis=bm_axis)
72607344
return self._constructor(new_data)

0 commit comments

Comments
 (0)