Polars 数据框中的条件分配

我想知道是否有一种方法可以在不使用 numpy 相关的情况下处理 Polars 数据框中的条件分配。

import polars as pl
import numpy as np

df = pl.DataFrame({'team': ['A', 'A', 'A', 'B', 'B', 'C'],
                   'conference': ['East', 'East', 'East', 'West', 'West', 'East'],
                   'points': [11, 8, 10, 6, 6, 5],
                   'rebounds': [7, 7, 6, 9, 12, 8]})

shape: (6, 4)
┌──────┬────────────┬────────┬──────────┐
│ team ┆ conference ┆ points ┆ rebounds │
│ ---  ┆ ---        ┆ ---    ┆ ---      │
│ str  ┆ str        ┆ i64    ┆ i64      │
╞══════╪════════════╪════════╪══════════╡
│ A    ┆ East       ┆ 11     ┆ 7        │
│ A    ┆ East       ┆ 8      ┆ 7        │
│ A    ┆ East       ┆ 10     ┆ 6        │
│ B    ┆ West       ┆ 6      ┆ 9        │
│ B    ┆ West       ┆ 6      ┆ 12       │
│ C    ┆ East       ┆ 5      ┆ 8        │
└──────┴────────────┴────────┴──────────┘

使用 numpy，我们可以这样做：

conditions = [
    df['points'].le(6) & df['rebounds'].le(9),
    df['points'].gt(10) & df['rebounds'].gt(6)
]
choicelist = ['Bad','Good']

df.with_columns(rating = np.select(conditions, choicelist, 'Aveg'))

shape: (6, 5)
┌──────┬────────────┬────────┬──────────┬────────┐
│ team ┆ conference ┆ points ┆ rebounds ┆ rating │
│ ---  ┆ ---        ┆ ---    ┆ ---      ┆ ---    │
│ str  ┆ str        ┆ i64    ┆ i64      ┆ str    │
╞══════╪════════════╪════════╪══════════╪════════╡
│ A    ┆ East       ┆ 11     ┆ 7        ┆ Good   │
│ A    ┆ East       ┆ 8      ┆ 7        ┆ Aveg   │
│ A    ┆ East       ┆ 10     ┆ 6        ┆ Aveg   │
│ B    ┆ West       ┆ 6      ┆ 9        ┆ Bad    │
│ B    ┆ West       ┆ 6      ┆ 12       ┆ Aveg   │
│ C    ┆ East       ┆ 5      ┆ 8        ┆ Bad    │
└──────┴────────────┴────────┴──────────┴────────┘

可以使用 when.then.otherwise 语句来使用 Polars 执行条件分配。以下是实现请求的代码示例：

import polars as pl

df = pl.DataFrame({'team': ['A', 'A', 'A', 'B', 'B', 'C'],
                   'conference': ['East', 'East', 'East', 'West', 'West', 'East'],
                   'points': [11, 8, 10, 6, 6, 5],
                   'rebounds': [7, 7, 6, 9, 12, 8]})

df = df.with_column(
    pl.when((pl.col("points") <= 6) & (pl.col("rebounds") <= 9))
    .then("Bad")
    .when((pl.col("points") > 10) & (pl.col("rebounds") > 6))
    .then("Good")
    .otherwise("Aveg")
    .alias("rating")
)

print(df)

此代码通过执行以下操作来工作：

1. with_column ： 这用于基于现有列添加一个新列。
2. pl.when.then.otherwise ： 此结构用于基于条件逻辑创建新列“评级”。
3. pl.when((pl.col("points") <= 6) & (pl.col("rebounds") <= 9)) ： 这定义了第一个条件，其中如果“分数”小于或等于 6 且“篮板”小于或等于 9，则条件为真。
4. .then("Bad") ： 如果第一个条件为真，则在“评级”列中分配值“差”。
5. .when((pl.col("points") > 10) & (pl.col("rebounds") > 6)) ： 这定义了第二个条件，其中如果“分数”大于 10 且“篮板”大于 6，则条件为真。
6. .then("Good") ： 如果第二个条件为真，则在“评级”列中分配值“好”。
7. .otherwise("Aveg") ： 如果前两个条件都不为真，则为“评级”列分配默认值“平均”。
8. .alias("rating") ： 这会将结果表达式命名为“评级”。

此方法利用了 Polars 的表达语言，为操作数据框提供了一种简洁且高效的方式。