CLN: always dispatch-to-series (#34286)

* comments

* CLN: always dispatch-to-series

* More accurate assertions

* clarify comment, annotations

Author: jbrockmendel

pandas/_libs/tslibs/offsets.pyx

@@ -1299,7 +1299,7 @@ cdef class QuarterOffset(SingleConstructorOffset):
         return type(dtindex)._simple_new(shifted, dtype=dtindex.dtype)
 
 
-cdef class MonthOffset(BaseOffset):
+cdef class MonthOffset(SingleConstructorOffset):
     def is_on_offset(self, dt) -> bool:
         if self.normalize and not is_normalized(dt):
             return False
@@ -1316,13 +1316,6 @@ cdef class MonthOffset(BaseOffset):
         shifted = shift_months(dtindex.asi8, self.n, self._day_opt)
         return type(dtindex)._simple_new(shifted, dtype=dtindex.dtype)
 
-    @classmethod
-    def _from_name(cls, suffix=None):
-        # default _from_name calls cls with no args
-        if suffix:
-            raise ValueError(f"Bad freq suffix {suffix}")
-        return cls()
-
 
 # ----------------------------------------------------------------------
 # RelativeDelta Arithmetic

pandas/core/frame.py

@@ -135,7 +135,6 @@
     sanitize_index,
     to_arrays,
 )
-from pandas.core.ops.missing import dispatch_fill_zeros
 from pandas.core.series import Series
 from pandas.core.sorting import ensure_key_mapped
 
@@ -5734,14 +5733,7 @@ def _arith_op(left, right):
                 left, right = ops.fill_binop(left, right, fill_value)
                 return func(left, right)
 
-        if ops.should_series_dispatch(self, other, func):
-            # iterate over columns
-            new_data = ops.dispatch_to_series(self, other, _arith_op)
-        else:
-            with np.errstate(all="ignore"):
-                res_values = _arith_op(self.values, other.values)
-            new_data = dispatch_fill_zeros(func, self.values, other.values, res_values)
-
+        new_data = ops.dispatch_to_series(self, other, _arith_op)
         return new_data
 
     def _construct_result(self, result) -> "DataFrame":

pandas/core/ops/__init__.py

@@ -4,7 +4,7 @@
 This is not a public API.
 """
 import operator
-from typing import TYPE_CHECKING, Optional, Set
+from typing import TYPE_CHECKING, Optional, Set, Type
 
 import numpy as np
 
@@ -21,13 +21,11 @@
 from pandas.core.ops.array_ops import (
     arithmetic_op,
     comparison_op,
-    define_na_arithmetic_op,
     get_array_op,
     logical_op,
 )
 from pandas.core.ops.array_ops import comp_method_OBJECT_ARRAY  # noqa:F401
 from pandas.core.ops.common import unpack_zerodim_and_defer
-from pandas.core.ops.dispatch import should_series_dispatch
 from pandas.core.ops.docstrings import (
     _arith_doc_FRAME,
     _flex_comp_doc_FRAME,
@@ -154,7 +152,7 @@ def _maybe_match_name(a, b):
 # -----------------------------------------------------------------------------
 
 
-def _get_frame_op_default_axis(name):
+def _get_frame_op_default_axis(name: str) -> Optional[str]:
     """
     Only DataFrame cares about default_axis, specifically:
     special methods have default_axis=None and flex methods
@@ -277,7 +275,11 @@ def dispatch_to_series(left, right, func, axis=None):
         return type(left)(bm)
 
     elif isinstance(right, ABCDataFrame):
-        assert right._indexed_same(left)
+        assert left.index.equals(right.index)
+        assert left.columns.equals(right.columns)
+        # TODO: The previous assertion `assert right._indexed_same(left)`
+        #  fails in cases with empty columns reached via
+        #  _frame_arith_method_with_reindex
 
         array_op = get_array_op(func)
         bm = left._mgr.operate_blockwise(right._mgr, array_op)
@@ -345,6 +347,7 @@ def _arith_method_SERIES(cls, op, special):
     Wrapper function for Series arithmetic operations, to avoid
     code duplication.
     """
+    assert special  # non-special uses _flex_method_SERIES
     op_name = _get_op_name(op, special)
 
     @unpack_zerodim_and_defer(op_name)
@@ -368,6 +371,7 @@ def _comp_method_SERIES(cls, op, special):
     Wrapper function for Series arithmetic operations, to avoid
     code duplication.
     """
+    assert special  # non-special uses _flex_method_SERIES
     op_name = _get_op_name(op, special)
 
     @unpack_zerodim_and_defer(op_name)
