大规模数据集下Pandas按组差异化随机采样的性能优化问询

问题背景 #

需要处理9000万条规模的数据集，按列a（含10万个唯一值）分组后，每组采样不同数量的样本：

• 若组内数据量≥目标采样数m，执行无放回采样
• 若组内数据量<目标采样数m，执行有放回采样

原方案通过循环遍历每个唯一值、过滤分组后采样再合并，功能正常但性能极差，需优化为高效分组/向量化实现。

优化思路 #

核心是避免逐组过滤全量数据集的低效操作，改为基于分组索引批量生成采样结果：

1. 先将目标采样数关联到原始数据集的每条记录
2. 计算每个分组的实际大小，判断采样模式（有放回/无放回）
3. 针对每个分组，直接从组内索引列表中批量采样，最后合并所有采样结果

高效实现方案 #

方案1：利用groupby.apply优化分组采样 #

import pandas as pd
import numpy as np

# 设置随机种子保证结果可复现
np.random.seed(6)

# 假设：
# df = 你的9000万条原始数据集
# sample_counts = 采样数量表（包含a列和对应的count列）

# 1. 关联目标采样数到原始数据
df = df.merge(sample_counts, on='a', how='left').rename(columns={'count': 'target_m'})

# 2. 计算每个分组的大小
group_sizes = df.groupby('a').size().reset_index(name='group_size')
df = df.merge(group_sizes, on='a', how='left')

# 3. 定义分组采样逻辑
def sample_single_group(group):
    target_m = group['target_m'].iloc[0]
    group_size = group['group_size'].iloc[0]
    # 判断是否需要有放回采样
    replace_mode = group_size < target_m
    # 从组内索引中采样
    sampled_idx = np.random.choice(group.index, size=target_m, replace=replace_mode)
    return df.loc[sampled_idx]

# 4. 执行分组采样并合并结果
final_sampled = df.groupby('a').apply(sample_single_group).reset_index(drop=True)
# 保留需要的列（根据实际需求调整）
final_sampled = final_sampled[['a', 'x']]

方案2：基于分组索引批量生成采样结果（性能更优） #

该方案跳过groupby.apply的中间包装，直接操作分组索引列表，适合超大规模数据集：

import pandas as pd
import numpy as np

np.random.seed(6)

# 1. 关联采样数并计算分组大小
df = df.merge(sample_counts, on='a', how='left').rename(columns={'count': 'target_m'})
group_info = df.groupby('a').agg(
    target_m=('target_m', 'first'),
    group_size=('a', 'size'),
    group_indices=('a', lambda x: x.index.tolist())
).reset_index()

# 2. 批量生成所有采样索引
sampled_indices = []
for _, row in group_info.iterrows():
    replace = row['group_size'] < row['target_m']
    # 从当前组的索引列表中采样
    idx = np.random.choice(row['group_indices'], size=row['target_m'], replace=replace)
    sampled_indices.extend(idx)

# 3. 提取最终采样结果
final_sampled = df.loc[sampled_indices][['a', 'x']].reset_index(drop=True)

性能优势说明 #

• 避免了原方案中逐组过滤全量数据集的重复IO和内存开销
• 利用numpy的向量采样逻辑，比pandas的sample方法在批量操作时更快
• 方案2直接操作索引列表，减少了pandas分组操作的中间对象创建，在10万级分组规模下性能提升更明显

示例验证 #

用你提供的测试数据验证：

• 采样数量表：

• 原始数据集：

执行方案后，会得到符合预期的结果：

• 组a=1：无放回采样1条
• 组a=2：有放回采样3条（会出现重复值）
• 组a=3：无放回采样2条

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