如何在Polars中按Group ID进行有放回抽样N-1样本？

要实现按group_id分组的有放回抽样（每组抽取M=N-1个样本），可以利用Polars的分组聚合结合sample方法完成，步骤如下：

核心代码实现 #

import polars as pl

# 加载原始数据
water_data = {
    'group_id': [1,1,1,1,2,2,2,3,3,3,4,4,4,4,5,5,5],
    'obs_id_within_group': [1,2,3,4,1,2,3,1,2,3,1,2,3,4,1,2,3],
    'N': [4,4,4,4,3,3,3,3,3,3,4,4,4,4,3,3,3],
    'M': [3,3,3,3,2,2,2,2,2,2,3,3,3,3,2,2,2],
    'water_gallons': [12,23,21,11,10,10,10,23,24,25,27,30,17,12,11,14,20],
    'water_source': ['lake','lake','pond','river','lake','glacier','glacier','lake','pond','river','lake','lake','pond','river','river','lake','glacier'],
    'water_acidity': [3,4,5,1,2,4,3,2,3,3,4,6,7,8,8,3,1]
}
df = pl.DataFrame(water_data)

# 执行按组有放回抽样
df_sampled = df.group_by('group_id').agg(
    pl.all().sample(
        n=pl.col('M').first(),
        replace=True,
        shuffle=True
    ).explode()
)

print(df_sampled)

代码说明 #

• 分组逻辑：通过group_by('group_id')将数据按组划分，确保每组独立处理。
• 抽样参数：
  • pl.all().sample()：对每组的所有列执行抽样操作，保留原数据的全部字段。
  • n=pl.col('M').first()：取每组第一个M值作为抽样数量（因同一组内M值统一，无需额外处理）。
  • replace=True：开启有放回抽样，满足Bootstrapping的需求。
  • shuffle=True：保证抽样结果的随机性。
• 展开结果：.explode()将抽样得到的列表型数据展开为多行，恢复标准DataFrame结构。

注意事项 #

如果你的场景中存在组内M值不一致的情况，需要先确保每组抽样数量统一（比如取组内M的唯一值），否则会引发抽样数量不匹配的错误。

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