按组填充Python Polars的lazyframe中的空值，以每组中唯一值的数量为条件

Question

我有一个很大的（~300M 行 x 44 列）数据框，我需要根据每个组的特征以某些方式填充空值。

例如，假设我们有

lf = pl.LazyFrame(
    {'group':(1,1,1,2,2,2,3,3,3),
     'val':('yes', None, 'no', '2', '2', '2', 'answer', None, 'answer')
     }
)

┌───────┬────────┐
│ group ┆ val    │
│ ---   ┆ ---    │
│ i64   ┆ str    │
╞═══════╪════════╡
│ 1     ┆ yes    │
│ 1     ┆ null   │
│ 1     ┆ no     │
│ 2     ┆ 2      │
│ 2     ┆ 2      │
│ 2     ┆ 2      │
│ 3     ┆ answer │
│ 3     ┆ null   │
│ 3     ┆ answer │
└───────┴────────┘

当且仅当该组在其他单元格中包含单个非空唯一值时，我才想填充空值，因为在我的上下文中，这是数据的期望以及存在多个唯一值（或所有空值）组中的信号表明另一个问题将以不同的方式处理。

我可以使用以下内容填充每个组的空值：

filled_lf = (
    lf
    .with_columns(
        pl.col('val')
        .fill_null(pl.col('val').unique().first().over('group')).alias('filled_val')
        )
    )

但是，首先，pl.col('val').unique() 似乎包含“null”作为其中一个值，并且排序是随机的，因此选择列表中的第一个值会产生不一致的结果。其次，它不包含我需要的条件。

想要的结果：

┌───────┬────────┬────────────┐
│ group ┆ val    ┆ filled_val │
│ ---   ┆ ---    ┆ ---        │
│ i64   ┆ str    ┆ str        │
╞═══════╪════════╪════════════╡
│ 1     ┆ yes    ┆ yes        │
│ 1     ┆ null   ┆ null       │
│ 1     ┆ no     ┆ no         │
│ 2     ┆ 2      ┆ 2          │
│ 2     ┆ 2      ┆ 2          │
│ 2     ┆ 2      ┆ 2          │
│ 3     ┆ answer ┆ answer     │
│ 3     ┆ null   ┆ answer     │
│ 3     ┆ answer ┆ answer     │
└───────┴────────┴────────────┘

熊猫3.12 极地0.20.1

提前感谢您的建议！

Answer 1

您可以添加：

• .drop_nulls()

• .unique()

  与

  maintain_order=True

• when/then/otherwise

  实现条件计数/长度逻辑

unique = pl.col("val").drop_nulls().unique(maintain_order=True)

df.with_columns(
   pl.when(unique.len().over("group") == 1)
     .then(pl.col("val").fill_null(unique.first().over("group")))
     .otherwise(pl.col("val"))
     .alias("filled")
)

shape: (9, 3)
┌───────┬────────┬────────┐
│ group ┆ val    ┆ filled │
│ ---   ┆ ---    ┆ ---    │
│ i64   ┆ str    ┆ str    │
╞═══════╪════════╪════════╡
│ 1     ┆ yes    ┆ yes    │
│ 1     ┆ null   ┆ null   │
│ 1     ┆ no     ┆ no     │
│ 2     ┆ 2      ┆ 2      │
│ 2     ┆ 2      ┆ 2      │
│ 2     ┆ 2      ┆ 2      │
│ 3     ┆ answer ┆ answer │
│ 3     ┆ null   ┆ answer │
│ 3     ┆ answer ┆ answer │
└───────┴────────┴────────┘