如何用Polars原生表达式API替代map_elements生成序列列表列？

我们需要生成一个新列，其中每个元素是包含从1到Qty列对应数值的所有整数的列表。已知map_elements的实现速度远慢于Polars原生表达式API，因此需要改用原生API完成需求。

输入DataFrame：

shape: (3, 1)
┌─────┐
│ Qty │
│ --- │
│ i64 │
╞═════╡
│ 1   │
│ 2   │
│ 3   │
└─────┘

期望输出DataFrame：

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

原map_elements实现 #

import polars as pl

# 创建DataFrame
df = pl.DataFrame({'Qty': [1, 2, 3]})

# 添加List列
df = df.with_columns(
    pl.col('Qty').map_elements(lambda x: [i+1 for i in range(x)]).alias('List')
)

原生表达式优化实现 #

Polars提供的pl.int_range是原生矢量化函数，能直接生成符合要求的整数列表，性能远优于逐行处理的map_elements。

import polars as pl

df = pl.DataFrame({'Qty': [1, 2, 3]})

df = df.with_columns(
    pl.int_range(1, pl.col('Qty') + 1).alias('List')
)

代码说明 #

• pl.int_range(start, end)会生成从start到end-1的整数序列，因此将end设为pl.col('Qty') + 1，就能精准生成从1到Qty的完整整数列表。
• 该操作是Polars原生支持的矢量化计算，在处理大数据集时，性能优势会非常明显。

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