Polars如何基于另一列的列名动态选取对应列的值？

现有两个Polars DataFrame结构如下：

import polars as pl

dfA = pl.DataFrame(
    {
       "year": [2000, 2001, 2002, 2003],
       "comp": ["1A", "1A", "1B", "1C"]
    }
)

dfB = pl.DataFrame(
   {
       "year": [2000, 2001, 2002, 2003],
       "1A": [10, 12, 14, 12],
       "1B": [20, 30, 40, 20],
       "1C": [1, 2, 3, 4]
    }
)

需要在dfA中新增一列N，取值规则为：匹配dfA与dfB的year，同时选取dfB中列名等于dfA的comp列值的单元格数据。预期输出如下：

┌──────┬──────┬─────┐
│ year ┆ comp ┆ N   │
│ ---  ┆ ---  ┆ --- │
│ i64  ┆ str  ┆ i64 │
╞══════╪══════╪═════╡
│ 2000 ┆ 1A   ┆ 10  │
│ 2001 ┆ 1A   ┆ 12  │
│ 2002 ┆ 1B   ┆ 40  │
│ 2003 ┆ 1C   ┆ 4   │
└──────┴──────┴─────┘

尝试过按year左连接后通过comp列选取目标列，但无法将pl.col("comp")转换为外层pl.col()可识别的有效字符串；且因1A、1B这类列数量较多，无法使用pl.when().then().otherwise()方案，寻求可行解决方法。

推荐将dfB重塑为**长格式（long format）**后再与dfA连接，这是Polars中最符合向量化操作逻辑的高效方案：

# 将dfB从宽表转成窄表
dfB_melted = dfB.melt(
    id_vars="year",  # 保留year作为关联标识列
    variable_name="comp",  # 原列名映射为comp列
    value_name="N"  # 原列值映射为目标N列
)

# 与dfA按year和comp做左连接
result = dfA.join(dfB_melted, on=["year", "comp"], how="left")
print(result)

方案说明 #

• melt重塑表结构：把dfB中以公司代码命名的列（1A/1B/1C）转换成comp列，对应的值存入N列，让dfB_melted的结构与dfA完全匹配，直接通过year+comp即可精准关联目标值。
• 高效向量化操作：melt和join都是Polars原生的向量化操作，处理大数据集时性能远高于逐行判断方案，且无需关心comp的具体取值数量。

如果需要保留原连接后的宽表结构再取值，也可以用逐行映射的方式（性能略低于上述方案）：

result = dfA.join(dfB, on="year", how="left").with_columns(
    pl.row(return_dtype=pl.Int64).map_elements(lambda x: x[x["comp"]], return_dtype=pl.Int64).alias("N")
).drop(dfB.columns[1:])  # 移除不需要的宽表列

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