如何在Polars中按条件随机丢弃指定比例的行？

问题背景 #

现有如下Polars DataFrame：

需求：保留所有Labels == 1的行，对Labels == 0的行以50%的概率随机丢弃（即保留50%的随机样本）。

尝试了以下代码但无法运行：

df = df.drop(df.filter(pl.col("Labels") == 0).sample(frac=0.5))

解决方案 #

方法1：拆分处理后合并（直观易懂） #

df.drop()方法仅接受列名或行索引集合，不能直接传入DataFrame，因此原代码报错。正确思路是拆分两组分别处理后合并：

import polars as pl

# 构造原始DataFrame
df = pl.DataFrame({
    "特征1": [100, 150, 200, 230, 120, 130, 150, 180],
    "特征2": [25, 18, 15, 28, 12, 34, 23, 25],
    "Labels": [1, 0, 0, 0, 1, 1, 1, 0]
})

# 处理逻辑：保留Labels=1的所有行，对Labels=0的行抽样50%
df_processed = pl.concat([
    df.filter(pl.col("Labels") == 1),
    df.filter(pl.col("Labels") == 0).sample(frac=0.5, seed=42)  # seed参数可选，用于固定抽样结果
])

# 可选：重置索引
df_processed = df_processed.reset_index(drop=True)
print(df_processed)

方法2：分组聚合抽样（简洁高效） #

通过group_by按Labels分组，对不同组应用不同抽样比例：

df_processed = df.group_by("Labels").agg(
    pl.all().sample(frac=pl.when(pl.col("Labels") == 0).then(0.5).otherwise(1.0))
).explode(pl.all())

两种方法都能实现需求，第一种逻辑清晰适合新手，第二种代码更简洁适合批量处理场景。

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