Pandas处理10亿行数据按id1分组取最小score前100的内存优化方法

核心原因说明 #

你遇到的内存报错是因为groupby.nsmallest实现逻辑导致的：该方法会在内存中为每个id1分组缓存全量对应行的索引、数值，10亿行数据的临时缓存会占用远超原数据本身的内存，才会触发8.53G的内存分配失败。

以下是从易到难的低内存实现方案：

方案1：分块迭代处理（最推荐，无需额外依赖） #

不需要把全量数据一次性读进内存，逐块处理后再汇总，内存占用可控制在几百MB级别：

import pandas as pd

# 按你的内存大小调整单块行数，1000万行约占100-200MB内存
CHUNK_SIZE = 10_000_000
temp_result = []

# 读文件时直接指定更紧凑的数据类型，进一步降低内存占用
# 3万多唯一id用int32完全足够，float32可满足大部分距离精度需求，内存直接减少50%+
for chunk in pd.read_csv(
    "你的数据文件路径.csv",
    chunksize=CHUNK_SIZE,
    dtype={"permid_1": "int32", "permid_2": "int32", "distance": "float32"}
):
    # 单块内先按id1、距离升序排序，取每个id1的前100条
    chunk_top = chunk.sort_values(["permid_1", "distance"]).groupby("permid_1").head(100)
    temp_result.append(chunk_top)

# 合并所有块的中间结果，总数据量最多只有4万*100=400万行，内存完全无压力
all_temp = pd.concat(temp_result, ignore_index=True)
# 全局再取一次每个id1的前100，得到最终结果
final_df = all_temp.sort_values(["permid_1", "distance"]).groupby("permid_1").head(100)

方案2：先排序再取分组头部（适合数据已经加载进内存的场景） #

比原生nsmallest内存占用低60%以上，只需要一次全局排序，groupby阶段无需额外计算：

# 先按id1升序、距离升序全局排序，再取每个分组前100条
df = df.sort_values(["id1", "score"]).groupby("id1").head(100)

方案3：用Dask并行处理（适合高频处理此类大数据场景） #

Dask语法和Pandas几乎完全一致，自动分块调度，不需要手动处理分块逻辑：

import dask.dataframe as dd

# 自动分块读取数据
ddf = dd.read_csv(
    "你的数据文件路径.csv",
    dtype={"permid_1": "int32", "permid_2": "int32", "distance": "float32"}
)
# 执行和Pandas一致的逻辑
ddf = ddf.sort_values(["permid_1", "distance"]).groupby("permid_1").head(100, npartitions=-1)
# 触发计算得到最终Pandas DataFrame
final_df = ddf.compute()

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