如何基于另一列长度填充Polars DataFrame中的列表列？

问题背景 #

现有包含pad_size（目标长度）和pad_start（待补全列表）两列的Polars DataFrame，需要生成pad_end列，将pad_start中的列表追加0，直到其长度等于对应行的pad_size值。

原始数据定义 #

import polars as pl

# 调整显示配置，优化列表内容展示
pl.Config(fmt_table_cell_list_len=6, fmt_str_lengths=100)

# 示例DataFrame
data = {
    'pad_size': [5, 4, 3],
    'pad_start': [[1, 2], [4, 5, 6], [7]]
}
df = pl.DataFrame(data)

原始DataFrame输出 #

┌──────────┬───────────┐
│ pad_size ┆ pad_start │
│ ---      ┆ ---       │
│ i64      ┆ list[i64] │
╞══════════╪═══════════╡
│ 5        ┆ [1, 2]    │
│ 4        ┆ [4, 5, 6] │
│ 3        ┆ [7]       │
└──────────┴───────────┘

期望结果 #

生成pad_end列后的DataFrame如下：

┌──────────┬─────────────────┐
│ pad_size ┆ pad_end         │
│ ---      ┆ ---             │
│ i64      ┆ list[i64]       │
╞══════════╪═════════════════╡
│ 5        ┆ [1, 2, 0, 0, 0] │
│ 4        ┆ [4, 5, 6, 0]    │
│ 3        ┆ [7, 0, 0]       │
└──────────┴─────────────────┘

优化解决方案 #

针对Polars 0.18.4版本，可通过以下代码直接生成目标列，无需保留中间临时列：

df = df.with_columns(
    pad_end=pl.col("pad_start")
    .list.concat(
        pl.lit(0).repeat_by(pl.col("pad_size") - pl.col("pad_start").list.lengths())
    )
)

逻辑说明 #

1. 计算需补充的0的数量：通过pl.col("pad_size") - pl.col("pad_start").list.lengths()得到每行列表需要追加的0的个数
2. 生成对应长度的0列表：使用pl.lit(0).repeat_by(...)将单个0重复指定次数，生成符合长度要求的0列表
3. 拼接列表：调用list.concat()将原pad_start列表与生成的0列表拼接，得到最终补全后的pad_end列

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