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

然而，这确实感觉像熊猫风格。我希望能够在

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_columns(
    pl.fold(0, lambda acc, s: acc + s, pl.exclude("id")).alias("sum_horizontal")
)

这将输出：

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

Answer 2

您也可以结合

with_columns

和

select

来执行此操作：

df.with_columns(
    df.select(pl.all().exclude("id")).sum(axis=1).alias("cat_total")
)

这会产生相同的结果：

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