@@ -252,47 +252,43 @@ def test_pivot_df_multiindex_index_none():
252
252
tm .assert_frame_equal (result , expected )
253
253
254
254
255
- def test_pivot_table_values_in_columns ():
255
+ @pytest .mark .parametrize (
256
+ argnames = ["index" , "columns" ],
257
+ argvalues = [("index" , ["col" , "value" ]), (["index" , "value" ], "col" )],
258
+ ids = ["values-and-columns" , "values-and-index" ],
259
+ )
260
+ def test_pivot_table_multiindex_values_as_two_params (index , columns , request ):
256
261
data = [
257
262
["A" , 1 , 50 , - 1 ],
258
263
["B" , 1 , 100 , - 2 ],
259
264
["A" , 2 , 100 , - 2 ],
260
265
["B" , 2 , 200 , - 4 ],
261
266
]
262
267
df = pd .DataFrame (data = data , columns = ["index" , "col" , "value" , "extra" ])
263
- result = df .pivot_table (values = "value" , index = " index" , columns = [ "col" , "value" ] )
268
+ result = df .pivot_table (values = "value" , index = index , columns = columns )
264
269
nan = np .nan
265
- e_data = [
266
- [ 50.0 , nan , 100.0 , nan ],
267
- [ nan , 100 .0 , nan , 200.0 ],
268
- ]
269
- e_index = Index ( data = [ "A" , "B" ], name = "index" )
270
- e_cols = MultiIndex . from_arrays (
271
- arrays = [[ 1 , 1 , 2 , 2 ], [ 50 , 100 , 100 , 200 ]], names = [ "col" , "value" ]
272
- )
273
- expected = pd . DataFrame ( data = e_data , index = e_index , columns = e_cols )
274
- tm . assert_frame_equal ( left = result , right = expected )
270
+ if request . node . callspec . id == "values-and-columns" :
271
+ e_data = [
272
+ [ 50 .0 , nan , 100.0 , nan ],
273
+ [ nan , 100.0 , nan , 200.0 ],
274
+ ]
275
+ e_index = Index ( data = [ "A" , "B" ], name = "index" )
276
+ e_cols = MultiIndex . from_arrays (
277
+ arrays = [[ 1 , 1 , 2 , 2 ], [ 50 , 100 , 100 , 200 ]], names = [ "col" , "value" ]
278
+ )
279
+ expected = pd . DataFrame ( data = e_data , index = e_index , columns = e_cols )
275
280
281
+ else :
282
+ e_data = [
283
+ [50.0 , nan ],
284
+ [nan , 100.0 ],
285
+ [100.0 , nan ],
286
+ [nan , 200.0 ],
287
+ ]
288
+ e_index = MultiIndex .from_arrays (
289
+ arrays = [["A" , "A" , "B" , "B" ], [50 , 100 , 100 , 200 ]], names = ["index" , "value" ]
290
+ )
291
+ e_cols = Index (data = [1 , 2 ], name = "col" )
292
+ expected = pd .DataFrame (data = e_data , index = e_index , columns = e_cols )
276
293
277
- def test_pivot_table_values_in_index ():
278
- data = [
279
- ["A" , 1 , 50 , - 1 ],
280
- ["B" , 1 , 100 , - 2 ],
281
- ["A" , 2 , 100 , - 2 ],
282
- ["B" , 2 , 200 , - 4 ],
283
- ]
284
- df = pd .DataFrame (data = data , columns = ["index" , "col" , "value" , "extra" ])
285
- result = df .pivot_table (values = "value" , index = ["index" , "value" ], columns = "col" )
286
- nan = np .nan
287
- e_data = [
288
- [50.0 , nan ],
289
- [nan , 100.0 ],
290
- [100.0 , nan ],
291
- [nan , 200.0 ],
292
- ]
293
- e_index = MultiIndex .from_arrays (
294
- arrays = [["A" , "A" , "B" , "B" ], [50 , 100 , 100 , 200 ]], names = ["index" , "value" ]
295
- )
296
- e_cols = Index (data = [1 , 2 ], name = "col" )
297
- expected = pd .DataFrame (data = e_data , index = e_index , columns = e_cols )
298
294
tm .assert_frame_equal (left = result , right = expected )
0 commit comments