如何在Polars中对列表元素进行重复或截断以固定长度

嘿，我懂你现在的需求——你的Polars LazyFrame里，points列的每个元素都是长短不一的坐标列表，想要把它们统一成固定长度：长的截断到目标长度，短的就用现有元素重复补够，对吧？刚好有个实用的解法，用Polars的列表表达式就能轻松搞定。

先看你的原始数据：

import polars as pl

lf = pl.LazyFrame(
    {
        "points": [
            [[1.0, 2.0]],
            [[3.0, 4.0], [5.0, 6.0]],
            [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]],
        ],
        "other": ["foo", "bar", "baz"]
    }
)

假设我们要把每个points子列表统一成3个元素，可以这么写代码：

target_length = 3

# 生成固定长度的points列
result = lf.with_columns(
    fixed_points=pl.col("points").list.eval(
        # 把原列表重复足够次数，展开后取前target_length个元素
        pl.element().repeat(pl.lit(target_length)).list.flatten().list.head(target_length)
    )
).collect()

print(result)

运行后你会得到这样的结果：

shape: (3, 3)
┌───────────────────────────────────┬───────┬───────────────────────────────────┐
│ points                            ┆ other ┆ fixed_points                      │
│ ---                               ┆ ---   ┆ ---                               │
│ list[list[f64]]                   ┆ str   ┆ list[list[f64]]                   │
╞═══════════════════════════════════╪═══════╪═══════════════════════════════════╡
│ [[1.0, 2.0]]                      ┆ foo   ┆ [[1.0, 2.0], [1.0, 2.0], [1.0, 2… │
│ [[3.0, 4.0], [5.0, 6.0]]          ┆ bar   ┆ [[3.0, 4.0], [5.0, 6.0], [3.0, 4… │
│ [[7.0, 8.0], [9.0, 10.0], [11.0…  ┆ baz   ┆ [[7.0, 8.0], [9.0, 10.0], [11.0…  │
└───────────────────────────────────┴───────┴───────────────────────────────────┘

我给你拆解下这段代码的逻辑：

• list.eval是核心，它能让我们对points列里的每一个子列表单独处理
• pl.element()指代当前正在处理的那个子列表，repeat(target_length)会把这个子列表重复3次，比如原来只有1个元素的列表，重复后就变成[[[1,2]], [[1,2]], [[1,2]]]
• list.flatten()把嵌套的结构展开成一维的坐标列表，再用list.head(target_length)取前3个，刚好就是我们要的固定长度

如果你想更高效一点（避免不必要的重复），也可以先判断原列表的长度，按需补充：

result = lf.with_columns(
    fixed_points=pl.col("points").list.eval(
        pl.when(pl.element().list.length() >= target_length)
        # 长度够的话直接截断到目标长度
        .then(pl.element().list.head(target_length))
        # 长度不够的话，计算需要补的数量，只重复必要的次数再拼接截断
        .otherwise(
            pl.element().concat(
                pl.element().repeat(
                    (target_length - pl.element().list.length() + pl.element().list.length() - 1) // pl.element().list.length()
                ).list.flatten()
            ).list.head(target_length)
        )
    )
).collect()

这个写法会先检查原列表长度，够长就直接截断；不够的话只重复刚好能补够长度的次数，比如原长度2、目标3，只需要重复1次（总共4个元素），再取前3个就行，比直接重复3次更省资源。

总之，Polars的列表操作API灵活性拉满，用这些组合就能轻松实现你要的“统一列表长度”的需求啦！

备注：内容来源于stack exchange，提问作者DJDuque