@@ -225,7 +225,6 @@ def data_for_grouping(dtype):
225
225
return pd .array ([B , B , None , None , A , A , B , C ], dtype = dtype )
226
226
227
227
228
- @pytest .fixture
229
228
def expected_inferred_result_dtype (dtype ):
230
229
"""
231
230
When the data pass through aggregate,
@@ -1144,6 +1143,27 @@ def test_comp_masked_numpy(self, masked_dtype, comparison_op):
1144
1143
expected = pd .Series (exp , dtype = ArrowDtype (pa .bool_ ()))
1145
1144
tm .assert_series_equal (result , expected )
1146
1145
1146
+ def test_groupby_agg_extension (self , data_for_grouping ):
1147
+ # GH#38980 groupby agg on extension type fails for non-numeric types
1148
+ df = pd .DataFrame ({"A" : [1 , 1 , 2 , 2 , 3 , 3 , 1 , 4 ], "B" : data_for_grouping })
1149
+
1150
+ expected_df = pd .DataFrame (
1151
+ {"A" : [1 , 1 , 2 , 2 , 3 , 3 , 1 , 4 ], "B" : data_for_grouping }
1152
+ )
1153
+ expected = expected_df .iloc [[0 , 2 , 4 , 7 ]]
1154
+ expected = expected .set_index ("A" )
1155
+ expected_dtype = expected_inferred_result_dtype (expected ["B" ].dtype )
1156
+ expected ["B" ] = expected ["B" ].astype (expected_dtype )
1157
+
1158
+ result = df .groupby ("A" ).agg ({"B" : "first" })
1159
+ tm .assert_frame_equal (result , expected )
1160
+
1161
+ result = df .groupby ("A" ).agg ("first" )
1162
+ tm .assert_frame_equal (result , expected )
1163
+
1164
+ result = df .groupby ("A" ).first ()
1165
+ tm .assert_frame_equal (result , expected )
1166
+
1147
1167
1148
1168
class TestLogicalOps :
1149
1169
"""Various Series and DataFrame logical ops methods."""
0 commit comments