Pandas groupby在重组组时更改返回类型

Question

我有一个数据帧:

df = pd.DataFrame({'c':[0,1,1,2,2,2],   'date':pd.to_datetime(['2016-01-01','2016-02-01','2016-03-01','2016-04-01','2016-05-01','2016-06-01'])})

对于每一行,我想得到一个数字=每个日期的月份数(Jan = 1,Feb = 2等)+该组的长度(第一组有1个成员,第2组有2个等):

所以它应该返回类似的东西:

c       date   num
0 2016-01-01    2
1 2016-02-01    4
1 2016-03-01    5
2 2016-04-01    7
2 2016-05-01    8
2 2016-06-01    9

我创建了一个函数:

def testlambda(x):
    print(x)
    return x.dt.month.astype('int') + len(x)

并使用groupby + transform:

df['num'] = df.groupby(['c'])['date'].transform(lambda x: testlambda(x))

但是返回的新列仍然是日期格式,即使我的lambda返回int.

该怎么办？

Answer 1

尝试使用DataFrameGroupBy.transform()而不是SeriesGroupBy.transform()后者尝试将结果转换为源dtype:

In [131]: def testlambda(x):
     ...:     #print(x)
     ...:     return x.dt.month.astype('int') + len(x)
     ...:

In [132]: df
Out[132]:
   c       date
0  0 2016-01-01
1  1 2016-02-01
2  1 2016-03-01
3  2 2016-04-01
4  2 2016-05-01
5  2 2016-06-01

#                                      v        v - thats's the only difference    
In [133]: df['num'] = df.groupby(['c'])[['date']].transform(lambda x: testlambda(x))

In [134]: df
Out[134]:
   c       date  num
0  0 2016-01-01    2
1  1 2016-02-01    4
2  1 2016-03-01    5
3  2 2016-04-01    7
4  2 2016-05-01    8
5  2 2016-06-01    9