Python按变量类型批量处理大数据集的性能优化求助

你的思路非常准确：批量处理所有数值（NUM）列 + 仅针对类别（CAT）列做必要处理，能从根源上解决逐列读取CSV带来的巨大IO开销——这正是你当前方案耗时超12小时的核心原因。下面是具体实现步骤和代码：

核心优化逻辑 #

原来的逐列读取方案需要重复读取CSV文件3400次，IO操作是性能瓶颈。优化后：

• 仅读取CSV文件1次，减少99.9%的IO开销
• 数值列用Pandas向量化操作批量填充缺失值，效率远高于逐列处理
• 仅对少量类别列做独热编码（且过滤无效列），避免不必要的计算

完整实现代码 #

import pandas as pd
import numpy as np

# 1. 一次性读取完整数据集（彻底避免重复IO）
df_full = pd.read_csv("df_seg_sample.csv")

# 2. 拆分NUM和CAT列，提前分类减少后续判断
num_cols = [col for col, dtype in attribute_dict.items() if dtype == 'NUM']
cat_cols = [col for col, dtype in attribute_dict.items() if dtype == 'CAT']

# 3. 批量处理NUM列：用向量化操作统一填充NA为0
df_num_processed = df_full[num_cols].fillna(0)

# 4. 处理CAT列：仅对unique值≤100的列做独热编码，过滤无效列
df_cat_processed = pd.DataFrame()
for col in cat_cols:
    unique_count = df_full[col].unique().size
    if unique_count <= 100:
        # 生成包含NA的独热编码列
        dummies = pd.get_dummies(df_full[col], dummy_na=True, prefix=col)
        # 移除可能出现的重复列
        dummies = dummies.loc[:, ~dummies.columns.duplicated()]
        # 过滤方差为0的无效列（所有值完全相同的列无建模价值）
        valid_dummies = dummies.loc[:, dummies.var() != 0]
        df_cat_processed = pd.concat([df_cat_processed, valid_dummies], axis=1)

# 5. 合并处理后的数值列和类别列，得到最终结果
df_final = pd.concat([df_num_processed, df_cat_processed], axis=1)

# 可选：保存最终预处理结果
df_final.to_csv("df_final_processed.csv", index=False)

内存紧张时的分块处理补充 #

如果500万行×3400列的数据集超出内存上限，可以用分块读取+逐块处理的方式：

chunk_size = 100000  # 每次处理10万行，可根据本机内存调整
df_final_list = []

for chunk in pd.read_csv("df_seg_sample.csv", chunksize=chunk_size):
    # 对当前数据块执行相同的处理逻辑
    chunk_num = chunk[num_cols].fillna(0)
    chunk_cat = pd.DataFrame()
    for col in cat_cols:
        unique_count = chunk[col].unique().size
        if unique_count <= 100:
            dummies = pd.get_dummies(chunk[col], dummy_na=True, prefix=col)
            dummies = dummies.loc[:, ~dummies.columns.duplicated()]
            valid_dummies = dummies.loc[:, dummies.var() != 0]
            chunk_cat = pd.concat([chunk_cat, valid_dummies], axis=1)
    # 合并当前块的处理结果并加入列表
    chunk_final = pd.concat([chunk_num, chunk_cat], axis=1)
    df_final_list.append(chunk_final)

# 合并所有块的结果得到最终数据集
df_final = pd.concat(df_final_list, axis=0)
df_final.to_csv("df_final_processed.csv", index=False)

预期性能提升 #

这种优化方案能把运行时间从12小时压缩到几十分钟甚至更短：

• IO操作从3400次减少到1次（或分块的几次），直接消除最大性能瓶颈
• 数值列的向量化填充比逐列处理快10~100倍
• 类别列的处理仅针对少数列，计算量大幅降低

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