如何在Polars中实现List[str]列的累积拼接？

问题描述 #

我有一个包含两列的pl.DataFrame()：第一列为Date，第二列为List[str]，具体数据如下：

import polars as pl

df = pl.DataFrame([
    pl.Series('Date', [2000, 2001, 2002]),
    pl.Series('Ids', [
        ['a'], 
        ['b', 'c'], 
        ['d'], 
    ])
])

对应表格：

需要对该List[str]列进行累积操作，使每行包含自身及之前所有行的列表，得到如下结果：

解决方案 #

方法一：使用cum_reduce（Polars 0.19.0+支持） #

这是最直观的实现方式，利用cum_reduce逐行累积拼接列表：

import polars as pl

df = pl.DataFrame([
    pl.Series('Date', [2000, 2001, 2002]),
    pl.Series('Ids', [['a'], ['b', 'c'], ['d']])
])

result = df.with_columns(
    pl.col("Ids").cum_reduce(lambda a, b: pl.concat([a, b]).alias("Ids"), dtype=pl.List(pl.String))
)

print(result)

方法二：兼容旧版本的实现 #

如果使用的Polars版本低于0.19.0，可以通过结构体转换配合cum_sum实现：

result = df.with_columns(
    pl.col("Ids")
    .arr.to_struct()
    .cum_sum()
    .arr.from_struct()
    .arr.flatten()
    .alias("Ids")
)

print(result)

输出结果 #

运行上述代码后，会得到目标数据框：

shape: (3, 2)
┌──────┬───────────────────┐
│ Date ┆ Ids               │
│ ---  ┆ ---               │
│ i64  ┆ list[str]         │
╞══════╪═══════════════════╡
│ 2000 ┆ ["a"]             │
│ 2001 ┆ ["a", "b", "c"]   │
│ 2002 ┆ ["a", "b", "c", "d"] │
└──────┴───────────────────┘

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