Início / coding / Perguntas / 78842466
Accepted
Andi
Asked: 2024-08-07 15:52:59 +0800 CST

Divida um DataFrame polar em vários pedaços com groupby

  • 772

Considere o seguinte pl.DataFrame:

import datetime

import polars as pl

df_orig = pl.DataFrame(
    {
        "symbol": [*["A"] * 10, *["B"] * 8],
        "date": [
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 10),
                eager=True,
            ),
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 8),
                eager=True,
            ),
        ],
        "data": [*range(10), *range(8)],
    }
)

df_helper = pl.DataFrame({"symbol": ["A", "B"], "start_idx": [[0, 5], [0, 4]]})

chunk_size = 5

with pl.Config(tbl_rows=30):
    print(df_orig)
    print(df_helper)


shape: (18, 3)
┌────────┬─────────────────────┬──────┐
│ symbol ┆ date                ┆ data │
│ ---    ┆ ---                 ┆ ---  │
│ str    ┆ datetime[μs]        ┆ i64  │
╞════════╪═════════════════════╪══════╡
│ A      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ A      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ A      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ A      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ A      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ A      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ A      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ A      ┆ 2024-01-08 00:00:00 ┆ 7    │
│ A      ┆ 2024-01-09 00:00:00 ┆ 8    │
│ A      ┆ 2024-01-10 00:00:00 ┆ 9    │
│ B      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ B      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ B      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ B      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ B      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ B      ┆ 2024-01-08 00:00:00 ┆ 7    │
└────────┴─────────────────────┴──────┘

shape: (2, 2)
┌────────┬───────────┐
│ symbol ┆ start_idx │
│ ---    ┆ ---       │
│ str    ┆ list[i64] │
╞════════╪═══════════╡
│ A      ┆ [0, 5]    │
│ B      ┆ [0, 3]    │
└────────┴───────────┘

Agora, preciso dividir o dataframe em dois pedaços de comprimento 5 ( chunk_size) agrupados pela symbolcoluna. A coluna start_idxindica as linhas para iniciar o bloco em cada grupo. Ou seja, o grupo A será dividido em dois pedaços de comprimento 5 começando nas linhas 0 e 5, enquanto os pedaços do grupo B começam nas linhas 0 e 3. Finalmente, todos os pedaços precisam ser concatenados axis=0, onde uma nova coluna split_idxindica onde a divisão está vindo.

Aqui está o que estou procurando:

shape: (20, 4)
┌────────────────────┬─────────────────────┬──────┐
│ split_idx ┆ symbol ┆ date                ┆ data │
│           ┆ ---    ┆ ---                 ┆ ---  │
│ i64       ┆ str    ┆ datetime[μs]        ┆ i64  │
╞═══════════╪════════╪═════════════════════╪══════╡
│ 0         ┆ A      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ 0         ┆ A      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ 0         ┆ A      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ 0         ┆ A      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 0         ┆ A      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 0         ┆ B      ┆ 2024-01-01 00:00:00 ┆ 0    │
│ 0         ┆ B      ┆ 2024-01-02 00:00:00 ┆ 1    │
│ 0         ┆ B      ┆ 2024-01-03 00:00:00 ┆ 2    │
│ 0         ┆ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 0         ┆ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 1         ┆ A      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ 1         ┆ A      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ 1         ┆ A      ┆ 2024-01-08 00:00:00 ┆ 7    │
│ 1         ┆ A      ┆ 2024-01-09 00:00:00 ┆ 8    │
│ 1         ┆ A      ┆ 2024-01-10 00:00:00 ┆ 9    │
│ 1         ┆ B      ┆ 2024-01-04 00:00:00 ┆ 3    │
│ 1         ┆ B      ┆ 2024-01-05 00:00:00 ┆ 4    │
│ 1         ┆ B      ┆ 2024-01-06 00:00:00 ┆ 5    │
│ 1         ┆ B      ┆ 2024-01-07 00:00:00 ┆ 6    │
│ 1         ┆ B      ┆ 2024-01-08 00:00:00 ┆ 7    │
└───────────┴────────┴─────────────────────┴──────┘

Tenha em mente que a lista na coluna start_idxpode ter comprimento variável para cada linha individual. O comprimento de cada lista determina o número de pedaços para cada grupo.