@@ -394,6 +398,7 @@ def _bool_method_SERIES(cls, op, special):
     Wrapper function for Series arithmetic operations, to avoid
     code duplication.
     """
+    assert special  # non-special uses _flex_method_SERIES
     op_name = _get_op_name(op, special)
 
     @unpack_zerodim_and_defer(op_name)
@@ -412,6 +417,7 @@ def wrapper(self, other):
 
 
 def _flex_method_SERIES(cls, op, special):
+    assert not special  # "special" also means "not flex"
     name = _get_op_name(op, special)
     doc = _make_flex_doc(name, "series")
 
@@ -574,7 +580,7 @@ def to_series(right):
 
 
 def _should_reindex_frame_op(
-    left: "DataFrame", right, op, axis, default_axis: int, fill_value, level
+    left: "DataFrame", right, op, axis, default_axis, fill_value, level
 ) -> bool:
     """
     Check if this is an operation between DataFrames that will need to reindex.
@@ -629,11 +635,12 @@ def _frame_arith_method_with_reindex(
     return result.reindex(join_columns, axis=1)
 
 
-def _arith_method_FRAME(cls, op, special):
+def _arith_method_FRAME(cls: Type["DataFrame"], op, special: bool):
+    # This is the only function where `special` can be either True or False
     op_name = _get_op_name(op, special)
     default_axis = _get_frame_op_default_axis(op_name)
 
-    na_op = define_na_arithmetic_op(op)
+    na_op = get_array_op(op)
     is_logical = op.__name__.strip("_").lstrip("_") in ["and", "or", "xor"]
 
     if op_name in _op_descriptions:
@@ -650,18 +657,19 @@ def f(self, other, axis=default_axis, level=None, fill_value=None):
         ):
             return _frame_arith_method_with_reindex(self, other, op)
 
+        # TODO: why are we passing flex=True instead of flex=not special?
+        #  15 tests fail if we pass flex=not special instead
         self, other = _align_method_FRAME(self, other, axis, flex=True, level=level)
 
         if isinstance(other, ABCDataFrame):
             # Another DataFrame
-            pass_op = op if should_series_dispatch(self, other, op) else na_op
-            pass_op = pass_op if not is_logical else op
-
-            new_data = self._combine_frame(other, pass_op, fill_value)
+            new_data = self._combine_frame(other, na_op, fill_value)
 
         elif isinstance(other, ABCSeries):
             # For these values of `axis`, we end up dispatching to Series op,
             # so do not want the masked op.
+            # TODO: the above comment is no longer accurate since we now
+            #  operate blockwise if other._values is an ndarray
             pass_op = op if axis in [0, "columns", None] else na_op
             pass_op = pass_op if not is_logical else op
 
@@ -684,9 +692,11 @@ def f(self, other, axis=default_axis, level=None, fill_value=None):
     return f
 
 
-def _flex_comp_method_FRAME(cls, op, special):
+def _flex_comp_method_FRAME(cls: Type["DataFrame"], op, special: bool):
+    assert not special  # "special" also means "not flex"
     op_name = _get_op_name(op, special)
     default_axis = _get_frame_op_default_axis(op_name)
