Polars DataFrame中按列名分组跨行求和的实现方法

生成目标DataFrame的代码：

import polars as pl

df = pl.DataFrame({
    'id': ['CHECK.US1', 'CHECK.US2', 'CHECK.CA9'],
    'libor.M2': [99, 332, 934],
    'libor.Y5': [11, -10, 904],
    'estr.M2':  [99, 271, 741],
    'estr.Y3':  [-8, -24, 183],
    'estr.Y5':  [88, 771, 455]
})

原始DataFrame结构：

┌───────────┬──────────┬──────────┬─────────┬─────────┬─────────┐
│ id        ┆ libor.M2 ┆ libor.Y5 ┆ estr.M2 ┆ estr.Y3 ┆ estr.Y5 │
│ ---       ┆ ---      ┆ ---      ┆ ---     ┆ ---     ┆ ---     │
│ str       ┆ i64      ┆ i64      ┆ i64     ┆ i64     ┆ i64     │
╞═══════════╪══════════╪══════════╪═════════╪═════════╪═════════╡
│ CHECK.US1 ┆ 99       ┆ 11       ┆ 99      ┆ -8      ┆ 88      │
│ CHECK.US2 ┆ 332      ┆ -10      ┆ 271     ┆ -24     ┆ 771     │
│ CHECK.CA9 ┆ 934      ┆ 904      ┆ 741     ┆ 183     ┆ 455     │
└───────────┴──────────┴──────────┴─────────┴─────────┴─────────┘

需求：将列名简化为后缀（如M2、Y5），对每行中同后缀的列求和，得到目标结果：

┌───────────┬──────┬──────┬──────┐
│ id        ┆ M2   ┆ Y5   ┆ Y3   │
│ ---       ┆ ---  ┆ ---  ┆ ---  │
│ str       ┆ i64  ┆ i64  ┆ i64  │
╞═══════════╪══════╪══════╪══════╡
│ CHECK.US1 ┆ 198  ┆ 99   ┆ -8   │
│ CHECK.US2 ┆ 603  ┆ 761  ┆ -24  │
│ CHECK.CA9 ┆ 1675 ┆ 1359 ┆ 183  │
└───────────┴──────┴──────┴──────┘

尝试直接重命名列触发错误：
polars.exceptions.DuplicateError: the name 'M2' is duplicate

尝试的未成功代码：

rename_func = lambda col: col.split('.')[-1]
new_cols = set([rename_func(c) for c in df.columns])

df.with_columns([
  pl.sum_horizontal(pl.all().map(rename_func) == c).alias(c) for c in new_cols
])

方法一：正则匹配后缀直接求和 #

针对已知的后缀，用正则筛选对应列后横向求和：

result = df.select(
    pl.col('id'),
    pl.sum_horizontal(pl.col(r'.*\.M2$')).alias('M2'),
    pl.sum_horizontal(pl.col(r'.*\.Y5$')).alias('Y5'),
    pl.sum_horizontal(pl.col(r'.*\.Y3$')).alias('Y3')
)

方法二：动态生成求和表达式（适配多后缀场景） #

自动提取所有后缀，批量生成求和逻辑：

# 提取所有非id列的后缀
suffixes = {col.split('.')[-1] for col in df.columns if col != 'id'}

# 批量生成求和表达式
sum_exprs = [
    pl.sum_horizontal(pl.col(f'*.{suffix}')).alias(suffix)
    for suffix in suffixes
]

# 生成结果
result = df.select('id', *sum_exprs)

方法三：melt + pivot 转表实现 #

通过宽表转长表分组求和，再转回宽表：

result = df.melt(id_vars='id') \
           .with_columns(
               pl.col('variable').str.split('.').list.get(-1).alias('suffix')
           ) \
           .group_by(['id', 'suffix']) \
           .agg(pl.sum('value')) \
           .pivot(index='id', columns='suffix', values='value')

以上三种方法均可得到目标DataFrame结构。

内容的提问来源于stack exchange，提问作者Phil-ZXX