如何将分组变换后的结果通过df.join合并回原始DataFrame？

原始数据 #

import pandas as pd

d = {'bank': [1, 1, 2], 'account_money': [100, 300, 500]}
df = pd.DataFrame(data=d)
print(df)

输出：

bank  account_money
0     1            100
1     1            300
2     2            500

分组滚动求和操作 #

def transform(grp_obj):
    result = grp_obj['account_money'].rolling(2, min_periods=0).sum()
    return result

transformed = df.groupby('bank').pipe(transform)
print(transformed)

输出：

bank   
1     0    100.0
      1    400.0
2     2    500.0
Name: account_money, dtype: float64

用join合并结果 #

由于transformed是带多层索引（bank + 原始行号）的Series，只需移除bank层级索引，将其转为与原始df索引匹配的DataFrame，再通过join按索引对齐合并即可：

# 将滚动求和结果转为DataFrame，保留原始行索引，移除bank层级
transformed_df = transformed.reset_index(level=0, drop=True).to_frame(name='rolling_sum')

# 合并回原始DataFrame
merged_df = df.join(transformed_df)
print(merged_df)

输出：

bank  account_money  rolling_sum
0     1            100        100.0
1     1            300        400.0
2     2            500        500.0

操作说明 #

• reset_index(level=0, drop=True)：移除transformed的第一层索引（bank），仅保留原始行号作为索引，确保和原始df的索引完全匹配
• to_frame(name='rolling_sum')：将Series转为DataFrame并给新列命名，方便合并后识别含义
• df.join(transformed_df)：默认按行索引对齐，直接将滚动求和列追加到原始df中

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