如何在Polars中根据另一列的索引提取列表列中的元素

要实现根据第二列的索引值提取第一列列表对应位置的元素并新增为第三列，直接使用Polars的表达式即可完成，无需手动构造Series。

核心解法 #

使用pl.col("lst").get(pl.col("ind"))表达式，该表达式会逐行匹配：取当前行lst列的列表，用同一行ind列的值作为索引提取对应元素，再通过with_columns将结果作为新列加入DataFrame。

完整代码示例 #

import polars as pl

# 输入DataFrame
df = pl.DataFrame({
    "lst": [[1, 2, 3], [4, 5, 6]], 
    "ind": [1, 2]
})

# 提取元素并新增列
res = df.with_columns(
    pl.col("lst").get(pl.col("ind")).alias("list[ind]")
)

print(res)

输出结果 #

┌───────────┬─────┬───────────┐
│ lst       ┆ ind ┆ list[ind] │
│ ---       ┆ --- ┆ ---       │
│ list[i64] ┆ i64 ┆ i64       │
╞═══════════╪═════╪═══════════╡
│ [1, 2, 3] ┆ 1   ┆ 2         │
│ [4, 5, 6] ┆ 2   ┆ 6         │
└───────────┴─────┴───────────┘

补充说明 #

• 如果ind列的索引值超出对应lst列列表的长度，Polars会返回null，不会抛出索引越界错误。
• 该方法是向量化操作，性能远优于逐行遍历处理。

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