如何优化超大规模DataFrame按簇条件抽样的效率？

问题背景 #

我有一个约3500万行的大型DataFrame，包含embeddings、cluster_ID、label三列，其中有180万个唯一cluster_ID。需要生成新DataFrame，规则如下：

• 若簇内同时存在label 0和label 1，各随机抽取1行
• 若簇内仅存在label 1，随机抽取1行
  原代码运行耗时超过15小时，急需优化。

原代码的性能瓶颈 #

原代码通过循环遍历每个唯一cluster_ID，每次过滤子DataFrame再抽样，存在两大低效点：

• 循环过滤开销极大：每次result_df_30[result_df_30['cluster_ID'] == cluster_id]都是O(n)操作，180万次循环的总复杂度达到天文量级
• 小DataFrame拼接开销大：每次抽样后将小DataFrame加入列表，最后pd.concat需要合并大量小对象，内存和时间成本极高

优化方案：使用pandas分组向量化操作 #

利用pandas的groupby实现向量化处理，底层为C优化，性能比Python循环提升几个数量级：

优化代码 #

import pandas as pd

# 步骤1：先过滤出label为0或1的行（如果存在其他label值）
filtered_df = result_df_30[result_df_30['label'].isin([0, 1])]

# 步骤2：按cluster_ID和label分组，每组随机抽取1行
sampled = filtered_df.groupby(['cluster_ID', 'label'], group_keys=False).sample(n=1, random_state=42)

# 步骤3：筛选出符合需求的簇（仅保留包含label 1的簇，排除仅含label 0的簇）
# 先获取每个簇的label集合
cluster_label_set = filtered_df.groupby('cluster_ID')['label'].unique()
# 只保留包含label 1的簇ID
valid_clusters = cluster_label_set[cluster_label_set.apply(lambda x: 1 in x)].index

# 步骤4：生成最终结果
selected_rows_df = sampled[sampled['cluster_ID'].isin(valid_clusters)]

额外优化建议 #

• 内存优化：如果内存紧张，可将embeddings列转为更节省空间的类型（如numpy数组而非列表），或分块处理（但groupby本身已做内存优化）
• 并行处理：若仍有性能需求，可使用dask.dataframe替代pandas，实现分布式并行分组抽样，适合超大数据量场景

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