Início / coding / Perguntas / 79242541
Accepted
Samuel Allain
Asked: 2024-12-02 05:44:44 +0800 CST

No Polars, como você cria um contador de grupo / ID de grupo?

  • 772

Como você obtém uma group_idcoluna como esta, agrupando por colunas col1e col2?

col1 col2 id_do_grupo
UM Z 1
UM E 2
UM Z 1
B Z 3

com base em tal DataFrame:

df = pl.DataFrame({
  'col1': ['A', 'A', 'A', 'B'],
  'col2': ['Z', 'Y', 'Z', 'Z']}
)

Em outras palavras, estou procurando polares equivalentes a R data.table .GRP ( df[, group_id:=.GRP, by = .(col1, col2)]). Obrigado!

Contexto: Quero construir um ID de evento porque, em meus dados, tenho muitas linhas detalhadas para um evento. Depois que o ID do evento for criado, usarei-o para executar várias operações de janela. Prefiro ter esse ID de evento em vez de manter uma lista de variáveis ​​de agrupamento.

python

2 respostas

  1. Best Answer

    Se você não se importa com "primeira ocorrência tem classificação mais baixa", então você pode simplesmente rankcombinar as colunas col1e col2.

    df.with_columns(group_id = pl.struct("col1","col2").rank("dense"))

    shape: (4, 3)
┌──────┬──────┬──────────┐
│ col1 ┆ col2 ┆ group_id │
│ ---  ┆ ---  ┆ ---      │
│ str  ┆ str  ┆ u32      │
╞══════╪══════╪══════════╡
│ A    ┆ Z    ┆ 2        │
│ A    ┆ Y    ┆ 1        │
│ A    ┆ Z    ┆ 2        │
│ B    ┆ Z    ┆ 3        │
└──────┴──────┴──────────┘

    Caso contrário, será um pouco mais complicado - você pode criar um índice para rastrear a ordem atual das linhas e então atribuir uma classificação com base nesse índice.

    (
    df.with_row_index("group_id")
    .with_columns(
        pl.col.group_id.min().over("col1","col2").rank("dense")
    )
)

    shape: (4, 3)
┌──────────┬──────┬──────┐
│ group_id ┆ col1 ┆ col2 │
│ ---      ┆ ---  ┆ ---  │
│ u32      ┆ str  ┆ str  │
╞══════════╪══════╪══════╡
│ 1        ┆ A    ┆ Z    │
│ 2        ┆ A    ┆ Y    │
│ 1        ┆ A    ┆ Z    │
│ 3        ┆ B    ┆ Z    │
└──────────┴──────┴──────┘
    • 2

  2. Uma maneira de fazer isso seria obter uma lista de combinações únicas de col1e col2e com um índice de linha inteiro. Essa seria a coluna de agrupamento que você está procurando.

    Por exemplo:

    import polars as pl

# your example
df = pl.DataFrame({
    'col1': ['A', 'A', 'A', 'B'],
    'col2': ['Z', 'Y', 'Z', 'Z'],
})

unique_pairs = df.select(['col1', 'col2']).unique().with_row_index(name="group")
df_with_group = df.join(unique_pairs, on=['col1', 'col2'])

print(df_with_group)

    A junção emparelha cada linha com a linha apropriada unique_pairse adiciona as outras colunas dessa linha (apenas groupneste caso) ao resultado.

    Isso pode resultar em uma numeração diferente. Se você precisar que a numeração siga a ordem dos dados originais, como no seu exemplo, você pode alterá-la para:

    unique_pairs = (
    df.select(['col1', 'col2'])
    .unique(maintain_order=True)
    .with_row_index(name="group")
)

    Agora, eles unique_pairsterão uma numeração de pares exclusivos que corresponde à ordem em que aparecem pela primeira vez, em df.

    Saída:

    shape: (4, 3)
┌──────┬──────┬───────┐
│ col1 ┆ col2 ┆ group │
│ ---  ┆ ---  ┆ ---   │
│ str  ┆ str  ┆ u32   │
╞══════╪══════╪═══════╡
│ A    ┆ Z    ┆ 0     │
│ A    ┆ Y    ┆ 1     │
│ A    ┆ Z    ┆ 0     │
│ B    ┆ Z    ┆ 2     │
└──────┴──────┴───────┘

    No entanto, sua pergunta me faz pensar se esse é um Problema XY . Por que exatamente você precisa desse agrupamento? Porque se o agrupamento for gerado apenas para executar operações de grupo mais tarde, você pode simplesmente usar a funcionalidade de agrupamento do Polars.

    Por exemplo:

    import polars as pl

df = pl.DataFrame({
    'col1': ['A', 'A', 'A', 'B'],
    'col2': ['Z', 'Y', 'Z', 'Z'],
    'values': [1, 2, 3, 4],       # added some data to operate on
})

grouped_df = df.group_by(['col1', 'col2'], maintain_order=True).agg([
    pl.col('values').sum().alias('sum_values'),
    pl.col('values').mean().alias('mean_values'),
    pl.col('values').count().alias('count_values')
])

print(grouped_df)

    Resultados em:

    shape: (3, 5)
┌──────┬──────┬────────────┬─────────────┬──────────────┐
│ col1 ┆ col2 ┆ sum_values ┆ mean_values ┆ count_values │
│ ---  ┆ ---  ┆ ---        ┆ ---         ┆ ---          │
│ str  ┆ str  ┆ i64        ┆ f64         ┆ u32          │
╞══════╪══════╪════════════╪═════════════╪══════════════╡
│ A    ┆ Z    ┆ 4          ┆ 2.0         ┆ 2            │
│ A    ┆ Y    ┆ 2          ┆ 2.0         ┆ 1            │
│ B    ┆ Z    ┆ 4          ┆ 4.0         ┆ 1            │
└──────┴──────┴────────────┴─────────────┴──────────────┘
    • 1

relate perguntas