Python如何高效处理百万级百列DataFrame/CSV并划分训练测试集

• 加载阶段优先做裁剪和类型优化：读取CSV时不要默认加载全列，通过usecols参数指定实际需要用到的列；提前通过dtype参数配置列类型，将默认的int64/float64按数值范围下探到int32/int16、float32，低基数字符串列直接指定为category类型，整体内存占用可降低60%-80%。读取引擎优先选pyarrow，或者直接用Polars替代Pandas做加载，读取速度是Pandas默认引擎的3-10倍，内存占用低50%左右。
• 超内存文件用分块/懒加载方案：单文件大小超过可用内存时，用Pandas的chunksize参数按固定行数分块迭代处理，或者用Dask DataFrame做并行懒计算，不需要加载全量数据即可完成过滤、聚合等常规操作。
• 加载后第一时间做内存诊断：跑df.info(memory_usage='deep')逐列排查内存占用，对加载后自动识别为object类型的列做类型转换，避免混合类型、冗余字符串占用内存。
• 计算阶段规避低效率操作：绝对不要用iterrows做逐行遍历，优先用Pandas/Polars自带的矢量化接口做计算；复杂自定义函数可以用swifter自动适配矢量化或并行计算，多核CPU下处理速度能提升数倍。
• 中间结果别存CSV：处理过程中的临时结果、最终结果优先存为带ZSTD压缩的Parquet/Feather格式，读写速度是CSV的10倍以上，还能自动保留列类型，下次加载不需要重新做类型转换。
• 内存及时回收：处理过程中不需要的中间DataFrame直接用del删除，配合gc.collect()主动回收内存，避免无意义的内存占用。

针对你已经生成好的400万行规模数据集，核心优化思路是避免全量数据复制、减少不必要的整表操作，按优先级推荐以下实现：

方案1：Numpy索引切分（Pandas环境下最优，零额外大内存开销） #

不要直接打乱整表、不要用默认参数的train_test_split（默认会复制全量数据，额外占用数GB内存，速度慢），先生成乱序整数索引，再按位置切分，全程只占用几十MB的索引内存：

import numpy as np
import pandas as pd
import gc

# 假设df为你已经生成好的目标DataFrame
total_len = len(df)
test_size = 0.2  # 测试集比例按需求调整
random_seed = 42

# 生成乱序索引，仅占几十MB内存
rng = np.random.default_rng(random_seed)
shuffled_idx = rng.permutation(total_len)
split_pos = int(total_len * (1 - test_size))

# 按位置切分，不触发全量数据复制
train_df = df.iloc[shuffled_idx[:split_pos]]
test_df = df.iloc[shuffled_idx[split_pos:]]

# 切分完成后删除原DataFrame，回收内存
del df, shuffled_idx
gc.collect()

方案2：Polars原生切分（速度最快，适合后续高性能计算场景） #

如果可以换用Polars做数据处理，切分速度是Pandas方案的5-8倍，内存控制更优，切分后的数据可以直接对接LightGBM、XGBoost等框架训练，不需要额外格式转换：

import polars as pl

# 假设df为Polars格式的DataFrame（Pandas转Polars仅需df = pl.from_pandas(pd_df)，零拷贝）
train_df = df.sample(fraction=0.8, seed=42, shuffle=True)
test_df = df.join(train_df, on=df.columns, how="anti")  # 取差集保证无重叠

切分避坑提示 #

• 非必要不要用sklearn的train_test_split处理该规模数据，默认参数下会触发全量数据复制，内存占用翻倍，切分耗时是索引切分的3倍以上；如果必须用sklearn接口，记得传入shuffle=False，提前做好索引打乱再切分。
• 切分好的训练集、测试集直接存为ZSTD压缩的Parquet格式，后续加载仅需数秒，不要存CSV或者pickle格式。
• 如果后续要做多次交叉验证，直接基于索引做切分即可，不需要每次都复制数据。

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