Início / coding / Perguntas / 78992045
Accepted
usdn
Asked: 2024-09-17 06:40:08 +0800 CST

Modo de rolamento em Polars

  • 772

Tenho um data frame de ~100M linhas contendo IDs em diferentes grupos. Alguns deles estão errados (indicados pelo 99). Estou tentando corrigi-los com uma janela de modo de rolagem, semelhante ao exemplo de código abaixo. Existe uma maneira melhor de fazer isso, já que rolling_map() é super lento?

import polars as pl
from scipy import stats

def dummy(input):
    return stats.mode(input)[0]

df = pl.DataFrame({'group': [10, 10, 10, 10, 10, 10, 10, 20, 20, 20, 20],
                   'id': [1, 1, 99, 1, 1, 2, 2, 3, 3, 99, 3]})

df.with_columns(pl.col('id')
                 .rolling_map(function=dummy,
                              window_size=3,
                              min_periods=1,
                              center=True)
                 .over('group')
                 .alias('id_mode'))

shape: (11, 3)
╭───────┬─────┬─────────╮
│ group ┆  id ┆ id_mode │
│   i64 ┆ i64 ┆     i64 │
╞═══════╪═════╪═════════╡
│    10 ┆   1 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆  99 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆   1 ┆       1 │
│    10 ┆   2 ┆       2 │
│    10 ┆   2 ┆       2 │
│    20 ┆   3 ┆       3 │
│    20 ┆   3 ┆       3 │
│    20 ┆  99 ┆       3 │
│    20 ┆   3 ┆       3 │
╰───────┴─────┴─────────╯
python-polars

1 respostas

  1. Best Answer

    Existem .rolling()métodos de nível de quadro que permanecem na "terra da expressão" .

    (df.with_row_index()
   .rolling(group_by="group", index_column="index", period="3i")
   .agg(pl.col("id").mode().first())
)

    shape: (11, 3)
┌───────┬───────┬─────┐
│ group ┆ index ┆ id  │
│ ---   ┆ ---   ┆ --- │
│ i64   ┆ u32   ┆ i64 │
╞═══════╪═══════╪═════╡
│ 10    ┆ 0     ┆ 1   │
│ 10    ┆ 1     ┆ 1   │
│ 10    ┆ 2     ┆ 1   │
│ 10    ┆ 3     ┆ 1   │
│ 10    ┆ 4     ┆ 1   │
│ 10    ┆ 5     ┆ 1   │
│ 10    ┆ 6     ┆ 2   │
│ 20    ┆ 7     ┆ 3   │
│ 20    ┆ 8     ┆ 3   │
│ 20    ┆ 9     ┆ 3   │
│ 20    ┆ 10    ┆ 3   │
└───────┴───────┴─────┘
    • 1

relate perguntas