如何在Python Polars中为列表列添加最大值+2的元素？

问题场景 #

现有包含列表列的Polars DataFrame：

import polars as pl

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

需要为每个列表追加该列表最大值加2的元素，得到目标DataFrame：

expected_df = pl.DataFrame({'lists': [[1,2,3,5], [2,3,4,6], [3,4,5,7]]})

原写法问题分析 #

你尝试的写法无法正确运行，原因是list.add_element()仅支持传入单个标量值，不能直接使用列表达式来动态计算每行的追加值。

正确实现方式 #

方法一：向量化拼接（推荐，性能更优）

利用list.concat()将原列表与计算出的单元素列表拼接，需要先把标量结果转为列表类型：

result_df = df.with_columns(
    pl.col("lists").list.concat(
        (pl.col("lists").list.max() + 2).cast(pl.List(pl.Int64))
    ).alias("lists")
)

方法二：逐行映射处理

适合逻辑更复杂的场景，通过map_elements逐行处理每个列表：

result_df = df.with_columns(
    pl.col("lists").map_elements(
        lambda lst: lst + [max(lst) + 2],
        return_dtype=pl.List(pl.Int64)
    ).alias("lists")
)

结果验证 #

运行上述任意一种方法，result_df的结构和数据都会与expected_df完全一致。

内容的提问来源于stack exchange，提问作者sports-modelling