基于组内非空唯一值数量，条件填充Polars LazyFrame空值

针对3亿行×44列的大型LazyFrame，我们可以通过分组统计+条件判断实现需求：仅当组内非空单元格唯一值数量为1时填充空值，否则保留原空值。

核心思路 #

1. 按组统计两个关键指标：
   • 组内非空值的唯一数量（排除null干扰）
   • 组内的唯一非空值（仅当数量为1时有效）
2. 基于上述指标，通过条件判断实现精准填充：仅当原值为null且组内非空唯一值数量为1时，用组内唯一值填充；其他情况保留原值。

可行代码实现 #

方式一：窗口函数（推荐，代码简洁）

import polars as pl

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

# 执行填充逻辑
result = lf.with_columns(
    # 统计组内非空值的唯一数量
    non_null_unique_count=pl.col('val').n_unique(drop_nulls=True).over('group'),
    # 获取组内的唯一非空值（仅当数量为1时有效）
    group_unique_val=pl.col('val').drop_nulls().first().over('group')
).with_columns(
    filled_val=pl.when(
        # 满足条件：原值为空，且组内非空唯一值数量为1
        pl.col('val').is_null() & (pl.col('non_null_unique_count') == 1)
    ).then(pl.col('group_unique_val'))
    # 其他情况保留原值
    .otherwise(pl.col('val'))
).drop(['non_null_unique_count', 'group_unique_val'])

# 查看结果
print(result.collect())

方式二：分组聚合+Join（适合极端大数据量场景）

如果担心窗口函数的内存占用，也可以先分组聚合再关联：

import polars as pl

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

# 先分组计算指标
agg_df = lf.group_by('group').agg(
    non_null_unique_count=pl.col('val').n_unique(drop_nulls=True),
    group_unique_val=pl.col('val').drop_nulls().first()
)

# 关联后执行填充
result = lf.join(agg_df, on='group', how='left').with_columns(
    filled_val=pl.when(
        pl.col('val').is_null() & (pl.col('non_null_unique_count') == 1)
    ).then(pl.col('group_unique_val'))
    .otherwise(pl.col('val'))
).drop(['non_null_unique_count', 'group_unique_val'])

print(result.collect())

关键细节说明 #

• 解决null干扰问题：使用n_unique(drop_nulls=True)而非unique()，精准统计组内非空值的唯一数量，避免null被计入。
• 条件判断逻辑：通过pl.when(...).then(...).otherwise(...)实现严格的填充规则：
  • 组内非空唯一值>1（如group1的yes/no）：空值保留
  • 组内全为空：非空唯一值数量为0，空值保留
  • 组内非空唯一值=1（如group2的2、group3的answer）：空值填充为该唯一值
• 性能适配：基于Polars LazyFrame的惰性执行机制，两种方式都会被优化为高效的查询计划，适配3亿行的大数据量处理。

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