如何基于列名是否为None使用polars.when实现指定功能？

核心问题原因 #

你遇到的TypeError是因为直接在Polars表达式中传入None作为列名时，Polars无法将其识别为合法的列标识符；返回形状为(1,1)则是因为未生成匹配DataFrame行数的序列，仅创建了单行常量值。

一行表达式实现代码 #

import polars as pl

def get_series(df: pl.DataFrame, col_name: str | None = None, default_value: int = 0) -> pl.Series:
    return df.select(pl.col(col_name) if col_name is not None else pl.lit(default_value).repeat(df.height)).to_series()

代码逻辑说明 #

• 当col_name不为None时，pl.col(col_name)直接选取DataFrame中对应列，自然匹配原数据行数
• 当col_name为None时，pl.lit(default_value).repeat(df.height)生成与原DataFrame行数完全一致的默认值序列，df.height获取数据行数确保长度匹配
• 最后通过to_series()将单列DataFrame转换为Polars Series类型

测试验证示例 #

# 测试用DataFrame
df = pl.DataFrame({"x": [10, 20, 30], "y": [40, 50, 60]})

# 测试指定列名的情况
print(get_series(df, "x"))
# 输出：
# shape: (3,)
# Series: 'x' [i64]
# [
#  10
#  20
#  30
# ]

# 测试列名为None的情况
print(get_series(df, None, 99))
# 输出：
# shape: (3,)
# Series: 'literal' [i64]
# [
#  99
#  99
#  99
# ]

原方案问题解析 #

polars.when是用于列内的逐行条件判断，并非用来切换整个列的选择逻辑。直接传入None作为列名会触发Polars的类型校验错误；若仅使用pl.lit(default_value)而未调用repeat，会生成单行常量，导致返回结果行数不匹配。

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