Início / coding / Perguntas / 77978208
Accepted
kejtos
Asked: 2024-02-12 03:54:38 +0800 CST

Como contar linhas acima da atual com base nas condições polares?

  • 772

Vamos ter um polars df:

df = pl.DataFrame(
    {
    'date': ['2022-01-01', '2022-01-02', '2022-01-07', '2022-01-17', '2022-03-02', '2022-06-05', '2022-06-07', '2022-07-02'],
    'col1': [4, 4, 2, 2, 2, 3, 2, 1],
    'col2': [1, 2, 3, 4, 1, 3, 3, 4],
    'col3': [2, 3, 4, 4, 3, 2, 2, 1]
    }
)
data col1 col2 col3
01/01/2022 1 1 2
2022-01-02 1 2 3
07/01/2022 2 3 4
17/01/2022 2 4 1
02/03/2022 3 1 3
05/06/2022 3 2 2
07/06/2022 4 3 4
02/07/2022 4 4 1

O df é classificado por data. Gostaria de criar uma coluna que me desse uma contagem de todas as linhas anteriores (data inferior) com todas as 3 colunas tendo um valor maior ou igual ao valor da linha atual. Ou em outras palavras:

Count rows
    where row_index < current_row_index &
          col1[row_index] >= col1[current_row_index] &
          col2[row_index] >= col2[current_row_index] &
          col3[row_index] >= col3[current_row_index]
)

O resultado deve ficar assim:

data col1 col2 col3 ge
01/01/2022 4 1 2 0
2022-01-02 4 2 3 0
07/01/2022 2 3 4 0
17/01/2022 2 4 4 0
02/03/2022 2 1 3 3
05/06/2022 3 3 2 0
07/06/2022 2 3 2 3
02/07/2022 1 4 1 1

Eu tentei várias combinações de ,,,, etc. shift, mas não consegui descobrir. Também não consegui encontrar uma pergunta semelhante o suficiente para adotar sua resposta com êxito. Existe uma maneira de conseguir isso usando Polars? Agradeço antecipadamente.qeoverwhencum_count

python

2 respostas

  1. Best Answer

    Essencialmente, a ideia é reunir todas as informações necessárias cumulative_evalem listas e, em seguida, explodir as listas resultantes em linhas e levar em consideração apenas as linhas onde todos os valores da lista são verdadeiros:

    import polars.selectors as cs

(
    df
    .with_columns(
        cs.starts_with('col').cumulative_eval(
            (pl.element() >= pl.element().last()).implode()
        ).name.prefix('_')
    )
    .explode(cs.starts_with('_col'))
    .with_columns(
        pl.all_horizontal(cs.starts_with('_col')).alias('ge')
    )
    .group_by('date','col1','col2','col3').agg(pl.col('ge').sum() - 1).sort('date')    
)

┌────────────┬──────┬──────┬──────┬─────┐
│ date       ┆ col1 ┆ col2 ┆ col3 ┆ ge  │
│ ---        ┆ ---  ┆ ---  ┆ ---  ┆ --- │
│ str        ┆ i64  ┆ i64  ┆ i64  ┆ u32 │
╞════════════╪══════╪══════╪══════╪═════╡
│ 2022-01-01 ┆ 4    ┆ 1    ┆ 2    ┆ 0   │
│ 2022-01-02 ┆ 4    ┆ 2    ┆ 3    ┆ 0   │
│ 2022-01-07 ┆ 2    ┆ 3    ┆ 4    ┆ 0   │
│ 2022-01-17 ┆ 2    ┆ 4    ┆ 4    ┆ 0   │
│ 2022-03-02 ┆ 2    ┆ 1    ┆ 3    ┆ 3   │
│ 2022-06-05 ┆ 3    ┆ 3    ┆ 2    ┆ 0   │
│ 2022-06-07 ┆ 2    ┆ 3    ┆ 2    ┆ 3   │
│ 2022-07-02 ┆ 1    ┆ 4    ┆ 1    ┆ 1   │
└────────────┴──────┴──────┴──────┴─────┘

    Você também pode usar duckdbpara fazer isso, graças à integração perfeita com polares :

    import duckdb

duckdb.sql("""
    select
        d.*,
        (
            select count(*)
            from df as tt
            where
                tt.date < d.date and
                tt.col1 >= d.col1 and
                tt.col2 >= d.col2 and
                tt.col3 >= d.col3
        ) as ge
    from df as d
""").pl()

┌────────────┬──────┬──────┬──────┬─────┐
│ date       ┆ col1 ┆ col2 ┆ col3 ┆ ge  │
│ ---        ┆ ---  ┆ ---  ┆ ---  ┆ --- │
│ str        ┆ i64  ┆ i64  ┆ i64  ┆ i64 │
╞════════════╪══════╪══════╪══════╪═════╡
│ 2022-01-01 ┆ 4    ┆ 1    ┆ 2    ┆ 0   │
│ 2022-01-02 ┆ 4    ┆ 2    ┆ 3    ┆ 0   │
│ 2022-01-07 ┆ 2    ┆ 3    ┆ 4    ┆ 0   │
│ 2022-01-17 ┆ 2    ┆ 4    ┆ 4    ┆ 0   │
│ 2022-03-02 ┆ 2    ┆ 1    ┆ 3    ┆ 3   │
│ 2022-06-05 ┆ 3    ┆ 3    ┆ 2    ┆ 0   │
│ 2022-06-07 ┆ 2    ┆ 3    ┆ 2    ┆ 3   │
│ 2022-07-02 ┆ 1    ┆ 4    ┆ 1    ┆ 1   │
└────────────┴──────┴──────┴──────┴─────┘
    • 1

  2. Você também pode usar uma struct para colocar toda a lógica condicional dentro.cumulative_eval()

    cols = "col1", "col2", "col3"

df.with_columns(ge =
   pl.struct(cols).cumulative_eval(
       pl.all_horizontal(
          pl.element().struct[col] >= pl.element().struct[col].last()
          for col in cols
       )
       .sum() - 1 # subtract 1 as we compare each row against itself
   )
)

    shape: (8, 5)
┌────────────┬──────┬──────┬──────┬─────┐
│ date       ┆ col1 ┆ col2 ┆ col3 ┆ ge  │
│ ---        ┆ ---  ┆ ---  ┆ ---  ┆ --- │
│ date       ┆ i64  ┆ i64  ┆ i64  ┆ u32 │
╞════════════╪══════╪══════╪══════╪═════╡
│ 2022-01-01 ┆ 4    ┆ 1    ┆ 2    ┆ 0   │
│ 2022-01-02 ┆ 4    ┆ 2    ┆ 3    ┆ 0   │
│ 2022-01-07 ┆ 2    ┆ 3    ┆ 4    ┆ 0   │
│ 2022-01-17 ┆ 2    ┆ 4    ┆ 4    ┆ 0   │
│ 2022-03-02 ┆ 2    ┆ 1    ┆ 3    ┆ 3   │
│ 2022-06-05 ┆ 3    ┆ 3    ┆ 2    ┆ 0   │
│ 2022-06-07 ┆ 2    ┆ 3    ┆ 2    ┆ 3   │
│ 2022-07-02 ┆ 1    ┆ 4    ┆ 1    ┆ 1   │
└────────────┴──────┴──────┴──────┴─────┘
    • 1

relate perguntas