+    assert default_axis == "columns", default_axis  # because we are not "special"
 
     doc = _flex_comp_doc_FRAME.format(
         op_name=op_name, desc=_op_descriptions[op_name]["desc"]
@@ -715,7 +725,8 @@ def f(self, other, axis=default_axis, level=None):
     return f
 
 
-def _comp_method_FRAME(cls, op, special):
+def _comp_method_FRAME(cls: Type["DataFrame"], op, special: bool):
+    assert special  # "special" also means "not flex"
     op_name = _get_op_name(op, special)
 
     @Appender(f"Wrapper for comparison method {op_name}")

pandas/core/ops/array_ops.py

@@ -121,13 +121,6 @@ def masked_arith_op(x: np.ndarray, y, op):
     return result
 
 
-def define_na_arithmetic_op(op):
-    def na_op(x, y):
-        return na_arithmetic_op(x, y, op)
-
-    return na_op
-
-
 def na_arithmetic_op(left, right, op, is_cmp: bool = False):
     """
     Return the result of evaluating op on the passed in values.
@@ -378,8 +371,13 @@ def get_array_op(op):
 
     Returns
     -------
-    function
+    functools.partial
     """
+    if isinstance(op, partial):
+        # We get here via dispatch_to_series in DataFrame case
+        # TODO: avoid getting here
+        return op
+
     op_name = op.__name__.strip("_")
     if op_name in {"eq", "ne", "lt", "le", "gt", "ge"}:
         return partial(comparison_op, op=op)

pandas/core/ops/dispatch.py

@@ -5,12 +5,6 @@
 
 from pandas._typing import ArrayLike
 
-from pandas.core.dtypes.common import (
-    is_datetime64_dtype,
-    is_integer_dtype,
-    is_object_dtype,
-    is_timedelta64_dtype,
-)
 from pandas.core.dtypes.generic import ABCExtensionArray
 
 
@@ -28,57 +22,3 @@ def should_extension_dispatch(left: ArrayLike, right: Any) -> bool:
     bool
     """
     return isinstance(left, ABCExtensionArray) or isinstance(right, ABCExtensionArray)
-
-
-def should_series_dispatch(left, right, op):
-    """
-    Identify cases where a DataFrame operation should dispatch to its
-    Series counterpart.
-
-    Parameters
-    ----------
-    left : DataFrame
-    right : DataFrame or Series
-    op : binary operator
-
-    Returns
-    -------
-    override : bool
-    """
-    if left._is_mixed_type or right._is_mixed_type:
-        return True
-
-    if op.__name__.strip("_") in ["and", "or", "xor", "rand", "ror", "rxor"]:
-        # TODO: GH references for what this fixes
-        # Note: this check must come before the check for nonempty columns.
-        return True
-
-    if right.ndim == 1:
-        # operating with Series, short-circuit checks that would fail
-        #  with AttributeError.
-        return False
-
-    if not len(left.columns) or not len(right.columns):
-        # ensure obj.dtypes[0] exists for each obj
-        return False
-
-    ldtype = left.dtypes.iloc[0]
-    rdtype = right.dtypes.iloc[0]
-
-    if (
-        is_timedelta64_dtype(ldtype)
-        and (is_integer_dtype(rdtype) or is_object_dtype(rdtype))
-    ) or (
-        is_timedelta64_dtype(rdtype)
-        and (is_integer_dtype(ldtype) or is_object_dtype(ldtype))
-    ):
-        # numpy integer dtypes as timedelta64 dtypes in this scenario
-        return True
-
-    if (is_datetime64_dtype(ldtype) and is_object_dtype(rdtype)) or (
-        is_datetime64_dtype(rdtype) and is_object_dtype(ldtype)
-    ):
-        # in particular case where one is an array of DateOffsets
-        return True
-
-    return False

pandas/core/ops/methods.py

@@ -207,7 +207,6 @@ def _create_methods(cls, arith_method, comp_method, bool_method, special):
             dict(
                 and_=bool_method(cls, operator.and_, special),
                 or_=bool_method(cls, operator.or_, special),
-                # For some reason ``^`` wasn't used in original.
                 xor=bool_method(cls, operator.xor, special),
                 rand_=bool_method(cls, rand_, special),
                 ror_=bool_method(cls, ror_, special),

pandas/tests/arithmetic/test_timedelta64.py

@@ -552,7 +552,10 @@ def test_tda_add_dt64_object_array(self, box_with_array, tz_naive_fixture):
         obj = tm.box_expected(tdi, box)
         other = tm.box_expected(dti, box)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = None
+        if box is not pd.DataFrame or tz_naive_fixture is None:
+            warn = PerformanceWarning
+        with tm.assert_produces_warning(warn):
             result = obj + other.astype(object)
         tm.assert_equal(result, other)
 

pandas/tseries/offsets.py

@@ -790,7 +790,7 @@ def __init__(
 # Month-Based Offset Classes
 
 
-class MonthEnd(SingleConstructorMixin, liboffsets.MonthOffset):
+class MonthEnd(liboffsets.MonthOffset):
     """
     DateOffset of one month end.
     """
@@ -799,7 +799,7 @@ class MonthEnd(SingleConstructorMixin, liboffsets.MonthOffset):
     _day_opt = "end"
 
 
-class MonthBegin(SingleConstructorMixin, liboffsets.MonthOffset):
+class MonthBegin(liboffsets.MonthOffset):
     """
     DateOffset of one month at beginning.
     """
@@ -808,7 +808,7 @@ class MonthBegin(SingleConstructorMixin, liboffsets.MonthOffset):
     _day_opt = "start"
 
 
-class BusinessMonthEnd(SingleConstructorMixin, liboffsets.MonthOffset):
+class BusinessMonthEnd(liboffsets.MonthOffset):
     """
     DateOffset increments between business EOM dates.
     """
@@ -817,7 +817,7 @@ class BusinessMonthEnd(SingleConstructorMixin, liboffsets.MonthOffset):
     _day_opt = "business_end"
 
 
-class BusinessMonthBegin(SingleConstructorMixin, liboffsets.MonthOffset):
+class BusinessMonthBegin(liboffsets.MonthOffset):
     """
     DateOffset of one business month at beginning.
     """
@@ -827,9 +827,7 @@ class BusinessMonthBegin(SingleConstructorMixin, liboffsets.MonthOffset):
 
 
 @doc(bound="bound")
-class _CustomBusinessMonth(
-    CustomMixin, BusinessMixin, SingleConstructorMixin, liboffsets.MonthOffset
-):
+class _CustomBusinessMonth(CustomMixin, BusinessMixin, liboffsets.MonthOffset):
     """
     DateOffset subclass representing custom business month(s).