如何高效在Pandas中按指定列分组并保留size_col最大值对应的行（处理500万行数据集）

问题场景 #

我手头有个500万行、7列的大型DataFrame df1：

>>> df1
col_1 col_2 size_col other_col
0    aaa    abc        4       zxc
1    aaa    abc        3       xcv
2    aaa    abc        1       cvb
3    bbb    bbc        7       vbn
4    bbb    bbc        3       bnm
5    ccc    cbc        1       asd
6    ddd    dbc        9       sdf
7    ccc    cbc        3       dfg
8    ccc    cbc        1       fgh

需求很明确：按col_1和col_2分组，每组只保留size_col值最大的那一行，最终得到df2：

>>> df2
col_1 col_2 size_col other_col
0    aaa    abc        4       zxc
3    bbb    bbc        7       vbn
6    ddd    dbc        9       sdf
7    ccc    cbc        3       dfg

想请教下，针对这种大数据量的场景，怎么处理最高效？

高效解决方案 #

对于500万行这种量级的数据，咱们得优先选性能好、内存占用可控的方法，下面两种都是经过实践验证的靠谱方案：

方法1：排序后去重（首推，性能最优）

思路是先把数据按分组列排序，再按size_col降序排列，这样每组里最大值的行就排在最前面，然后用drop_duplicates保留每组的第一行就行：

# 先按分组列升序排，再按size_col降序排，确保每组最大值行在最前面
df_sorted = df1.sort_values(by=['col_1', 'col_2', 'size_col'], ascending=[True, True, False])
# 按col_1和col_2去重，只保留每组第一行（也就是size_col最大的行）
df2 = df_sorted.drop_duplicates(subset=['col_1', 'col_2'], keep='first')
# 要是需要恢复原数据的索引顺序，加这一行就行（可选）
# df2 = df2.sort_index()

这个方法的优势在于：排序和去重都是pandas底层优化的矢量化操作，处理百万级数据时速度比groupby方法更快，内存压力也更小。

方法2：GroupBy + idxmax（逻辑更直观）

如果觉得排序的思路绕，也可以用分组找最大值索引的方法，逻辑更直接：

# 按col_1和col_2分组，找出每组size_col最大的行的索引
max_row_indices = df1.groupby(['col_1', 'col_2'])['size_col'].idxmax()
# 根据索引提取对应的行
df2 = df1.loc[max_row_indices]

这个方法容易理解，但在处理超大数据量时，groupby的内存开销会比方法1略高一点，不过对于500万行的场景，完全能hold住。

额外提醒 #

• 如果同一组里有多个行的size_col都是最大值，上面两种方法都会保留第一个出现的那一行。要是想保留所有最大值行，可以把drop_duplicates里的keep='first'改成keep='all'，或者用transform配合布尔索引实现。
• 建议用pandas 1.3以上的版本，新版本对大数据量的处理性能有不少优化。
• 要是内存实在紧张，可以考虑分块处理，但500万行7列的数据，只要你内存有16G以上，用上面的方法就能快速跑完。

内容的提问来源于stack exchange，提问作者Naveen Reddy Marthala