BUG: flex op with DataFrame, Series and ea vs ndarray

doc/source/whatsnew/v1.1.0.rst

@@ -686,6 +686,7 @@ Numeric
 - Bug in DataFrame reductions using ``numeric_only=True`` and ExtensionArrays (:issue:`33256`).
 - Bug in :meth:`DataFrame.corr` and :meth:`DataFrame.cov` raising when handling nullable integer columns with ``pandas.NA`` (:issue:`33803`)
 - Bug in :class:`DataFrame` and :class:`Series` addition and subtraction between object-dtype objects and ``datetime64`` dtype objects (:issue:`33824`)
+- Bug in :class:`DataFrame` flex arithmetic methods with ExtensionDtypes (:issue:`34277`)
 
 Conversion
 ^^^^^^^^^^

pandas/core/ops/__init__.py

@@ -57,7 +57,7 @@
 )
 
 if TYPE_CHECKING:
-    from pandas import DataFrame  # noqa:F401
+    from pandas import DataFrame, Series  # noqa:F401
 
 # -----------------------------------------------------------------------------
 # constants
@@ -511,19 +511,7 @@ def _combine_series_frame(left, right, func, axis: int, str_rep: str):
     # We assume that self.align(other, ...) has already been called
 
     rvalues = right._values
-    if isinstance(rvalues, np.ndarray):
-        # TODO(EA2D): no need to special-case with 2D EAs
-        # We can operate block-wise
-        if axis == 0:
-            rvalues = rvalues.reshape(-1, 1)
-        else:
-            rvalues = rvalues.reshape(1, -1)
-
-        rvalues = np.broadcast_to(rvalues, left.shape)
-
-        array_op = get_array_op(func, str_rep=str_rep)
-        bm = left._mgr.apply(array_op, right=rvalues.T, align_keys=["right"])
-        return type(left)(bm)
+    assert not isinstance(rvalues, np.ndarray)  # handled by align_series_as_frame
 
     if axis == 0:
         new_data = dispatch_to_series(left, right, func)
@@ -619,6 +607,7 @@ def to_series(right):
         left, right = left.align(
             right, join="outer", axis=axis, level=level, copy=False
         )
+        right = _align_series_as_frame(left, right, axis)
 
     return left, right
 
@@ -679,6 +668,25 @@ def _frame_arith_method_with_reindex(
     return result.reindex(join_columns, axis=1)
 
 
+def _align_series_as_frame(frame: "DataFrame", series: "Series", axis: int):
+    """
+    If the Series operand is not EA-dtype, we can broadcast to 2D and operate
+    blockwise.
+    """
+    rvalues = series._values
+    if not isinstance(rvalues, np.ndarray):
+        # TODO(EA2D): no need to special-case with 2D EAs
+        return series
+
+    if axis == 0:
+        rvalues = rvalues.reshape(-1, 1)
+    else:
+        rvalues = rvalues.reshape(1, -1)
+
+    rvalues = np.broadcast_to(rvalues, frame.shape)
+    return type(frame)(rvalues, index=frame.index, columns=frame.columns)
+
+
 def _arith_method_FRAME(cls, op, special):
     str_rep = _get_opstr(op)
     op_name = _get_op_name(op, special)
@@ -701,6 +709,10 @@ def f(self, other, axis=default_axis, level=None, fill_value=None):
         ):
             return _frame_arith_method_with_reindex(self, other, op)
 
+        if isinstance(other, ABCSeries) and fill_value is not None:
+            # TODO: We could allow this in cases where we end up going
+            #  through the DataFrame path
+            raise NotImplementedError(f"fill_value {fill_value} not supported.")
         self, other = _align_method_FRAME(self, other, axis, flex=True, level=level)
 
         if isinstance(other, ABCDataFrame):
@@ -716,9 +728,6 @@ def f(self, other, axis=default_axis, level=None, fill_value=None):
             pass_op = op if axis in [0, "columns", None] else na_op
             pass_op = pass_op if not is_logical else op
 
