如何通过采样列的其余部分来使用Polars DataFrame中的`fill_null`?

Question

你想知道怎么在一个叫做 DataFrame 的数据表里，用其他列的值来填补空值（null），每次填补的时候都要随机选一个独特的值。

举个例子：假设我有这样一个 DataFrame：

>>> daf=pl.DataFrame({"a":[1,2, None, None, 5, 6, None, None, None, 10], "b":[None, "two", "three", None, "five", "six", "seven", None, "nine", "ten"]})
>>> daf
shape: (10, 2)
┌──────┬───────┐
│ a    ┆ b     │
│ ---  ┆ ---   │
│ i64  ┆ str   │
╞══════╪═══════╡
│ 1    ┆ null  │
│ 2    ┆ two   │
│ null ┆ three │
│ null ┆ null  │
│ 5    ┆ five  │
│ 6    ┆ six   │
│ null ┆ seven │
│ null ┆ null  │
│ null ┆ nine  │
│ 10   ┆ ten   │
└──────┴───────┘

我想把所有列里的空值填上，填的值是从该列其他值中随机选的。比如，"a" 列里的空值应该从 [1,2,5,6,10] 中随机选一个，而 "b" 列里的空值应该从 ["two", "three", "five", "six", "seven", "nine", "ten"] 中随机选一个。但要注意，"a" 列里不能有字符串（str），而 "b" 列里不能有整数（i64）。而且，同一列里的每个空值填的值不能重复。

我之前是通过一个普通的 Python for 循环来实现这个功能的：

>>> for c in daf.columns:
...   repl=daf[c].drop_nulls().sample(daf[c].len(), with_replacement=True)
...   daf=daf.with_columns(pl.col(c).fill_null(repl))
...
>>> daf
shape: (10, 2)
┌─────┬───────┐
│ a   ┆ b     │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ two   │
│ 2   ┆ two   │
│ 1   ┆ three │
│ 2   ┆ six   │
│ 5   ┆ five  │
│ 6   ┆ six   │
│ 5   ┆ seven │
│ 10  ┆ two   │
│ 5   ┆ nine  │
│ 10  ┆ ten   │
└─────┴───────┘

但我想知道有没有更简洁的方法，能用 Polars 的表达式语法来实现这个功能？

Answer 1

解决方案

这个方法看起来有效。它是基于你最初的想法，只不过是用表达式对每一列进行相同的操作。

daf.select(
    pl.col("*").fill_null(
        pl.col("*").drop_nulls().sample(pl.len(), with_replacement=True)
    )
)

┌─────┬───────┐
│ a   ┆ b     │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ two   │
│ 2   ┆ two   │
│ 6   ┆ three │
│ 10  ┆ six   │
│ 5   ┆ five  │
│ 6   ┆ six   │
│ 1   ┆ seven │
│ 2   ┆ two   │
│ 2   ┆ nine  │
│ 10  ┆ ten   │
└─────┴───────┘

解释

pl.col("*")（或者pl.all()）就是单独选择每一列，然后对它进行后续的操作。之后它保持了相同的结构，因为例如这样做是有效的：

>>> df = pl.DataFrame({"a": [1, 2, 3], "b": [100, 200, 300]})
>>> df.select(pl.col("*").mul(pl.col("*").mean()))
shape: (3, 2)
┌─────┬─────────┐
│ a   ┆ b       │
│ --- ┆ ---     │
│ f64 ┆ f64     │
╞═════╪═════════╡
│ 2.0 ┆ 20000.0 │
│ 4.0 ┆ 40000.0 │
│ 6.0 ┆ 60000.0 │
└─────┴─────────┘

如果我使用LazyFrame和explain，它在后台做的事情是这样的：

>>> daf.select(pl.col("*").fill_null(pl.col("*").drop_nulls().sample(pl.len(), with_replacement=True))).explain()
 SELECT [col("a").fill_null([col("a").drop_nulls().sample([len()])]), col("b").fill_null([col("b").drop_nulls().sample([len()])])] FROM
  DF ["a", "b"]; PROJECT 2/2 COLUMNS; SELECTION: "None"

让我有点惊讶的是，fill_null的部分实际上是有效的。关于fill_null的文档并没有很清楚地说明，你可以传递一个返回其他值而不是单一值的表达式，它就是这样工作的。

另外，人们可能会认为用1或100代替pl.len()会得到相同的结果，但实际上并不是。1会为每个空值替换返回相同的样本，而任何其他值（除了1或pl.len()）都会导致形状错误。

看起来它执行的是一个表大小的or操作，把每个空值替换为另一侧对应位置的值（那边的值也可能是空的），这点还挺有意思的。

>>> daf.select(pl.col("*").fill_null(pl.col("*").reverse()))
shape: (10, 2)
┌──────┬───────┐
│ a    ┆ b     │
│ ---  ┆ ---   │
│ i64  ┆ str   │
╞══════╪═══════╡
│ 1    ┆ ten   │
│ 2    ┆ two   │
│ null ┆ three │
│ null ┆ seven │
│ 5    ┆ five  │
│ 6    ┆ six   │
│ null ┆ seven │
│ null ┆ three │
│ 2    ┆ nine  │
│ 10   ┆ ten   │
└──────┴───────┘