Polars是否有替代numpy select的纯Polars条件赋值实现？

完全可以不依赖Numpy，用Polars原生语法实现条件赋值操作，效果和你用Numpy的np.select完全一致。

首先看原始的DataFrame：

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]})

输出结果：

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的实现方式：

import numpy as np

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    │
└──────┴────────────┴────────┴──────────┴────────┘

Polars原生实现方法 #

用Polars的when-then-otherwise链式语法，完全替代Numpy的select：

df.with_columns(
    rating=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")
)

这段代码会得到和Numpy版本完全相同的输出结果，全程不需要依赖Numpy，完全利用Polars的向量化操作，性能更贴合Polars的优化逻辑。

如果有大量条件需要处理，也可以把条件和选项打包成列表，用循环构建链式调用：

conditions_choices = [
    ((pl.col("points") <= 6) & (pl.col("rebounds") <= 9), "Bad"),
    ((pl.col("points") > 10) & (pl.col("rebounds") > 6), "Good")
]

# 构建when-then链
expr = pl.when(conditions_choices[0][0]).then(conditions_choices[0][1])
for cond, choice in conditions_choices[1:]:
    expr = expr.when(cond).then(choice)
expr = expr.otherwise("Aveg")

df.with_columns(rating=expr)

这种方式适合条件较多的场景，代码更易维护。

内容的提问来源于stack exchange，提问作者codedancer