Polars中基于条件添加列的错误排查与最佳实践

问题场景 #

给定如下Polars DataFrame：

df = pl.DataFrame({
    'a': [0.3, 0.7, 0.5, 0.1, 0.9]
})

需要添加新列new，规则为：当a的值大于阈值（如0.5）时赋值0，否则赋值1。在Pandas中可通过np.where实现：

import numpy as np

THRESHOLD = 0.5
df['new'] = np.where(df.a > THRESHOLD, 0, 1)

用户尝试了两种Polars写法：

1. 依赖NumPy的可运行写法（非最佳实践）：

df = df.with_columns(
    pl.lit(np.where(df.select('a').to_numpy() > THRESHOLD, 0, 1).ravel())
    .alias('new')
)

2. 触发错误的写法：

df = df.with_columns(
    pl.lit(df.filter(pl.col('a') > THRESHOLD).select([0, 1]))
    .alias('new')
)

错误信息：

DuplicateError                            Traceback (most recent call last)
Cell In[47], line 5
      1 THRESHOLD = 0.5
      2 DELAY_TOLERANCE = 10
      4 df = df.with_columns(
----> 5     pl.lit(df.filter(pl.col('a') > THRESHOLD).select([0, 1]))
      6     .alias('new')
      7 )
      8 df.head()

DuplicateError: column with name 'literal' has more than one occurrences

用户疑问：

1. 第二种写法错误的原因是什么？
2. Polars中进行此类条件赋值的最佳实践是什么？

错误原因分析 #

第二种写法出错的核心问题：

• df.filter(pl.col('a') > THRESHOLD).select([0,1])返回的是包含两行的Polars DataFrame，而非一维数组或单列Series，长度和原DataFrame不匹配，无法对应到原表的每一行。
• 用pl.lit()包裹这个多列DataFrame时，Polars会尝试将其作为字面量插入，过程中会产生重复的literal列名，最终触发DuplicateError。

Polars条件赋值的最佳实践 #

Polars提供原生条件表达式支持，无需依赖NumPy，推荐两种写法：

方式1：使用pl.when().then().otherwise()（直观对应np.where） #

这是Polars处理条件逻辑的标准写法，语义清晰：

THRESHOLD = 0.5
df = df.with_columns(
    pl.when(pl.col('a') > THRESHOLD)
      .then(0)
      .otherwise(1)
      .alias('new')
)

方式2：布尔值转整数（更简洁） #

Polars的布尔列可直接通过cast转为整数（True→1，False→0），结合取反操作实现需求：

THRESHOLD = 0.5
df = df.with_columns(
    (pl.col('a') > THRESHOLD).not_().cast(pl.Int32).alias('new')
)

解释：pl.col('a') > THRESHOLD生成布尔列（大于阈值为True），not_()取反后，转整数就得到0和1的结果，完全匹配需求。

最终结果 #

两种写法运行后，得到的DataFrame均为：

shape: (5, 2)
┌─────┬─────┐
│ a   ┆ new │
│ --- ┆ --- │
│ f64 ┆ i32 │
╞═════╪═════╡
│ 0.3 ┆ 1   │
│ 0.7 ┆ 0   │
│ 0.5 ┆ 1   │
│ 0.1 ┆ 1   │
│ 0.9 ┆ 0   │
└─────┴─────┘

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