如何在Polars数据框中生成某列的累积列表新列

要给现有数据框新增c4列，让该列每个元素成为c1列的累积列表，可使用Polars的cumulative_eval方法配合collect()表达式实现，代码如下：

import polars as pl

df = pl.from_repr("""
┌─────┬─────┬─────┐
│ c1  ┆ c2  ┆ c3  │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 2   ┆ 3   │
│ 3   ┆ 4   ┆ 5   │
│ 5   ┆ 6   ┆ 7   │
│ 7   ┆ 8   ┆ 9   │
└─────┴─────┴─────┘
""")

# 新增累积列表列c4
result_df = df.with_columns(
    c4=pl.col("c1").cumulative_eval(pl.collect())
)

# 查看c4列结果
print(result_df.select("c4"))

执行后c4列的输出结果如下：

shape: (4, 1)
┌──────────────┐
│ c4           │
│ ---          │
│ list[i64]    │
╞══════════════╡
│ [1]          │
│ [1, 3]       │
│ [1, 3, 5]    │
│ [1, 3, 5, 7] │
└──────────────┘

代码说明 #

• cumulative_eval会遍历c1列的每个位置，对从列开头到当前位置的所有元素组成的子集执行传入的表达式。
• pl.collect()的作用是将当前子集的所有元素打包为一个列表，恰好生成我们需要的累积序列。

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