Pandas分组执行自定义回归函数后新列全为NaN的问题排查

问题场景 #

需要对按gvkey分组后的每个子数据集执行OLS回归，将回归系数（乘以100）存入新列b，但运行代码后b列全部显示NaN。相关代码如下：

自定义回归函数 #

def simplereg(g, y, x):
    try:
        xvar = sm.add_constant(g[x])
        yvar = g[y]
        model = sm.OLS(yvar, xvar, missing='drop').fit()
        b = model.params[x]
        return pd.Series([b*100]*len(g))
    except Exception as e:
        return pd.Series([np.NaN]*len(g))

样本数据生成 #

import pandas as pd
import numpy as np
import statsmodels.api as sm

# 设置参数
gvkeys = ['A', 'B', 'C', 'D']  # gvkey的可选取值
years = np.arange(2000, 2020)  # year的可选取值

# 每个gvkey的行数，确保每个组有5-7条观测值
num_rows_per_gvkey = np.random.randint(5, 8, size=len(gvkeys))
total_rows = sum(num_rows_per_gvkey)

# 创建DataFrame
np.random.seed(0)  # 保证结果可复现
df = pd.DataFrame({
    'gvkey': np.repeat(gvkeys, num_rows_per_gvkey),
    'year': np.random.choice(years, size=total_rows),
    'y': np.random.rand(total_rows),
    'x': np.random.rand(total_rows)
})

df.sort_values(by='year', ignore_index=True, inplace=True) # 验证代码是否能处理未排序的数据

分组执行代码 #

df['b'] = df.groupby('gvkey').apply(simplereg, y='y', x='x')

核心原因：索引不匹配 #

问题出在返回的Series索引与原DataFrame索引不匹配：

• 原DataFrame执行sort_values(..., ignore_index=True)后，索引被重置为连续的0、1、2...
• 自定义函数返回的Series使用的是分组内部的局部索引（每个组内从0开始计数），和原df的全局索引完全不对应
• 赋值时Pandas按索引匹配，找不到对应值就填充NaN

修复方法 #

有两种简洁的修复方案：

方案1：让返回的Series使用组的原始索引 #

修改自定义函数的返回语句，明确指定Series的索引为当前组的索引（g.index），确保和原df索引对齐：

def simplereg(g, y, x):
    try:
        xvar = sm.add_constant(g[x])
        yvar = g[y]
        model = sm.OLS(yvar, xvar, missing='drop').fit()
        b = model.params[x]
        # 关键：指定索引为当前组的索引
        return pd.Series([b*100]*len(g), index=g.index)
    except Exception as e:
        return pd.Series([np.NaN]*len(g), index=g.index)

方案2：使用transform替代apply #

transform会自动处理索引对齐，更适合这种返回与组长度一致结果的场景。可以简化函数返回标量（Pandas会自动将标量广播为组内所有行的值）：

def simplereg(g, y, x):
    try:
        xvar = sm.add_constant(g[x])
        yvar = g[y]
        model = sm.OLS(yvar, xvar, missing='drop').fit()
        return model.params[x] * 100
    except Exception as e:
        return np.NaN

# 调用时用transform
df['b'] = df.groupby('gvkey').transform(simplereg, y='y', x='x')

验证结果 #

运行修复后的代码，df['b']会显示每个gvkey组对应的回归系数（乘以100），不再是NaN。以方案2为例，输出的df中同一gvkey的所有行b值相同，符合预期。

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