Início / coding / Perguntas / 79573564
Accepted
Phil-ZXX
Asked: 2025-04-14 23:57:58 +0800 CST

Coluna Group-By em Polars DataFrame dentro de with_columns

  • 772

Tenho o seguinte dataframe:

import polars as pl

df = pl.DataFrame({
    'ID': [1, 1, 5, 5, 7, 7, 7],
    'YEAR': [2025, 2025, 2023, 2024, 2020, 2021, 2021]
})
shape: (7, 2)
┌─────┬──────┐
│ ID  ┆ YEAR │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 1   ┆ 2025 │
│ 1   ┆ 2025 │
│ 5   ┆ 2023 │
│ 5   ┆ 2024 │
│ 7   ┆ 2020 │
│ 7   ┆ 2021 │
│ 7   ┆ 2021 │
└─────┴──────┘

Agora eu gostaria de obter o número único de anos por ID, ou seja

shape: (7, 3)
┌─────┬──────┬──────────────┐
│ ID  ┆ YEAR ┆ UNIQUE_YEARS │
│ --- ┆ ---  ┆ ---          │
│ i64 ┆ i64  ┆ u32          │
╞═════╪══════╪══════════════╡
│ 1   ┆ 2025 ┆ 1            │
│ 1   ┆ 2025 ┆ 1            │
│ 5   ┆ 2023 ┆ 2            │
│ 5   ┆ 2024 ┆ 2            │
│ 7   ┆ 2020 ┆ 2            │
│ 7   ┆ 2021 ┆ 2            │
│ 7   ┆ 2021 ┆ 2            │
└─────┴──────┴──────────────┘

Então eu tentei df.with_columns(pl.col('YEAR').over('ID').alias('UNIQUE_YEARS')), mas isso dá o resultado errado. Então eu pensei em

df.join(df.group_by('ID').agg(pl.col('YEAR').unique().len().alias('UNIQUE_YEARS')), on='ID', how='left')

O que dá o resultado correto! Mas parece um pouco desajeitado, e eu me pergunto se existe uma maneira mais natural de usar with_columnse over?

python

1 respostas

  1. Best Answer

    Você pode usar Expr.n_unique:

    out = df.with_columns(
    pl.col('YEAR').n_unique().over('ID').alias('UNIQUE_YEARS')
    )

    Saída:

    shape: (7, 3)
┌─────┬──────┬──────────────┐
│ ID  ┆ YEAR ┆ UNIQUE_YEARS │
│ --- ┆ ---  ┆ ---          │
│ i64 ┆ i64  ┆ u32          │
╞═════╪══════╪══════════════╡
│ 1   ┆ 2025 ┆ 1            │
│ 1   ┆ 2025 ┆ 1            │
│ 5   ┆ 2023 ┆ 2            │
│ 5   ┆ 2024 ┆ 2            │
│ 7   ┆ 2020 ┆ 2            │
│ 7   ┆ 2021 ┆ 2            │
│ 7   ┆ 2021 ┆ 2            │
└─────┴──────┴──────────────┘

    Da mesma forma, groupbypode tomar .n_unique()em vez do seu .unique().len().

    • 3

relate perguntas