SageMaker ml.t3.large实例中pd.merge大表致Jupyter内核崩溃的优化方案

针对你在ml.t3.large实例上遇到的pandas合并内存溢出问题，结合df1（300万行）和df2（5万行）的规模差异，以下是具体优化方案：

1. 压缩数据类型（最直接的内存节省手段） #

原数据的int64、object、float64类型大多存在压缩空间，可大幅降低内存占用：

• 整数类型：检查ID的取值范围，将int64降级为int32或更小的类型（如int16，若ID最大值不超过32767）
• 浮点类型：将float64转为float32（若精度要求允许）
• 字符串类型：将object类型的Name转为category（若Name的唯一值占比低，比如小于10%）

示例代码：

import pandas as pd

# 压缩df1的类型
df1['ID'] = df1['ID'].astype('int32')
# 遍历浮点列转float32
for col in df1.select_dtypes(include=['float64']).columns:
    df1[col] = df1[col].astype('float32')
# 将Name转为category
df1['Name'] = df1['Name'].astype('category')

# 对df2做同样的压缩
df2['ID'] = df2['ID'].astype('int32')
for col in df2.select_dtypes(include=['float64']).columns:
    df2[col] = df2[col].astype('float32')
df2['Name'] = df2['Name'].astype('category')

# 再次查看内存占用
df1.info(memory_usage="deep")

2. 预处理df2减少重复连接键 #

如果df2中存在重复的ID+Name组合，左连接会导致结果行数暴增，内存占用翻倍。先对df2去重：

# 按连接键去重，保留首次出现的行（或根据业务需求保留特定行）
df2 = df2.drop_duplicates(subset=['ID', 'Name'], keep='first')

3. 分块合并（Chunking） #

将大的df1拆分为多个小块，分别与df2合并后再拼接结果，避免一次性加载所有数据到内存：

import numpy as np
import gc

# 将df1拆分为10个块（可根据内存情况调整数量）
chunks = np.array_split(df1, 10)
merged_chunks = []

for chunk in chunks:
    # 单块与df2合并
    merged_chunk = pd.merge(chunk, df2, on=['ID', 'Name'], how='left')
    merged_chunks.append(merged_chunk)
    # 释放当前块的内存
    del chunk
    gc.collect()

# 拼接所有块
df_merge = pd.concat(merged_chunks, ignore_index=True)

4. 使用Dask替代Pandas处理超内存数据 #

Dask支持并行计算和分块处理，可轻松处理超出单节点内存的数据集：

import dask.dataframe as dd

# 若数据已在内存，转为dask dataframe
df1_dask = dd.from_pandas(df1, npartitions=10)
df2_dask = dd.from_pandas(df2, npartitions=1)

# 执行合并
df_merge_dask = df1_dask.merge(df2_dask, on=['ID', 'Name'], how='left')

# 计算结果（按需转为pandas或导出到文件）
df_merge = df_merge_dask.compute()

5. 合并后及时释放内存 #

合并完成后立即删除原数据集并触发垃圾回收，避免内存占用过高：

# 执行合并
df_merge = pd.merge(df1, df2, on=['ID', 'Name'], how='left')

# 释放原数据内存
del df1, df2
gc.collect()

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