Início / coding / Perguntas / 79077774
Accepted
Kazdegotepu
Asked: 2024-10-11 18:13:29 +0800 CST

Polars.col() aninhado [duplicado]

  • 772

Gostaria de obter valor de uma coluna cujo nome está em outra coluna.

Na minha tabela há muitas colunas, então fazer isso pl.when().then()não é uma opção.

Como entrada temos este dataframe:

foo contagem_de_foo bar contagem_de_barras base contagem_base maior
1 23 4 43 5 64 base
2 45 6 45 1 43 bar
3 234 9 453 15 231 base
4 55 2 67 3 94 foo

e gostaria de transformar isso com with_columns():

foo contagem_de_foo bar contagem_de_barras base contagem_base maior maior_contagem
1 23 4 43 5 64 base 64
2 45 6 45 1 43 bar 45
3 234 9 453 15 231 base 231
4 55 2 67 3 94 foo 4

Este (pseudo)código ilustra o que tenho em mente (mas é claro que isso não funciona)

df = pl.DataFrame({"foo" : [1, 2, 3, 4],
                   "foo_count" : [23, 45 ,234, 55],
                   "bar" : [4 ,6 ,9, 2],
                   "bar_count" : [43, 45 ,453, 67],
                   "baz": [5,1,15, 3],
                   "baz_count" : [64, 43 ,231, 94],
                   "largest" : ["baz", "bar", "baz", "foo"]})

df.with_columns(
    pl.col(f"{pl.col('largest')}_count").alias("largest_count")
)
python-polars

2 respostas

  1. Best Answer

    Aqui está uma abordagem simples. Primeiro, criamos expressões para colunas mascaradas usando um gerador de expressões when-then . Então, confiamos em pl.coalescepara combinar as colunas mascaradas. Ainda há um loop python ingênuo sobre os elementos exclusivos em largest .

    df.with_columns(
    pl.coalesce(
        pl.when(pl.col("largest") == col).then(f"{col}_count")
        for col in df.get_column("largest").unique()
    ).alias("largest_count")
)

    shape: (4, 8)
┌─────┬───────────┬─────┬───────────┬─────┬───────────┬─────────┬───────────────┐
│ foo ┆ foo_count ┆ bar ┆ bar_count ┆ baz ┆ baz_count ┆ largest ┆ largest_count │
│ --- ┆ ---       ┆ --- ┆ ---       ┆ --- ┆ ---       ┆ ---     ┆ ---           │
│ i64 ┆ i64       ┆ i64 ┆ i64       ┆ i64 ┆ i64       ┆ str     ┆ i64           │
╞═════╪═══════════╪═════╪═══════════╪═════╪═══════════╪═════════╪═══════════════╡
│ 1   ┆ 23        ┆ 4   ┆ 43        ┆ 5   ┆ 64        ┆ baz     ┆ 64            │
│ 2   ┆ 45        ┆ 6   ┆ 45        ┆ 1   ┆ 43        ┆ bar     ┆ 45            │
│ 3   ┆ 234       ┆ 9   ┆ 453       ┆ 15  ┆ 231       ┆ baz     ┆ 231           │
│ 4   ┆ 55        ┆ 2   ┆ 67        ┆ 3   ┆ 94        ┆ foo     ┆ 55            │
└─────┴───────────┴─────┴───────────┴─────┴───────────┴─────────┴───────────────┘
    • 3

  2. Aqui está uma abordagem que escala mesmo com uma quantidade muito grande de colunas. Podemos extrair apenas as colunas de contagem:

    counts = df.select(pl.col("^.*_count$").name.map(lambda c: c.removesuffix("_count")))

    shape: (4, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ baz │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 23  ┆ 43  ┆ 64  │
│ 45  ┆ 45  ┆ 43  │
│ 234 ┆ 453 ┆ 231 │
│ 55  ┆ 67  ┆ 94  │
└─────┴─────┴─────┘

    Em seguida, crie uma tabela normalizada com um índice:

    norm_counts = counts.with_row_index().unpivot(index="index")

    shape: (12, 3)
┌───────┬──────────┬───────┐
│ index ┆ variable ┆ value │
│ ---   ┆ ---      ┆ ---   │
│ u32   ┆ str      ┆ i64   │
╞═══════╪══════════╪═══════╡
│ 0     ┆ foo      ┆ 23    │
│ 1     ┆ foo      ┆ 45    │
│ 2     ┆ foo      ┆ 234   │
│ 3     ┆ foo      ┆ 55    │
│ 0     ┆ bar      ┆ 43    │
│ …     ┆ …        ┆ …     │
│ 3     ┆ bar      ┆ 67    │
│ 0     ┆ baz      ┆ 64    │
│ 1     ┆ baz      ┆ 43    │
│ 2     ┆ baz      ┆ 231   │
│ 3     ┆ baz      ┆ 94    │
└───────┴──────────┴───────┘

    Então podemos juntar as contagens normalizadas ao dataframe original usando um índice temporário:

    out = (
    df
    .with_row_index()
    .join(norm_counts,
          left_on=["index", "largest"],
          right_on=["index", "variable"],
          how="left")
    .rename({"value": "largest_count"})
    .drop("index")
)

    shape: (4, 8)
┌─────┬───────────┬─────┬───────────┬─────┬───────────┬─────────┬───────────────┐
│ foo ┆ foo_count ┆ bar ┆ bar_count ┆ baz ┆ baz_count ┆ largest ┆ largest_count │
│ --- ┆ ---       ┆ --- ┆ ---       ┆ --- ┆ ---       ┆ ---     ┆ ---           │
│ i64 ┆ i64       ┆ i64 ┆ i64       ┆ i64 ┆ i64       ┆ str     ┆ i64           │
╞═════╪═══════════╪═════╪═══════════╪═════╪═══════════╪═════════╪═══════════════╡
│ 1   ┆ 23        ┆ 4   ┆ 43        ┆ 5   ┆ 64        ┆ baz     ┆ 64            │
│ 2   ┆ 45        ┆ 6   ┆ 45        ┆ 1   ┆ 43        ┆ bar     ┆ 45            │
│ 3   ┆ 234       ┆ 9   ┆ 453       ┆ 15  ┆ 231       ┆ baz     ┆ 231           │
│ 4   ┆ 55        ┆ 2   ┆ 67        ┆ 3   ┆ 94        ┆ foo     ┆ 55            │
└─────┴───────────┴─────┴───────────┴─────┴───────────┴─────────┴───────────────┘
    • 1

relate perguntas