如何在Polars中利用两个列表列实现类似numpy.take的取元素操作？

直接使用Polars内置的gather方法即可满足需求，无需依赖性能低下的map_elements，也能避开list.eval无法跨列引用的限制：

import polars as pl

df = pl.DataFrame({"0": [[1,0], [0, 1]], "1": [[55, 101], [3, 9]]})

result = df.with_columns(
    wanted=pl.col("1").gather(pl.col("0"))
)

print(result)

代码说明 #

• pl.col("1").gather(pl.col("0"))：对每行的"1"列列表，使用对应行"0"列的索引列表提取元素，完全等效于numpy.take的逻辑
• 该操作是Polars的向量化实现，性能远高于map_elements这类逐行处理的方式

输出结果 #

shape: (2, 3)
┌───────────┬───────────┬───────────┐
│ 0         ┆ 1         ┆ wanted    │
│ ---       ┆ ---       ┆ ---       │
│ list[i64] ┆ list[i64] ┆ list[i64] │
╞═══════════╪═══════════╪═══════════╡
│ [1, 0]    ┆ [55, 101] ┆ [101, 55] │
│ [0, 1]    ┆ [3, 9]    ┆ [3, 9]    │
└───────────┴───────────┴───────────┘

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