Polars中when-then-otherwise条件替换列值出错及正确实现

初始DataFrame #

创建代码：

import polars as pl

df = pl.DataFrame({'r_num':['Yes', '', 'Yes'], 'pin': ['Yes','',''],'fin':['','','']})

对应的DataFrame结构：

shape: (3, 3)
┌───────┬─────┬─────┐
│ r_num ┆ pin ┆ fin │
│ ---   ┆ --- ┆ --- │
│ str   ┆ str ┆ str │
╞═══════╪═════╪═════╡
│ Yes   ┆ Yes ┆     │
│       ┆     ┆     │
│ Yes   ┆     ┆     │
└───────┴─────┴─────┘

需求说明 #

需要筛选出满足以下条件的行：

• r_num等于Yes
• pin等于Yes
• fin为空字符串

将这些行的r_num和pin字段设置为空字符串。

尝试的代码及错误结果 #

我试了这段代码：

df.with_columns(
    pl.when((pl.col('r_num')=='Yes') & (pl.col('pin')=='Yes') & (pl.col('fin') !='Yes'))
    .then(pl.col('r_num')=='')
    .otherwise(pl.col('r_num'))
)

得到错误结果：

shape: (3, 3)
┌───────┬─────┬─────┐
│ r_num ┆ pin ┆ fin │
│ ---   ┆ --- ┆ --- │
│ str   ┆ str ┆ str │
╞═══════╪═════╪═════╡
│ false ┆ Yes ┆     │
│       ┆     ┆     │
│ Yes   ┆     ┆     │
└───────┴─────┴─────┘

为什么r_num会被填充成false？

另外，我在Pandas里是这么实现的，能得到正确结果：

df_pd = df.to_pandas()
df_pd.loc[(df_pd['r_num']=='Yes') & (df_pd['pin']=='Yes') & (df_pd['fin']!='Yes'),['r_num','pin']] = ''

期望的正确结果：

shape: (3, 3)
┌───────┬─────┬─────┐
│ r_num ┆ pin ┆ fin │
│ ---   ┆ --- ┆ --- │
│ str   ┆ str ┆ str │
╞═══════╪═════╪═════╡
│       ┆     ┆     │
│       ┆     ┆     │
│ Yes   ┆     ┆     │
└───────┴─────┴─────┘

错误原因 #

你写的.then(pl.col('r_num')=='')是一个布尔比较表达式，它会判断r_num是否等于空字符串，返回的是布尔值（True或False），而不是把字段设置为空字符串。第一行的r_num是Yes，和空字符串比较结果是False，所以就被填充成了false。

另外，你条件里的pl.col('fin') !='Yes'其实不符合需求，应该写成pl.col('fin') == ''来匹配空字符串的条件。

正确的Polars实现方法 #

方法1：单独更新每个字段 #

直接在then里返回空字符串，而非做比较：

df = df.with_columns(
    # 更新r_num
    pl.when((pl.col('r_num') == 'Yes') & (pl.col('pin') == 'Yes') & (pl.col('fin') == ''))
    .then('')
    .otherwise(pl.col('r_num'))
    .alias('r_num'),
    # 更新pin
    pl.when((pl.col('r_num') == 'Yes') & (pl.col('pin') == 'Yes') & (pl.col('fin') == ''))
    .then('')
    .otherwise(pl.col('pin'))
    .alias('pin')
)

方法2：批量处理多个字段（更简洁） #

用列表推导式批量处理需要更新的字段，避免重复写条件：

condition = (pl.col('r_num') == 'Yes') & (pl.col('pin') == 'Yes') & (pl.col('fin') == '')

df = df.with_columns(
    [
        pl.when(condition).then('').otherwise(pl.col(col)).alias(col)
        for col in ['r_num', 'pin']
    ]
)

方法3：使用where反向筛选 #

利用where方法，当条件不满足时保留原字段，满足时替换为空：

condition = (pl.col('r_num') == 'Yes') & (pl.col('pin') == 'Yes') & (pl.col('fin') == '')

df = df.with_columns(
    pl.col(['r_num', 'pin']).where(~condition, '')
)

运行以上任意一种方法，都能得到你期望的结果。

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