python-polars

2 respostas

  1. Best Answer

    Aqui está uma solução que permanece totalmente dentro da API de expressão polars.

    A ideia principal é pré-processar o dataframe auxiliar em um dataframe de symbol, split_idxe row_idx. Aqui, row_idxestá o índice de uma linha dentro de um grupo definido por símbolo e índice de divisão.

    Ele pode servir como um "esqueleto" e podemos (depois de adicionar esse índice de linha a df_orig) usá-lo facilmente para uma mesclagem à esquerda com df_orig.

    pl.Config().set_tbl_rows(-1)

def preprocess_helper(df_helper: pl.DataFrame) -> pl.DataFrame:
    return (
        df_helper
        .explode("start_idx")
        .with_columns(
            pl.int_range(pl.len()).over("symbol").alias("split_idx"),
            pl.int_ranges(pl.col("start_idx"), pl.col("start_idx") + chunk_size).alias("row_idx"),
        )
        .explode("row_idx")
    )

(
    preprocess_helper(df_helper)
    .join(
        df_orig.with_columns(pl.int_range(pl.len()).over("symbol").alias("row_idx")),
        on=["symbol", "row_idx"],
        how="left",
    )
    .drop("row_idx", "start_idx")
    .sort("split_idx", "symbol")
)

    Observação.pl.DataFrame.drop O / final pl.DataFrame.sortpode ser omitido se as colunas/ordem exatas da saída não importarem.

    shape: (20, 4)
┌────────┬───────────┬─────────────────────┬──────┐
│ symbol ┆ split_idx ┆ date                ┆ data │
│ ---    ┆ ---       ┆ ---                 ┆ ---  │
│ str    ┆ i64       ┆ datetime[μs]        ┆ i64  │
╞════════╪═══════════╪═════════════════════╪══════╡
│ A      ┆ 0         ┆ 2024-01-01 00:00:00 ┆ 0    │
│ A      ┆ 0         ┆ 2024-01-02 00:00:00 ┆ 1    │
│ A      ┆ 0         ┆ 2024-01-03 00:00:00 ┆ 2    │
│ A      ┆ 0         ┆ 2024-01-04 00:00:00 ┆ 3    │
│ A      ┆ 0         ┆ 2024-01-05 00:00:00 ┆ 4    │
│ B      ┆ 0         ┆ 2024-01-01 00:00:00 ┆ 0    │
│ B      ┆ 0         ┆ 2024-01-02 00:00:00 ┆ 1    │
│ B      ┆ 0         ┆ 2024-01-03 00:00:00 ┆ 2    │
│ B      ┆ 0         ┆ 2024-01-04 00:00:00 ┆ 3    │
│ B      ┆ 0         ┆ 2024-01-05 00:00:00 ┆ 4    │
│ A      ┆ 1         ┆ 2024-01-06 00:00:00 ┆ 5    │
│ A      ┆ 1         ┆ 2024-01-07 00:00:00 ┆ 6    │
│ A      ┆ 1         ┆ 2024-01-08 00:00:00 ┆ 7    │
│ A      ┆ 1         ┆ 2024-01-09 00:00:00 ┆ 8    │
│ A      ┆ 1         ┆ 2024-01-10 00:00:00 ┆ 9    │
│ B      ┆ 1         ┆ 2024-01-04 00:00:00 ┆ 3    │
│ B      ┆ 1         ┆ 2024-01-05 00:00:00 ┆ 4    │
│ B      ┆ 1         ┆ 2024-01-06 00:00:00 ┆ 5    │
│ B      ┆ 1         ┆ 2024-01-07 00:00:00 ┆ 6    │
│ B      ┆ 1         ┆ 2024-01-08 00:00:00 ┆ 7    │
└────────┴───────────┴─────────────────────┴──────┘
    • 2

  2. Não consegui encontrar uma solução totalmente nativa, pois polars.Expr.sliceparece não suportar valores diferentes para cada linha.

    Aqui está uma maneira de fazer isso com algum código Python misturado:

    Primeiro, agrupe df_origpor símbolo e coloque-o em um ditado.

    Em seguida, itere df_helpere extraia as fatias e coloque-as em uma lista.

    Por fim, concatene os dataframes em um único dataframe e classifique-o.

    df_orig_by_symbol = {
    k: v for ((k,), v) in df_orig.group_by("symbol", maintain_order=True)
}

