Polars筛选数据后如何随机选取指定行数填充目标值

需求说明 #

针对包含列a（仅由值1、2构成）的Polars DataFrame，新增列b需满足以下规则：

• 若a列值为1，对应行b取值为True
• 若a列值为2，从中无放回随机抽取数量与a列中1的总数相等的行，对应b取值为True，抽样需保证无偏
• b列其余所有行取值为False

测试样例构造 #

import polars as pl
import numpy as np

df = pl.DataFrame({
    'a': [2, 2, 2, 1, 2, 1]
})

样例初始结构：

shape: (6, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 2   │
│ 2   │
│ 2   │
│ 1   │
│ 2   │
│ 1   │
└─────┘

原生Polars实现代码 #

无需转换为pandas对象，全程使用Polars原生表达式即可完成逻辑，无额外格式转换开销：

# 统计a=1的总行数，作为a=2分组的抽样样本量
n_1 = df.filter(pl.col("a") == 1).height

df_res = df.with_columns(
    pl.when(pl.col("a") == 1)
    .then(True)
    .when(
        (pl.col("a") == 2) & (pl.int_range(pl.len()).shuffle().over("a") < n_1)
    )
    .then(True)
    .otherwise(False)
    .alias("b")
)

实现逻辑说明 #

• 第一步先统计a=1的样本总数n_1，作为a=2分组需要抽取的正样本数量
• 用pl.when/then/otherwise做条件分支赋值：
  • 第一分支直接匹配所有a=1的行，固定赋值True
  • 第二分支针对a=2的分组做窗口计算：通过pl.int_range(pl.len()).shuffle().over("a")在a=2的分组内生成打乱顺序的连续整数序列，序列值小于n_1的行即为随机无放回抽中的样本，赋值True
  • 剩余所有行统一赋值False
• 该实现和pandas版本逻辑完全等价，抽样结果无偏，性能优于跨库转换的实现方式。

运行结果示例 #

每次运行因随机打乱顺序会得到不同的抽样结果，符合随机要求，示例输出：

shape: (6, 2)
┌─────┬───────┐
│ a   ┆ b     │
│ --- ┆ ---   │
│ i64 ┆ bool  │
╞═════╪═══════╡
│ 2   ┆ false │
│ 2   ┆ true  │
│ 2   ┆ false │
│ 1   ┆ true  │
│ 2   ┆ true  │
│ 1   ┆ true  │
└─────┴───────┘

内容的提问来源于stack exchange，提问作者Piotr Pasza Storożenko