为何groupby.transform对第一组最后一次调用传入整个DataFrame？

为了明确groupby.transform方法调用时自定义函数的输入内容，我做了以下测试：

• 编写get_data函数生成堆叠格式的DataFrame；
• 定义函数f，打印每次调用的输入参数；
• 执行dfs.groupby('date').transform(f)后发现：第一组（2010）前两次输入是该组的symbol、ret列（Series类型），第三次却传入了整个组的DataFrame，总调用次数为9次，而非预期的8次（4组×2列），对此现象存在疑问。

测试代码 #

# 生成测试数据框
def get_data(stack=False):
    dic = {'APPL':[10,5,-4,8],'MSFT':[20,-9,12,4],'TSLA':[-8,15,-7,14]}
    df = pd.DataFrame(dic,index=pd.Series([2010,2011,2012,2013],name='date'))
    if stack:
        df = pd.DataFrame(df.stack()).reset_index()
        df.columns =['date','symbol','ret']
    return df

# 示例数据
dfs = get_data(stack=True).copy()


def f(x):
    global i
    print("-----------------------")
    print("函数调用次数:", i,'\n')
    print("函数输入内容\n",x,'\n')
    i = i + 1

i = 1
dfs.groupby('date').transform(f)

执行结果 #

---- 堆叠后的数据 ---- 
     date symbol  ret
0   2010   APPL   10
1   2010   MSFT   20
2   2010   TSLA   -8
3   2011   APPL    5
4   2011   MSFT   -9
5   2011   TSLA   15
6   2012   APPL   -4
7   2012   MSFT   12
8   2012   TSLA   -7
9   2013   APPL    8
10  2013   MSFT    4
11  2013   TSLA   14

======= Transform方法执行日志 ========
-----------------------
函数调用次数: 1 

函数输入内容
 0    APPL
1    MSFT
2    TSLA
Name: symbol, dtype: object 

-----------------------
函数调用次数: 2 

函数输入内容
 0    10
1    20
2    -8
Name: ret, dtype: int64 

-----------------------
函数调用次数: 3 

函数输入内容
   symbol  ret
0   APPL   10
1   MSFT   20
2   TSLA   -8 

-----------------------
函数调用次数: 4 

函数输入内容
 3    APPL
4    MSFT
5    TSLA
Name: symbol, dtype: object 

-----------------------
函数调用次数: 5 

函数输入内容
 3     5
4    -9
5    15
Name: ret, dtype: int64 

-----------------------
函数调用次数: 6 

函数输入内容
 6    APPL
7    MSFT
8    TSLA
Name: symbol, dtype: object 

-----------------------
函数调用次数: 7 

函数输入内容
 6    -4
7    12
8    -7
Name: ret, dtype: int64 

-----------------------
函数调用次数: 8 

函数输入内容
 9     APPL
10    MSFT
11    TSLA
Name: symbol, dtype: object 

-----------------------
函数调用次数: 9 

函数输入内容
 9      8
10     4
11    14
Name: ret, dtype: int64 

原因说明 #

这是pandas内部优化机制导致的：当transform检测到自定义函数的返回值无法匹配原数据形状时，会先尝试传入单列Series，若验证不通过，就会退回到传入整个组的DataFrame重新执行。

在本次测试中，函数f没有返回值（默认返回None），pandas无法判断该函数是否适配单列输入，因此在第一组尝试单列调用后，又传入整个DataFrame进行验证，多了一次调用。后续组因为已经确认函数可处理单列输入，所以不再重复传入整个DataFrame。

如果给函数f添加符合要求的返回值（比如返回与输入等长的Series），调用次数就会变成预期的8次：

def f(x):
    global i
    print("-----------------------")
    print("函数调用次数:", i,'\n')
    print("函数输入内容\n",x,'\n')
    i = i + 1
    return x  # 添加返回值，匹配输入形状

内容的提问来源于stack exchange，提问作者jay lee