Início / coding / Perguntas / 78914882
Accepted
Andi
Asked: 2024-08-26 22:10:19 +0800 CST

Substitua uma célula em uma coluna com base em uma célula em outra coluna em um DataFrame polar

  • 772

Considere o seguinte pl.DataFrame:

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["s1", "s1", "s2", "s2"],
        "signal": [0, 1, 2, 0],
        "trade": [None, 1, None, -1],
    }
)

shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ null  │
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘

Agora preciso agrupar o dataframe symbole verificar se a primeira linha de cada grupo da coluna signalnão é igual a 0 (zero). Se for igual a True, preciso substituir a célula correspondente na coluna tradepelo valor na célula em signal.

Aqui está o que estou realmente procurando:

shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ 2     │ <- copy value from the ``signal`` column
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘
python

3 respostas

  1. Best Answer

    Para isso, uma construção quando-então-caso contrário pode ser usada.

    • Criamos uma condição que avalia Trueexatamente para as primeiras linhas (crie o índice instantaneamente usando pl.int_range) em cada grupo com signalvalor diferente de 0.
    • Com base nessa condição, selecionamos o valor na coluna signalou trade.
    df.with_columns(
    trade=pl.when(
        pl.col("signal") != 0,
        pl.int_range(pl.len()) == 0,
    ).then("signal").otherwise("trade").over("symbol")
)

    shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ 2     │
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘
    • 4

  2. Você pode usar .is_first_distinct()para identificar as linhas desejadas.

    df.with_columns(pl.col.symbol.is_first_distinct().alias("checkme"))

    shape: (4, 4)
┌────────┬────────┬───────┬─────────┐
│ symbol ┆ signal ┆ trade ┆ checkme │
│ ---    ┆ ---    ┆ ---   ┆ ---     │
│ str    ┆ i64    ┆ i64   ┆ bool    │
╞════════╪════════╪═══════╪═════════╡
│ s1     ┆ 0      ┆ null  ┆ true    │
│ s1     ┆ 1      ┆ 1     ┆ false   │
│ s2     ┆ 2      ┆ null  ┆ true    │
│ s2     ┆ 0      ┆ -1    ┆ false   │
└────────┴────────┴───────┴─────────┘

    E use como parte da lógica quando/então:

    df.with_columns(
    pl.when(
       pl.col.symbol.is_first_distinct(),
       pl.col.signal != 0
    )
    .then("signal")
    .otherwise("trade")
    .alias("trade")
)

    shape: (4, 3)
┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ 2     │
│ s2     ┆ 0      ┆ -1    │
└────────┴────────┴───────┘
    • 2

  3. Usando o símbolo e :cum_count overwhen

    df.with_columns(
    trade=pl.when(
        pl.cum_count('symbol').over('symbol').eq(1),
        pl.col('signal') != 0,
    )
    .then('signal')
    .otherwise('trade')
)

    Saída:

    ┌────────┬────────┬───────┐
│ symbol ┆ signal ┆ trade │
│ ---    ┆ ---    ┆ ---   │
│ str    ┆ i64    ┆ i64   │
╞════════╪════════╪═══════╡
│ s1     ┆ 0      ┆ null  │
│ s1     ┆ 1      ┆ 1     │
│ s2     ┆ 2      ┆ 2     │
│ s2     ┆ 2      ┆ -2    │
└────────┴────────┴───────┘
    • 0

relate perguntas