REF: avoid broadcasting of Series to DataFrame in ops for ArrayManager

pandas/core/frame.py

@@ -7061,8 +7061,11 @@ def _dispatch_frame_op(self, right, func: Callable, axis: int | None = None):
             assert right.index.equals(self.columns)
 
             right = right._values
-            # maybe_align_as_frame ensures we do not have an ndarray here
-            assert not isinstance(right, np.ndarray)
+            # BlockManager only: maybe_align_as_frame ensures we do not have
+            # an ndarray here (`assert not isinstance(right, np.ndarray)`)
+
+            if isinstance(right, TimedeltaArray):
+                right = right._data
 
             with np.errstate(all="ignore"):
                 arrays = [

pandas/core/generic.py

@@ -321,6 +321,10 @@ def _as_manager(self: NDFrameT, typ: str, copy: bool_t = True) -> NDFrameT:
         # fastpath of passing a manager doesn't check the option/manager class
         return self._constructor(new_mgr).__finalize__(self)
 
+    @property
+    def _using_array_manager(self):
+        return isinstance(self._mgr, ArrayManager)
+
     # ----------------------------------------------------------------------
     # attrs and flags
 

pandas/core/ops/__init__.py

@@ -307,7 +307,10 @@ def to_series(right):
         left, right = left.align(
             right, join="outer", axis=axis, level=level, copy=False
         )
-        right = _maybe_align_series_as_frame(left, right, axis)
+
+        if not left._using_array_manager:
+            # for BlockManager maybe broadcast Series to a DataFrame
+            right = _maybe_align_series_as_frame(left, right, axis)
 
     return left, right
 

pandas/core/ops/array_ops.py

@@ -120,7 +120,10 @@ def _masked_arith_op(x: np.ndarray, y, op):
 
         # mask is only meaningful for x
         result = np.empty(x.size, dtype=x.dtype)
-        mask = notna(xrav)
+        if isna(y):
+            mask = np.zeros(x.size, dtype=bool)
+        else:
+            mask = notna(xrav)
 
         # 1 ** np.nan is 1. So we have to unmask those.
         if op is pow:
@@ -309,6 +312,7 @@ def na_logical_op(x: np.ndarray, y, op):
             y = ensure_object(y)
             result = libops.vec_binop(x.ravel(), y.ravel(), op)
         else:
+            x = ensure_object(x)
             # let null fall thru
             assert lib.is_scalar(y)
             if not isna(y):

pandas/tests/arithmetic/test_datetime64.py

@@ -1477,7 +1477,13 @@ def test_dt64arr_add_sub_DateOffset(self, box_with_array):
     @pytest.mark.parametrize("op", [operator.add, roperator.radd, operator.sub])
     @pytest.mark.parametrize("box_other", [True, False])
     def test_dt64arr_add_sub_offset_array(
-        self, tz_naive_fixture, box_with_array, box_other, op, other
+        self,
+        tz_naive_fixture,
+        box_with_array,
+        box_other,
+        op,
+        other,
+        using_array_manager,
     ):
         # GH#18849
         # GH#10699 array of offsets
@@ -1493,7 +1499,10 @@ def test_dt64arr_add_sub_offset_array(
         if box_other:
             other = tm.box_expected(other, box_with_array)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box_with_array is pd.DataFrame and not box_other and using_array_manager:
+            warn = None
+        with tm.assert_produces_warning(warn):
             res = op(dtarr, other)
 
         tm.assert_equal(res, expected)
@@ -2425,7 +2434,7 @@ def test_dti_addsub_offset_arraylike(
 
     @pytest.mark.parametrize("other_box", [pd.Index, np.array])
     def test_dti_addsub_object_arraylike(
-        self, tz_naive_fixture, box_with_array, other_box
+        self, tz_naive_fixture, box_with_array, other_box, using_array_manager
     ):
         tz = tz_naive_fixture
 
@@ -2437,14 +2446,18 @@ def test_dti_addsub_object_arraylike(
         expected = DatetimeIndex(["2017-01-31", "2017-01-06"], tz=tz_naive_fixture)
         expected = tm.box_expected(expected, xbox)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box_with_array is pd.DataFrame and using_array_manager:
+            warn = None
+
+        with tm.assert_produces_warning(warn):
             result = dtarr + other
         tm.assert_equal(result, expected)
 
         expected = DatetimeIndex(["2016-12-31", "2016-12-29"], tz=tz_naive_fixture)
         expected = tm.box_expected(expected, xbox)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             result = dtarr - other
         tm.assert_equal(result, expected)
 

pandas/tests/arithmetic/test_timedelta64.py

@@ -1304,7 +1304,9 @@ def test_td64arr_sub_timedeltalike(self, two_hours, box_with_array):
     # ------------------------------------------------------------------
     # __add__/__sub__ with DateOffsets and arrays of DateOffsets
 
-    def test_td64arr_add_sub_offset_index(self, names, box_with_array):
+    def test_td64arr_add_sub_offset_index(
+        self, names, box_with_array, using_array_manager
+    ):
         # GH#18849, GH#19744
         box = box_with_array
         exname = get_expected_name(box, names)
@@ -1324,19 +1326,23 @@ def test_td64arr_add_sub_offset_index(self, names, box_with_array):
         expected = tm.box_expected(expected, box)
         expected_sub = tm.box_expected(expected_sub, box)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box is DataFrame and using_array_manager:
+            warn = None
+
+        with tm.assert_produces_warning(warn):
             res = tdi + other
         tm.assert_equal(res, expected)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             res2 = other + tdi
         tm.assert_equal(res2, expected)
 
         with tm.assert_produces_warning(PerformanceWarning):
             res_sub = tdi - other
         tm.assert_equal(res_sub, expected_sub)
 
-    def test_td64arr_add_sub_offset_array(self, box_with_array):
+    def test_td64arr_add_sub_offset_array(self, box_with_array, using_array_manager):
         # GH#18849, GH#18824
         box = box_with_array
         tdi = TimedeltaIndex(["1 days 00:00:00", "3 days 04:00:00"])
@@ -1352,11 +1358,15 @@ def test_td64arr_add_sub_offset_array(self, box_with_array):
         tdi = tm.box_expected(tdi, box)
         expected = tm.box_expected(expected, box)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box is DataFrame and using_array_manager:
+            warn = None
+
+        with tm.assert_produces_warning(warn):
             res = tdi + other
         tm.assert_equal(res, expected)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             res2 = other + tdi
         tm.assert_equal(res2, expected)
 
@@ -1365,7 +1375,9 @@ def test_td64arr_add_sub_offset_array(self, box_with_array):
             res_sub = tdi - other
         tm.assert_equal(res_sub, expected_sub)
 
-    def test_td64arr_with_offset_series(self, names, box_with_array):
+    def test_td64arr_with_offset_series(
+        self, names, box_with_array, using_array_manager
+    ):
         # GH#18849
         box = box_with_array
         box2 = Series if box in [pd.Index, tm.to_array, pd.array] else box
@@ -1378,18 +1390,22 @@ def test_td64arr_with_offset_series(self, names, box_with_array):
         obj = tm.box_expected(tdi, box)
         expected_add = tm.box_expected(expected_add, box2)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box is DataFrame and using_array_manager:
+            warn = None
+
+        with tm.assert_produces_warning(warn):
             res = obj + other
         tm.assert_equal(res, expected_add)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             res2 = other + obj
         tm.assert_equal(res2, expected_add)
 
         expected_sub = Series([tdi[n] - other[n] for n in range(len(tdi))], name=exname)
         expected_sub = tm.box_expected(expected_sub, box2)
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             res3 = obj - other
         tm.assert_equal(res3, expected_sub)
 
@@ -1420,7 +1436,7 @@ def test_td64arr_addsub_anchored_offset_arraylike(self, obox, box_with_array):
     # ------------------------------------------------------------------
     # Unsorted
 
-    def test_td64arr_add_sub_object_array(self, box_with_array):
+    def test_td64arr_add_sub_object_array(self, box_with_array, using_array_manager):
         box = box_with_array
         xbox = np.ndarray if box is pd.array else box
 
@@ -1429,7 +1445,11 @@ def test_td64arr_add_sub_object_array(self, box_with_array):
 
         other = np.array([Timedelta(days=1), offsets.Day(2), Timestamp("2000-01-04")])
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        warn = PerformanceWarning
+        if box is DataFrame and using_array_manager:
+            warn = None
+
+        with tm.assert_produces_warning(warn):
             result = tdarr + other
 
         expected = pd.Index(
@@ -1440,10 +1460,10 @@ def test_td64arr_add_sub_object_array(self, box_with_array):
 
         msg = "unsupported operand type|cannot subtract a datelike"
         with pytest.raises(TypeError, match=msg):
-            with tm.assert_produces_warning(PerformanceWarning):
+            with tm.assert_produces_warning(warn):
                 tdarr - other
 
-        with tm.assert_produces_warning(PerformanceWarning):
+        with tm.assert_produces_warning(warn):
             result = other - tdarr
 
         expected = pd.Index([Timedelta(0), Timedelta(0), Timestamp("2000-01-01")])
@@ -2056,7 +2076,6 @@ def test_td64arr_div_numeric_array(
             expected = pd.Index(expected)  # do dtype inference
             expected = tm.box_expected(expected, xbox)
             assert tm.get_dtype(expected) == "m8[ns]"
-
             tm.assert_equal(result, expected)
 
         with pytest.raises(TypeError, match=pattern):

pandas/tests/frame/test_arithmetic.py

@@ -1241,7 +1241,12 @@ def test_combineFrame(self, float_frame, mixed_float_frame, mixed_int_frame):
         _check_mixed_float(added, dtype="float64")
 
     def test_combine_series(
-        self, float_frame, mixed_float_frame, mixed_int_frame, datetime_frame
+        self,
+        float_frame,
+        mixed_float_frame,
+        mixed_int_frame,
+        datetime_frame,
+        using_array_manager,
     ):
 
         # Series
@@ -1264,7 +1269,14 @@ def test_combine_series(
 
         # no upcast needed
         added = mixed_float_frame + series
-        assert np.all(added.dtypes == series.dtype)
+        if using_array_manager:
+            # INFO(ArrayManager) addition is performed column-wise, and
+            # apparently in numpy the dtype of array gets priority in
+            # arithmetic casting rules (array[float32] + scalar[float64] gives
+            # array[float32] and not array[float64])
+            assert np.all(added.dtypes == mixed_float_frame.dtypes)
+        else:
+            assert np.all(added.dtypes == series.dtype)
 
         # vs mix (upcast) as needed
         added = mixed_float_frame + series.astype("float32")
@@ -1624,7 +1636,7 @@ def test_inplace_ops_identity2(self, op):
         expected = id(df)
         assert id(df) == expected
 
-    def test_alignment_non_pandas(self):
+    def test_alignment_non_pandas(self, using_array_manager):
         index = ["A", "B", "C"]
         columns = ["X", "Y", "Z"]
         df = DataFrame(np.random.randn(3, 3), index=index, columns=columns)
@@ -1637,13 +1649,21 @@ def test_alignment_non_pandas(self):
             range(1, 4),
         ]:
 
-            expected = DataFrame({"X": val, "Y": val, "Z": val}, index=df.index)
-            tm.assert_frame_equal(align(df, val, "index")[1], expected)
+            if not using_array_manager:
+                expected = DataFrame({"X": val, "Y": val, "Z": val}, index=df.index)
+                tm.assert_frame_equal(align(df, val, "index")[1], expected)
+            else:
+                expected = Series(val, index=index)
+                tm.assert_series_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)
+            if not using_array_manager:
+                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)
+            else:
+                expected = Series(val, index=columns)
+                tm.assert_series_equal(align(df, val, "columns")[1], expected)
 
         # length mismatch
         msg = "Unable to coerce to Series, length must be 3: given 2"