@@ -1179,32 +1179,32 @@ def test_apply_empty_string_nan_coerce_bug():
1179
1179
tm .assert_frame_equal (result , expected )
1180
1180
1181
1181
1182
- def test_apply_index_key_error_bug ():
1182
+ @pytest .mark .parametrize ("is_float" , [True , False ])
1183
+ def test_apply_index_key_error_bug (is_float ):
1183
1184
# GH 44310
1184
- # Tests for a key error while applying in a grouped DF
1185
-
1186
- result = DataFrame (
1187
- {
1188
- "a" : ["aa" , "a2" , "a3" ],
1189
- "b" : [1 , 2 , 3 ],
1190
- },
1191
- index = Index ([1.0 , 2.0 , 3.0 ]),
1192
- )
1193
-
1185
+ if is_float :
1186
+ result = DataFrame (
1187
+ {
1188
+ "a" : ["aa" , "a2" , "a3" ],
1189
+ "b" : [1 , 2 , 3 ],
1190
+ },
1191
+ index = Index ([1.0 , 2.0 , 3.0 ]),
1192
+ )
1193
+ else :
1194
+ result = DataFrame (
1195
+ {
1196
+ "a" : ["aa" , "a2" , "a3" ],
1197
+ "b" : [1 , 2 , 3 ],
1198
+ },
1199
+ index = Index ([1 , 2 , 3 ]),
1200
+ )
1194
1201
expected = DataFrame (
1195
1202
{
1196
- "a" : ["aa" , "a2" , "a3" ],
1197
- "b" : [1 , 2 , 3 ],
1203
+ "b_mean" : [2.0 , 3.0 , 1.0 ],
1198
1204
},
1199
- index = Index ([1 , 2 , 3 ] ),
1205
+ index = Index (["a2" , "a3" , "aa" ], name = "a" ),
1200
1206
)
1201
- try :
1202
- result = result .groupby ("a" ).apply (
1203
- lambda df : Series ([df ["b" ].mean ()], index = ["b_mean" ])
1204
- )
1205
- expected = expected .groupby ("a" ).apply (
1206
- lambda df : Series ([df ["b" ].mean ()], index = ["b_mean" ])
1207
- )
1208
- tm .assert_frame_equal (result , expected )
1209
- except KeyError :
1210
- assert False
1207
+ result = result .groupby ("a" ).apply (
1208
+ lambda df : Series ([df ["b" ].mean ()], index = ["b_mean" ])
1209
+ )
1210
+ tm .assert_frame_equal (result , expected )
0 commit comments