Python DataFrame按分组提取某列最小值对应行的实现方法

问题背景 #

我有一个包含7700万行的遗传变异与基因关联DataFrame，每个基因对应多个变异，需要获取每个基因pval_nominal值最小的行。想知道最佳实现方式，是否需要先生成所有唯一基因列表再筛选？

数据样例 #

EnsemblID   VariantID   tss_distance    CAF_eQTL    ma_samples  ma_count    pval_nominal    Beta    SE  chromosome  ... Source  AverageMaximumPosteriorCall Info    AA_N    AB_N    BB_N    TotalN  MAF MissingDataProportion   HWE_P
0   ENSG00000187634 1:693731    -240519 0.134185    153 168 0.881492    0.005953    0.039917    1   ... HRCr11  0.902046    0.624501    1637.01 451.103 35.569  2124    0.122957    0.000075    0.544841
1   ENSG00000187634 1:714596    -219654 0.038339    48  48  0.834746    -0.015369   0.073635    1   ... HRCr11  0.968278    0.577564    1987.73 135.495 0.677   2124    0.032217    0.000024    0.272504
2   ENSG00000187634 1:715367    -218883 0.038339    48  48  0.834746    -0.015369   0.073635    1   ... HRCr11  0.975207    0.671960    1976.85 146.512 0.556   2124    0.034753    0.000019    0.176551
...
51006816    ENSG00000176383 12:121614338    -591979 0.398562    392 499 0.188786    0.035313    0.026839    12  ... 1000Gp3v5   0.998417    0.997017    782.742 992.043 349.155 2124    0.397929    0.000014    0.257526

期望结果 #

每个基因仅保留pval_nominal最小的一行：

EnsemblID   VariantID   tss_distance    CAF_eQTL    ma_samples  ma_count    pval_nominal    Beta    SE  chromosome  ... Source  AverageMaximumPosteriorCall Info    AA_N    AB_N    BB_N    TotalN  MAF MissingDataProportion   HWE_P
1   ENSG00000187634 1:714596    -219654 0.038339    48  48  0.834746    -0.015369   0.073635    1   ... HRCr11  0.968278    0.577564    1987.73 135.495 0.677   2124    0.032217    0.000024    0.272504
...
51006816    ENSG00000176383 12:121614338    -591979 0.398562    392 499 0.188786    0.035313    0.026839    12  ... 1000Gp3v5   0.998417    0.997017    782.742 992.043 349.155 2124    0.397929    0.000014    0.257526

最优解决方案 #

完全不需要手动生成唯一基因列表，Pandas内置的分组优化方法效率远高于循环遍历，以下是两种最适合大数据集的实现：

方法1：groupby() + idxmin()（内存友好，首选） #

这是处理该场景最高效的方式，idxmin()直接返回每个基因组中p值最小的行索引，再通过索引提取行，无需全量排序：

# 获取每个基因组中pval_nominal最小的行索引
min_pval_indices = df.groupby('EnsemblID')['pval_nominal'].idxmin()
# 根据索引提取目标行
result_df = df.loc[min_pval_indices]

• 优势：时间复杂度接近O(n)，内存占用低，适合7700万行的超大数据集
• 注意：如果同一基因存在多个p值相同的行，idxmin()会返回第一个出现的行，符合多数科研场景需求

方法2：sort_values() + drop_duplicates()（灵活适配额外规则） #

如果需要对p值相同的行增加额外筛选逻辑（比如优先保留tss_distance最小的行），可以先排序再去重：

# 先按基因分组，再按p值升序、距离升序排序，保留每个基因的第一行
result_df = df.sort_values(['EnsemblID', 'pval_nominal', 'tss_distance']).drop_duplicates('EnsemblID', keep='first')

• 注意：排序7700万行会消耗更多内存，建议仅在需要额外排序规则时使用

超大数据集的内存优化建议 #

如果机器内存不足以加载全量数据，推荐使用Dask DataFrame分块处理，语法和Pandas兼容：

import dask.dataframe as dd
# 分块读取大文件
ddf = dd.read_csv('your_large_data.csv')
# 按基因分组筛选p值最小的行
result_ddf = ddf.groupby('EnsemblID').apply(
    lambda x: x[x['pval_nominal'] == x['pval_nominal'].min()],
    meta=ddf.dtypes
)
# 将计算结果转为Pandas DataFrame
result_df = result_ddf.compute()

另外，还可以提前过滤掉不需要的列，只保留EnsemblID、pval_nominal及需要输出的字段，进一步减少内存占用。

为什么不建议手动遍历唯一基因？ #

手动循环每个基因并筛选的方法（比如for gene in df['EnsemblID'].unique(): ...）会重复遍历整个数据集，时间复杂度为O(n*m)（n是行数，m是基因数），效率极低，对于7700万行的数据集来说，运行时间会是内置方法的几十甚至上百倍，完全不可取。

内容的提问来源于stack exchange，提问作者Sander W. van der Laan