如何在Polars中为满足布尔条件的行用列表值创建新列？

我想在Polars DataFrame中创建一个新列，仅在满足特定条件的行中用现有列表填充值，同时保持总行数不变。

示例数据如下：

import polars as pl

df = pl.DataFrame({
    'cond1': [-1, 0, 1, 2, 2, 3, 4, 5],
    'cond2': ['No', 'No', 'No', 'Yes','Yes', 'Yes', 'Yes', 'Yes']
})

我知道在Pandas中可以这样实现（列表长度与筛选后的行数一致，output为新列）：

df_pd = df.to_pandas()

conditions = ((df_pd['cond1']>=1) & (df_pd['cond2'] == 'Yes'))
list_new = [1,2,3,4,5]

df_pd.loc[conditions, 'output'] = list_new

执行后结果：

#    cond1 cond2  output
# 0     -1    No     NaN
# 1      0    No     NaN
# 2      1    No     NaN
# 3      2   Yes     1.0
# 4      2   Yes     2.0
# 5      3   Yes     3.0
# 6      4   Yes     4.0
# 7      5   Yes     5.0

请问如何在Polars中实现相同功能？我尝试过以下代码的多种变体但均未成功：

df = df.with_columns(pl.when(conditions).then(list_new).alias('output'))

解决方案 #

Polars中不能直接在pl.when(...).then(...)里传入Python列表（then要求接收列表达式，而非逐行赋值的列表），可以通过以下两种方式实现需求：

方法一：索引匹配+左连接 #

先筛选出符合条件的行并赋值，再通过原表索引与筛选表左连接，保留所有行：

import polars as pl

df = pl.DataFrame({
    'cond1': [-1, 0, 1, 2, 2, 3, 4, 5],
    'cond2': ['No', 'No', 'No', 'Yes','Yes', 'Yes', 'Yes', 'Yes']
})

list_new = [1,2,3,4,5]
cond = (pl.col('cond1') >= 1) & (pl.col('cond2') == 'Yes')

# 给符合条件的行添加output列，并保留原索引
filtered_df = df.filter(cond).with_columns(
    pl.Series(list_new).alias('output')
).with_row_index('idx')

# 原表添加索引后左连接，再移除索引列
result = df.with_row_index('idx').join(filtered_df, on='idx', how='left').drop('idx')
print(result)

方法二：直接操作Series替换指定位置 #

先创建与原表长度一致的空Series，再替换符合条件的位置的值：

import polars as pl

df = pl.DataFrame({
    'cond1': [-1, 0, 1, 2, 2, 3, 4, 5],
    'cond2': ['No', 'No', 'No', 'Yes','Yes', 'Yes', 'Yes', 'Yes']
})

list_new = [1,2,3,4,5]
cond = (pl.col('cond1') >= 1) & (pl.col('cond2') == 'Yes')

# 创建全为None的空Series，类型与列表匹配
output_series = pl.Series([None]*len(df), dtype=pl.Int64)
# 获取符合条件的行索引
target_indices = df.select(pl.int_range(0, pl.count()).filter(cond)).to_series().to_list()
# 替换指定索引位置的值
output_series = output_series.set(target_indices, list_new)

# 添加新列到原表
result = df.with_columns(output_series.alias('output'))
print(result)

两种方法执行后都会得到与Pandas示例一致的结果：

shape: (8, 3)
┌───────┬───────┬────────┐
│ cond1 ┆ cond2 ┆ output │
│ ---   ┆ ---   ┆ ---    │
│ i64   ┆ str   ┆ i64    │
╞═══════╪═══════╪════════╡
│ -1    ┆ No    ┆ null   │
│ 0     ┆ No    ┆ null   │
│ 1     ┆ No    ┆ null   │
│ 2     ┆ Yes   ┆ 1      │
│ 2     ┆ Yes   ┆ 2      │
│ 3     ┆ Yes   ┆ 3      │
│ 4     ┆ Yes   ┆ 4      │
│ 5     ┆ Yes   ┆ 5      │
└───────┴───────┴────────┘

原代码失败原因 #

你之前尝试的pl.when(conditions).then(list_new)无法工作，是因为Polars的then方法要求传入列表达式（比如pl.col()、pl.Series()等），而不是直接传入Python列表。直接传列表会导致Polars把整个列表作为单一值，赋值给所有符合条件的行，而非逐行分配列表中的元素，因此达不到预期效果。

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