Polars中如何按索引列表提取列表列元素？

直接解决方案 #

使用Polars内置的list.take方法，它支持接受列表类型的索引列，直接提取对应位置的元素，完全符合Polars的矢量化惯用写法：

import polars as pl

df = pl.DataFrame({
    'idxs': [[0], [1], [0, 2]], 
     'prices': [[0.0, 3.5], [4.6, 0.0], [0.0, 7.8, 0.0]]
})

result = df.with_columns(
    pl.col('prices').list.take(pl.col('idxs')).alias('zero_prices')
)

print(result)

执行结果：

shape: (3, 3)
┌───────────┬─────────────────┬─────────────┐
│ idxs      ┆ prices          ┆ zero_prices │
│ ---       ┆ ---             ┆ ---         │
│ list[i64] ┆ list[f64]       ┆ list[f64]   │
╞═══════════╪═════════════════╪═════════════╡
│ [0]       ┆ [0.0, 3.5]      ┆ [0.0]       │
│ [1]       ┆ [4.6, 0.0]      ┆ [0.0]       │
│ [0, 2]    ┆ [0.0, 7.8, 0.0] ┆ [0.0, 0.0]  │
└───────────┴─────────────────┴─────────────┘

原方法报错原因 #

你之前使用list.get报错，是因为list.get的设计目的是提取单个元素，它要求传入的索引是标量（单个数值）或者返回标量的表达式，而你传入的是idxs列（列表类型），类型不匹配，因此抛出cannot cast List type to Int64的错误。

与自定义UDF方案的对比 #

你用map_elements结合自定义函数的方案虽然能实现功能，但它属于逐行的Python-level循环，在处理大数据集时性能远不如Polars内置的矢量化方法（比如list.take）。内置方法是在Rust层面执行的，避免了Python与Rust之间的上下文切换，效率提升非常明显。

内容的提问来源于stack exchange，提问作者Some lazy dude