如何降低numpy where方法处理大数据集时的内存消耗？

问题核心原因 #

• 新增的10列均为object类型字符串，5000万行×10列的字符串会占用极高内存
• data.copy(deep=True)深拷贝直接让内存占用翻倍
• 拆分DataFrame时用globals()存储所有子DataFrame，等于保留了全量数据的多份副本，完全没起到内存优化作用
• 嵌套np.where会生成多个临时布尔数组，额外占用大量内存

具体解决方案 #

1. 用分类数据类型（Categorical）替代字符串 #

因为标记只有Above/Below/Around三个固定值，转为category类型后，内存占用可降至原字符串类型的1/10甚至更低。

# 提前定义分类类型
category_type = pd.CategoricalDtype(categories=['Above', 'Below', 'Around'], ordered=False)

# 改用浅拷贝，若无需保留原数据可直接操作原df
data_percnt = data.copy(deep=False)
b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]

for j in b:
    col_name = f'PCT{j*100}'
    # 计算条件，用np.select替代嵌套where，减少中间变量
    cond1 = data_percnt['PREDICTED'] > data_percnt['TARGET'] * (1 + j)
    cond2 = data_percnt['PREDICTED'] < data_percnt['TARGET'] * (1 - j)
    data_percnt[col_name] = np.select(
        [cond1, cond2],
        ['Above', 'Below'],
        default='Around'
    ).astype(category_type)

2. 优化拆分处理逻辑，避免全量数据驻留内存 #

原拆分方式会把所有子DataFrame存于内存，应处理一个子块就写入磁盘，再释放内存：

import gc

category_type = pd.CategoricalDtype(categories=['Above', 'Below', 'Around'], ordered=False)
b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]

# 按日期分组处理，不保留所有子df
for date in sorted(data['DATA_DATE'].unique()):
    df_chunk = data[data['DATA_DATE'] == date].copy(deep=False)
    
    for j in b:
        col_name = f'PCT{j*100}'
        cond1 = df_chunk['PREDICTED'] > df_chunk['TARGET'] * (1 + j)
        cond2 = df_chunk['PREDICTED'] < df_chunk['TARGET'] * (1 - j)
        df_chunk[col_name] = np.select(
            [cond1, cond2],
            ['Above', 'Below'],
            default='Around'
        ).astype(category_type)
    
    # 处理完直接写入文件（推荐用Parquet格式，比CSV更省空间）
    df_chunk.to_parquet(f'data_processed_{date}.parquet', index=False)
    # 手动释放内存
    del df_chunk
    gc.collect()

3. 避免不必要的深拷贝 #

如果不需要保留原始data DataFrame，直接在原数据上操作，跳过深拷贝步骤：

# 直接操作原数据，省掉一份全量数据的内存占用
data_percnt = data
# 后续处理同上

4. 分块读取数据（若数据来自文件） #

如果数据是从CSV/Parquet等文件读取的，直接用分块读取处理，无需一次性加载全量数据：

import gc

category_type = pd.CategoricalDtype(categories=['Above', 'Below', 'Around'], ordered=False)
b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]

# 分块读取CSV，chunksize根据内存情况调整
for chunk in pd.read_csv('your_raw_data.csv', chunksize=1_000_000):
    for j in b:
        col_name = f'PCT{j*100}'
        cond1 = chunk['PREDICTED'] > chunk['TARGET'] * (1 + j)
        cond2 = chunk['PREDICTED'] < chunk['TARGET'] * (1 - j)
        chunk[col_name] = np.select(
            [cond1, cond2],
            ['Above', 'Below'],
            default='Around'
        ).astype(category_type)
    
    # 追加写入处理后的文件
    chunk.to_csv('processed_data.csv', mode='a', header=False, index=False)
    del chunk
    gc.collect()

5. 提前计算目标值，减少重复计算 #

预先缓存TARGET和PREDICTED列，提前计算倍数阈值，避免循环中重复生成中间数组：

data_percnt = data.copy(deep=False)
b = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1]
category_type = pd.CategoricalDtype(categories=['Above', 'Below', 'Around'], ordered=False)

# 提前缓存列，减少索引开销
target = data_percnt['TARGET']
predicted = data_percnt['PREDICTED']

for j in b:
    col_name = f'PCT{j*100}'
    upper_threshold = target * (1 + j)
    lower_threshold = target * (1 - j)
    cond1 = predicted > upper_threshold
    cond2 = predicted < lower_threshold
    data_percnt[col_name] = np.select(
        [cond1, cond2],
        ['Above', 'Below'],
        default='Around'
    ).astype(category_type)

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