Polars中如何用列值作为列名在when-then表达式生成新列？

问题描述 #

在Polars DataFrame中，存在一列id_column_name，其值为DataFrame中其他列的名称（如id_column1、id_column2）。需要根据id_column_name的取值，从对应列中获取值，生成新列id，每行的id值可来自DataFrame中不同的列。

初始DataFrame代码 #

# 初始DataFrame
df = pl.DataFrame({
        'id_column1': [123, 456],
        'id_column2': ['abc', 'def'],
        'id_column_name': ['id_column1', 'id_column2']
})

期望输出 #

# 目标输出DataFrame
df_out = pl.DataFrame({
    'id_column1': [123, 456],
    'id_column2': ['abc', 'def'],
    'id_column_name': ['id_column1', 'id_column2'],
    'id': ['123', 'def'] 
})

尝试的代码及报错 #

尝试用when-then结合pl.col()动态取列，但报错：

# 尝试的代码
df = df.with_columns(
    pl.when(pl.col('id_column_name').is_not_null())
    .then(pl.col(pl.col('id_column_name')))
    .otherwise(None)
    .cast(pl.String)
    .alias('id')
)

报错信息：

Expected str or DataType, got 'Expr'.

尝试用str()或.str转换表达式为字符串时，出现另一个错误：

Expected str or DataType, got 'ExprStringNameSpace'.

解决方案 #

问题出在pl.col()无法直接接受表达式作为参数，它需要传入字符串类型的列名。我们可以通过将目标列打包成Struct，再动态索引Struct的方式实现需求：

方法1：指定目标列（适用于列数少的情况） #

df = df.with_columns(
    # 将需要提取的列打包成Struct
    pl.struct(['id_column1', 'id_column2'])
    # 根据id_column_name的值动态获取Struct中的对应字段
    .get(pl.col('id_column_name'))
    .cast(pl.String)
    .alias('id')
)

方法2：自动匹配所有其他列（通用方案） #

如果需要提取的列较多，可以自动排除id_column_name列，打包剩余所有列：

df = df.with_columns(
    pl.struct(pl.all().exclude('id_column_name'))
    .get(pl.col('id_column_name'))
    .cast(pl.String)
    .alias('id')
)

执行上述代码后，即可得到符合预期的id列。

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