Polars中如何根据另一列列表长度重复特定值n次？

需求说明 #

需要为LazyFrame添加一列col3，规则如下：

• col3的列表长度必须和对应行col1的列表长度完全一致
• 如果col2和col1的列表长度相等，直接沿用col2的内容
• 否则生成一个长度与col1相同、所有元素为None的列表

测试场景与问题 #

先看测试数据和当前遇到的问题：

创建测试LazyFrame：

import polars as pl

lf = pl.LazyFrame(
    {
        "col1": [[1, 2, 3], [1, 2], [1]],
        "col2": [["A", "B", "C"], ["C"], ["D", "E"]],
    }
)

print(lf.collect())

输出：

shape: (3, 2)
┌───────────┬─────────────────┐
│ col1      ┆ col2            │
│ ---       ┆ ---             │
│ list[i64] ┆ list[str]       │
╞═══════════╪═════════════════╡
│ [1, 2, 3] ┆ ["A", "B", "C"] │
│ [1, 2]    ┆ ["C"]           │
│ [1]       ┆ ["D", "E"]      │
└───────────┴─────────────────┘

尝试了以下代码，但输出不符合预期：

condition = pl.col("col2").list.len().eq(pl.col("col1").list.len())

lf = lf.with_columns(
    pl.when(condition)
    .then(pl.col("col2"))
    .otherwise(
        pl.repeat(
            value=None,
            n=pl.col("col1").list.len(),
        )
    )
    .alias("col3")
).collect()

print(lf)

当前错误输出（col3直接为null，而非对应长度的None列表）：

shape: (3, 3)
┌───────────┬─────────────────┬─────────────────┐
│ col1      ┆ col2            ┆ col3            │
│ ---       ┆ ---             ┆ ---             │
│ list[i64] ┆ list[str]       ┆ list[str]       │
╞═══════════╪═════════════════╪═════════════════╡
│ [1, 2, 3] ┆ ["A", "B", "C"] ┆ ["A", "B", "C"] │
│ [1, 2]    ┆ ["C"]           ┆ null            │
│ [1]       ┆ ["D", "E"]      ┆ null            │
└───────────┴─────────────────┴─────────────────┘

我们需要的预期输出是：

shape: (3, 3)
┌───────────┬─────────────────┬─────────────────┐
│ col1      ┆ col2            ┆ col3            │
│ ---       ┆ ---             ┆ ---             │
│ list[i64] ┆ list[str]       ┆ list[str]       │
╞═══════════╪═════════════════╪═════════════════╡
│ [1, 2, 3] ┆ ["A", "B", "C"] ┆ ["A", "B", "C"] │
│ [1, 2]    ┆ ["C"]           ┆ [null, null]    │
│ [1]       ┆ ["D", "E"]      ┆ [null]          │
└───────────┴─────────────────┴─────────────────┘

问题原因 #

原来的pl.repeat在when/otherwise分支中使用列表达式作为n参数时，无法正确逐行生成对应长度的列表。因为pl.repeat默认接受标量参数，当传入列表达式时，上下文无法正确解析为逐行操作，最终返回了null。

解决方案 #

可以通过Polars的列表操作实现逐行生成对应长度的None列表，推荐两种简单有效的方法：

方法1：用list.eval替换原列表元素为None

利用col1.list.eval(pl.lit(None))将col1中的每个元素替换为None，得到的列表长度自然和col1一致：

import polars as pl

lf = pl.LazyFrame(
    {
        "col1": [[1, 2, 3], [1, 2], [1]],
        "col2": [["A", "B", "C"], ["C"], ["D", "E"]],
    }
)

condition = pl.col("col2").list.len().eq(pl.col("col1").list.len())

lf = lf.with_columns(
    pl.when(condition)
    .then(pl.col("col2"))
    .otherwise(pl.col("col1").list.eval(pl.lit(None)))
    .alias("col3")
).collect()

print(lf)

方法2：用list.set结合索引范围生成None列表

先生成与col1长度一致的索引范围，再将所有位置的元素设置为None：

import polars as pl

lf = pl.LazyFrame(
    {
        "col1": [[1, 2, 3], [1, 2], [1]],
        "col2": [["A", "B", "C"], ["C"], ["D", "E"]],
    }
)

condition = pl.col("col2").list.len().eq(pl.col("col1").list.len())

lf = lf.with_columns(
    pl.when(condition)
    .then(pl.col("col2"))
    .otherwise(pl.col("col1").list.set(pl.int_range(0, pl.col("col1").list.len()), None))
    .alias("col3")
).collect()

print(lf)

两种方法都能精准生成符合要求的col3列，得到预期输出。

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