如何在Polars DataFrame中根据指定列名查找对应行值

原始数据 #

import polars as pl

df = pl.DataFrame({
    'a': [1, 2, 1],
    'b': [2, 1, 2],
    'c': [3, 3, 2],
    'column': ['a', 'c', 'b'],
})

原始数据结构：

shape: (3, 4)
┌─────┬─────┬─────┬────────┐
│ a   ┆ b   ┆ c   ┆ column │
│ --- ┆ --- ┆ --- ┆ ---    │
│ i64 ┆ i64 ┆ i64 ┆ str    │
╞═════╪═════╪═════╪════════╡
│ 1   ┆ 2   ┆ 3   ┆ a      │
│ 2   ┆ 1   ┆ 3   ┆ c      │
│ 1   ┆ 2   ┆ 2   ┆ b      │
└─────┴─────┴─────┴────────┘

需求 #

新增一列lookup，每行取值为该行中column列指定列名对应的数值，最终结果如下：

shape: (3, 5)
┌─────┬─────┬─────┬────────┬────────┐
│ a   ┆ b   ┆ c   ┆ column ┆ lookup │
│ --- ┆ --- ┆ --- ┆ ---    ┆ ---    │
│ i64 ┆ i64 ┆ i64 ┆ str    ┆ i64    │
╞═════╪═════╪═════╪════════╪════════╡
│ 1   ┆ 2   ┆ 3   ┆ a      ┆ 1      │
│ 2   ┆ 1   ┆ 3   ┆ c      ┆ 3      │
│ 1   ┆ 2   ┆ 2   ┆ b      ┆ 2      │
└─────┴─────┴─────┴────────┴────────┘

解决方案 #

使用Polars的row()方法将每行转换为结构体，再通过struct.field()动态提取对应列的值：

df = df.with_columns(
    lookup=pl.row(by_name=True).struct.field(pl.col("column"))
)

说明 #

• pl.row(by_name=True)：将每行数据转换为以列名为键的结构体
• struct.field(pl.col("column"))：根据column列的动态值，提取结构体中对应字段的内容

这种方法属于行内操作，无需数据重塑或表连接，性能高效，适合处理大规模数据集。

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