本文介绍如何对包含3500万行的大型dataframe进行高性能条件采样：为每个cluster_id最多抽取2行（label=0和label=1各1行），自动跳过缺失标签的簇，全程避免低效循环，将耗时从15小时级降至数分钟。
本文介绍如何对包含3500万行的大型dataframe进行高性能条件采样：为每个cluster_id最多抽取2行（label=0和label=1各1行），自动跳过缺失标签的簇，全程避免低效循环，将耗时从15小时级降至数分钟。
在处理超大规模结构化数据（如聚类后的嵌入向量）时，基于分组逻辑的随机采样极易因不当实现成为性能瓶颈。原始代码中对每个唯一 cluster_ID 进行显式 Python 循环、多次布尔索引过滤及重复 sample() 调用，导致约 180 万次 DataFrame 子集构建——这不仅引发大量内存拷贝，还严重阻碍 Pandas 的向量化优势。
核心优化原则：用向量化操作替代显式循环，用一次全局打散 + 分组聚合替代逐簇采样。
✅ 推荐方案一：采样每个「cluster_ID + label」组合的首行（最简高效，满足多数场景）
该方法等价于“对每个存在 label=0 或 label=1 的 cluster_ID，各取 1 行（若存在）”，天然支持不均衡标签分布，且性能最优：
# 确保 label 列仅含 0/1（如题所述）
assert df['label'].isin([0, 1]).all()
# 一步到位：全局随机打散 → 按 (cluster_ID, label) 分组 → 取每组首行
sampled = (
df.sample(frac=1, random_state=42) # 全局 shuffle，保证随机性
.groupby(['cluster_ID', 'label'], as_index=False, dropna=False)
.first() # 每个 (cluster_ID, label) 组取第 1 行（已随机）
)
✅ 优势：单次 sample() + 单次 groupby().first()，时间复杂度 O(n)，实测 3500 万行可在 2–4 分钟内完成（取决于硬件）。
⚠️ 注意：groupby(...).first() 依赖 sample(frac=1) 的全局随机顺序，无需 random_state 重复设置；若需复现结果，请固定 sample 的 random_state。
✅ 推荐方案二：严格只采「同时含 label=0 和 label=1」的簇（原需求精确版）
当业务要求必须成对出现（即仅保留双标签簇，并从中各取 1 行）时，采用索引映射法，避免 apply 或循环：
# 1. 预筛选并重置索引（关键！确保 iloc 可靠）
df_clean = df[df['label'].isin([0, 1])].reset_index(drop=True)
# 2. 全局 shuffle 后，按 cluster_ID & label 记录首个出现的原始位置
pivot_idx = (
df_clean.sample(frac=1, random_state=42)
.reset_index()
.pivot_table(
index='cluster_ID',
columns='label',
values='index', # 原始重置后的整数索引
aggfunc='first' # 每组取第一个随机位置
)
)
# 3. 筛出同时有 0 和 1 的簇，并提取对应两行索引
valid_clusters = pivot_idx.dropna(subset=[0, 1])
row_indices = valid_clusters[[0, 1]].values.flatten().astype(int)
# 4. 一次性切片获取最终样本
sampled = df_clean.iloc[row_indices].copy()
✅ 优势：完全避免分组内采样开销，pivot_table 高度优化，iloc 索引访问为 O(1)；结果严格满足“每簇至多 2 行，且仅来自双标签簇”。
⚠️ 注意：pivot_table 默认会丢弃无某标签的簇（符合需求）；dropna(subset=[0,1]) 精确过滤双标签簇；最终 iloc 返回视图或副本需根据后续操作决定是否 .copy()。
? 为什么原始代码极慢？
- for cluster_id in df['cluster_ID'].unique()：触发 180 万次 Python 解释器迭代；
- df[df['cluster_ID'] == cluster_id]：每次执行全量布尔索引（O(n) × 1.8M ≈ O(63e12) 次比较）；
- sample(n=1) 在小 DataFrame 上开销显著（初始化随机引擎、校验参数等）；
- selected_rows.append(...) 导致频繁内存分配，pd.concat 在末尾集中合并进一步放大开销。
? 最终建议与验证
- 优先使用方案一：简洁、快速、可读性强，且统计意义上等价于“每个有效簇-标签对独立随机抽 1 行”；
- 若需严格双标签约束，选用方案二，并通过 len(sampled) // 2 验证双标签簇数量；
- 对于超大数据，可考虑 dask.dataframe 或 polars 进一步加速（后者在 groupby sampling 场景下快 3–5×）；
- 始终在采样前执行 df = df.reset_index(drop=True)，避免索引碎片影响性能。
通过向量化重构，您不仅能将运行时间压缩至分钟级，更能获得可维护、可扩展、符合 Pandas 最佳实践的生产就绪代码。
