@@ -99,6 +99,7 @@ def test_df_flex_cmp_constant_return_types_empty(self, opname):
99
99
# Arithmetic
100
100
101
101
class TestFrameFlexArithmetic (object ):
102
+
102
103
def test_df_add_td64_columnwise (self ):
103
104
# GH#22534 Check that column-wise addition broadcasts correctly
104
105
dti = pd .date_range ('2016-01-01' , periods = 10 )
@@ -252,6 +253,99 @@ def test_arith_flex_zero_len_raises(self):
252
253
253
254
254
255
class TestFrameArithmetic (object ):
256
+ def test_df_add_2d_array_rowlike_broadcasts (self ):
257
+ # GH#23000
258
+ arr = np .arange (6 ).reshape (3 , 2 )
259
+ df = pd .DataFrame (arr , columns = [True , False ], index = ['A' , 'B' , 'C' ])
260
+
261
+ rowlike = arr [[1 ], :] # shape --> (1, ncols)
262
+ assert rowlike .shape == (1 , df .shape [1 ])
263
+
264
+ expected = pd .DataFrame ([[2 , 4 ],
265
+ [4 , 6 ],
266
+ [6 , 8 ]],
267
+ columns = df .columns , index = df .index ,
268
+ # specify dtype explicitly to avoid failing
269
+ # on 32bit builds
270
+ dtype = arr .dtype )
271
+ result = df + rowlike
272
+ tm .assert_frame_equal (result , expected )
273
+ result = rowlike + df
274
+ tm .assert_frame_equal (result , expected )
275
+
276
+ def test_df_add_2d_array_collike_broadcasts (self ):
277
+ # GH#23000
278
+ arr = np .arange (6 ).reshape (3 , 2 )
279
+ df = pd .DataFrame (arr , columns = [True , False ], index = ['A' , 'B' , 'C' ])
280
+
281
+ collike = arr [:, [1 ]] # shape --> (nrows, 1)
282
+ assert collike .shape == (df .shape [0 ], 1 )
283
+
284
+ expected = pd .DataFrame ([[1 , 2 ],
285
+ [5 , 6 ],
286
+ [9 , 10 ]],
287
+ columns = df .columns , index = df .index ,
288
+ # specify dtype explicitly to avoid failing
289
+ # on 32bit builds
290
+ dtype = arr .dtype )
291
+ result = df + collike
292
+ tm .assert_frame_equal (result , expected )
293
+ result = collike + df
294
+ tm .assert_frame_equal (result , expected )
295
+
296
+ def test_df_arith_2d_array_rowlike_broadcasts (self ,
297
+ all_arithmetic_operators ):
298
+ # GH#23000
299
+ opname = all_arithmetic_operators
300
+
301
+ arr = np .arange (6 ).reshape (3 , 2 )
302
+ df = pd .DataFrame (arr , columns = [True , False ], index = ['A' , 'B' , 'C' ])
303
+
304
+ rowlike = arr [[1 ], :] # shape --> (1, ncols)
305
+ assert rowlike .shape == (1 , df .shape [1 ])
306
+
307
+ exvals = [getattr (df .loc ['A' ], opname )(rowlike .squeeze ()),
308
+ getattr (df .loc ['B' ], opname )(rowlike .squeeze ()),
309
+ getattr (df .loc ['C' ], opname )(rowlike .squeeze ())]
310
+
311
+ expected = pd .DataFrame (exvals , columns = df .columns , index = df .index )
312
+
313
+ if opname in ['__rmod__' , '__rfloordiv__' ]:
314
+ # exvals will have dtypes [f8, i8, i8] so expected will be
315
+ # all-f8, but the DataFrame operation will return mixed dtypes
316
+ # use exvals[-1].dtype instead of "i8" for compat with 32-bit
317
+ # systems/pythons
318
+ expected [False ] = expected [False ].astype (exvals [- 1 ].dtype )
319
+
320
+ result = getattr (df , opname )(rowlike )
321
+ tm .assert_frame_equal (result , expected )
322
+
323
+ def test_df_arith_2d_array_collike_broadcasts (self ,
324
+ all_arithmetic_operators ):
325
+ # GH#23000
326
+ opname = all_arithmetic_operators
327
+
328
+ arr = np .arange (6 ).reshape (3 , 2 )
329
+ df = pd .DataFrame (arr , columns = [True , False ], index = ['A' , 'B' , 'C' ])
330
+
331
+ collike = arr [:, [1 ]] # shape --> (nrows, 1)
332
+ assert collike .shape == (df .shape [0 ], 1 )
333
+
334
+ exvals = {True : getattr (df [True ], opname )(collike .squeeze ()),
335
+ False : getattr (df [False ], opname )(collike .squeeze ())}
336
+
337
+ dtype = None
338
+ if opname in ['__rmod__' , '__rfloordiv__' ]:
339
+ # Series ops may return mixed int/float dtypes in cases where
340
+ # DataFrame op will return all-float. So we upcast `expected`
341
+ dtype = np .common_type (* [x .values for x in exvals .values ()])
342
+
343
+ expected = pd .DataFrame (exvals , columns = df .columns , index = df .index ,
344
+ dtype = dtype )
345
+
346
+ result = getattr (df , opname )(collike )
347
+ tm .assert_frame_equal (result , expected )
348
+
255
349
def test_df_bool_mul_int (self ):
256
350
# GH#22047, GH#22163 multiplication by 1 should result in int dtype,
257
351
# not object dtype
0 commit comments