分组Pandas DataFrame上SKLearn线性回归（非聚合）报错排查

问题重现 #

测试代码 #

import pandas as pd
from sklearn.linear_model import LinearRegression

# 构造测试数据
df = pd.DataFrame({
    'group': ['A', 'A', 'A', 'B', 'B', 'B'],
    'x': [1, 2, 3, 4, 5, 6],
    'y': [2, 4, 6, 8, 10, 12]
})

# 错误的自定义回归函数
def add_regression_coeffs(group):
    X = group[['x']]
    y = group['y']
    model = LinearRegression().fit(X, y)
    # 错误：返回标量组而非完整分组
    return model.coef_[0], model.intercept_

# 调用apply触发报错
df.groupby('group').apply(add_regression_coeffs)

实际报错 #

ValueError: Length of values (2) does not match length of index (3)

预期输出 #

R tidyverse等价实现 #

library(tidyverse)
library(broom)

df <- tibble(
  group = c('A', 'A', 'A', 'B', 'B', 'B'),
  x = c(1,2,3,4,5,6),
  y = c(2,4,6,8,10,12)
)

df %>%
  group_by(group) %>%
  mutate(
    model = list(lm(y ~ x, data = cur_data())),
    coef = map_dbl(model, ~ coef(.)[['x']]),
    intercept = map_dbl(model, ~ coef(.)[['(Intercept)']])
  ) %>%
  select(-model)

问题原因 #

groupby.apply()的返回值处理逻辑分为两种：

1. 如果函数返回标量/短向量，Pandas会将其视为该分组的聚合结果，尝试映射到分组索引，导致长度不匹配（示例中每个分组有3行，但返回2个值）。
2. 如果函数返回与原分组行数一致的DataFrame/Series，Pandas会自动将所有分组的结果按行拼接，保留原数据的行结构。

而for循环可行，是因为循环中直接对每个分组的行赋值并合并，本质是保留了原分组的完整行结构。

解决办法 #

方法1：让apply函数返回完整分组 #

修改自定义函数，添加新列后返回整个分组DataFrame，apply会自动拼接所有分组：

def add_regression_coeffs(group):
    X = group[['x']]
    y = group['y']
    model = LinearRegression().fit(X, y)
    group['coef'] = model.coef_[0]
    group['intercept'] = model.intercept_
    return group

result = df.groupby('group').apply(add_regression_coeffs).reset_index(drop=True)

方法2：使用transform广播标量 #

transform专门用于将分组计算的标量广播到每行，可拆分系数和截距分别处理：

def get_regression_params(group):
    model = LinearRegression().fit(group[['x']], group['y'])
    return pd.Series([model.coef_[0], model.intercept_], index=['coef', 'intercept'])

# 先计算分组参数，再合并到原表
params = df.groupby('group').apply(get_regression_params).reset_index()
result = df.merge(params, on='group', how='left')

方法3：结合assign与apply（简洁写法） #

直接在分组内用assign添加新列，返回完整分组：

result = df.groupby('group').apply(lambda g: g.assign(
    coef=LinearRegression().fit(g[['x']], g['y']).coef_[0],
    intercept=LinearRegression().fit(g[['x']], g['y']).intercept_
)).reset_index(drop=True)

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