Início / coding / Perguntas / 77878728
Accepted
FredMaster
Asked: 2024-01-25 17:15:53 +0800 CST

Dataframe Polars: grupos sobrepostos

  • 772

Atualmente estou "convertendo" de pandas para polares porque gosto muito da API. Esta pergunta é uma pergunta mais geral para uma pergunta anterior minha (veja aqui )

Eu tenho o seguinte dataframe


# Dummy data
df = pl.DataFrame({
    "Buy_Signal": [1, 0, 1, 0, 1, 0, 0], 
    "Returns": [0.01, 0.02, 0.03, 0.02, 0.01, 0.00, -0.01],
})

Em última análise, quero fazer agregações na coluna Returnscondicionais a intervalos diferentes - que são dados por column Buy_Signal. No caso acima, o comprimento é de 1 até o final do dataframe. O dataframe resultante deve, portanto, ter a seguinte aparência

| group     | Returns   |
|------:    |--------:  |
|   u32     |     f64   |
|     1     |    0.01   |
|     1     |    0.02   |
|     1     |    0.03   |
|     1     |    0.02   |
|     1     |    0.01   |
|     1     |     0.0   |
|     1     |   -0.01   |
|     2     |    0.03   |
|     2     |    0.02   |
|     2     |    0.01   |
|     2     |     0.0   |
|     2     |   -0.01   |
|     3     |    0.01   |
|     3     |     0.0   |
|     3     |   -0.01   |

Uma abordagem postada como resposta à minha pergunta anterior é a seguinte:


# Build overlapping group index
idx = df.select(index=
          pl.when(pl.col("Buy_Signal") == 1)
          .then(pl.int_ranges(pl.int_range(pl.len()), pl.len()  ))
).explode(pl.col("index")).drop_nulls().cast(pl.UInt32)

# Join index with original data
df = (df.with_row_index()
    .join(idx, on="index")
    .with_columns(group = (pl.col("index") == pl.col("index").max())
                .shift().cum_sum().backward_fill() + 1)
    .select(["group", "Returns"])
)
df

Pergunta: existem outras boas soluções para este problema.

Com bom quero dizer (i) legível e/ou (ii) rápido.

Meu problema real contém conjuntos de dados muito maiores.

Obrigado

python

3 respostas

  1. Best Answer

    Para completar, aqui está uma solução alternativa que não depende de funcionalidade experimental.

    (
    df
    .with_columns(
        pl.col("Buy_Signal").cum_sum().alias("group")
    )
    .with_columns(
        pl.int_ranges(pl.col("group").min(), pl.col("group")+1)
    )
    .explode("group")
    .sort("group")
)

    Saída.

    shape: (15, 3)
┌────────────┬─────────┬───────┐
│ Buy_Signal ┆ Returns ┆ group │
│ ---        ┆ ---     ┆ ---   │
│ i64        ┆ f64     ┆ i64   │
╞════════════╪═════════╪═══════╡
│ 1          ┆ 0.01    ┆ 1     │
│ 0          ┆ 0.02    ┆ 1     │
│ 1          ┆ 0.03    ┆ 1     │
│ 0          ┆ 0.02    ┆ 1     │
│ 1          ┆ 0.01    ┆ 1     │
│ …          ┆ …       ┆ …     │
│ 0          ┆ 0.0     ┆ 2     │
│ 0          ┆ -0.01   ┆ 2     │
│ 1          ┆ 0.01    ┆ 3     │
│ 0          ┆ 0.0     ┆ 3     │
│ 0          ┆ -0.01   ┆ 3     │
└────────────┴─────────┴───────┘
    • 2

  2. Eu acho que é um bom caso para uso de cumulative_eval().

    À primeira vista, você precisa ter 2 passagens no seu DataFrame. Primeiro, você precisa calcular cum_sumpara acertar o agrupamento:

    df = df.with_columns(
        Holdings = pl.col('Buy_Signal').cum_sum()
    ).group_by("Holdings").agg(
        pl.col(["Returns"])).sort("Holdings", descending=True
    )

