Python Pandas大表高效随机抽样：替代循环实现方案求助

这确实是Pandas处理百万级数据时的典型性能坑——iterrows逐行循环在数据量上去后，速度会慢到让人崩溃。咱们可以利用Pandas的分组+向量化抽样能力，彻底抛弃Python层面的循环，把速度提升几个数量级。

先回顾下你的表格结构：

table1示例： #

ID;COUNT;FOREIGN_ID;OTHER_DATA
1;3;xyz1
2;1;xyz2
3;1;xyz3

table2示例： #

ID;FOREIGN_ID;OTHER_DATA
1;xyz1;000001
2;xyz1;000002
3;xyz1;000003
4;xyz1;000004
5;xyz1;000005
6;xyz2;000000
7;xyz2;000000
8;xyz3;000000
9;xyz3;000000

完整高效代码 #

import pandas as pd

# 加载数据（注意分隔符是分号）
df_table1 = pd.read_csv('table1.csv', delimiter=';')
df_table2 = pd.read_csv('table2.csv', delimiter=';')

# 构建FOREIGN_ID到需要抽取数量的映射字典
count_map = df_table1.set_index('FOREIGN_ID')['COUNT'].to_dict()

# 按FOREIGN_ID分组，对每个组按指定数量随机抽样
# 加个min判断，避免COUNT大于组内行数时报错
sampled_groups = df_table2.groupby('FOREIGN_ID').apply(
    lambda group: group.sample(n=min(count_map.get(group.name, 0), len(group)), replace=False)
)

# 整理结果：重置索引，去掉分组产生的额外层级
df_result = sampled_groups.reset_index(drop=True)

代码解释 #

1. 构建count_map字典：把table1的FOREIGN_ID作为键，COUNT作为值，这样分组时可以快速获取当前组需要抽取的数量，比每次循环查询要高效得多。
2. 分组抽样：groupby('FOREIGN_ID')把table2按外键分组，然后用apply对每个组执行抽样操作。这里用min(count_map.get(group.name, 0), len(group))做了安全处理——如果某个FOREIGN_ID在table1里没有对应记录，或者COUNT超过了组内的行数，就取组内所有数据（或者0，根据你的需求调整）。
3. 整理结果：分组抽样后会产生多级索引，用reset_index(drop=True)把它转换成普通的DataFrame结构，和你原来的结果格式一致。

为什么比原方案快？ #

• 原方案的iterrows是Python层面的逐行循环，每一次循环都要执行布尔索引、创建小DataFrame、append（append本身每次都会复制整个结果DataFrame，数据量越大越慢），这些操作在百万级数据下会产生大量的性能开销。
• 新方案的groupby和sample都是Pandas底层优化过的C级别操作，避免了Python循环的 overhead，所有数据处理都是批量进行的，速度能提升几十甚至上百倍。

额外注意事项 #

• 如果table1中存在table2里没有的FOREIGN_ID，这些记录不会产生任何抽样结果，如果你需要保留这类记录（比如合并table1的OTHER_DATA），可以在最后把df_result和df_table1做左连接。
• 如果你的COUNT可能大于对应组的行数，且允许重复抽样，可以把replace=False改成replace=True，同时去掉min判断（根据你的业务需求调整）。

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