Início / coding / Perguntas / 79429531
Accepted
Quinten
Asked: 2025-02-11 17:57:15 +0800 CST

Selecione a primeira e a última linha por grupo no dataframe Polars

  • 772

Estou tentando usar polarsdataframe onde gostaria de selecionar a linha firste lastpor grupo. Aqui está um exemplo simples selecionando a primeira linha por grupo:

import polars as pl

df = pl.DataFrame(
    {
        "a": [1, 2, 2, 3, 4, 5],
        "b": [0.5, 0.5, 4, 10, 14, 13],
        "c": [True, True, True, False, False, True],
        "d": ["Apple", "Apple", "Apple", "Banana", "Banana", "Banana"],
    }
)
result = df.group_by("d", maintain_order=True).first()
print(result)

Saída:

shape: (2, 4)
┌────────┬─────┬──────┬───────┐
│ d      ┆ a   ┆ b    ┆ c     │
│ ---    ┆ --- ┆ ---  ┆ ---   │
│ str    ┆ i64 ┆ f64  ┆ bool  │
╞════════╪═════╪══════╪═══════╡
│ Apple  ┆ 1   ┆ 0.5  ┆ true  │
│ Banana ┆ 3   ┆ 10.0 ┆ false │
└────────┴─────┴──────┴───────┘

Isso funciona bem e podemos usar .lastpara fazer isso na última fileira. Mas como podemos combinar essas duas em uma group_by?

python

3 respostas

  1. Best Answer

    Como colunas

    Você pode usar agg, você terá que adicionar um suffix(ou prefix) para diferenciar os nomes das colunas:

    result = (df.group_by('d', maintain_order=True)
            .agg(pl.all().first().name.suffix('_first'),
                 pl.all().last().name.suffix('_last'))
         )

    Saída:

    ┌────────┬─────────┬─────────┬─────────┬────────┬────────┬────────┐
│ d      ┆ a_first ┆ b_first ┆ c_first ┆ a_last ┆ b_last ┆ c_last │
│ ---    ┆ ---     ┆ ---     ┆ ---     ┆ ---    ┆ ---    ┆ ---    │
│ str    ┆ i64     ┆ f64     ┆ bool    ┆ i64    ┆ f64    ┆ bool   │
╞════════╪═════════╪═════════╪═════════╪════════╪════════╪════════╡
│ Apple  ┆ 1       ┆ 0.5     ┆ true    ┆ 2      ┆ 4.0    ┆ true   │
│ Banana ┆ 3       ┆ 10.0    ┆ false   ┆ 5      ┆ 13.0   ┆ true   │
└────────┴─────────┴─────────┴─────────┴────────┴────────┴────────┘

    Como linhas

    Se você quiser várias linhas, precisará concat:

    g = df.group_by('d', maintain_order=True)

result = pl.concat([g.first(), g.last()]).sort(by='d', maintain_order=True)

    Saída:

    ┌────────┬─────┬──────┬───────┐
│ d      ┆ a   ┆ b    ┆ c     │
│ ---    ┆ --- ┆ ---  ┆ ---   │
│ str    ┆ i64 ┆ f64  ┆ bool  │
╞════════╪═════╪══════╪═══════╡
│ Apple  ┆ 1   ┆ 0.5  ┆ true  │
│ Apple  ┆ 2   ┆ 4.0  ┆ true  │
│ Banana ┆ 3   ┆ 10.0 ┆ false │
│ Banana ┆ 5   ┆ 13.0 ┆ true  │
└────────┴─────┴──────┴───────┘

    Ou usando filtercom int_range+ over:

    result = df.filter((pl.int_range(pl.len()).over('d') == 0)
                  |(pl.int_range(pl.len(), 0, -1).over('d') == 1)
                  )

    Saída:

    ┌─────┬──────┬───────┬────────┐
│ a   ┆ b    ┆ c     ┆ d      │
│ --- ┆ ---  ┆ ---   ┆ ---    │
│ i64 ┆ f64  ┆ bool  ┆ str    │
╞═════╪══════╪═══════╪════════╡
│ 1   ┆ 0.5  ┆ true  ┆ Apple  │
│ 2   ┆ 4.0  ┆ true  ┆ Apple  │
│ 3   ┆ 10.0 ┆ false ┆ Banana │
│ 5   ┆ 13.0 ┆ true  ┆ Banana │
└─────┴──────┴───────┴────────┘
    • 2

  2. As soluções do @mozway funcionam bem! Para completar, eu também queria compartilhar duas soluções que dependem de pl.Expr.gather.

    Em um contexto selecionado

    df.select(
    pl.all().gather([0, -1]).over("d", mapping_strategy="explode")
)

    Em um contexto de agrupamento por

    (
    df
    .group_by("d", maintain_order=True)
    .agg(
        pl.all().gather([0, -1])
    )
    .explode(pl.exclude("d"))
)

    Considerações sobre desempenho

    Também executei cronometragem preliminar desses métodos (no pequeno conjunto de dados de exemplo).

    Método Tempos (média ± desvio padrão de 7 execuções, 1.000 loops cada)
    group_by+concat 452 μs ± 7,34 μs por loop
    filter 396 μs ± 10,2 μs por loop
    group_by+gather 255 μs ± 4,09 μs por loop
    select+gather 172 μs ± 1,29 μs por loop
    • 2

  3. Existem métodos dedicados de primeiro/último.

    df.filter(
    pl.any_horizontal(
        pl.col("d").is_first_distinct(),
        pl.col("d").is_last_distinct()
    )
)

    shape: (4, 4)
┌─────┬──────┬───────┬────────┐
│ a   ┆ b    ┆ c     ┆ d      │
│ --- ┆ ---  ┆ ---   ┆ ---    │
│ i64 ┆ f64  ┆ bool  ┆ str    │
╞═════╪══════╪═══════╪════════╡
│ 1   ┆ 0.5  ┆ true  ┆ Apple  │
│ 2   ┆ 4.0  ┆ true  ┆ Apple  │
│ 3   ┆ 10.0 ┆ false ┆ Banana │
│ 5   ┆ 13.0 ┆ true  ┆ Banana │
└─────┴──────┴───────┴────────┘

    Você pode usar uma struct se o identificador do grupo tiver várias colunas.

    pl.struct("c", "d").is_first_distinct()
    • 1

relate perguntas