对极坐标数据框中同名（或“键”）的列求和

Question

我有这个代码

import polars as pl

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

这给出了

┌───────────┬──────────┬──────────┬─────────┬─────────┬─────────┐
│ 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     │
└───────────┴──────────┴──────────┴─────────┴─────────┴─────────┘

现在我想做的是将列重命名为更短的名称，例如

┌───────────┬──────┬──────┬─────┬─────┬─────┐
│ id        ┆ M2   ┆ Y5   ┆ M2  ┆ Y3  ┆ Y5  │
│ ---       ┆ ---  ┆ ---  ┆ --- ┆ --- ┆ --- │
│ str       ┆ i64  ┆ i64  ┆ i64 ┆ i64 ┆ i64 │
╞═══════════╪══════╪══════╪═════╪═════╪═════╡
or
┌───────────┬──────┬──────┬─────┬─────┬─────┐
│ id        ┆ libor┆ libor┆ estr┆ estr┆ estr│
│ ---       ┆ ---  ┆ ---  ┆ --- ┆ --- ┆ --- │
│ str       ┆ i64  ┆ i64  ┆ i64 ┆ i64 ┆ i64 │
╞═══════════╪══════╪══════╪═════╪═════╪═════╡

然后折叠（= sum）具有相同名称的列，这样我就得到了例如

┌───────────┬──────┬──────┬──────┐
│ 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
])

但是不太管用。

Answer 1

您可以构建要聚合的列的字典，然后使用

pl.sum_horizontal

:

cols = {}
for col in df.columns:
    suffix = col.split('.')[-1]
    cols.setdefault(suffix, []).append(col)
# {'id': ['id'], 'M2': ['libor.M2', 'estr.M2'],
#  'Y5': ['libor.Y5', 'estr.Y5'], 'Y3': ['estr.Y3']}

df.select([pl.sum_horizontal(pl.col(c) for c in l).alias(k)
           for k, l in cols.items()
          ])

输出：

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