Python大规模数据集内存高效处理优化策略问询

针对你处理百万级CSV数据集时遇到的内存过高问题，结合你的代码和使用场景（Python 3.8、Windows），可以从以下几个核心方向优化：

1. 避免缓存全部处理结果（核心问题） #

你的当前代码将所有处理后的数据存入processed_data列表，这直接导致内存被全部结果占满——哪怕用了生成器加载原始数据，最终还是把所有数据加载到了内存。解决方法是让处理函数本身成为生成器，不存储全部结果：

修改后的代码：

def process_large_dataset():
    data_generator = load_large_dataset()
    for batch in data_generator:
        processed_batch = expensive_processing(batch)
        # 逐个返回处理后的元素，而非缓存到列表
        yield from processed_batch

如果需要输出结果，直接边处理边写入磁盘，避免内存堆积：

import csv

with open('processed_output.csv', 'w', newline='') as output_file:
    writer = csv.writer(output_file)
    # 写入表头（如果需要）
    writer.writerow(['col1', 'col2', 'col3'])
    # 迭代处理结果并写入
    for processed_item in process_large_dataset():
        writer.writerow(processed_item)

2. 优化expensive_processing的内存开销 #

（1）使用高效数据类型 #

处理CSV时，默认的数据类型往往冗余，比如用int64存储小范围整数、用object存储重复类别字符串。指定合适的 dtype 能大幅降低单条记录的内存占用：

import pandas as pd

def expensive_processing(batch_df):
    # 将类别列转为category类型，减少重复字符串内存消耗
    batch_df['category_col'] = batch_df['category_col'].astype('category')
    # 将整数列转为更小的数值类型（如int16替代int64，范围足够时）
    batch_df['small_int_col'] = batch_df['small_int_col'].astype('int16')
    # 删除不需要的列，即时释放内存
    batch_df.drop(columns=['unused_col'], inplace=True)
    # 返回迭代器而非完整DataFrame，减少内存持有
    return batch_df.itertuples(index=False, name=None)

（2）避免不必要的数据副本 #

处理数据时尽量使用原地操作（inplace=True），或用视图替代副本。比如Pandas中df[col] = df[col].apply(func)比创建新列更节省内存；NumPy中避免用.copy()，除非必要。

（3）即时清理中间变量 #

在expensive_processing中，处理完的临时变量及时删除，并手动触发垃圾回收（注意不要频繁调用，避免影响速度）：

def expensive_processing(batch):
    temp_data = some_heavy_operation(batch)
    processed = transform(temp_data)
    # 删除临时变量，释放内存
    del temp_data
    import gc
    gc.collect()
    return processed

3. 优化生成器的加载逻辑 #

（1）调整批次大小 #

当前的load_large_dataset如果批次过大，会导致单批次占用内存过高；批次太小则会增加IO次数拖慢速度。建议测试不同批次大小（比如1000、5000、10000条），找到内存和速度的平衡点：

import csv

def load_large_dataset(batch_size=5000):
    with open('large_data.csv', 'r') as f:
        reader = csv.DictReader(f)
        batch = []
        for row in reader:
            batch.append(row)
            if len(batch) == batch_size:
                yield batch
                batch = []
        # 处理最后一批不足batch_size的数据
        if batch:
            yield batch

（2）用流式读取替代批次加载 #

如果数据集极大，甚至单批次都占内存，可以用逐行读取的方式处理，避免批次缓存：

def load_large_dataset():
    with open('large_data.csv', 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            # 逐行返回，而非批次
            yield row

4. 选择内存友好的库 #

（1）Polars替代Pandas #

Polars是专为大数据设计的库，内存效率远高于Pandas，支持流式处理，无需加载全部数据到内存：

import polars as pl

def process_large_dataset():
    # 流式扫描CSV，不加载全部数据
    df = pl.scan_csv('large_data.csv')
    # 定义处理逻辑（过滤、转换、新增列等）
    processed_df = df.with_columns(
        pl.col('value').cast(pl.Int16),
        pl.col('category').cast(pl.Categorical)
    ).filter(pl.col('value') > 0)
    # 将结果写入Parquet（比CSV更省空间，支持列存储）
    processed_df.sink_parquet('processed_data.parquet')

（2）Vaex实现懒加载 #

Vaex支持对TB级数据集进行懒加载，只在需要时处理数据，内存占用几乎可以忽略：

import vaex

# 懒加载CSV，不占内存
df = vaex.read_csv('large_data.csv')
# 处理数据（如计算、过滤），实际执行时才加载必要数据
processed_df = df[df['value'] > 0]
processed_df['new_col'] = df['col1'] + df['col2']
# 写入结果
processed_df.export_csv('processed_output.csv')

（3）正确使用Dask #

如果之前用Dask效果不佳，可能是因为调用了compute()将全部数据加载到内存。正确的做法是用Dask的分块处理和延迟执行，直接将结果写入磁盘：

import dask.dataframe as dd

# 分块读取CSV
df = dd.read_csv('large_data.csv', blocksize='64MB')
# 处理逻辑
processed_df = df.assign(
    new_col=df['col1'] + df['col2']
).drop(columns=['unused_col'])
# 写入结果，不加载全部数据到内存
processed_df.to_csv('processed_output_*.csv')

5. 并行处理的内存优化 #

如果用多进程并行，避免用multiprocessing.Pool.map()（会缓存全部结果），改用imap()迭代获取结果，边处理边写入：

from multiprocessing import Pool

def process_single_batch(batch):
    return expensive_processing(batch)

def process_large_dataset():
    data_generator = load_large_dataset()
    with Pool(processes=4) as pool:
        # 迭代获取处理后的批次，避免缓存全部
        for processed_batch in pool.imap(process_single_batch, data_generator):
            yield from processed_batch

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