@@ -253,11 +253,41 @@ def test_pivot_df_multiindex_index_none():
253
253
254
254
255
255
@pytest .mark .parametrize (
256
- "index, columns" ,
257
- [("index" , ["col" , "value" ]), (["index" , "value" ], "col" )],
256
+ "index, columns, e_data, e_index, e_cols" ,
257
+ [
258
+ (
259
+ "index" ,
260
+ ["col" , "value" ],
261
+ [
262
+ [50.0 , (nan := np .nan ), 100.0 , nan ],
263
+ [nan , 100.0 , nan , 200.0 ],
264
+ ],
265
+ Index (data = ["A" , "B" ], name = "index" ),
266
+ MultiIndex .from_arrays (
267
+ arrays = [[1 , 1 , 2 , 2 ], [50 , 100 , 100 , 200 ]], names = ["col" , "value" ]
268
+ ),
269
+ ),
270
+ (
271
+ ["index" , "value" ],
272
+ "col" ,
273
+ [
274
+ [50.0 , nan ],
275
+ [nan , 100.0 ],
276
+ [100.0 , nan ],
277
+ [nan , 200.0 ],
278
+ ],
279
+ MultiIndex .from_arrays (
280
+ arrays = [["A" , "A" , "B" , "B" ], [50 , 100 , 100 , 200 ]],
281
+ names = ["index" , "value" ],
282
+ ),
283
+ Index (data = [1 , 2 ], name = "col" ),
284
+ ),
285
+ ],
258
286
ids = ["values-and-columns" , "values-and-index" ],
259
287
)
260
- def test_pivot_table_multiindex_values_as_two_params (index , columns , request ):
288
+ def test_pivot_table_multiindex_values_as_two_params (
289
+ index , columns , e_data , e_index , e_cols
290
+ ):
261
291
# GH#61292
262
292
data = [
263
293
["A" , 1 , 50 , - 1 ],
@@ -267,28 +297,5 @@ def test_pivot_table_multiindex_values_as_two_params(index, columns, request):
267
297
]
268
298
df = pd .DataFrame (data = data , columns = ["index" , "col" , "value" , "extra" ])
269
299
result = df .pivot_table (values = "value" , index = index , columns = columns )
270
- nan = np .nan
271
- if request .node .callspec .id == "values-and-columns" :
272
- e_data = [
273
- [50.0 , nan , 100.0 , nan ],
274
- [nan , 100.0 , nan , 200.0 ],
275
- ]
276
- e_index = Index (data = ["A" , "B" ], name = "index" )
277
- e_cols = MultiIndex .from_arrays (
278
- arrays = [[1 , 1 , 2 , 2 ], [50 , 100 , 100 , 200 ]], names = ["col" , "value" ]
279
- )
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
-
293
300
expected = pd .DataFrame (data = e_data , index = e_index , columns = e_cols )
294
301
tm .assert_frame_equal (result , expected )
0 commit comments