print(df)

shape: (3, 2)
┌──────────┬────────────────────┐
│ Holdings ┆ Returns            │
│ ---      ┆ ---                │
│ i64      ┆ list[f64]          │
╞══════════╪════════════════════╡
│ 3        ┆ [0.01, 0.0, -0.01] │
│ 2        ┆ [0.03, 0.02]       │
│ 1        ┆ [0.01, 0.02]       │
└──────────┴────────────────────┘

    Observe que atualmente nosso DataFrame está classificado com descending=True, isso é importante, pois queremos que nossas listas sejam agregadas do maior grupo para o menor.

    E agora cumulative_eval(), junto com implode():

    df = df.with_columns(
        pl.col("Returns").cumulative_eval(pl.element().explode().implode())
    )

print(df.explode("Returns"))

shape: (15, 2)
┌──────────┬─────────┐
│ Holdings ┆ Returns │
│ ---      ┆ ---     │
│ i64      ┆ f64     │
╞══════════╪═════════╡
│ 3        ┆ 0.01    │
│ 3        ┆ 0.0     │
│ 3        ┆ -0.01   │
│ 2        ┆ 0.01    │
│ …        ┆ …       │
│ 1        ┆ 0.03    │
│ 1        ┆ 0.02    │
│ 1        ┆ 0.01    │
│ 1        ┆ 0.02    │
└──────────┴─────────┘

    Você pode ver que o resultado corresponde ao que você esperaria. No entanto, você deve ter cuidado com o desempenho, pois há uma observação na cumulative_evaldocumentação:

    Aviso

    Esta funcionalidade é experimental e pode mudar sem ser considerada uma alteração significativa.

    Isso pode ser muito lento, pois pode ter complexidade O(n^2). Não use isto para operações que visitam todos os elementos.

    • 1

  3. Apenas para adicionar mais algumas informações, pois acho que o título da pergunta é baseado na minha redação anterior:

    Encontrei uma definição existente do problema (qual pode ser um título melhor para a pergunta?)

    Um Polars equivalente ao ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING enquadramento de janela do SQL.

    ┌────────────┬─────────┐
│ Buy_Signal ┆ Returns │
│ ---        ┆ ---     │
│ i64        ┆ f64     │ # New "group/window" starts when Buy_Signal == 1
╞════════════╪═════════╡ # 
│ 1          ┆ 0.01    │ # ┌─ group 1 ─┐
│ 0          ┆ 0.02    │ # │           │ 
│ 1          ┆ 0.03    │ # │           │┌─ group 2 ─┐
│ 0          ┆ 0.02    │ # │           ││           │ 
│ 1          ┆ 0.01    │ # │           ││           │┌─ group 3 ─┐
│ 0          ┆ 0.0     │ # │           ││           ││           │
│ 0          ┆ -0.01   │ # └───────────┘└───────────┘└───────────┘
└────────────┴─────────┘

    por exemplo, usando DuckDB para demonstrar:

    duckdb.sql("""
from df
select 
   buy_signal, 
   returns, 
   case when buy_signal = 1
   then
      list(returns) over (rows between current row and unbounded following)
   end as group
""")

    ┌────────────┬─────────┬────────────────────────────────────────────┐
│ Buy_Signal │ Returns │                   group                    │
│   int64    │ double  │                  double[]                  │
├────────────┼─────────┼────────────────────────────────────────────┤
│          1 │    0.01 │ [0.01, 0.02, 0.03, 0.02, 0.01, 0.0, -0.01] │
│          0 │    0.02 │ NULL                                       │
│          1 │    0.03 │ [0.03, 0.02, 0.01, 0.0, -0.01]             │
│          0 │    0.02 │ NULL                                       │
│          1 │    0.01 │ [0.01, 0.0, -0.01]                         │
│          0 │     0.0 │ NULL                                       │
│          0 │   -0.01 │ NULL                                       │
└────────────┴─────────┴────────────────────────────────────────────┘
    • 0

relate perguntas