Polars分组后对多列应用含关联列的自定义函数的最优方法

更简洁高效的加权平均实现 #

你当前的实现逻辑是可行的，但可以利用Polars的批量列选择特性简化代码，同时保持甚至提升执行效率——Polars会自动对选中的列进行向量化处理，无需手动循环生成单个列的表达式：

import polars as pl

df = pl.DataFrame({
    'group': [1,1,2,2],
    'other': ['a', 'b', 'a', 'b'],
    'num_obs': [10, 5, 20, 10],
    'x': [1,2,3,4],
    'y': [5,6,7,8],
})

variables = ['x', 'y']

# 优化后的加权平均计算
result = df.group_by('group').agg(
    # 批量处理所有目标列，自动生成每个列的加权平均
    ((pl.col(variables) * pl.col('num_obs')).sum() / pl.col('num_obs').sum()).alias(f'{var}_weighted_mean' for var in variables),
)

print(result)

这段代码和你原来的逻辑完全一致，但通过pl.col(variables)一次性选中所有需要计算的列，避免了手动循环生成表达式，代码更简洁，也更符合Polars的向量化设计理念。

添加额外聚合操作（如sum(num_obs)） #

要在同一个分组聚合中加入其他操作（比如计算每组的总num_obs），只需在agg方法中追加对应的聚合表达式即可，多个表达式用逗号分隔：

result_with_total = df.group_by('group').agg(
    # 批量计算加权平均
    ((pl.col(variables) * pl.col('num_obs')).sum() / pl.col('num_obs').sum()).alias(f'{var}_weighted_mean' for var in variables),
    # 添加总观测数的聚合
    pl.sum('num_obs').alias('total_obs'),
    # 还可以按需添加其他聚合，比如取other列的唯一值集合
    pl.col('other').unique().alias('unique_others')
)

print(result_with_total)

这样就能在一次分组聚合中完成所有需要的计算，无需多次分组，效率更高。

封装为自定义函数（可选） #

如果需要重复使用这个加权平均逻辑，可以把它封装成一个自定义聚合函数，方便复用：

def weighted_mean(cols, weight_col):
    return (pl.col(cols) * pl.col(weight_col)).sum() / pl.col(weight_col).sum()

# 使用自定义函数
result = df.group_by('group').agg(
    weighted_mean(variables, 'num_obs').alias(f'{var}_weighted_mean' for var in variables),
    pl.sum('num_obs').alias('total_obs')
)

这种方式让代码的可读性更好，也便于后续维护。

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