使用 pl.when 创建一个元素数量相同但内容不同的列表

Question

我已经对 Pandas 提出了一个类似的问题，但正如我在其他地方提到的，我正在转向 Polars。所以，现在我需要一个解决方案，让 Polars 做同样的事情：

我有一个 Polars DataFrame，其中的值有时会变为 NA（对于 Polars，它是 null (=None)，如果我没猜错的话）。我想用与另一列长度相同的字符串列表填充此列：

import polars as pl

df = pl.DataFrame({"a": ["one", "two"],
                   "b": ["three", "four"],
                   "c": [[1, 2], [3, 4]],
                   "d": [[5, 6], None]})

a	b	c	d
一个	三	[1, 2]	[5, 6]
两个	四	[3, 4]	NaN

我希望这成为

a	b	c	d
一个	三	[1, 2]	[5, 6]
两个	四	[3, 4]	[无值，无值]

我试过了

df = df.with_columns(d = pl.when(pl.col('d').is_null())
                        #    .then(pl.Series([['no_value'] * len(lst) for lst in pl.col('c')])) # 'Expr' object is not iterable
                        #    .then(pl.Series([['no_value'] * pl.col('c').list.len()])) # failed to determine supertype of object and list[i64]
                        #    .then([['no_value'] * pl.col('c').list.len()]) # not yet implemented: Nested object types
                           .otherwise(pl.col('d')))

Answer 1

一般来说，使用 when-then-otherwise 结构的方法很好。您可以使用

pl.Expr.repeat_by

将标量列重复到列表列中。

df.with_columns(
    pl.when(
        pl.col("d").is_null()
    ).then(
        pl.lit("no_value").repeat_by(pl.col("c").list.len())
    ).otherwise(
        pl.col("d")
    ).alias("d")
)

shape: (2, 4)
┌─────┬───────┬───────────┬──────────────────────────┐
│ a   ┆ b     ┆ c         ┆ d                        │
│ --- ┆ ---   ┆ ---       ┆ ---                      │
│ str ┆ str   ┆ list[i64] ┆ list[str]                │
╞═════╪═══════╪═══════════╪══════════════════════════╡
│ one ┆ three ┆ [1, 2]    ┆ ["5", "6"]               │
│ two ┆ four  ┆ [3, 4]    ┆ ["no_value", "no_value"] │
└─────┴───────┴───────────┴──────────────────────────┘