dfs = []

for symbol, start_idxs in df_helper.select("symbol", "start_idx").iter_rows():
    for split_idx, start_idx in enumerate(start_idxs):
        dfs.append(
            df_orig_by_symbol[symbol]
            .slice(start_idx, chunk_size)
            .with_columns(split_idx=split_idx)
        )

final_df = pl.concat(dfs).sort("split_idx", "symbol")

    Saída:

    shape: (20, 4)
┌────────┬─────────────────────┬──────┬───────────┐
│ symbol ┆ date                ┆ data ┆ split_idx │
│ ---    ┆ ---                 ┆ ---  ┆ ---       │
│ str    ┆ datetime[μs]        ┆ i64  ┆ i32       │
╞════════╪═════════════════════╪══════╪═══════════╡
│ A      ┆ 2024-01-01 00:00:00 ┆ 0    ┆ 0         │
│ A      ┆ 2024-01-02 00:00:00 ┆ 1    ┆ 0         │
│ A      ┆ 2024-01-03 00:00:00 ┆ 2    ┆ 0         │
│ A      ┆ 2024-01-04 00:00:00 ┆ 3    ┆ 0         │
│ A      ┆ 2024-01-05 00:00:00 ┆ 4    ┆ 0         │
│ B      ┆ 2024-01-01 00:00:00 ┆ 0    ┆ 0         │
│ B      ┆ 2024-01-02 00:00:00 ┆ 1    ┆ 0         │
│ B      ┆ 2024-01-03 00:00:00 ┆ 2    ┆ 0         │
│ B      ┆ 2024-01-04 00:00:00 ┆ 3    ┆ 0         │
│ B      ┆ 2024-01-05 00:00:00 ┆ 4    ┆ 0         │
│ A      ┆ 2024-01-06 00:00:00 ┆ 5    ┆ 1         │
│ A      ┆ 2024-01-07 00:00:00 ┆ 6    ┆ 1         │
│ A      ┆ 2024-01-08 00:00:00 ┆ 7    ┆ 1         │
│ A      ┆ 2024-01-09 00:00:00 ┆ 8    ┆ 1         │
│ A      ┆ 2024-01-10 00:00:00 ┆ 9    ┆ 1         │
│ B      ┆ 2024-01-04 00:00:00 ┆ 3    ┆ 1         │
│ B      ┆ 2024-01-05 00:00:00 ┆ 4    ┆ 1         │
│ B      ┆ 2024-01-06 00:00:00 ┆ 5    ┆ 1         │
│ B      ┆ 2024-01-07 00:00:00 ┆ 6    ┆ 1         │
│ B      ┆ 2024-01-08 00:00:00 ┆ 7    ┆ 1         │
└────────┴─────────────────────┴──────┴───────────┘

    Código completo:

    import datetime

import polars as pl

df_orig = pl.DataFrame(
    {
        "symbol": [*["A"] * 10, *["B"] * 8],
        "date": [
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 10),
                eager=True,
            ),
            *pl.datetime_range(
                start=datetime.date(2024, 1, 1),
                end=datetime.date(2024, 1, 8),
                eager=True,
            ),
        ],
        "data": [*range(10), *range(8)],
    }
)

df_helper = pl.DataFrame({"symbol": ["A", "B"], "start_idx": [[0, 5], [0, 3]]})


chunk_size = 5

df_orig_by_symbol = {
    k: v for ((k,), v) in df_orig.group_by("symbol", maintain_order=True)
}

dfs = []

for symbol, start_idxs in df_helper.select("symbol", "start_idx").iter_rows():
    for split_idx, start_idx in enumerate(start_idxs):
        dfs.append(
            df_orig_by_symbol[symbol]
            .slice(start_idx, chunk_size)
            .with_columns(split_idx=split_idx)
        )

final_df = pl.concat(dfs).sort("split_idx", "symbol")

with pl.Config(tbl_rows=100):
    print(final_df)
    • 0

relate perguntas