Python Polars将相同的自定义函数应用于组中的多个列，

发布于03月29日

将自定义函数应用于Polars中的多个列的最佳方法是什么？具体地说，我需要函数引用另一个列的框架.假设我有以下内容:

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],
})

我想按group分组，计算x和y的平均值，再乘以num_obs.我可以做这样的事

variables = ['x', 'y']
df.group_by('group').agg((pl.col(var) * pl.col('num_obs')).sum()/pl.col('num_obs').sum() for var in variables)

但我想知道有没有更好的办法此外，我不知道如何向这种方法添加其他聚合，但有没有一种方法可以添加pl.sum('n_obs')个聚合？谢谢！

推荐答案

您可以将列列表传递到pl.col():

df.group_by('group').agg(
    (pl.col('x','y') * pl.col('num_obs')).sum() / pl.col('num_obs').sum(),
    pl.col('num_obs').sum()
)

┌───────┬──────────┬──────────┬─────────┐
│ group ┆ x        ┆ y        ┆ num_obs │
│ ---   ┆ ---      ┆ ---      ┆ ---     │
│ i64   ┆ f64      ┆ f64      ┆ i64     │
╞═══════╪══════════╪══════════╪═════════╡
│ 1     ┆ 1.333333 ┆ 5.333333 ┆ 15      │
│ 2     ┆ 3.333333 ┆ 7.333333 ┆ 30      │
└───────┴──────────┴──────────┴─────────┘