导出大型CSV时触发ArrayMemoryError，求内存优化方案

处理大型CSV文件时触发内存错误：
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 4.85 GiB for an array with shape (651219361,) and data type object

导入的CSV约有140万行，处理后数据量会大幅膨胀——比如导入100行时，输出能达到2400行。求不崩溃导出CSV的方法。

程序代码 #

import pandas as pd
import numpy as np
import csv

patents = pd.read_csv('E:/Work/final.csv')

patents = patents[['publn_nr', 'permno_adj', 'publn_year', 'IPC1']].dropna().drop_duplicates().reset_index(drop=True)
patents = patents[(patents['publn_year'] >= 1980) & (patents['publn_year'] < 2016)].reset_index(drop=True)
patents['permno_adj'] = patents['permno_adj'].astype(str) + patents['publn_year'].astype(str)
print(patents.head())

patents = patents.groupby(['permno_adj', 'IPC1'])['publn_nr'].nunique().reset_index()
patents.columns = ['permno_adj', 'IPC1', 'ipc_patents']

patents['total_patents'] = patents.groupby(['permno_adj'])['ipc_patents'].transform('sum')
patents['share'] = patents['ipc_patents'] / patents['total_patents']

for v in ['IPC1']:

    temp = patents.copy()

    T = temp.pivot(index='%s' % v, columns='permno_adj', values='share')
    T = T.fillna(0)

    X_t = temp.pivot(index='permno_adj', columns='%s' % v, values='share')
    X_t = X_t.fillna(0)

    T_t = T.copy()
    for column in list(T_t):
        T_t[column] = T_t[column] / (np.sqrt(np.dot(T_t[[column]].values.T, T_t[[column]].values)[0][0],))

    om_f = X_t.T.dot(X_t)
    om_s = X_t.T.dot(X_t)

    for sic1 in list(om_s):
        for sic2 in list(om_s):
            om_s.loc[sic1][sic2] = om_s.loc[sic1][sic2] / (np.sqrt(om_f[sic1][sic1]) * np.sqrt(om_f[sic2][sic2]))

    tech = T_t.T.dot(om_s).dot(T_t)

    tech = tech.unstack().reset_index(level=1)

    if 'IPC1' in v:
        tech.columns = ['permno_adj_pat', 'tech_mahal_sim']

    tech = tech.reset_index()

    tech = tech[(tech['permno_adj'] != tech['permno_adj_pat'])].sort_values(
        ['permno_adj', 'permno_adj_pat']).reset_index(drop=True)

    if 'IPC1' in v:
        tech.to_csv('gajuf.csv')
        print(done)

完整错误信息 #

Traceback (most recent call last):
  File "E:\Work\SMS London PPT by Gaju Masare\python work\sic_tech_only.py", line 41, in <module>
    tech = tech.unstack().reset_index(level=1)
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\util\_decorators.py", line 317, in wrapper
    return func(*args, **kwargs)
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\core\series.py", line 1584, in reset_index
    return df.reset_index(
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\util\_decorators.py", line 317, in wrapper
    return func(*args, **kwargs)
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\core\frame.py", line 6349, in reset_index
    level_values = algorithms.take(
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\core\algorithms.py", line 1565, in take
    result = take_nd(
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\core\array_algos\take.py", line 117, in take_nd
    return _take_nd_ndarray(arr, indexer, axis, fill_value, allow_fill)
  File "C:\Users\Gaju\PycharmProjects\first project\venv\lib\site-packages\pandas\core\array_algos\take.py", line 156, in _take_nd_ndarray
    out = np.empty(out_shape, dtype=dtype, order="F")
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 4.85 GiB for an array with shape (651219361,) and data type object

输出格式示例 #

1. 压缩内存占用，避免超大中间数组 #

错误出现在tech.unstack()步骤，本质是tech矩阵规模过大（6.5亿元素），直接unstack会瞬间占满内存。可以从这些环节优化：

（1）降低数值精度

如果业务允许，把float64类型改成float32，直接减少一半内存占用：

# 处理share列时指定类型
patents['share'] = (patents['ipc_patents'] / patents['total_patents']).astype('float32')
# 矩阵运算后强制转成float32
om_f = X_t.T.dot(X_t).astype('float32')
om_s = om_f.copy()  # 去掉重复计算，直接复制
tech = T_t.T.dot(om_s).dot(T_t).astype('float32')

（2）分批次计算+增量写入

如果permno_adj的唯一值数量极多，全量生成N×N的tech矩阵必然内存溢出。可以把permno_adj拆分成小批次，逐个计算并追加写入CSV：

# 获取所有唯一permno_adj值
unique_permnos = patents['permno_adj'].unique()
batch_size = 1000  # 根据自身内存调整批次大小

# 先写入CSV表头
with open('gajuf.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    writer.writerow(['permno_adj', 'permno_adj_pat', 'tech_mahal_sim'])

# 分批次处理
for i in range(0, len(unique_permnos), batch_size):
    batch_permnos = unique_permnos[i:i+batch_size]
    # 只取当前批次的T_t子集
    T_t_batch = T_t[batch_permnos]
    # 计算当前批次的结果
    tech_batch = T_t_batch.T.dot(om_s).dot(T_t)
    # 转换格式并过滤自身匹配
    tech_batch = tech_batch.unstack().reset_index(level=1)
    tech_batch.columns = ['permno_adj_pat', 'tech_mahal_sim']
    tech_batch = tech_batch.reset_index()
    tech_batch = tech_batch[tech_batch['permno_adj'] != tech_batch['permno_adj_pat']]
    # 追加写入CSV
    tech_batch.to_csv('gajuf.csv', mode='a', header=False, index=False)
    print(f'完成批次 {i//batch_size + 1}')

2. 用Dask替代Pandas处理超大数据 #

Dask会自动把数据拆分成小块，在磁盘上并行处理，不用全量加载到内存。核心修改：

• 用dask.dataframe替代pandas读取和预处理数据
• 矩阵运算部分结合Dask数组模块，或者沿用上面的分块逻辑

3. 清理冗余计算 #

原代码中om_f = X_t.T.dot(X_t)和om_s = X_t.T.dot(X_t)是重复计算，直接改成om_s = om_f.copy()，节省一次矩阵运算的时间和内存。

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