Início / coding / Perguntas / 78977874
Accepted
Phil-ZXX
Asked: 2024-09-12 19:17:14 +0800 CST

Soma entre colunas do mesmo nome (ou "chave") no dataframe polars

  • 772

Eu tenho esse código

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

o que dá

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

Agora o que estou tentando fazer é renomear as colunas para nomes mais curtos, como

┌───────────┬──────┬──────┬─────┬─────┬─────┐
│ 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 │
╞═══════════╪══════╪══════╪═════╪═════╪═════╡

e então recolher (= soma) nas colunas com o mesmo nome, para que eu obtenha, por exemplo

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

Tentei renomeá-los primeiro, mas não obtive sucesso polars.exceptions.DuplicateError: the name 'M2' is duplicate.

Existe uma maneira de conseguir o que estou tentando fazer?

editar: Eu também tentei algo como

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

mas não funciona muito bem.

sum

2 respostas

  1. Best Answer

    Você pode criar um dicionário das colunas a serem agregadas e então usar 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()
          ])

    Saída:

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

  2. Você pode criar seus títulos usando setentão usar regex pl.colpara selecionar essas colunas

    headings = set([y[1] for x in df.columns if len(y:=x.split('.'))==2])
df.select('id',
    *[pl.sum_horizontal(pl.col(f"^.+?{y}$")).alias(y) for y in headings]
)

    Alternativamente, você pode fazer um unpivotre pivotcomo este

    (
    df
    .unpivot(index='id')
    .with_columns(pl.col('variable').str.split('.').list.last())
    .pivot('variable', index='id', aggregate_function='sum')
)

    O unpivot/pivot provavelmente é mais lento, mas para algumas estruturas de dados a primeira abordagem não funciona.

    • 0

relate perguntas