如何在Polars中对DataFrame按分组抽取随机样本

你需要的每组抽取固定数量随机样本的功能，在Polars中可以直接按以下方式实现，完全适配LazyFrame场景，不需要加载全量数据：

直接实现目标写法（Polars 0.19+ 版本支持） #

你最开始预期的group_by + agg(sample(n))逻辑本身就原生支持，之前的近似方案出错是因为你传入的是浮点型的采样比例，换成固定采样数参数n即可：

# 聚合结构输出：每个分组1行，采样结果以列表形式存储在列中
sample_agg_lf = lf.group_by("column").agg(
    pl.all().sample(n=10)
)

# 如果需要摊平为行级结构（每个样本占1行），追加explode即可
sample_flat_lf = sample_agg_lf.explode(pl.exclude("column"))

这个写法会在每个分组内独立做无放回随机抽样，每个分组最多返回10条样本（分组本身行数不足10时返回该组全部数据），不会出现和全局采样等价的问题，查询全程可以在Lazy模式下执行，Polars会自动下推算子，不会加载全量数据到内存。

低版本兼容写法（所有Polars版本通用） #

如果使用的是较早的不支持agg内传sample(n)的版本，可以用窗口函数组内打随机序号再过滤的方式实现，性能和原生写法基本一致，同样支持Lazy模式：

SAMPLE_PER_GROUP = 10
sample_flat_lf = lf.with_columns(
    # 给每个分组内的行生成随机打乱的行序号
    pl.int_range(pl.len()).shuffle().over("column").alias("_group_rand_id")
).filter(
    # 过滤出每组序号前10的行
    pl.col("_group_rand_id") < SAMPLE_PER_GROUP
).drop("_group_rand_id")

如果需要有放回抽样，只需要把shuffle逻辑替换为随机数排序即可：

sample_flat_lf = lf.with_columns(
    pl.int_range(pl.len()).sort_by(pl.rand(seed=42)).over("column").alias("_group_rand_id")
).filter(
    pl.col("_group_rand_id") < SAMPLE_PER_GROUP
).drop("_group_rand_id")

旧近似方案的问题说明 #

之前用pl.all().sample(.001)的写法，传入的是采样比例参数，会按每个分组行数的0.1%比例抽样本，当采样比例极低时，抽样分布和全局按比例采样后再分组的结果几乎一致，既没法保证每个分组固定返回10个样本，也没法保证小分组一定能抽到样本，不符合需求。

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