-            if fill_value is not None:
-                raise NotImplementedError(f"fill_value {fill_value} not supported.")
-
             axis = self._get_axis_number(axis) if axis is not None else 1
             new_data = _combine_series_frame(
                 self, other, pass_op, axis=axis, str_rep=str_rep

pandas/core/ops/array_ops.py

@@ -386,8 +386,13 @@ def get_array_op(op, str_rep: Optional[str] = None):
 
     Returns
     -------
-    function
+    functools.partial
     """
+    if isinstance(op, partial):
+        # We get here via dispatch_to_series in DataFrame case
+        # TODO: try to 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, str_rep=str_rep)

pandas/tests/arithmetic/test_timedelta64.py

@@ -1467,8 +1467,6 @@ def test_td64arr_add_sub_object_array(self, box_with_array):
             [pd.Timedelta(days=2), pd.Timedelta(days=4), pd.Timestamp("2000-01-07")]
         )
         expected = tm.box_expected(expected, box_with_array)
-        if box_with_array is pd.DataFrame:
-            expected = expected.astype(object)
         tm.assert_equal(result, expected)
 
         msg = "unsupported operand type|cannot subtract a datelike"
@@ -1483,8 +1481,6 @@ def test_td64arr_add_sub_object_array(self, box_with_array):
             [pd.Timedelta(0), pd.Timedelta(0), pd.Timestamp("2000-01-01")]
         )
         expected = tm.box_expected(expected, box_with_array)
-        if box_with_array is pd.DataFrame:
-            expected = expected.astype(object)
         tm.assert_equal(result, expected)
 
 
@@ -2009,7 +2005,7 @@ def test_td64arr_div_numeric_array(self, box_with_array, vector, any_real_dtype)
         tm.assert_equal(result, expected)
 
         pattern = (
-            "true_divide cannot use operands|"
+            "true_divide'? cannot use operands|"
             "cannot perform __div__|"
             "cannot perform __truediv__|"
             "unsupported operand|"

pandas/tests/frame/test_arithmetic.py

@@ -339,6 +339,20 @@ def test_df_flex_cmp_constant_return_types_empty(self, opname):
         result = getattr(empty, opname)(const).dtypes.value_counts()
         tm.assert_series_equal(result, pd.Series([2], index=[np.dtype(bool)]))
 
+    def test_df_flex_cmp_ea_dtype_with_ndarray_series(self):
+        ii = pd.IntervalIndex.from_breaks([1, 2, 3])
+        df = pd.DataFrame({"A": ii, "B": ii})
+
+        ser = pd.Series([0, 0])
+        res = df.eq(ser, axis=0)
+
+        expected = pd.DataFrame({"A": [False, False], "B": [False, False]})
+        tm.assert_frame_equal(res, expected)
+
+        ser2 = pd.Series([1, 2], index=["A", "B"])
+        res2 = df.eq(ser2, axis=1)
+        tm.assert_frame_equal(res2, expected)
+
 
 # -------------------------------------------------------------------
 # Arithmetic
@@ -1410,12 +1424,13 @@ def test_alignment_non_pandas(self):
             range(1, 4),
         ]:
 
-            tm.assert_series_equal(
-                align(df, val, "index")[1], Series([1, 2, 3], index=df.index)
-            )
-            tm.assert_series_equal(
-                align(df, val, "columns")[1], Series([1, 2, 3], index=df.columns)
+            expected = DataFrame({"X": val, "Y": val, "Z": val}, index=df.index)
+            tm.assert_frame_equal(align(df, val, "index")[1], expected)
+
+            expected = DataFrame(
+                {"X": [1, 1, 1], "Y": [2, 2, 2], "Z": [3, 3, 3]}, index=df.index
             )
+            tm.assert_frame_equal(align(df, val, "columns")[1], expected)
 
         # length mismatch
         msg = "Unable to coerce to Series, length must be 3: given 2"