Polars：在 select/with_column 调用中添加一些列的总和

Question

我想添加一列，它是所有列的总和，但一些带有极坐标的 id 列除外。这可以使用 polars.DataFrame.sum(axis=1):

来完成

import polars as pl
df = pl.DataFrame(
    {
        "id": [1, 2],
        "cat_a": [2, 7],
        "cat_b": [5, 1],
        "cat_c": [0, 3]
    }
)
df["cat_total"] = df.select(pl.all().exclude("id")).sum(axis=1)
df

不过，这确实是pandas风格的感觉。我希望能够在 select 或 with_column 调用中的更长的调用序列中进行此操作：

# Throws TypeError: sum() got an unexpected keyword argument 'axis'
# because polars.Expr.sum does not support choosing an axis
(df
     # [...]
    .with_column(pl.all().exclude("id").sum(axis=1).alias("cat_total"))
     # [...]
)

如何做到这一点（不明确标识列名）？

Answer 1

您可以使用 fold 表达式，它接受一个累加器：acc、一个二元函数 Fn(acc, Series) -> Series 和一个或多个表达式来应用折叠。

df.with_column(
    pl.fold(0, lambda acc, s: acc + s, pl.all().exclude("id")).alias("horizontal_sum")
)

这将输出：

shape: (2, 5)
┌─────┬───────┬───────┬───────┬────────────────┐
│ id  ┆ cat_a ┆ cat_b ┆ cat_c ┆ horizontal_sum │
│ --- ┆ ---   ┆ ---   ┆ ---   ┆ ---            │
│ i64 ┆ i64   ┆ i64   ┆ i64   ┆ i64            │
╞═════╪═══════╪═══════╪═══════╪════════════════╡
│ 1   ┆ 2     ┆ 5     ┆ 0     ┆ 7              │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 7     ┆ 1     ┆ 3     ┆ 11             │
└─────┴───────┴───────┴───────┴────────────────┘

Polars：在 select/with_column 调用中添加一些列的总和

Polars: add the sum of some columns inside select/with_column call

python-polars