Início / coding / Perguntas / 79576828
Accepted
gernophil
Asked: 2025-04-16 17:29:12 +0800 CST

Obtenha uma soma agrupada em polares, mas mantenha todas as linhas individuais

  • 772

Estou quebrando a cabeça com essa pergunta, provavelmente bem simples, e não consigo encontrar a resposta em lugar nenhum. Quero criar uma nova coluna com a soma agrupada de outra coluna, mas quero manter todas as linhas individuais. Então, aqui está o que a documentação diz:

import polars as pl

df = pl.DataFrame(
    {
        "a": ["a", "b", "a", "b", "c"],
        "b": [1, 2, 1, 3, 3],
    }
)

df.group_by("a").agg(pl.col("b").sum())

O resultado disso seria:

shape: (3, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ i64 │
╞═════╪═════╡
│ a   ┆ 2   │
│ c   ┆ 3   │
│ b   ┆ 5   │
└─────┴─────┘

No entanto, o que eu preciso é disto:

shape: (5, 3)
┌─────┬─────┬────────┐
│ a   ┆ b   ┆ sum(b) │
│ --- ┆ --- ┆ ------ │
│ str ┆ i64 ┆ i64    │
╞═════╪═════╪════════╡
│ a   ┆ 1   ┆ 2      │
│ b   ┆ 2   ┆ 5      │
│ a   ┆ 1   ┆ 2      │
│ b   ┆ 3   ┆ 5      │
│ c   ┆ 3   ┆ 3      │
└─────┴─────┴────────┘

Eu poderia criar a soma em um df separado e depois uni-la ao original, mas tenho certeza de que há uma solução mais fácil.

python

1 respostas

  1. Best Answer

    Tudo o que você precisa é de uma função de janela:

    df.with_columns(
    b_sum=pl.col("b").sum().over(pl.col("a"))
)


shape: (5, 3)
┌─────┬─────┬───────┐
│ a   ┆ b   ┆ b_sum │
│ --- ┆ --- ┆ ---   │
│ str ┆ i64 ┆ i64   │
╞═════╪═════╪═══════╡
│ a   ┆ 1   ┆ 2     │
│ b   ┆ 2   ┆ 5     │
│ a   ┆ 1   ┆ 2     │
│ b   ┆ 3   ┆ 5     │
│ c   ┆ 3   ┆ 3     │
└─────┴─────┴───────┘
    • 3

relate perguntas