Polars在分组中对多个列应用相同自定义函数

Question

在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')呢？谢谢！

Answer 1

你只需要把想要的列名放进 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      │
└───────┴──────────┴──────────┴─────────┘